From 5a9391c12fe350b851291203e682ce05ea8e14d2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Gustavo=20I=C3=B1iguez=20Goia?= <gooffy1@gmail.com>
Date: Tue, 29 Mar 2022 18:00:11 +0200
Subject: [PATCH] Added support for quota expression (#149)

---
 expr/quota.go    | 76 ++++++++++++++++++++++++++++++++++++++++++++++++
 nftables_test.go | 71 ++++++++++++++++++++++++++++++++++++++++++++
 rule.go          |  2 ++
 3 files changed, 149 insertions(+)
 create mode 100644 expr/quota.go

diff --git a/expr/quota.go b/expr/quota.go
new file mode 100644
index 0000000..91db028
--- /dev/null
+++ b/expr/quota.go
@@ -0,0 +1,76 @@
+// Copyright 2018 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 expr
+
+import (
+	"encoding/binary"
+
+	"github.com/google/nftables/binaryutil"
+	"github.com/mdlayher/netlink"
+	"golang.org/x/sys/unix"
+)
+
+// Quota defines a threshold against a number of bytes.
+type Quota struct {
+	Bytes    uint64
+	Consumed uint64
+	Over     bool
+}
+
+func (q *Quota) marshal() ([]byte, error) {
+	attrs := []netlink.Attribute{
+		{Type: unix.NFTA_QUOTA_BYTES, Data: binaryutil.BigEndian.PutUint64(q.Bytes)},
+		{Type: unix.NFTA_QUOTA_CONSUMED, Data: binaryutil.BigEndian.PutUint64(q.Consumed)},
+	}
+
+	flags := uint32(0)
+	if q.Over {
+		flags = unix.NFT_QUOTA_F_INV
+	}
+	attrs = append(attrs, netlink.Attribute{
+		Type: unix.NFTA_QUOTA_FLAGS,
+		Data: binaryutil.BigEndian.PutUint32(flags),
+	})
+
+	data, err := netlink.MarshalAttributes(attrs)
+	if err != nil {
+		return nil, err
+	}
+
+	return netlink.MarshalAttributes([]netlink.Attribute{
+		{Type: unix.NFTA_EXPR_NAME, Data: []byte("quota\x00")},
+		{Type: unix.NLA_F_NESTED | unix.NFTA_EXPR_DATA, Data: data},
+	})
+}
+
+func (q *Quota) unmarshal(data []byte) error {
+	ad, err := netlink.NewAttributeDecoder(data)
+	if err != nil {
+		return err
+	}
+
+	ad.ByteOrder = binary.BigEndian
+	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 ad.Err()
+}
diff --git a/nftables_test.go b/nftables_test.go
index 87d009f..25b0184 100644
--- a/nftables_test.go
+++ b/nftables_test.go
@@ -4768,6 +4768,77 @@ func TestNotrack(t *testing.T) {
 	}
 }
 
+func TestQuota(t *testing.T) {
+
+	want := [][]byte{
+		// batch begin
+		[]byte("\x00\x00\x00\x0a"),
+		// nft flush ruleset
+		[]byte("\x00\x00\x00\x00"),
+		// nft add table ip filter
+		[]byte("\x02\x00\x00\x00\x0b\x00\x01\x00\x66\x69\x6c\x74\x65\x72\x00\x00\x08\x00\x02\x00\x00\x00\x00\x00"),
+		// nft add chain ip filter regular-chain
+		[]byte("\x02\x00\x00\x00\x0b\x00\x01\x00\x66\x69\x6c\x74\x65\x72\x00\x00\x12\x00\x03\x00\x72\x65\x67\x75\x6c\x61\x72\x2d\x63\x68\x61\x69\x6e\x00\x00\x00"),
+		// nft add rule ip filter regular-chain quota over 6 bytes
+		[]byte("\x02\x00\x00\x00\x0b\x00\x01\x00\x66\x69\x6c\x74\x65\x72\x00\x00\x12\x00\x02\x00\x72\x65\x67\x75\x6c\x61\x72\x2d\x63\x68\x61\x69\x6e\x00\x00\x00\x38\x00\x04\x80\x34\x00\x01\x80\x0a\x00\x01\x00\x71\x75\x6f\x74\x61\x00\x00\x00\x24\x00\x02\x80\x0c\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x06\x0c\x00\x04\x00\x00\x00\x00\x00\x00\x00\x00\x00\x08\x00\x02\x00\x00\x00\x00\x01"),
+		// batch end
+		[]byte("\x00\x00\x00\x0a"),
+	}
+
+	c := &nftables.Conn{
+		TestDial: func(req []netlink.Message) ([]netlink.Message, error) {
+			for idx, msg := range req {
+				b, err := msg.MarshalBinary()
+				if err != nil {
+					t.Fatal(err)
+				}
+				if len(b) < 16 {
+					continue
+				}
+				b = b[16:]
+				if len(want[idx]) == 0 {
+					t.Errorf("no want entry for message %d: %x", idx, b)
+					continue
+				}
+				got := b
+				if !bytes.Equal(got, want[idx]) {
+					t.Errorf("message %d: %s", idx, linediff(nfdump(got), nfdump(want[idx])))
+				}
+			}
+			return req, nil
+		},
+	}
+
+	c.FlushRuleset()
+
+	filter := c.AddTable(&nftables.Table{
+		Family: nftables.TableFamilyIPv4,
+		Name:   "filter",
+	})
+
+	output := c.AddChain(&nftables.Chain{
+		Name:  "regular-chain",
+		Table: filter,
+	})
+
+	c.AddRule(&nftables.Rule{
+		Table: filter,
+		Chain: output,
+		Exprs: []expr.Any{
+			// [ quota bytes 6 consumed 0 flags 1 ]
+			&expr.Quota{
+				Bytes:    6,
+				Consumed: 0,
+				Over:     true,
+			},
+		},
+	})
+
+	if err := c.Flush(); err != nil {
+		t.Fatal(err)
+	}
+}
+
 func TestStatelessNAT(t *testing.T) {
 	// The want byte sequences come from stracing nft(8), e.g.:
 	// strace -f -v -x -s 2048 -eraw=sendto nft add rule filter prerouting mark set jhash ip saddr mod 2
diff --git a/rule.go b/rule.go
index cf7e422..275b03b 100644
--- a/rule.go
+++ b/rule.go
@@ -251,6 +251,8 @@ func exprsFromMsg(b []byte) ([]expr.Any, error) {
 						e = &expr.NAT{}
 					case "limit":
 						e = &expr.Limit{}
+					case "quota":
+						e = &expr.Quota{}
 					case "dynset":
 						e = &expr.Dynset{}
 					case "log":