From 71337b220cc888d7b3b00a33806a5ee94054c678 Mon Sep 17 00:00:00 2001
From: Maxime Demode <maxime.demode@sewan.fr>
Date: Tue, 15 Oct 2019 18:03:13 +0200
Subject: [PATCH] [expr] Add SourceRegister field and modify marshal function
 to work with it.

---
 expr/ct.go | 37 +++++++++++++++++++++++++++++--------
 1 file changed, 29 insertions(+), 8 deletions(-)

diff --git a/expr/ct.go b/expr/ct.go
index 0ac5207..5808463 100644
--- a/expr/ct.go
+++ b/expr/ct.go
@@ -51,21 +51,42 @@ const (
 
 // Ct defines type for NFT connection tracking
 type Ct struct {
-	Register uint32
-	Key      CtKey
+	Register       uint32
+	SourceRegister bool
+	Key            CtKey
 }
 
 func (e *Ct) marshal() ([]byte, error) {
-	data, err := netlink.MarshalAttributes([]netlink.Attribute{
-		{Type: unix.NFTA_CT_KEY, Data: binaryutil.BigEndian.PutUint32(uint32(e.Key))},
-		{Type: unix.NFTA_CT_DREG, Data: binaryutil.BigEndian.PutUint32(e.Register)},
-	})
+	regData := []byte{}
+	exprData, err := netlink.MarshalAttributes(
+		[]netlink.Attribute{
+			{Type: unix.NFTA_CT_KEY, Data: binaryutil.BigEndian.PutUint32(uint32(e.Key))},
+		},
+	)
 	if err != nil {
 		return nil, err
 	}
+	if e.SourceRegister {
+		regData, err = netlink.MarshalAttributes(
+			[]netlink.Attribute{
+				{Type: unix.NFTA_CT_SREG, Data: binaryutil.BigEndian.PutUint32(e.Register)},
+			},
+		)
+	} else {
+		regData, err = netlink.MarshalAttributes(
+			[]netlink.Attribute{
+				{Type: unix.NFTA_CT_DREG, Data: binaryutil.BigEndian.PutUint32(e.Register)},
+			},
+		)
+	}
+	if err != nil {
+		return nil, err
+	}
+	exprData = append(exprData, regData...)
+
 	return netlink.MarshalAttributes([]netlink.Attribute{
 		{Type: unix.NFTA_EXPR_NAME, Data: []byte("ct\x00")},
-		{Type: unix.NLA_F_NESTED | unix.NFTA_EXPR_DATA, Data: data},
+		{Type: unix.NLA_F_NESTED | unix.NFTA_EXPR_DATA, Data: exprData},
 	})
 }
 
@@ -78,7 +99,7 @@ func (e *Ct) unmarshal(data []byte) error {
 	for ad.Next() {
 		switch ad.Type() {
 		case unix.NFTA_CT_KEY:
-			e.Key = ad.Uint32()
+			e.Key = CtKey(ad.Uint32())
 		case unix.NFTA_CT_DREG:
 			e.Register = ad.Uint32()
 		}