Add stream reuse fixes and performance improvements

This commit is contained in:
Alec
2022-01-30 21:46:21 -07:00
parent c750bc1d92
commit 8da27edd39

View File

@@ -30,7 +30,7 @@ var (
tunDev *tun.TUN tunDev *tun.TUN
// RevLookup allow quick lookups of an incoming stream // RevLookup allow quick lookups of an incoming stream
// for security before accepting or responding to any data. // for security before accepting or responding to any data.
RevLookup map[string]bool RevLookup map[string]string
// activeStreams is a map of active streams to a peer // activeStreams is a map of active streams to a peer
activeStreams map[string]network.Stream activeStreams map[string]network.Stream
) )
@@ -91,9 +91,9 @@ func UpRun(r *cmd.Root, c *cmd.Sub) {
} }
// Setup reverse lookup hash map for authentication. // Setup reverse lookup hash map for authentication.
RevLookup = make(map[string]bool, len(cfg.Peers)) RevLookup = make(map[string]string, len(cfg.Peers))
for _, id := range cfg.Peers { for ip, id := range cfg.Peers {
RevLookup[id.ID] = true RevLookup[id.ID] = ip
} }
fmt.Println("[+] Creating TUN Device") fmt.Println("[+] Creating TUN Device")
@@ -200,9 +200,10 @@ func UpRun(r *cmd.Root, c *cmd.Sub) {
continue continue
} }
stream.Close() stream.Close()
delete(activeStreams, dst)
ok = false ok = false
} }
if !ok { if _, ok := peerTable[dst]; ok {
stream, err = host.NewStream(ctx, peerTable[dst], p2p.Protocol) stream, err = host.NewStream(ctx, peerTable[dst], p2p.Protocol)
if err != nil { if err != nil {
log.Println(err) log.Println(err)
@@ -257,8 +258,8 @@ func streamHandler(stream network.Stream) {
for { for {
plen, err = stream.Read(packet) plen, err = stream.Read(packet)
if err != nil { if err != nil {
log.Println(err)
stream.Close() stream.Close()
delete(activeStreams, RevLookup[stream.Conn().RemotePeer().Pretty()])
return return
} }
tunDev.Iface.Write(packet[:plen]) tunDev.Iface.Write(packet[:plen])