Merge pull request #150 from coroot/fast_fqdn_normalization

optimize FQDNs normalization
coroot · Dec 4, 2024 · 42a8e5c · 42a8e5c
2 parents 98808a0 + bcf6e0b
commit 42a8e5c
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 5 deletions.
diff --git a/common/net.go b/common/net.go
@@ -1,6 +1,7 @@
 package common
 
 import (
+	"bytes"
 	"fmt"
 	"net"
 	"regexp"
@@ -210,14 +211,30 @@ func NewDestinationKey(dst, actualDst netaddr.IPPort, fqdn string) DestinationKe
 	}
 }
 
-var ec2NodeRegex = regexp.MustCompile(`ip-\d+-\d+-\d+-\d+\.ec2`)
-var externalDomainWithSuffix = regexp.MustCompile(`(.+\.(com|net|org|io))\..+`)
-
 func NormalizeFQDN(fqdn string, requestType string) string {
 	if requestType == "TypePTR" {
 		return "IP.in-addr.arpa"
 	}
-	fqdn = ec2NodeRegex.ReplaceAllLiteralString(fqdn, "IP.ec2")
-	fqdn = externalDomainWithSuffix.ReplaceAllString(fqdn, "$1.search_path_suffix")
+	if strings.HasPrefix(fqdn, "ip-") {
+		if idx := strings.Index(fqdn, "."); idx > 0 && strings.HasPrefix(fqdn[idx+1:], "ec2") {
+			return "IP.ec2" + fqdn[idx+4:]
+		}
+	}
+	buf := bytes.NewBuffer(nil)
+	partsCount := 0
+	for i, r := range fqdn {
+		if r != '.' {
+			buf.WriteRune(r)
+		} else {
+			if partsCount > 0 && len(fqdn) > i {
+				switch string(buf.Bytes()) {
+				case "com", "net", "org", "io":
+					return fqdn[:i] + ".search_path_suffix"
+				}
+			}
+			buf.Reset()
+			partsCount++
+		}
+	}
 	return fqdn
 }
diff --git a/common/net_test.go b/common/net_test.go
@@ -62,6 +62,7 @@ func TestNormalizeFQDN(t *testing.T) {
 	assert.Equal(t, "IP.in-addr.arpa", NormalizeFQDN("4.3.2.1.in-addr.arpa", "TypePTR"))
 	assert.Equal(t, "coroot.com", NormalizeFQDN("coroot.com", "TypeA"))
 	assert.Equal(t, "IP.ec2.internal", NormalizeFQDN("ip-172-1-2-3.ec2.internal", "TypeA"))
+	assert.Equal(t, "IP.ec2", NormalizeFQDN("ip-172-1-2-3.ec2", "TypeA"))
 
 	assert.Equal(t, "example.com", NormalizeFQDN("example.com", "TypeA"))
 	assert.Equal(t, "example.com.search_path_suffix", NormalizeFQDN("example.com.cluster.local", "TypeA"))
@@ -72,3 +73,10 @@ func TestNormalizeFQDN(t *testing.T) {
 	assert.Equal(t, "example.org.search_path_suffix", NormalizeFQDN("example.org.svc.default.cluster.local", "TypeA"))
 	assert.Equal(t, "example.io.search_path_suffix", NormalizeFQDN("example.io.svc.default.cluster.local", "TypeA"))
 }
+
+func BenchmarkNormalizeFQDN(b *testing.B) {
+	for i := 0; i < b.N; i++ {
+		NormalizeFQDN("ip-172-1-2-3.ec2.internal", "TypeA")
+		NormalizeFQDN("example.io.svc.default.cluster.local", "TypeA")
+	}
+}