From 508bb1ffd4a2dc62c6c5b59152ff94e5ab9dc22d Mon Sep 17 00:00:00 2001
From: Nikita Vorontsov <vorontsov.nstu@yandex.ru>
Date: Mon, 14 Jul 2025 14:43:15 +0300
Subject: [PATCH] fix set.KeyType overwrite for vmap (#319)

---
 set.go      |  7 +++----
 set_test.go | 24 ++++++++++++++++++++++++
 2 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/set.go b/set.go
index 5818ca7..e320117 100644
--- a/set.go
+++ b/set.go
@@ -572,7 +572,7 @@ func (cc *Conn) AddSet(s *Set, vals []SetElement) error {
 	if s.IsMap {
 		// Check if it is vmap case
 		if s.DataType.nftMagic == 1 {
-			// For Verdict data type, the expected magic is 0xfffff0
+			// For Verdict data type, the expected magic is 0xffffff00
 			tableInfo = append(tableInfo, netlink.Attribute{Type: unix.NFTA_SET_DATA_TYPE, Data: binaryutil.BigEndian.PutUint32(uint32(unix.NFT_DATA_VERDICT))},
 				netlink.Attribute{Type: unix.NFTA_SET_DATA_LEN, Data: binaryutil.BigEndian.PutUint32(s.DataType.Bytes)})
 		} else {
@@ -772,9 +772,8 @@ func setsFromMsg(msg netlink.Message) (*Set, error) {
 		case unix.NFTA_SET_DATA_TYPE:
 			nftMagic := ad.Uint32()
 			// Special case for the data type verdict, in the message it is stored as 0xffffff00 but it is defined as 1
-			if nftMagic == 0xffffff00 {
-				set.KeyType = TypeVerdict
-				break
+			if nftMagic == unix.NFT_DATA_VERDICT {
+				nftMagic = 1
 			}
 			dt, err := parseSetDatatype(nftMagic)
 			if err != nil {
diff --git a/set_test.go b/set_test.go
index dd30f45..df40e57 100644
--- a/set_test.go
+++ b/set_test.go
@@ -240,6 +240,30 @@ func TestMarshalSet(t *testing.T) {
 				Timeout:    30 * time.Second,
 			},
 		},
+		{
+			name: "Map ip-ip", // generic case
+			set: Set{
+				Name:     "test-map",
+				ID:       uint32(3),
+				Table:    tbl,
+				KeyType:  TypeIPAddr,
+				DataType: TypeIPAddr,
+				IsMap:    true,
+			},
+		},
+		{
+			// special case, see
+			// sets.go:setsFromMsg:(case unix.NFTA_SET_DATA_TYPE) and sets.go:AddSet:(if s.DataType.nftMagic == 1)
+			name: "Vedict map",
+			set: Set{
+				Name:     "test-map",
+				ID:       uint32(3),
+				Table:    tbl,
+				KeyType:  TypeIPAddr,
+				DataType: TypeVerdict,
+				IsMap:    true,
+			},
+		},
 	}
 
 	for i, tt := range tests {