From 73c788181ee16c9d1839181d1a1d4ca0667509ff Mon Sep 17 00:00:00 2001
From: Georg Pfuetzenreuter <mail@georg-pfuetzenreuter.net>
Date: Fri, 30 Aug 2024 18:26:57 +0200
Subject: [PATCH] Correctly return set elements on GET

Signed-off-by: Georg Pfuetzenreuter <mail@georg-pfuetzenreuter.net>
---
 nft.go               |  9 ++++++---
 nftables-http-api.go |  3 +--
 utils.go             | 20 ++++++++++++++++++++
 3 files changed, 27 insertions(+), 5 deletions(-)

diff --git a/nft.go b/nft.go
index 87f62eb..e228985 100644
--- a/nft.go
+++ b/nft.go
@@ -14,6 +14,7 @@ package main
 import (
 	"github.com/google/nftables"
 	"log"
+	"net"
 )
 
 type nftError struct {
@@ -79,6 +80,7 @@ func getNftSet(nft *nftables.Conn, setName string) (*nftables.Set, error) {
 		log.Printf("Set lookup for %s failed, cannot proceed: %s", setName, err)
 		return nil, err
 	}
+	log.Printf("Found set %s", foundSet.Name)
 
 	return foundSet, nil
 }
@@ -97,9 +99,10 @@ func getNftSetElements(nft *nftables.Conn, setName string) ([]string, error) {
 
 	var returnElements []string
 
-	for _, element := range setElements {
-		log.Printf("element: %s", element.Key)
-		returnElements = append(returnElements, string(element.Key))
+	for i, element := range setElements {
+		ip := net.IP(element.Key)
+		log.Printf("Element %d: %s", i, ip)
+		returnElements = append(returnElements, ip.String())
 	}
 
 	return returnElements, nil
diff --git a/nftables-http-api.go b/nftables-http-api.go
index 10aec3c..0d8ca37 100644
--- a/nftables-http-api.go
+++ b/nftables-http-api.go
@@ -18,7 +18,6 @@ import (
 	"log"
 	"net/http"
 	"os"
-	"strings"
 )
 
 var config Config
@@ -108,7 +107,7 @@ func handleSetRoute(w http.ResponseWriter, r *http.Request) {
 			doReturn(w, http.StatusInternalServerError, "nftables failure")
 		}
 		if nftResult != nil {
-			doReturn(w, http.StatusOK, strings.Join(nftResult.([]string), ""))
+			doReturnSet(w, http.StatusOK, "", nftResult.([]string))
 		}
 	}
 }
diff --git a/utils.go b/utils.go
index b28ac87..4ff1a19 100644
--- a/utils.go
+++ b/utils.go
@@ -25,6 +25,11 @@ type Response struct {
 	RResult string `json:"result,omitempty"`
 }
 
+type ResponseSet struct {
+	RError  string   `json:"error,omitempty"`
+	RResult []string `json:"result,omitempty"`
+}
+
 func doReturn(w http.ResponseWriter, status int, text string) {
 	var response any
 	if status == http.StatusOK {
@@ -40,6 +45,21 @@ func doReturn(w http.ResponseWriter, status int, text string) {
 	w.Write(j)
 }
 
+func doReturnSet(w http.ResponseWriter, status int, text string, elements []string) {
+	var response any
+	if status == http.StatusOK {
+		response = ResponseSet{RResult: elements}
+	} else {
+		response = Response{RError: text}
+	}
+	j, err := json.Marshal(response)
+	if err != nil {
+		log.Fatalf("Failed to marshal JSON: %s", err)
+	}
+	w.WriteHeader(status)
+	w.Write(j)
+}
+
 func doCheckToken(token string, hash string) bool {
 	err := bcrypt.CompareHashAndPassword([]byte(hash), []byte(token))
 	if err == nil {