From 5bdeb6a4166f66750ad1dac3e5477e21ddf2dcc1 Mon Sep 17 00:00:00 2001
From: jarppiko <23549925+jarppiko@users.noreply.github.com>
Date: Mon, 29 May 2023 00:05:00 +0300
Subject: [PATCH] Add metrics collection support to set-only mode

---
 pkg/nftables/metrics.go | 17 +++++++++++++----
 1 file changed, 13 insertions(+), 4 deletions(-)

diff --git a/pkg/nftables/metrics.go b/pkg/nftables/metrics.go
index f3733f28..d6bf1ff0 100644
--- a/pkg/nftables/metrics.go
+++ b/pkg/nftables/metrics.go
@@ -38,8 +38,8 @@ type Set struct {
 	} `json:"nftables"`
 }
 
-func (c *nftContext) collectDroppedPackets(path string, hook string) (int, int, error) {
-	cmd := exec.Command(path, "-j", "list", "chain", c.ipFamily(), c.tableName, c.chainName+"-"+hook)
+func (c *nftContext) collectDroppedPackets(path string, chain string) (int, int, error) {
+	cmd := exec.Command(path, "-j", "list", "chain", c.ipFamily(), c.tableName, chain)
 	out, err := cmd.CombinedOutput()
 
 	if err != nil {
@@ -98,13 +98,22 @@ func (c *nftContext) collectDropped(path string, hooks []string) (int, int, int)
 
 	var droppedPackets, droppedBytes, banned int
 
-	for _, hook := range hooks {
-		pkt, byt, err := c.collectDroppedPackets(path, hook)
+	if c.setOnly {
+		pkt, byt, err := c.collectDroppedPackets(path, c.chainName)
 		if err != nil {
 			log.Errorf("can't collect dropped packets for ip%s from nft: %s", c.version, err)
 		}
 		droppedPackets += pkt
 		droppedBytes += byt
+	} else {
+		for _, hook := range hooks {
+			pkt, byt, err := c.collectDroppedPackets(path, c.chainName+"-"+hook)
+			if err != nil {
+				log.Errorf("can't collect dropped packets for ip%s from nft: %s", c.version, err)
+			}
+			droppedPackets += pkt
+			droppedBytes += byt
+		}
 	}
 
 	banned, err := c.collectActiveBannedIPs(path)