Make packet streams bidirectional (#121)

node/node.go

@@ -40,17 +40,18 @@ type SharedStream struct {
 }
 
 type Node struct {
-	cfg           *config.Config
-	p2p           host.Host
-	dht           *dht.IpfsDHT
-	tunDev        *tun.TUN
-	activeStreams map[peer.ID]SharedStream
-	ctx           context.Context
-	cancel        func()
-	lockPath      string
-	configPath    string
-	interfaceName string
-	wg            *sync.WaitGroup
+	cfg               *config.Config
+	p2p               host.Host
+	dht               *dht.IpfsDHT
+	tunDev            *tun.TUN
+	activeStreams     map[peer.ID]SharedStream
+	activeStreamsLock sync.RWMutex
+	ctx               context.Context
+	cancel            func()
+	lockPath          string
+	configPath        string
+	interfaceName     string
+	wg                *sync.WaitGroup
 }
 
 func New(ctx context.Context, configPath string, ifName string) Node {
@@ -60,7 +61,6 @@ func New(ctx context.Context, configPath string, ifName string) Node {
 		cfg:           &config.Config{},
 		p2p:           nil,
 		tunDev:        &tun.TUN{},
-		activeStreams: map[peer.ID]SharedStream{},
 		ctx:           innerCtx,
 		cancel:        ctxCancel,
 		configPath:    configPath,
@@ -311,12 +311,51 @@ func (node *Node) Run() error {
 	return nil
 }
 
+func (node *Node) getActiveStream(pid peer.ID) (SharedStream, bool) {
+	node.activeStreamsLock.RLock()
+	defer node.activeStreamsLock.RUnlock()
+	s, ok := node.activeStreams[pid]
+	return s, ok
+}
+
+func (node *Node) insertActiveStream(pid peer.ID, ss SharedStream) bool {
+	node.activeStreamsLock.Lock()
+	defer node.activeStreamsLock.Unlock()
+	if _, exists := node.activeStreams[pid]; exists {
+		return false
+	}
+	node.activeStreams[pid] = ss
+	return true
+}
+
+func (node *Node) expireActiveStream(pid peer.ID) {
+	node.activeStreamsLock.Lock()
+	defer node.activeStreamsLock.Unlock()
+	delete(node.activeStreams, pid)
+}
+
 func (node *Node) streamHandler(stream network.Stream) {
+	remotePeerID := stream.Conn().RemotePeer()
+
 	// If the remote node ID isn't in the list of known nodes don't respond.
-	if _, ok := config.FindPeer(node.cfg.Peers, stream.Conn().RemotePeer()); !ok {
+	if _, ok := config.FindPeer(node.cfg.Peers, remotePeerID); !ok {
 		stream.Reset()
 		return
 	}
+
+	var streamLock = new(sync.Mutex)
+
+	// Version 0 nodes don't read from this stream, so we can't reuse it.
+	if stream.Protocol() != p2p.ProtocolV0 {
+		inserted := node.insertActiveStream(remotePeerID, SharedStream{
+			Stream: &stream,
+			Lock:   streamLock,
+		})
+		if inserted {
+			defer node.expireActiveStream(remotePeerID)
+		}
+	}
+
 	var packet = make([]byte, 1420)
 	var packetSize = make([]byte, 2)
 	for {
@@ -352,7 +391,7 @@ func (node *Node) streamHandler(stream network.Stream) {
 
 func (node *Node) sendPacket(dst peer.ID, packet []byte, plen int) {
 	// Check if we already have an open connection to the destination peer.
-	ms, ok := node.activeStreams[dst]
+	ms, ok := node.getActiveStream(dst)
 	if ok {
 		if func() bool {
 			ms.Lock.Lock()
@@ -374,14 +413,14 @@ func (node *Node) sendPacket(dst peer.ID, packet []byte, plen int) {
 			// If we encounter an error when writing to a stream we should
 			// close that stream and delete it from the active stream map.
 			(*ms.Stream).Close()
-			delete(node.activeStreams, dst)
+			node.expireActiveStream(dst)
 			return false
 		}() {
 			return
 		}
 	}
 
-	stream, err := node.p2p.NewStream(node.ctx, dst, p2p.Protocol)
+	stream, err := node.p2p.NewStream(node.ctx, dst, p2p.Protocols...)
 	if err != nil {
 		logger.With(zap.String("destination", dst.String()), zap.Error(err)).Error("Failed to open stream")
 		go p2p.Rediscover()
@@ -406,12 +445,7 @@ func (node *Node) sendPacket(dst peer.ID, packet []byte, plen int) {
 		return
 	}
 
-	// If all succeeds when writing the packet to the stream
-	// we should reuse this stream by adding it active streams map.
-	node.activeStreams[dst] = SharedStream{
-		Stream: &stream,
-		Lock:   &sync.Mutex{},
-	}
+	go node.streamHandler(stream)
 }
 
 func (node *Node) eventLogger(ctx context.Context, host host.Host) error {

p2p/node.go

@@ -21,6 +21,7 @@ import (
 	"github.com/libp2p/go-libp2p/core/network"
 	"github.com/libp2p/go-libp2p/core/peer"
 	"github.com/libp2p/go-libp2p/core/pnet"
+	"github.com/libp2p/go-libp2p/core/protocol"
 	"github.com/libp2p/go-libp2p/core/routing"
 	"github.com/libp2p/go-libp2p/p2p/discovery/backoff"
 	"github.com/libp2p/go-libp2p/p2p/host/autorelay"
@@ -48,8 +49,16 @@ func (c *httpRoutingWrapper) Bootstrap(ctx context.Context) error {
 	return nil
 }
 
-// Protocol is a descriptor for the Hyprspace P2P Protocol.
-const Protocol = "/hyprspace/0.0.1"
+// Version 0
+const ProtocolV0 = "/hyprspace/0.0.1"
+
+// Version 1: Bidirectional streams
+const ProtocolV1 = "/hyprspace/1"
+
+var Protocols = []protocol.ID{
+	ProtocolV1,
+	ProtocolV0,
+}
 
 func getExtraPeers(addr ma.Multiaddr) (nodesList []string) {
 	nodesList = []string{}
@@ -281,7 +290,9 @@ func CreateNode(ctx context.Context, privateKey crypto.PrivKey, listenAddreses [
 	node = routedhost.Wrap(basicHost, pr)
 
 	// Setup Hyprspace Stream Handler
-	node.SetStreamHandler(Protocol, handler)
+	for _, proto := range Protocols {
+		node.SetStreamHandler(proto, handler)
+	}
 
 	if err != nil {
 		return node, nil, err

rpc/server.go

@@ -7,6 +7,7 @@ import (
 	"net"
 	"net/rpc"
 	"os"
+	"slices"
 	"sync"
 	"syscall"
 
@@ -82,7 +83,7 @@ func (hsr *HyprspaceRPC) Route(args *RouteArgs, reply *RouteReply) error {
 			ConnLoop:
 				for _, c := range hsr.host.Network().ConnsToPeer(rte.Target.ID) {
 					for _, s := range c.GetStreams() {
-						if s.Protocol() == p2p.Protocol {
+						if slices.Contains(p2p.Protocols, s.Protocol()) {
 							if _, err := c.RemoteMultiaddr().ValueForProtocol(multiaddr.P_CIRCUIT); err == nil {
 								relay = true
 								if ra, err := c.RemoteMultiaddr().ValueForProtocol(multiaddr.P_P2P); err == nil {