Add UpdateBatch and UpateSingle

joyent · Jun 9, 2022 · 6b680b0 · 6b680b0
1 parent 51c81e4
commit 6b680b0
Show file tree

Hide file tree

Showing 5 changed files with 239 additions and 50 deletions.
diff --git a/conntrack.go b/conntrack.go
@@ -41,6 +41,11 @@ const (
 	ipctnlMsgExpGetStatsCPU
 )
 
+const (
+	MaxNetlinkMessageSize = 16384
+	NetlinkHeaderSize     = 16
+)
+
 // for detailes see https://github.com/tensorflow/tensorflow/blob/master/tensorflow/go/tensor.go#L488-L505
 var nativeEndian binary.ByteOrder
 
@@ -235,6 +240,110 @@ func (nfct *Nfct) Update(t Table, f Family, attributes Con) error {
 	return nfct.execute(req)
 }
 
+// send a message with no reply
+func (nfct *Nfct) UpdateSingle(t Table, f Family, attrs []*Con) error {
+	if t != Conntrack {
+		return ErrUnknownCtTable
+	}
+
+	for _, attr := range attrs {
+		query, err := nestAttributes(nfct.logger, attr)
+		if err != nil {
+			return err
+		}
+
+		data := putExtraHeader(uint8(f), unix.NFNETLINK_V0, unix.NFNL_SUBSYS_CTNETLINK)
+		data = append(data, query...)
+
+		req := netlink.Message{
+			Header: netlink.Header{
+				Type: netlink.HeaderType(t << 8),
+				//Flags: netlink.Request | netlink.Acknowledge,
+				// no reply
+				Flags: netlink.Request,
+			},
+			Data: data,
+		}
+
+		if t == Conntrack {
+			req.Header.Type |= netlink.HeaderType(ipctnlMsgCtNew)
+		} else if t == Expected {
+			req.Header.Type |= netlink.HeaderType(ipctnlMsgExpNew)
+		} else {
+			return ErrUnknownCtTable
+		}
+
+		_, e := nfct.Con.Send(req)
+		if e != nil {
+			return e
+		}
+	}
+
+	return nil
+}
+
+// send messages with no reply
+func (nfct *Nfct) UpdateBatch(t Table, f Family, attrs []*Con) error {
+	if t != Conntrack {
+		return ErrUnknownCtTable
+	}
+
+	var err error
+	var dataLen, l int
+
+	msgs := make([]netlink.Message, 0)
+
+	for _, attr := range attrs {
+		query, err := nestAttributes(nfct.logger, attr)
+		if err != nil {
+			return err
+		}
+		data := putExtraHeader(uint8(f), unix.NFNETLINK_V0, unix.NFNL_SUBSYS_CTNETLINK)
+		data = append(data, query...)
+
+		req := netlink.Message{
+			Header: netlink.Header{
+				Type: netlink.HeaderType(t << 8),
+				//Flags: netlink.Request | netlink.Acknowledge,
+				// no reply
+				Flags: netlink.Request,
+			},
+			Data: data,
+		}
+
+		if t == Conntrack {
+			req.Header.Type |= netlink.HeaderType(ipctnlMsgCtNew)
+		} else if t == Expected {
+			req.Header.Type |= netlink.HeaderType(ipctnlMsgExpNew)
+		} else {
+			return ErrUnknownCtTable
+		}
+
+		l = len(data) + NetlinkHeaderSize
+
+		// exceed max size
+		if dataLen+l > MaxNetlinkMessageSize {
+			_, err = nfct.Con.SendMessages(msgs)
+			if err != nil {
+				return err
+			}
+
+			msgs = make([]netlink.Message, 0)
+			dataLen = 0
+		}
+
+		dataLen += l
+		msgs = append(msgs, req)
+	}
+
+	err = nil
+	if len(msgs) > 0 {
+		_, err = nfct.Con.SendMessages(msgs)
+	}
+
+	return err
+}
+
 // Delete elements from the conntrack subsystem with certain attributes
 func (nfct *Nfct) Delete(t Table, f Family, filters Con) error {
 	query, err := nestAttributes(nfct.logger, &filters)

diff --git a/example/example b/example/example
diff --git a/example/main.go b/example/main.go
@@ -2,12 +2,14 @@ package main
 
 import (
 	"context"
+	"encoding/binary"
 	"fmt"
 	"net"
 	"time"
 
 	ct "github.com/florianl/go-conntrack"
 	"github.com/mdlayher/netlink"
+	"github.com/openlyinc/pointy"
 )
 
 func main() {
@@ -78,71 +80,138 @@ func ExampleUpdate() {
 	}
 	defer nfct.Close()
 
-	var filter ct.Con
-	src := net.ParseIP("172.30.1.60")
-	dst := net.ParseIP("172.30.1.72")
-	proto := uint8(6)
-	sp := uint16(50965)
-	dp := uint16(22)
-
-	label := make([]byte, 16)
-	label[0] = 0x11
-	label[1] = 0x99
-	/*
-		label[2] = 22
-		label[3] = 33
-	*/
-
-	labelMask := make([]byte, 16)
-	labelMask[0] = 0xff
-	labelMask[1] = 0xff
-	/*
-		labelMask[2] = 0xff
-		labelMask[1] = 0xff
-	*/
-
-	filter.Origin = &ct.IPTuple{
-		Src: &src,
-		Dst: &dst,
-		Proto: &ct.ProtoTuple{
-			Number:  &proto,
-			SrcPort: &sp,
-			DstPort: &dp,
-		},
+	var filter []*ct.Con
+
+	timestamp := uint32(time.Now().Unix())
+
+	if true {
+		f1 := ct.Con{}
+		src := net.ParseIP("172.30.1.60")
+		dst := net.ParseIP("172.30.1.72")
+		proto := uint8(6)
+		sp := uint16(51137)
+		dp := uint16(22)
+
+		label := make([]byte, 16)
+		binary.LittleEndian.PutUint32(label[0:4], timestamp)
+		//label[0] = 11
+		//label[1] = 22
+
+		labelMask := make([]byte, 16)
+		binary.LittleEndian.PutUint32(labelMask[0:4], ^uint32(0))
+		//labelMask[0] = 0xff
+		//labelMask[1] = 0xff
+
+		f1.Origin = &ct.IPTuple{
+			Src: &src,
+			Dst: &dst,
+			Proto: &ct.ProtoTuple{
+				Number:  &proto,
+				SrcPort: &sp,
+				DstPort: &dp,
+			},
+		}
+
+		f1.Label = &label
+		f1.LabelMask = &labelMask
+
+		filter = append(filter, &f1)
 	}
 
-	filter.Label = &label
-	filter.LabelMask = &labelMask
+	if false {
+		f2 := ct.Con{}
+		src := net.ParseIP("172.30.1.60")
+		dst := net.ParseIP("172.30.1.72")
+		proto := uint8(6)
+		sp := uint16(53044)
+		dp := uint16(22)
+
+		label := make([]byte, 16)
+		binary.LittleEndian.PutUint32(label[1:5], timestamp)
+		label[0] = 33
+		//label[1] = 22
+
+		labelMask := make([]byte, 16)
+		binary.LittleEndian.PutUint32(labelMask[1:5], ^uint32(0))
+		labelMask[0] = 0xff
+		//labelMask[1] = 0xff
+
+		f2.Origin = &ct.IPTuple{
+			Src: &src,
+			Dst: &dst,
+			Proto: &ct.ProtoTuple{
+				Number:  &proto,
+				SrcPort: &sp,
+				DstPort: &dp,
+			},
+		}
+
+		f2.Label = &label
+		f2.LabelMask = &labelMask
+
+		filter = append(filter, &f2)
+	}
 
 	//fmt.Printf("### Update: %#v\n", filter)
 
-	err = nfct.Update(ct.Conntrack, ct.IPv4, filter)
+	cnt := 100
+
+	//////////////////////////
+
+	for i := 1; i < cnt; i++ {
+		f := filter[0]
+		n := *f
+
+		n.Origin.Proto.SrcPort = pointy.Uint16(*n.Origin.Proto.SrcPort + 1)
+
+		filter = append(filter, &n)
+	}
+
+	start := time.Now()
+	err = nfct.UpdateBatch(ct.Conntrack, ct.IPv4, filter)
 	if err != nil {
-		fmt.Println("could not dump sessions:", err)
+		fmt.Println("error UpdateBatch:", err)
+		return
+	}
+	elapsed := time.Since(start)
+	fmt.Printf("### UpdateBatch(%d attrs) took %s \n", len(filter), elapsed)
+
+	///////////////////////////////////////
+
+	start = time.Now()
+	err = nfct.UpdateSingle(ct.Conntrack, ct.IPv4, filter)
+	if err != nil {
+		fmt.Println("error UpdateSingle:", err)
 		return
 	}
+	elapsed = time.Since(start)
+	fmt.Printf("### UpdateSingle(%d attrs) took %s \n", len(filter), elapsed)
 
-	//ExampleGet(&filter)
+	////////////////
+
+	//ExampleGet(filter)
 }
 
-func ExampleGet(filter *ct.Con) {
+func ExampleGet(filters []*ct.Con) {
 	nfct, err := ct.Open(&ct.Config{})
 	if err != nil {
 		fmt.Println("could not create nfct:", err)
 		return
 	}
 	defer nfct.Close()
 
-	sessions, err := nfct.Get(ct.Conntrack, ct.IPv4, *filter)
-	if err != nil {
-		fmt.Println("could not dump sessions:", err)
-		return
-	}
+	for _, filter := range filters {
+		sessions, err := nfct.Get(ct.Conntrack, ct.IPv4, *filter)
+		if err != nil {
+			fmt.Println("could not dump sessions:", err)
+			return
+		}
 
-	for i, session := range sessions {
-		fmt.Printf("### %d: %#v\n", i, session)
-		if session.Label != nil {
-			fmt.Printf("### Label: %+v \n", session.Label)
+		for i, session := range sessions {
+			fmt.Printf("### %d: %#v\n", i, session)
+			if session.Label != nil {
+				fmt.Printf("### Label: %+v \n", session.Label)
+			}
 		}
 	}
 }
diff --git a/go.mod b/go.mod
@@ -2,8 +2,9 @@ module github.com/florianl/go-conntrack
 
 require (
 	github.com/mdlayher/netlink v1.6.0
-	golang.org/x/net v0.0.0-20220127200216-cd36cc0744dd
-	golang.org/x/sys v0.0.0-20220128215802-99c3d69c2c27
+	github.com/openlyinc/pointy v1.1.2
+	golang.org/x/net v0.0.0-20220520000938-2e3eb7b945c2
+	golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a
 )
 
 go 1.13
diff --git a/go.sum b/go.sum
@@ -1,3 +1,5 @@
+github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
+github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/google/go-cmp v0.5.6/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
 github.com/google/go-cmp v0.5.7 h1:81/ik6ipDQS2aGcBfIN5dHDB36BwrStyeAQquSYCV4o=
 github.com/google/go-cmp v0.5.7/go.mod h1:n+brtR0CgQNWTVd5ZUFpTBC8YFBDLK/h/bpaJ8/DtOE=
@@ -7,18 +9,26 @@ github.com/mdlayher/netlink v1.6.0 h1:rOHX5yl7qnlpiVkFWoqccueppMtXzeziFjWAjLg6sz
 github.com/mdlayher/netlink v1.6.0/go.mod h1:0o3PlBmGst1xve7wQ7j/hwpNaFaH4qCRyWCdcZk8/vA=
 github.com/mdlayher/socket v0.1.1 h1:q3uOGirUPfAV2MUoaC7BavjQ154J7+JOkTWyiV+intI=
 github.com/mdlayher/socket v0.1.1/go.mod h1:mYV5YIZAfHh4dzDVzI8x8tWLWCliuX8Mon5Awbj+qDs=
+github.com/openlyinc/pointy v1.1.2 h1:LywVV2BWC5Sp5v7FoP4bUD+2Yn5k0VNeRbU5vq9jUMY=
+github.com/openlyinc/pointy v1.1.2/go.mod h1:w2Sytx+0FVuMKn37xpXIAyBNhFNBIJGR/v2m7ik1WtM=
+github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
+github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/stretchr/testify v1.2.2 h1:bSDNvY7ZPG5RlJ8otE/7V6gMiyenm9RtJ7IUVIAoJ1w=
+github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs=
 golang.org/x/net v0.0.0-20210928044308-7d9f5e0b762b/go.mod h1:9nx3DQGgdP8bBQD5qxJ1jj9UTztislL4KSBs9R2vV5Y=
-golang.org/x/net v0.0.0-20220127200216-cd36cc0744dd h1:O7DYs+zxREGLKzKoMQrtrEacpb0ZVXA5rIwylE2Xchk=
 golang.org/x/net v0.0.0-20220127200216-cd36cc0744dd/go.mod h1:CfG3xpIq0wQ8r1q4Su4UZFWDARRcnwPjda9FqA0JpMk=
+golang.org/x/net v0.0.0-20220520000938-2e3eb7b945c2 h1:NWy5+hlRbC7HK+PmcXVUmW1IMyFce7to56IUvhUFm7Y=
+golang.org/x/net v0.0.0-20220520000938-2e3eb7b945c2/go.mod h1:CfG3xpIq0wQ8r1q4Su4UZFWDARRcnwPjda9FqA0JpMk=
 golang.org/x/sync v0.0.0-20210220032951-036812b2e83c h1:5KslGYwFpkhGh+Q16bwMP3cOontH8FOep7tGV86Y7SQ=
 golang.org/x/sync v0.0.0-20210220032951-036812b2e83c/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20210423082822-04245dca01da/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20210927094055-39ccf1dd6fa6/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20211216021012-1d35b9e2eb4e/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
-golang.org/x/sys v0.0.0-20220128215802-99c3d69c2c27 h1:XDXtA5hveEEV8JB2l7nhMTp3t3cHp9ZpwcdjqyEWLlo=
 golang.org/x/sys v0.0.0-20220128215802-99c3d69c2c27/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a h1:dGzPydgVsqGcTRVwiLJ1jVbufYwmzD3LfVPLKsKg+0k=
+golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
 golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8=
 golang.org/x/text v0.3.6/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=