Support for quota as object (#244)

Fixes https://github.com/google/nftables/issues/238
This commit is contained in:
turekt 2023-12-12 21:49:45 +01:00 committed by GitHub
parent 32bfbb6627
commit 0f60df61a2
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 287 additions and 0 deletions

View File

@ -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.

15
obj.go
View File

@ -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 {

80
quota.go Normal file
View File

@ -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
}