Support for quota as object (#244)
Fixes https://github.com/google/nftables/issues/238
This commit is contained in:
parent
32bfbb6627
commit
0f60df61a2
192
nftables_test.go
192
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.
|
||||
|
|
15
obj.go
15
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 {
|
||||
|
|
|
@ -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
|
||||
}
|
Loading…
Reference in New Issue