Objects implementation refactor

Refactored obj.go to a more generic approach Added object support for already implemented expressions Added test for limit object Small lint changes Fixes #253
google · Apr 17, 2024 · 54b203b · 54b203b
1 parent 5e242ec
commit 54b203b
Show file tree

Hide file tree

Showing 46 changed files with 487 additions and 190 deletions.
diff --git a/counter.go b/counter.go
@@ -16,11 +16,12 @@ package nftables
 
 import (
 	"github.com/google/nftables/binaryutil"
+	"github.com/google/nftables/expr"
 	"github.com/mdlayher/netlink"
 	"golang.org/x/sys/unix"
 )
 
-// CounterObj implements Obj.
+// Deprecated: Use ObjAttr instead
 type CounterObj struct {
 	Table *Table
 	Name  string // e.g. “fwded”
@@ -41,6 +42,20 @@ func (c *CounterObj) unmarshal(ad *netlink.AttributeDecoder) error {
 	return ad.Err()
 }
 
+func (c *CounterObj) data() expr.Any {
+	return &expr.Counter{
+		Bytes:   c.Bytes,
+		Packets: c.Packets,
+	}
+}
+
+func (c *CounterObj) name() string {
+	return c.Name
+}
+func (c *CounterObj) objType() ObjType {
+	return ObjTypeCounter
+}
+
 func (c *CounterObj) table() *Table {
 	return c.Table
 }

diff --git a/expr/bitwise.go b/expr/bitwise.go
@@ -30,7 +30,7 @@ type Bitwise struct {
 	Xor            []byte
 }
 
-func (e *Bitwise) marshal(fam byte) ([]byte, error) {
+func (e *Bitwise) marshal(fam byte, dataOnly bool) ([]byte, error) {
 	mask, err := netlink.MarshalAttributes([]netlink.Attribute{
 		{Type: unix.NFTA_DATA_VALUE, Data: e.Mask},
 	})
@@ -54,6 +54,9 @@ func (e *Bitwise) marshal(fam byte) ([]byte, error) {
 	if err != nil {
 		return nil, err
 	}
+	if dataOnly {
+		return data, nil
+	}
 	return netlink.MarshalAttributes([]netlink.Attribute{
 		{Type: unix.NFTA_EXPR_NAME, Data: []byte("bitwise\x00")},
 		{Type: unix.NLA_F_NESTED | unix.NFTA_EXPR_DATA, Data: data},

diff --git a/expr/bitwise_test.go b/expr/bitwise_test.go
@@ -32,7 +32,7 @@ func TestBitwise(t *testing.T) {
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			nbw := Bitwise{}
-			data, err := tt.bw.marshal(0 /* don't care in this test */)
+			data, err := tt.bw.marshal(0 /* don't care in this test */, false)
 			if err != nil {
 				t.Fatalf("marshal error: %+v", err)
 

diff --git a/expr/byteorder.go b/expr/byteorder.go
@@ -37,7 +37,7 @@ type Byteorder struct {
 	Size           uint32
 }
 
-func (e *Byteorder) marshal(fam byte) ([]byte, error) {
+func (e *Byteorder) marshal(fam byte, dataOnly bool) ([]byte, error) {
 	data, err := netlink.MarshalAttributes([]netlink.Attribute{
 		{Type: unix.NFTA_BYTEORDER_SREG, Data: binaryutil.BigEndian.PutUint32(e.SourceRegister)},
 		{Type: unix.NFTA_BYTEORDER_DREG, Data: binaryutil.BigEndian.PutUint32(e.DestRegister)},
@@ -48,6 +48,9 @@ func (e *Byteorder) marshal(fam byte) ([]byte, error) {
 	if err != nil {
 		return nil, err
 	}
+	if dataOnly {
+		return data, nil
+	}
 	return netlink.MarshalAttributes([]netlink.Attribute{
 		{Type: unix.NFTA_EXPR_NAME, Data: []byte("byteorder\x00")},
 		{Type: unix.NLA_F_NESTED | unix.NFTA_EXPR_DATA, Data: data},

diff --git a/expr/connlimit.go b/expr/connlimit.go
@@ -36,14 +36,17 @@ type Connlimit struct {
 	Flags uint32
 }
 
-func (e *Connlimit) marshal(fam byte) ([]byte, error) {
+func (e *Connlimit) marshal(fam byte, dataOnly bool) ([]byte, error) {
 	data, err := netlink.MarshalAttributes([]netlink.Attribute{
 		{Type: NFTA_CONNLIMIT_COUNT, Data: binaryutil.BigEndian.PutUint32(e.Count)},
 		{Type: NFTA_CONNLIMIT_FLAGS, Data: binaryutil.BigEndian.PutUint32(e.Flags)},
 	})
 	if err != nil {
 		return nil, err
 	}
+	if dataOnly {
+		return data, nil
+	}
 
 	return netlink.MarshalAttributes([]netlink.Attribute{
 		{Type: unix.NFTA_EXPR_NAME, Data: []byte("connlimit\x00")},

diff --git a/expr/counter.go b/expr/counter.go
@@ -27,14 +27,17 @@ type Counter struct {
 	Packets uint64
 }
 
-func (e *Counter) marshal(fam byte) ([]byte, error) {
+func (e *Counter) marshal(fam byte, dataOnly bool) ([]byte, error) {
 	data, err := netlink.MarshalAttributes([]netlink.Attribute{
 		{Type: unix.NFTA_COUNTER_BYTES, Data: binaryutil.BigEndian.PutUint64(e.Bytes)},
 		{Type: unix.NFTA_COUNTER_PACKETS, Data: binaryutil.BigEndian.PutUint64(e.Packets)},
 	})
 	if err != nil {
 		return nil, err
 	}
+	if dataOnly {
+		return data, nil
+	}
 
 	return netlink.MarshalAttributes([]netlink.Attribute{
 		{Type: unix.NFTA_EXPR_NAME, Data: []byte("counter\x00")},

diff --git a/expr/ct.go b/expr/ct.go
@@ -63,8 +63,8 @@ type Ct struct {
 	Key            CtKey
 }
 
-func (e *Ct) marshal(fam byte) ([]byte, error) {
-	regData := []byte{}
+func (e *Ct) marshal(fam byte, dataOnly bool) ([]byte, error) {
+	var regData []byte
 	exprData, err := netlink.MarshalAttributes(
 		[]netlink.Attribute{
 			{Type: unix.NFTA_CT_KEY, Data: binaryutil.BigEndian.PutUint32(uint32(e.Key))},
@@ -90,6 +90,9 @@ func (e *Ct) marshal(fam byte) ([]byte, error) {
 		return nil, err
 	}
 	exprData = append(exprData, regData...)
+	if dataOnly {
+		return exprData, nil
+	}
 
 	return netlink.MarshalAttributes([]netlink.Attribute{
 		{Type: unix.NFTA_EXPR_NAME, Data: []byte("ct\x00")},

diff --git a/expr/dup.go b/expr/dup.go
@@ -28,7 +28,7 @@ type Dup struct {
 	IsRegDevSet bool
 }
 
-func (e *Dup) marshal(fam byte) ([]byte, error) {
+func (e *Dup) marshal(fam byte, dataOnly bool) ([]byte, error) {
 	attrs := []netlink.Attribute{
 		{Type: unix.NFTA_DUP_SREG_ADDR, Data: binaryutil.BigEndian.PutUint32(e.RegAddr)},
 	}
@@ -38,10 +38,12 @@ func (e *Dup) marshal(fam byte) ([]byte, error) {
 	}
 
 	data, err := netlink.MarshalAttributes(attrs)
-
 	if err != nil {
 		return nil, err
 	}
+	if dataOnly {
+		return data, nil
+	}
 
 	return netlink.MarshalAttributes([]netlink.Attribute{
 		{Type: unix.NFTA_EXPR_NAME, Data: []byte("dup\x00")},

diff --git a/expr/dynset.go b/expr/dynset.go
@@ -43,7 +43,7 @@ type Dynset struct {
 	Exprs      []Any
 }
 
-func (e *Dynset) marshal(fam byte) ([]byte, error) {
+func (e *Dynset) marshal(fam byte, dataOnly bool) ([]byte, error) {
 	// See: https://git.netfilter.org/libnftnl/tree/src/expr/dynset.c
 	var opAttrs []netlink.Attribute
 	opAttrs = append(opAttrs, netlink.Attribute{Type: unix.NFTA_DYNSET_SREG_KEY, Data: binaryutil.BigEndian.PutUint32(e.SrcRegKey)})
@@ -95,6 +95,9 @@ func (e *Dynset) marshal(fam byte) ([]byte, error) {
 	if err != nil {
 		return nil, err
 	}
+	if dataOnly {
+		return opData, nil
+	}
 
 	return netlink.MarshalAttributes([]netlink.Attribute{
 		{Type: unix.NFTA_EXPR_NAME, Data: []byte("dynset\x00")},
@@ -125,7 +128,7 @@ func (e *Dynset) unmarshal(fam byte, data []byte) error {
 		case unix.NFTA_DYNSET_FLAGS:
 			e.Invert = (ad.Uint32() & unix.NFT_DYNSET_F_INV) != 0
 		case unix.NFTA_DYNSET_EXPR:
-			exprs, err := parseexprfunc.ParseExprBytesFunc(fam, ad, ad.Bytes())
+			exprs, err := parseexprfunc.ParseExprBytesFunc(fam, ad)
 			if err != nil {
 				return err
 			}