From 0f60df61a2872af110f91fcd4f5fce6dd53972bf Mon Sep 17 00:00:00 2001 From: turekt <32360115+turekt@users.noreply.github.com> Date: Tue, 12 Dec 2023 21:49:45 +0100 Subject: [PATCH] Support for quota as object (#244) Fixes https://github.com/google/nftables/issues/238 --- nftables_test.go | 192 +++++++++++++++++++++++++++++++++++++++++++++++ obj.go | 15 ++++ quota.go | 80 ++++++++++++++++++++ 3 files changed, 287 insertions(+) create mode 100644 quota.go diff --git a/nftables_test.go b/nftables_test.go index 3158be3..160d12d 100644 --- a/nftables_test.go +++ b/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. diff --git a/obj.go b/obj.go index 08d43f4..e975bb8 100644 --- a/obj.go +++ b/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 { diff --git a/quota.go b/quota.go new file mode 100644 index 0000000..71cb9bb --- /dev/null +++ b/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 +}