feat: add connmark support

florianl · Dec 15, 2023 · 6fc2e97 · 6fc2e97
1 parent 39338a2
commit 6fc2e97
Show file tree

Hide file tree

Showing 2 changed files with 59 additions and 0 deletions.
diff --git a/nfqueue.go b/nfqueue.go
@@ -41,6 +41,27 @@ func (nfqueue *Nfqueue) SetVerdictWithMark(id uint32, verdict, mark int) error {
 	return nfqueue.setVerdict(id, verdict, false, attributes)
 }
 
+// SetVerdictWithConnMark signals the kernel the next action and the connmark for a specified package id
+func (nfqueue *Nfqueue) SetVerdictWithConnMark(id uint32, verdict, mark int) error {
+	buf := make([]byte, 4)
+	binary.BigEndian.PutUint32(buf, uint32(mark))
+	ctAttrs, err := netlink.MarshalAttributes([]netlink.Attribute{{
+		Type: ctaMark,
+		Data: buf,
+	}})
+	if err != nil {
+		return err
+	}
+	attributes, err := netlink.MarshalAttributes([]netlink.Attribute{{
+		Type: netlink.Nested | nfQaCt,
+		Data: ctAttrs,
+	}})
+	if err != nil {
+		return err
+	}
+	return nfqueue.setVerdict(id, verdict, false, attributes)
+}
+
 // SetVerdictModPacket signals the kernel the next action for an altered packet
 func (nfqueue *Nfqueue) SetVerdictModPacket(id uint32, verdict int, packet []byte) error {
 	data, err := netlink.MarshalAttributes([]netlink.Attribute{{
@@ -73,6 +94,33 @@ func (nfqueue *Nfqueue) SetVerdictModPacketWithMark(id uint32, verdict, mark int
 	return nfqueue.setVerdict(id, verdict, false, data)
 }
 
+// SetVerdictModPacketWithConnMark signals the kernel the next action and connmark for an altered packet
+func (nfqueue *Nfqueue) SetVerdictModPacketWithConnMark(id uint32, verdict, mark int, packet []byte) error {
+	buf := make([]byte, 4)
+	binary.BigEndian.PutUint32(buf, uint32(mark))
+	ctAttrs, err := netlink.MarshalAttributes([]netlink.Attribute{{
+		Type: ctaMark,
+		Data: buf,
+	}})
+	if err != nil {
+		return err
+	}
+	data, err := netlink.MarshalAttributes([]netlink.Attribute{
+		{
+			Type: nfQaPayload,
+			Data: packet,
+		},
+		{
+			Type: netlink.Nested | nfQaCt,
+			Data: ctAttrs,
+		},
+	})
+	if err != nil {
+		return err
+	}
+	return nfqueue.setVerdict(id, verdict, false, data)
+}
+
 // SetVerdict signals the kernel the next action for a specified package id
 func (nfqueue *Nfqueue) SetVerdict(id uint32, verdict int) error {
 	return nfqueue.setVerdict(id, verdict, false, []byte{})

diff --git a/types.go b/types.go
@@ -163,3 +163,14 @@ const (
 	NfQeueue
 	NfRepeat
 )
+
+// conntrack attributes
+const (
+	ctaTupleOrig  = 1
+	ctaTupleReply = 2
+	ctaStatus     = 3
+	ctaTimeout    = 7
+	ctaMark       = 8
+	ctaProtoInfo  = 4
+	ctaLabels     = 22
+)