tcp,udp: realips lack port to find active conns

Author: ignoramous

intra/common.go

@@ -193,12 +193,12 @@ func undoAlg(r dnsx.Resolver, algip netip.Addr) (realips, domains, probableDomai
 	return
 }
 
-func hasActiveConn(cm core.ConnMapper, ip, ips string) bool {
+func hasActiveConn(cm core.ConnMapper, ipp, ips, port string) bool {
 	if cm == nil {
 		return false
 	}
 	// TODO: filter by protocol (tcp/udp) when finding conns
-	return !hasSelfUid(cm.Find(ip), true) || !hasSelfUid(cm.FindAll(ips), true)
+	return !hasSelfUid(cm.Find(ipp), true) || !hasSelfUid(cm.FindAll(ips, port), true)
 }
 
 // returns proxy-id, conn-id, user-id

intra/core/connmap.go

@@ -21,7 +21,7 @@ type ConnMapper interface {
 	Clear() []string
 	Track(t ConnTuple, x ...net.Conn) int
 	Find(dst string) (t []ConnTuple)
-	FindAll(csvdst string) (t []ConnTuple)
+	FindAll(csvips, port string) (t []ConnTuple)
 	Get(cid string) []net.Conn
 	Untrack(cid string) int
 	UntrackBatch(cids []string) []string
@@ -72,6 +72,7 @@ func (h *cm) trackDstLocked(t ConnTuple, conns []net.Conn) {
 		}
 		dst := raddr.String()
 		if tups, ok := h.dsttracker[dst]; ok {
+			// TODO: do not add dup tuples (cid)
 			h.dsttracker[dst] = append(tups, t)
 		} else {
 			h.dsttracker[dst] = []ConnTuple{t}
@@ -158,18 +159,19 @@ func (h *cm) Find(dst string) (tups []ConnTuple) {
 	return
 }
 
-func (h *cm) FindAll(csvdst string) (out []ConnTuple) {
+func (h *cm) FindAll(csvips, port string) (out []ConnTuple) {
 	out = make([]ConnTuple, 0)
 
-	if len(csvdst) == 0 {
+	if len(csvips) == 0 {
 		return
 	}
 
 	h.RLock()
 	defer h.RUnlock()
 
-	dsts := strings.Split(csvdst, ",")
+	dsts := strings.Split(csvips, ",")
 	for _, dst := range dsts {
+		dst = net.JoinHostPort(dst, port)
 		if tups, ok := h.dsttracker[dst]; ok {
 			out = append(out, tups...)
 		}

intra/tcp.go

@@ -30,6 +30,7 @@ import (
 	"fmt"
 	"net"
 	"net/netip"
+	"strconv"
 	"time"
 
 	"github.com/celzero/firestack/intra/dnsx"
@@ -117,7 +118,8 @@ func (h *tcpHandler) onFlow(localaddr, target netip.AddrPort, realips, domains,
 	var proto int32 = 6 // tcp
 	src := localaddr.String()
 	dst := target.String()
-	dup := hasActiveConn(h.conntracker, dst, realips)
+	dport := strconv.Itoa(int(target.Port()))
+	dup := hasActiveConn(h.conntracker, dst, realips, dport)
 	res := h.listener.Flow(proto, uid, dup, src, dst, realips, domains, probableDomains, blocklists)
 
 	if res == nil {
@@ -278,7 +280,7 @@ func (h *tcpHandler) handle(px ipn.Proxy, src net.Conn, target netip.AddrPort, s
 			if r := recover(); r != nil {
 				log.W("tcp: forward: panic %v", r)
 			}
-			defer h.conntracker.Untrack(ct.CID)
+			h.conntracker.Untrack(ct.CID)
 		}()
 		forward(src, dst, l, smm) // src always *gonet.TCPConn
 	}()

intra/udp.go

@@ -30,6 +30,7 @@ import (
 	"io"
 	"net"
 	"net/netip"
+	"strconv"
 	"time"
 
 	"github.com/celzero/firestack/intra/dnsx"
@@ -118,9 +119,10 @@ func (h *udpHandler) onFlow(localaddr, target netip.AddrPort, realips, domains,
 	}
 
 	src := localaddr.String()
-	dst := "" // unconnected udp sockets may not have a valid target
+	dst, dport := "", "" // unconnected udp sockets may not have a valid target
 	if target.IsValid() {
 		dst = target.String()
+		dport = strconv.Itoa(int(target.Port()))
 	}
 	if len(realips) <= 0 || len(domains) <= 0 {
 		log.V("udp: onFlow: no realips(%s) or domains(%s + %s), for src=%s dst=%s", realips, domains, probableDomains, localaddr, dst)
@@ -136,7 +138,7 @@ func (h *udpHandler) onFlow(localaddr, target netip.AddrPort, realips, domains,
 	}
 
 	var proto int32 = 17 // udp
-	dup := hasActiveConn(h.conntracker, dst, realips)
+	dup := hasActiveConn(h.conntracker, dst, realips, dport)
 	res := h.listener.Flow(proto, uid, dup, src, dst, realips, domains, probableDomains, blocklists)
 
 	if res == nil {