Different matching strategies for TCP and UDP sockets

policy.go

@@ -181,17 +181,15 @@ func (fw *Firewall) filterPacket(pkt *nfqueue.Packet) {
 }
 
 func findProcessForPacket(pkt *nfqueue.Packet) *proc.ProcInfo {
-	proto := ""
 	switch pkt.Protocol {
 	case nfqueue.TCP:
-		proto = "tcp"
+		return proc.LookupTCPSocketProcess(pkt.SrcPort, pkt.Dst, pkt.DstPort)
 	case nfqueue.UDP:
-		proto = "udp"
+		return proc.LookupUDPSocketProcess(pkt.SrcPort)
 	default:
 		log.Warning("Packet has unknown protocol: %d", pkt.Protocol)
 		return nil
 	}
-	return proc.LookupSocketProcess(proto, pkt.SrcPort, pkt.Dst, pkt.DstPort)
 }
 
 func basicAllowPacket(pkt *nfqueue.Packet) bool {

proc/proc.go

@@ -9,7 +9,6 @@ import (
 	"strconv"
 	"strings"
 	"github.com/subgraph/fw-daemon/Godeps/_workspace/src/github.com/op/go-logging"
-	"github.com/subgraph/fw-daemon/nfqueue"
 )
 
 var log = logging.MustGetLogger("proc")
@@ -19,8 +18,17 @@ func SetLogger(logger *logging.Logger) {
 
 var pcache = &pidCache{}
 
-func LookupSocketProcess(proto string, srcPort uint16, dstAddr net.IP, dstPort uint16) *ProcInfo {
-	ss := findSocket(proto, srcPort, dstAddr, dstPort)
+
+func LookupUDPSocketProcess(srcPort uint16) *ProcInfo {
+	ss := findUDPSocket(srcPort)
+	if ss == nil {
+		return nil
+	}
+	return pcache.lookup(ss.inode)
+}
+
+func LookupTCPSocketProcess(srcPort uint16, dstAddr net.IP, dstPort uint16) *ProcInfo {
+	ss := findTCPSocket(srcPort, dstAddr, dstPort)
 	if ss == nil {
 		return nil
 	}
@@ -56,20 +64,6 @@ func (sa *socketAddr) parse(s string) error {
 }
 
 
-func printPacket(pkt *nfqueue.Packet) string {
-	proto := func() string {
-		switch pkt.Protocol {
-		case nfqueue.TCP:
-			return "TCP"
-		case nfqueue.UDP:
-			return "UDP"
-		default:
-			return "???"
-		}
-	}()
-	return fmt.Sprintf("(%s %s:%d --> %s:%d)", proto, pkt.Src, pkt.SrcPort, pkt.Dst.String(), pkt.DstPort)
-}
-
 func ParseIp(ip string) (net.IP, error) {
 	var result net.IP
 	dst, err := hex.DecodeString(ip)

proc/socket.go

@@ -29,15 +29,19 @@ func (ss *socketStatus) String() string {
 	return fmt.Sprintf("%s -> %s uid=%d inode=%d", ss.local, ss.remote, ss.uid, ss.inode)
 }
 
-func findUDPSocket(srcPort uint16, dstAddr net.IP, dstPort uint16) *socketStatus {
-	return findSocket("udp", srcPort, dstAddr, dstPort)
+func findUDPSocket(srcPort uint16) *socketStatus {
+	return findSocket("udp", func(ss socketStatus) bool {
+		return ss.local.port == srcPort
+	})
 }
 
 func findTCPSocket(srcPort uint16, dstAddr net.IP, dstPort uint16) *socketStatus {
-	return findSocket("tcp", srcPort, dstAddr, dstPort)
+	return findSocket("tcp", func(ss socketStatus) bool {
+		return ss.remote.port == dstPort && ss.remote.ip.Equal(dstAddr) && ss.local.port == srcPort
+	})
 }
 
-func findSocket(proto string, srcPort uint16, dstAddr net.IP, dstPort uint16) *socketStatus {
+func findSocket(proto string, matcher func(socketStatus) bool) *socketStatus {
 	var ss socketStatus
 	for _,line := range getSocketLines(proto) {
 		if len(line) == 0 {
@@ -47,7 +51,7 @@ func findSocket(proto string, srcPort uint16, dstAddr net.IP, dstPort uint16) *s
 			log.Warning("Unable to parse line from /proc/net/%s [%s]: %v", proto, line, err)
 			continue
 		}
-		if ss.remote.port == dstPort && ss.remote.ip.Equal(dstAddr) && ss.local.port == srcPort {
+		if matcher(ss) {
 			ss.line = line
 			return &ss
 		}