84 lines
1.6 KiB
Go
84 lines
1.6 KiB
Go
package rpctest
|
|
|
|
import (
|
|
"context"
|
|
"github.com/libp2p/go-libp2p/core/crypto"
|
|
"github.com/libp2p/go-libp2p/core/peer"
|
|
"net"
|
|
"storj.io/drpc"
|
|
"storj.io/drpc/drpcconn"
|
|
"storj.io/drpc/drpcmux"
|
|
"storj.io/drpc/drpcserver"
|
|
)
|
|
|
|
type SecConnMock struct {
|
|
net.Conn
|
|
localPrivKey crypto.PrivKey
|
|
remotePubKey crypto.PubKey
|
|
localId peer.ID
|
|
remoteId peer.ID
|
|
}
|
|
|
|
func (s *SecConnMock) LocalPeer() peer.ID {
|
|
return s.localId
|
|
}
|
|
|
|
func (s *SecConnMock) LocalPrivateKey() crypto.PrivKey {
|
|
return s.localPrivKey
|
|
}
|
|
|
|
func (s *SecConnMock) RemotePeer() peer.ID {
|
|
return s.remoteId
|
|
}
|
|
|
|
func (s *SecConnMock) RemotePublicKey() crypto.PubKey {
|
|
return s.remotePubKey
|
|
}
|
|
|
|
type ConnWrapper func(conn net.Conn) net.Conn
|
|
|
|
func NewSecConnWrapper(
|
|
localPrivKey crypto.PrivKey,
|
|
remotePubKey crypto.PubKey,
|
|
localId peer.ID,
|
|
remoteId peer.ID) ConnWrapper {
|
|
return func(conn net.Conn) net.Conn {
|
|
return &SecConnMock{
|
|
Conn: conn,
|
|
localPrivKey: localPrivKey,
|
|
remotePubKey: remotePubKey,
|
|
localId: localId,
|
|
remoteId: remoteId,
|
|
}
|
|
}
|
|
}
|
|
|
|
func NewTestServer() *TesServer {
|
|
ts := &TesServer{
|
|
Mux: drpcmux.New(),
|
|
}
|
|
ts.Server = drpcserver.New(ts.Mux)
|
|
return ts
|
|
}
|
|
|
|
type TesServer struct {
|
|
*drpcmux.Mux
|
|
*drpcserver.Server
|
|
}
|
|
|
|
func (ts *TesServer) Dial() drpc.Conn {
|
|
return ts.DialWrapConn(nil, nil)
|
|
}
|
|
|
|
func (ts *TesServer) DialWrapConn(serverWrapper ConnWrapper, clientWrapper ConnWrapper) drpc.Conn {
|
|
sc, cc := net.Pipe()
|
|
if serverWrapper != nil {
|
|
sc = serverWrapper(sc)
|
|
}
|
|
if clientWrapper != nil {
|
|
cc = clientWrapper(cc)
|
|
}
|
|
go ts.Server.ServeOne(context.Background(), sc)
|
|
return drpcconn.New(cc)
|
|
}
|