From 85c93f19db91fa601fb95721e9952a1f17d8d11e Mon Sep 17 00:00:00 2001
From: nickgarlis <nickgarlis@gmail.com>
Date: Tue, 25 Mar 2025 17:16:40 +0100
Subject: [PATCH] Add missing ct keys

The following keys were missing from the CtKey type:

 - `NFT_CT_SRC_IP`
 - `NFT_CT_DST_IP`
 - `NFT_CT_SRC_IP6`
 - `NFT_CT_DST_IP6`
 - `NFT_CT_ID`

Since they also seem to be missing from the unix package, their actual
values were added.
---
 expr/ct.go      | 10 +++++++++-
 expr/ct_test.go | 32 ++++++++++++++++++++++++++++++++
 2 files changed, 41 insertions(+), 1 deletion(-)

diff --git a/expr/ct.go b/expr/ct.go
index 0d58f2b..37d90ee 100644
--- a/expr/ct.go
+++ b/expr/ct.go
@@ -27,6 +27,7 @@ import (
 type CtKey uint32
 
 // Possible CtKey values.
+// Retrieved from https://git.netfilter.org/libnftnl/tree/include/linux/netfilter/nf_tables.h#n1121
 const (
 	CtKeySTATE      CtKey = unix.NFT_CT_STATE
 	CtKeyDIRECTION  CtKey = unix.NFT_CT_DIRECTION
@@ -48,6 +49,13 @@ const (
 	CtKeyZONE       CtKey = unix.NFT_CT_ZONE
 	CtKeyEVENTMASK  CtKey = unix.NFT_CT_EVENTMASK
 
+	// These values seem to be missing from the unix package
+	CtKeySRCIP  CtKey = 19
+	CtKeyDSTIP  CtKey = 20
+	CtKeySRCIP6 CtKey = 21
+	CtKeyDSTIP6 CtKey = 22
+	CtKeyID     CtKey = 23
+
 	// https://sources.debian.org/src//nftables/0.9.8-3/src/ct.c/?hl=39#L39
 	CtStateBitINVALID     uint32 = 1
 	CtStateBitESTABLISHED uint32 = 2
@@ -157,7 +165,7 @@ func (e *Ct) marshalData(fam byte) ([]byte, error) {
 	exprData = append(exprData, regData...)
 
 	switch e.Key {
-	case CtKeySRC, CtKeyDST, CtKeyPROTOSRC, CtKeyPROTODST:
+	case CtKeySRC, CtKeyDST, CtKeyPROTOSRC, CtKeyPROTODST, CtKeySRCIP, CtKeyDSTIP, CtKeySRCIP6, CtKeyDSTIP6:
 		regData, err = netlink.MarshalAttributes(
 			[]netlink.Attribute{
 				{Type: unix.NFTA_CT_DIRECTION, Data: binaryutil.BigEndian.PutUint32(e.Direction)},
diff --git a/expr/ct_test.go b/expr/ct_test.go
index 6b77c8a..8afa0be 100644
--- a/expr/ct_test.go
+++ b/expr/ct_test.go
@@ -46,6 +46,38 @@ func TestCt(t *testing.T) {
 				SourceRegister: true,
 			},
 		},
+		{
+			name: "Unmarshal Ct ip direction original case",
+			ct: Ct{
+				Register:  1,
+				Key:       CtKeySRCIP,
+				Direction: 0,
+			},
+		},
+		{
+			name: "Unmarshal Ct ip direction reply case",
+			ct: Ct{
+				Register:  1,
+				Key:       CtKeySRCIP,
+				Direction: 1,
+			},
+		},
+		{
+			name: "Unmarshal Ct ip6 direction original case",
+			ct: Ct{
+				Register:  1,
+				Key:       CtKeySRCIP6,
+				Direction: 0,
+			},
+		},
+		{
+			name: "Unmarshal Ct ip6 direction reply case",
+			ct: Ct{
+				Register:  1,
+				Key:       CtKeyDSTIP6,
+				Direction: 1,
+			},
+		},
 	}
 
 	for _, tt := range tests {