tcp,udp: track conn in same rountine to prevent race

Author: ignoramous

intra/common.go

@@ -66,13 +66,8 @@ func download(cid string, local net.Conn, remote net.Conn) (n int64, err error)
 
 // forward copies data between local and remote, and tracks the connection.
 // It also sends a summary to the listener when done. Always called in a goroutine.
-func forward(local net.Conn, remote net.Conn, t core.ConnMapper, l SocketListener, smm *SocketSummary) {
+func forward(local, remote net.Conn, l SocketListener, smm *SocketSummary) {
 	cid := smm.ID
-	uid := smm.UID
-	ct := core.ConnTuple{CID: cid, UID: uid}
-
-	t.Track(ct, local, remote)
-	defer t.Untrack(cid)
 
 	uploadch := make(chan ioinfo)
 

intra/tcp.go

@@ -269,15 +269,18 @@ func (h *tcpHandler) handle(px ipn.Proxy, src net.Conn, target netip.AddrPort, s
 		return err
 	}
 
+	ct := core.ConnTuple{CID: smm.ID, UID: smm.UID}
+
+	h.conntracker.Track(ct, src, dst)
 	go func() {
-		cm := h.conntracker
 		l := h.listener
 		defer func() {
 			if r := recover(); r != nil {
 				log.W("tcp: forward: panic %v", r)
 			}
+			defer h.conntracker.Untrack(ct.CID)
 		}()
-		forward(src, dst, cm, l, smm) // src always *gonet.TCPConn
+		forward(src, dst, l, smm) // src always *gonet.TCPConn
 	}()
 
 	log.I("tcp: new conn %s via proxy(%s); src(%s) -> dst(%s) for %s", smm.ID, px.ID(), src.LocalAddr(), target, smm.UID)

intra/udp.go

@@ -240,15 +240,17 @@ func (h *udpHandler) proxy(gconn net.Conn, src, dst netip.AddrPort) (ok bool) {
 		// no summary for dns queries
 		return true // ok
 	}
+
+	ct := core.ConnTuple{CID: smm.ID, UID: smm.UID}
+	h.conntracker.Track(ct, gconn, remote)
 	go func() {
-		cm := h.conntracker
 		defer func() {
 			if r := recover(); r != nil {
 				log.W("udp: forward: %s -> %s panic %v", src, dst, r)
 			}
+			h.conntracker.Untrack(ct.CID)
 		}()
-
-		forward(gconn, &rwext{remote}, cm, l, smm)
+		forward(gconn, &rwext{remote}, l, smm)
 	}()
 	return true // ok
 }