Use netip.Addr instead of net.IP

Using the net/netip package instead of the net package can help reduce the memory footprint of the library and help reduce the number of heap allocations. This is a breaking change for consumers of the library as exported types are updated to use fields of type netip.Addr instead of net.IP. We also remove the depedency on go-cmp and use assert.Equal instead in tests. Fixes #35 Signed-off-by: Antonin Bas <[email protected]>
ti-mo · Oct 16, 2023 · 7ac9c2a · 7ac9c2a
1 parent 59e136f
commit 7ac9c2a
Show file tree

Hide file tree

Showing 15 changed files with 141 additions and 209 deletions.
diff --git a/conn_test.go b/conn_test.go
@@ -3,7 +3,7 @@ package conntrack_test
 import (
 	"fmt"
 	"log"
-	"net"
+	"net/netip"
 	"testing"
 
 	"github.com/mdlayher/netlink"
@@ -34,8 +34,8 @@ func ExampleConn_createUpdateFlow() {
 	// Set up a new Flow object using a given set of attributes.
 	f := conntrack.NewFlow(
 		17, 0,
-		net.ParseIP("2a00:1450:400e:804::200e"),
-		net.ParseIP("2a00:1450:400e:804::200f"),
+		netip.MustParseAddr("2a00:1450:400e:804::200e"),
+		netip.MustParseAddr("2a00:1450:400e:804::200f"),
 		1234, 80, 120, 0,
 	)
 
@@ -72,12 +72,12 @@ func ExampleConn_dumpFilter() {
 	}
 
 	f1 := conntrack.NewFlow(
-		6, 0, net.IPv4(1, 2, 3, 4), net.IPv4(5, 6, 7, 8),
+		6, 0, netip.MustParseAddr("1.2.3.4"), netip.MustParseAddr("5.6.7.8"),
 		1234, 80, 120, 0x00ff, // Set a connection mark
 	)
 
 	f2 := conntrack.NewFlow(
-		17, 0, net.ParseIP("2a00:1450:400e:804::200e"), net.ParseIP("2a00:1450:400e:804::200f"),
+		17, 0, netip.MustParseAddr("2a00:1450:400e:804::200e"), netip.MustParseAddr("2a00:1450:400e:804::200f"),
 		1234, 80, 120, 0xff00, // Set a connection mark
 	)
 
@@ -116,12 +116,12 @@ func ExampleConn_flushFilter() {
 	}
 
 	f1 := conntrack.NewFlow(
-		6, 0, net.IPv4(1, 2, 3, 4), net.IPv4(5, 6, 7, 8),
+		6, 0, netip.MustParseAddr("1.2.3.4"), netip.MustParseAddr("5.6.7.8"),
 		1234, 80, 120, 0x00ff, // Set a connection mark
 	)
 
 	f2 := conntrack.NewFlow(
-		17, 0, net.ParseIP("2a00:1450:400e:804::200e"), net.ParseIP("2a00:1450:400e:804::200f"),
+		17, 0, netip.MustParseAddr("2a00:1450:400e:804::200e"), netip.MustParseAddr("2a00:1450:400e:804::200f"),
 		1234, 80, 120, 0xff00, // Set a connection mark
 	)
 
@@ -155,7 +155,7 @@ func ExampleConn_delete() {
 	}
 
 	f := conntrack.NewFlow(
-		6, 0, net.IPv4(1, 2, 3, 4), net.IPv4(5, 6, 7, 8),
+		6, 0, netip.MustParseAddr("1.2.3.4"), netip.MustParseAddr("5.6.7.8"),
 		1234, 80, 120, 0,
 	)
 

diff --git a/event_integration_test.go b/event_integration_test.go
@@ -3,7 +3,7 @@
 package conntrack
 
 import (
-	"net"
+	"net/netip"
 	"testing"
 
 	"github.com/mdlayher/netlink"
@@ -43,7 +43,7 @@ func TestConnListen(t *testing.T) {
 
 	var warn bool
 
-	ip := net.ParseIP("::f00")
+	ip := netip.MustParseAddr("::f00")
 	for _, proto := range []uint8{unix.IPPROTO_TCP, unix.IPPROTO_UDP, unix.IPPROTO_DCCP, unix.IPPROTO_SCTP} {
 		// Create the Flow.
 		f := NewFlow(

diff --git a/expect_integration_test.go b/expect_integration_test.go
@@ -3,7 +3,7 @@
 package conntrack
 
 import (
-	"net"
+	"net/netip"
 	"testing"
 
 	"golang.org/x/sys/unix"
@@ -27,7 +27,7 @@ func TestConnCreateExpect(t *testing.T) {
 	c, _, err := makeNSConn()
 	require.NoError(t, err)
 
-	f := NewFlow(6, 0, net.IPv4(1, 2, 3, 4), net.IPv4(5, 6, 7, 8), 42000, 21, 120, 0)
+	f := NewFlow(6, 0, netip.MustParseAddr("1.2.3.4"), netip.MustParseAddr("5.6.7.8"), 42000, 21, 120, 0)
 
 	err = c.Create(f)
 	require.NoError(t, err, "unexpected error creating flow", f)
@@ -37,8 +37,8 @@ func TestConnCreateExpect(t *testing.T) {
 		TupleMaster: f.TupleOrig,
 		Tuple: Tuple{
 			IP: IPTuple{
-				SourceAddress:      net.IPv4(1, 2, 3, 4),
-				DestinationAddress: net.IPv4(5, 6, 7, 8),
+				SourceAddress:      netip.MustParseAddr("1.2.3.4"),
+				DestinationAddress: netip.MustParseAddr("5.6.7.8"),
 			},
 			Proto: ProtoTuple{
 				Protocol:        6,
@@ -48,8 +48,8 @@ func TestConnCreateExpect(t *testing.T) {
 		},
 		Mask: Tuple{
 			IP: IPTuple{
-				SourceAddress:      net.IPv4(255, 255, 255, 255),
-				DestinationAddress: net.IPv4(255, 255, 255, 255),
+				SourceAddress:      netip.MustParseAddr("255.255.255.255"),
+				DestinationAddress: netip.MustParseAddr("255.255.255.255"),
 			},
 			Proto: ProtoTuple{
 				Protocol:        6,

diff --git a/expect_test.go b/expect_test.go
@@ -1,10 +1,9 @@
 package conntrack
 
 import (
-	"net"
+	"net/netip"
 	"testing"
 
-	"github.com/google/go-cmp/cmp"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"github.com/ti-mo/netfilter"
@@ -183,8 +182,8 @@ var corpusExpect = []struct {
 		exp: Expect{
 			TupleMaster: Tuple{
 				IP: IPTuple{
-					SourceAddress:      []byte{127, 0, 0, 1},
-					DestinationAddress: []byte{127, 0, 0, 2},
+					SourceAddress:      netip.MustParseAddr("127.0.0.1"),
+					DestinationAddress: netip.MustParseAddr("127.0.0.2"),
 				},
 				Proto: ProtoTuple{
 					Protocol:        6,
@@ -194,8 +193,8 @@ var corpusExpect = []struct {
 			},
 			Tuple: Tuple{
 				IP: IPTuple{
-					SourceAddress:      []byte{127, 0, 0, 1},
-					DestinationAddress: []byte{127, 0, 0, 2},
+					SourceAddress:      netip.MustParseAddr("127.0.0.1"),
+					DestinationAddress: netip.MustParseAddr("127.0.0.2"),
 				},
 				Proto: ProtoTuple{
 					Protocol:        6,
@@ -204,8 +203,8 @@ var corpusExpect = []struct {
 			},
 			Mask: Tuple{
 				IP: IPTuple{
-					SourceAddress:      []byte{255, 255, 255, 255},
-					DestinationAddress: []byte{255, 255, 255, 255},
+					SourceAddress:      netip.MustParseAddr("255.255.255.255"),
+					DestinationAddress: netip.MustParseAddr("255.255.255.255"),
 				},
 				Proto: ProtoTuple{
 					Protocol:        6,
@@ -263,11 +262,8 @@ func TestExpectUnmarshal(t *testing.T) {
 	for _, tt := range corpusExpect {
 		t.Run(tt.name, func(t *testing.T) {
 			var ex Expect
-			assert.NoError(t, ex.unmarshal(mustDecodeAttributes(tt.attrs)))
-
-			if diff := cmp.Diff(tt.exp, ex); diff != "" {
-				t.Fatalf("unexpected unmarshal (-want +got):\n%s", diff)
-			}
+			require.NoError(t, ex.unmarshal(mustDecodeAttributes(tt.attrs)))
+			assert.Equal(t, tt.exp, ex, "unexpected unmarshal")
 		})
 	}
 
@@ -355,9 +351,7 @@ func TestExpectMarshal(t *testing.T) {
 		},
 	}
 
-	if diff := cmp.Diff(want, exm); diff != "" {
-		t.Fatalf("unexpected Expect marshal (-want +got):\n%s", diff)
-	}
+	assert.Equal(t, want, exm, "unexpected Expect marshal")
 
 	// Cannot marshal without tuple/mask/master Tuples
 	_, err = Expect{}.marshal()
@@ -424,10 +418,7 @@ func TestExpectNATUnmarshal(t *testing.T) {
 			}
 
 			require.NoError(t, err)
-
-			if diff := cmp.Diff(tt.enat, enat); diff != "" {
-				t.Fatalf("unexpected unmarshal (-want +got):\n%s", diff)
-			}
+			assert.Equal(t, tt.enat, enat, "unexpected unmarshal")
 		})
 	}
 }
@@ -439,8 +430,8 @@ func TestExpectNATMarshal(t *testing.T) {
 		Direction: true,
 		Tuple: Tuple{
 			IP: IPTuple{
-				SourceAddress:      net.ParseIP("baa:baa::b"),
-				DestinationAddress: net.ParseIP("ef00:3f00::ba13"),
+				SourceAddress:      netip.MustParseAddr("baa:baa::b"),
+				DestinationAddress: netip.MustParseAddr("ef00:3f00::ba13"),
 			},
 			Proto: ProtoTuple{
 				Protocol:        13,
@@ -458,9 +449,7 @@ func TestExpectNATMarshal(t *testing.T) {
 
 	// Only verify first attribute (direction); Tuple marshal has its own tests
 	want := netfilter.Attribute{Type: uint16(ctaExpectNATDir), Data: []byte{0, 0, 0, 1}}
-	if diff := cmp.Diff(want, enm.Children[0]); diff != "" {
-		t.Fatalf("unexpected ExpectNAT marshal (-want +got):\n%s", diff)
-	}
+	assert.Equal(t, want, enm.Children[0], "unexpected ExpectNAT marshal")
 }
 
 func TestExpectTypeString(t *testing.T) {

diff --git a/filter_test.go b/filter_test.go
@@ -3,9 +3,9 @@ package conntrack
 import (
 	"testing"
 
-	"github.com/ti-mo/netfilter"
+	"github.com/stretchr/testify/assert"
 
-	"github.com/google/go-cmp/cmp"
+	"github.com/ti-mo/netfilter"
 )
 
 func TestFilterMarshal(t *testing.T) {
@@ -22,7 +22,5 @@ func TestFilterMarshal(t *testing.T) {
 		},
 	}
 
-	if diff := cmp.Diff(fm, f.marshal()); diff != "" {
-		t.Fatalf("unexpected Filter marshal (-want +got):\n%s", diff)
-	}
+	assert.Equal(t, fm, f.marshal(), "unexpected Filter marshal")
 }
diff --git a/flow.go b/flow.go
@@ -2,7 +2,7 @@ package conntrack
 
 import (
 	"fmt"
-	"net"
+	"net/netip"
 
 	"github.com/mdlayher/netlink"
 	"github.com/ti-mo/netfilter"
@@ -44,7 +44,7 @@ type Flow struct {
 // source and destination addresses. srcPort and dstPort are the source and
 // destination ports. timeout is the non-zero time-to-live of a connection in
 // seconds.
-func NewFlow(proto uint8, status StatusFlag, srcAddr, destAddr net.IP,
+func NewFlow(proto uint8, status StatusFlag, srcAddr, destAddr netip.Addr,
 	srcPort, destPort uint16, timeout, mark uint32) Flow {
 
 	var f Flow