Support for quota as object (#244)

Fixes #238

nftables_test.go

@@ -6077,6 +6077,198 @@ func TestGetRulesObjref(t *testing.T) {
 	}
 }
 
+func TestAddQuotaObj(t *testing.T) {
+	conn, newNS := nftest.OpenSystemConn(t, *enableSysTests)
+	defer nftest.CleanupSystemConn(t, newNS)
+	conn.FlushRuleset()
+	defer conn.FlushRuleset()
+
+	table := &nftables.Table{
+		Name:   "quota_demo",
+		Family: nftables.TableFamilyIPv4,
+	}
+	tr := conn.AddTable(table)
+
+	c := &nftables.Chain{
+		Name:  "filter",
+		Table: table,
+	}
+	conn.AddChain(c)
+
+	o := &nftables.QuotaObj{
+		Table:    tr,
+		Name:     "q_test",
+		Bytes:    0x06400000,
+		Consumed: 0,
+		Over:     true,
+	}
+	conn.AddObj(o)
+
+	if err := conn.Flush(); err != nil {
+		t.Errorf("conn.Flush() failed: %v", err)
+	}
+
+	obj, err := conn.GetObj(&nftables.QuotaObj{
+		Table: table,
+		Name:  "q_test",
+	})
+	if err != nil {
+		t.Fatalf("conn.GetObj() failed: %v", err)
+	}
+
+	if got, want := len(obj), 1; got != want {
+		t.Fatalf("unexpected object list length: got %d, want %d", got, want)
+	}
+
+	o1, ok := obj[0].(*nftables.QuotaObj)
+	if !ok {
+		t.Fatalf("unexpected type: got %T, want *QuotaObj", obj[0])
+	}
+	if got, want := o1.Name, o.Name; got != want {
+		t.Fatalf("quota name mismatch: got %s, want %s", got, want)
+	}
+	if got, want := o1.Bytes, o.Bytes; got != want {
+		t.Fatalf("quota bytes mismatch: got %d, want %d", got, want)
+	}
+	if got, want := o1.Consumed, o.Consumed; got != want {
+		t.Fatalf("quota consumed mismatch: got %d, want %d", got, want)
+	}
+	if got, want := o1.Over, o.Over; got != want {
+		t.Fatalf("quota over mismatch: got %v, want %v", got, want)
+	}
+}
+
+func TestAddQuotaObjRef(t *testing.T) {
+	conn, newNS := nftest.OpenSystemConn(t, *enableSysTests)
+	defer nftest.CleanupSystemConn(t, newNS)
+	conn.FlushRuleset()
+	defer conn.FlushRuleset()
+
+	table := &nftables.Table{
+		Name:   "quota_demo",
+		Family: nftables.TableFamilyIPv4,
+	}
+	tr := conn.AddTable(table)
+
+	c := &nftables.Chain{
+		Name:  "filter",
+		Table: table,
+	}
+	conn.AddChain(c)
+
+	o := &nftables.QuotaObj{
+		Table:    tr,
+		Name:     "q_test",
+		Bytes:    0x06400000,
+		Consumed: 0,
+		Over:     true,
+	}
+	conn.AddObj(o)
+
+	r := &nftables.Rule{
+		Table: table,
+		Chain: c,
+		Exprs: []expr.Any{
+			&expr.Objref{
+				Type: 2,
+				Name: "q_test",
+			},
+		},
+	}
+	conn.AddRule(r)
+	if err := conn.Flush(); err != nil {
+		t.Fatalf("failed to flush: %v", err)
+	}
+
+	rules, err := conn.GetRules(table, c)
+	if err != nil {
+		t.Fatalf("failed to get rules: %v", err)
+	}
+
+	if got, want := len(rules), 1; got != want {
+		t.Fatalf("unexpected number of rules: got %d, want %d", got, want)
+	}
+	if got, want := len(rules[0].Exprs), 1; got != want {
+		t.Fatalf("unexpected number of exprs: got %d, want %d", got, want)
+	}
+
+	objref, ok := rules[0].Exprs[0].(*expr.Objref)
+	if !ok {
+		t.Fatalf("Exprs[0] is type %T, want *expr.Objref", rules[0].Exprs[0])
+	}
+	if want := r.Exprs[0]; !reflect.DeepEqual(objref, want) {
+		t.Errorf("objref expr = %+v, wanted %+v", objref, want)
+	}
+}
+
+func TestDeleteQuotaObj(t *testing.T) {
+	conn, newNS := nftest.OpenSystemConn(t, *enableSysTests)
+	defer nftest.CleanupSystemConn(t, newNS)
+	conn.FlushRuleset()
+	defer conn.FlushRuleset()
+
+	table := &nftables.Table{
+		Name:   "quota_demo",
+		Family: nftables.TableFamilyIPv4,
+	}
+	tr := conn.AddTable(table)
+
+	c := &nftables.Chain{
+		Name:  "filter",
+		Table: table,
+	}
+	conn.AddChain(c)
+
+	o := &nftables.QuotaObj{
+		Table:    tr,
+		Name:     "q_test",
+		Bytes:    0x06400000,
+		Consumed: 0,
+		Over:     true,
+	}
+	conn.AddObj(o)
+
+	if err := conn.Flush(); err != nil {
+		t.Fatalf("conn.Flush() failed: %v", err)
+	}
+
+	obj, err := conn.GetObj(&nftables.QuotaObj{
+		Table: table,
+		Name:  "q_test",
+	})
+	if err != nil {
+		t.Fatalf("conn.GetObj() failed: %v", err)
+	}
+
+	if got, want := len(obj), 1; got != want {
+		t.Fatalf("unexpected number of objects: got %d, want %d", got, want)
+	}
+
+	if got, want := obj[0], o; !reflect.DeepEqual(got, want) {
+		t.Errorf("got = %+v, want = %+v", got, want)
+	}
+
+	conn.DeleteObject(&nftables.QuotaObj{
+		Table: tr,
+		Name:  "q_test",
+	})
+
+	if err := conn.Flush(); err != nil {
+		t.Fatalf("conn.Flush() failed: %v", err)
+	}
+
+	obj, err = conn.GetObj(&nftables.QuotaObj{
+		Table: table,
+		Name:  "q_test",
+	})
+	if err != nil {
+		t.Fatalf("conn.GetObj() failed: %v", err)
+	}
+	if got, want := len(obj), 0; got != want {
+		t.Fatalf("unexpected object list length: got %d, want %d", got, want)
+	}
+}
+
 func TestGetRulesQueue(t *testing.T) {
 	// Create a new network namespace to test these operations,
 	// and tear down the namespace at test completion.

obj.go

@@ -155,6 +155,21 @@ func objFromMsg(msg netlink.Message) (Obj, error) {
 					Name:  name,
 				}
 
+				ad.Do(func(b []byte) error {
+					ad, err := netlink.NewAttributeDecoder(b)
+					if err != nil {
+						return err
+					}
+					ad.ByteOrder = binary.BigEndian
+					return o.unmarshal(ad)
+				})
+				return &o, ad.Err()
+			case NFT_OBJECT_QUOTA:
+				o := QuotaObj{
+					Table: table,
+					Name:  name,
+				}
+
 				ad.Do(func(b []byte) error {
 					ad, err := netlink.NewAttributeDecoder(b)
 					if err != nil {

quota.go

@@ -0,0 +1,80 @@
+// Copyright 2023 Google LLC. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package nftables
+
+import (
+	"github.com/google/nftables/binaryutil"
+	"github.com/mdlayher/netlink"
+	"golang.org/x/sys/unix"
+)
+
+const (
+	NFTA_OBJ_USERDATA = 8
+	NFT_OBJECT_QUOTA  = 2
+)
+
+type QuotaObj struct {
+	Table    *Table
+	Name     string
+	Bytes    uint64
+	Consumed uint64
+	Over     bool
+}
+
+func (q *QuotaObj) unmarshal(ad *netlink.AttributeDecoder) error {
+	for ad.Next() {
+		switch ad.Type() {
+		case unix.NFTA_QUOTA_BYTES:
+			q.Bytes = ad.Uint64()
+		case unix.NFTA_QUOTA_CONSUMED:
+			q.Consumed = ad.Uint64()
+		case unix.NFTA_QUOTA_FLAGS:
+			q.Over = (ad.Uint32() & unix.NFT_QUOTA_F_INV) == 1
+		}
+	}
+	return nil
+}
+
+func (q *QuotaObj) marshal(data bool) ([]byte, error) {
+	flags := uint32(0)
+	if q.Over {
+		flags = unix.NFT_QUOTA_F_INV
+	}
+	obj, err := netlink.MarshalAttributes([]netlink.Attribute{
+		{Type: unix.NFTA_QUOTA_BYTES, Data: binaryutil.BigEndian.PutUint64(q.Bytes)},
+		{Type: unix.NFTA_QUOTA_CONSUMED, Data: binaryutil.BigEndian.PutUint64(q.Consumed)},
+		{Type: unix.NFTA_QUOTA_FLAGS, Data: binaryutil.BigEndian.PutUint32(flags)},
+	})
+	if err != nil {
+		return nil, err
+	}
+	attrs := []netlink.Attribute{
+		{Type: unix.NFTA_OBJ_TABLE, Data: []byte(q.Table.Name + "\x00")},
+		{Type: unix.NFTA_OBJ_NAME, Data: []byte(q.Name + "\x00")},
+		{Type: unix.NFTA_OBJ_TYPE, Data: binaryutil.BigEndian.PutUint32(NFT_OBJECT_QUOTA)},
+	}
+	if data {
+		attrs = append(attrs, netlink.Attribute{Type: unix.NLA_F_NESTED | unix.NFTA_OBJ_DATA, Data: obj})
+	}
+	return netlink.MarshalAttributes(attrs)
+}
+
+func (q *QuotaObj) table() *Table {
+	return q.Table
+}
+
+func (q *QuotaObj) family() TableFamily {
+	return q.Table.Family
+}