From 55407d5bb0af7f6061d6f59ad1fb72483c52de63 Mon Sep 17 00:00:00 2001
From: Paul Wells <paulwe@gmail.com>
Date: Fri, 19 Apr 2024 16:28:08 -0700
Subject: [PATCH] Prevent heap allocations in FiveTuple Fingerprint

---
 internal/allocation/allocation_manager.go |  4 +--
 internal/allocation/five_tuple.go         | 30 ++++++++++++++++++++---
 2 files changed, 29 insertions(+), 5 deletions(-)

diff --git a/internal/allocation/allocation_manager.go b/internal/allocation/allocation_manager.go
index 56126ea3..2b765921 100644
--- a/internal/allocation/allocation_manager.go
+++ b/internal/allocation/allocation_manager.go
@@ -30,7 +30,7 @@ type Manager struct {
 	lock sync.RWMutex
 	log  logging.LeveledLogger
 
-	allocations  map[string]*Allocation
+	allocations  map[FiveTupleFingerprint]*Allocation
 	reservations []*reservation
 
 	allocatePacketConn func(network string, requestedPort int) (net.PacketConn, net.Addr, error)
@@ -51,7 +51,7 @@ func NewManager(config ManagerConfig) (*Manager, error) {
 
 	return &Manager{
 		log:                config.LeveledLogger,
-		allocations:        make(map[string]*Allocation, 64),
+		allocations:        make(map[FiveTupleFingerprint]*Allocation, 64),
 		allocatePacketConn: config.AllocatePacketConn,
 		allocateConn:       config.AllocateConn,
 		permissionHandler:  config.PermissionHandler,
diff --git a/internal/allocation/five_tuple.go b/internal/allocation/five_tuple.go
index 99264623..6d812caf 100644
--- a/internal/allocation/five_tuple.go
+++ b/internal/allocation/five_tuple.go
@@ -4,7 +4,6 @@
 package allocation
 
 import (
-	"fmt"
 	"net"
 )
 
@@ -33,7 +32,32 @@ func (f *FiveTuple) Equal(b *FiveTuple) bool {
 	return f.Fingerprint() == b.Fingerprint()
 }
 
+// FiveTupleFingerprint is a comparable representation of a FiveTuple
+type FiveTupleFingerprint struct {
+	srcIP, dstIP     [16]byte
+	srcPort, dstPort uint16
+	protocol         Protocol
+}
+
 // Fingerprint is the identity of a FiveTuple
-func (f *FiveTuple) Fingerprint() string {
-	return fmt.Sprintf("%d_%s_%s", f.Protocol, f.SrcAddr.String(), f.DstAddr.String())
+func (f *FiveTuple) Fingerprint() (fp FiveTupleFingerprint) {
+	srcIP, srcPort := netAddrIPAndPort(f.SrcAddr)
+	copy(fp.srcIP[:], srcIP)
+	fp.srcPort = srcPort
+	dstIP, dstPort := netAddrIPAndPort(f.DstAddr)
+	copy(fp.dstIP[:], dstIP)
+	fp.dstPort = dstPort
+	fp.protocol = f.Protocol
+	return
+}
+
+func netAddrIPAndPort(addr net.Addr) (net.IP, uint16) {
+	switch a := addr.(type) {
+	case *net.UDPAddr:
+		return a.IP.To16(), uint16(a.Port)
+	case *net.TCPAddr:
+		return a.IP.To16(), uint16(a.Port)
+	default:
+		return nil, 0
+	}
 }