From 832360a02c3b6376a537d3e1640f42698f320f4b Mon Sep 17 00:00:00 2001
From: Paulo Nascimento <paulo.nascimento@medwing.com>
Date: Sat, 20 Mar 2021 19:25:25 -0300
Subject: [PATCH] Usignregex to determine if it os ipv6 or not

---
 pkg/iptables/iptables.go | 30 ++++++++++++++++++++----------
 1 file changed, 20 insertions(+), 10 deletions(-)

diff --git a/pkg/iptables/iptables.go b/pkg/iptables/iptables.go
index 1c0a399..1d83af5 100644
--- a/pkg/iptables/iptables.go
+++ b/pkg/iptables/iptables.go
@@ -73,6 +73,8 @@ type rule struct {
 	proto Protocol
 }
 
+var ipv6Regex,_ = regexp.Compile("[-]d\\s(.*:.*\\s[-]m\\scomment)")
+
 // NewRule creates a new iptables or ip6tables rule in the given table and chain
 // depending on the given protocol.
 func NewRule(proto Protocol, table, chain string, spec ...string) Rule {
@@ -125,7 +127,14 @@ func (r *rule) String() string {
 }
 
 func (r *rule) Proto() Protocol {
-	return r.proto
+	proto := ProtocolIPv4
+
+	ruleString := r.String()
+	if ipv6Regex.MatchString(ruleString) {
+		proto = ProtocolIPv6
+	}
+
+	return proto
 }
 
 // chain represents an iptables chain.
@@ -190,7 +199,14 @@ func (c *chain) String() string {
 }
 
 func (c *chain) Proto() Protocol {
-	return c.proto
+	proto := ProtocolIPv4
+
+	chainString := c.String()
+	if ipv6Regex.MatchString(chainString) {
+		proto = ProtocolIPv6
+	}
+
+	return proto
 }
 
 func chainToString(table, chain string) string {
@@ -355,8 +371,6 @@ func (c *Controller) Set(rules []Rule) error {
 	defer c.Unlock()
 	var i int
 
-	ipv6Regex,_ := regexp.Compile("[-]d\\s(.*:.*\\s[-]m\\scomment)")
-
 	for ; i < len(rules); i++ {
 		if i < len(c.rules) {
 			if rules[i].String() != c.rules[i].String() {
@@ -366,12 +380,7 @@ func (c *Controller) Set(rules []Rule) error {
 			}
 		}
 		if i >= len(c.rules) {
-			proto := ProtocolIPv4
-
-			ruleString := rules[i].String()
-			if ipv6Regex.MatchString(ruleString) {
-				proto = ProtocolIPv6
-			}
+			proto := rules[i].Proto()
 
 			protocolName := "ipv4"
 
@@ -379,6 +388,7 @@ func (c *Controller) Set(rules []Rule) error {
 				protocolName = "ipv6"
 			}
 
+			var ruleString = rules[i].String()
 			level.Debug(c.logger).Log("msg", "Applying Firewall Rule...", "Rule", ruleString, "Protocol", protocolName)
 			if err := rules[i].Add(c.v4); err != nil {
 				return fmt.Errorf("failed to add rule: %v", err)