Check TCP entries to make sure the connection is open

Signed-off-by: Matt Farina <[email protected]>

Author: mattfarina

pkg/guestagent/iptables/iptables.go

@@ -8,9 +8,11 @@ import (
 	"regexp"
 	"strconv"
 	"strings"
+	"time"
 )
 
 type Entry struct {
+	TCP  bool
 	IP   net.IP
 	Port int
 }
@@ -27,7 +29,7 @@ type Entry struct {
 // ipv4 IP address. We need to detect this IP.
 // --dport is the destination port. We need to detect this port
 // -j DNAT this tells us it's the line doing the port forwarding.
-var findPortRegex = regexp.MustCompile(`-A\s+CNI-DN-\w*\s+(?:-d ((?:\b25[0-5]|\b2[0-4][0-9]|\b[01]?[0-9][0-9]?)(?:\.(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)){3}))?(?:/32\s+)?-p .*--dport (\d+) -j DNAT`)
+var findPortRegex = regexp.MustCompile(`-A\s+CNI-DN-\w*\s+(?:-d ((?:\b25[0-5]|\b2[0-4][0-9]|\b[01]?[0-9][0-9]?)(?:\.(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)){3}))?(?:/32\s+)?-p (tcp)?.*--dport (\d+) -j DNAT`)
 
 func GetPorts() ([]Entry, error) {
 	// TODO: add support for ipv6
@@ -50,19 +52,29 @@ func GetPorts() ([]Entry, error) {
 		return nil, err
 	}
 
-	return parsePortsFromRules(res)
+	pts, err := parsePortsFromRules(res)
+	if err != nil {
+		return nil, err
+	}
+
+	return checkPortsOpen(pts)
 }
 
 func parsePortsFromRules(rules []string) ([]Entry, error) {
 	var entries []Entry
 	for _, rule := range rules {
 		if found := findPortRegex.FindStringSubmatch(rule); found != nil {
-			if len(found) == 3 {
-				port, err := strconv.Atoi(found[2])
+			if len(found) == 4 {
+				port, err := strconv.Atoi(found[3])
 				if err != nil {
 					return nil, err
 				}
 
+				istcp := false
+				if found[2] == "tcp" {
+					istcp = true
+				}
+
 				// if the IP is blank the port forwarding the portforwarding,
 				// which gets information from this, will skip it. When no IP
 				// is present localhost will work.
@@ -73,6 +85,7 @@ func parsePortsFromRules(rules []string) ([]Entry, error) {
 				ent := Entry{
 					IP:   net.ParseIP(ip),
 					Port: port,
+					TCP:  istcp,
 				}
 				entries = append(entries, ent)
 			}
@@ -109,3 +122,20 @@ func listNATRules(pth string) ([]string, error) {
 
 	return rules, nil
 }
+
+func checkPortsOpen(pts []Entry) ([]Entry, error) {
+	var entries []Entry
+	for _, pt := range pts {
+		if pt.TCP {
+			conn, err := net.DialTimeout("tcp", net.JoinHostPort(pt.IP.String(), strconv.Itoa(pt.Port)), time.Second)
+			if err == nil && conn != nil {
+				conn.Close()
+				entries = append(entries, pt)
+			}
+		} else {
+			entries = append(entries, pt)
+		}
+	}
+
+	return entries, nil
+}

pkg/guestagent/iptables/iptables_test.go

@@ -84,10 +84,10 @@ func TestParsePortsFromRules(t *testing.T) {
 		t.Fatalf("expected 2 ports parsed from iptables but parsed %d", l)
 	}
 
-	if res[0].IP.String() != "127.0.0.1" || res[0].Port != 8082 {
-		t.Errorf("expected port 8082 on IP 127.0.0.1 but go port %d on IP %s", res[0].Port, res[0].IP.String())
+	if res[0].IP.String() != "127.0.0.1" || res[0].Port != 8082 || res[0].TCP != true {
+		t.Errorf("expected port 8082 on IP 127.0.0.1 with TCP true but go port %d on IP %s with TCP %t", res[0].Port, res[0].IP.String(), res[0].TCP)
 	}
-	if res[1].IP.String() != "127.0.0.1" || res[1].Port != 8081 {
-		t.Errorf("expected port 8081 on IP 127.0.0.1 but go port %d on IP %s", res[1].Port, res[1].IP.String())
+	if res[1].IP.String() != "127.0.0.1" || res[1].Port != 8081 || res[1].TCP != true {
+		t.Errorf("expected port 8081 on IP 127.0.0.1 with TCP true but go port %d on IP %s with TCP %t", res[1].Port, res[1].IP.String(), res[1].TCP)
 	}
 }