From bee26f421c6f2e24362c2e045ece8c4efcc4b598 Mon Sep 17 00:00:00 2001
From: Lan Tian <xuyh0120@outlook.com>
Date: Sat, 31 Jul 2021 16:56:39 +0800
Subject: [PATCH] frontend: resolve asn in dns/whois/fail order & fix tests

---
 frontend/bgpmap.go      | 39 ++++++++++++++++++++++-----------------
 frontend/bgpmap_test.go | 13 ++++++++++++-
 frontend/whois.go       |  4 ++++
 frontend/whois_test.go  | 22 ++++++++++++++++++++++
 4 files changed, 60 insertions(+), 18 deletions(-)
 create mode 100644 frontend/whois_test.go

diff --git a/frontend/bgpmap.go b/frontend/bgpmap.go
index 8ad3d46..c332704 100644
--- a/frontend/bgpmap.go
+++ b/frontend/bgpmap.go
@@ -11,30 +11,35 @@ func getASNRepresentation(asn string) string {
 	if setting.dnsInterface != "" {
 		// get ASN representation using DNS
 		records, err := net.LookupTXT(fmt.Sprintf("AS%s.%s", asn, setting.dnsInterface))
-		if err != nil {
-			// DNS query failed, only use ASN as output
-			return fmt.Sprintf("AS%s", asn)
+		if err == nil {
+			result := strings.Join(records, " ")
+			if resultSplit := strings.Split(result, " | "); len(resultSplit) > 1 {
+				result = strings.Join(resultSplit[1:], "\\n")
+			}
+			return fmt.Sprintf("AS%s\\n%s", asn, result)
 		}
+	}
 
-		result := strings.Join(records, " ")
-		if resultSplit := strings.Split(result, " | "); len(resultSplit) > 1 {
-			result = strings.Join(resultSplit[1:], "\\n")
-		}
-		return fmt.Sprintf("AS%s\\n%s", asn, result)
-	} else {
+	if setting.whoisServer != "" {
 		// get ASN representation using WHOIS
 		records := whois(fmt.Sprintf("AS%s", asn))
-		recordsSplit := strings.Split(records, "\n")
-		result := "\\n"
-		for _, line := range recordsSplit {
-			if strings.Contains(line, "as-name:") || strings.Contains(line, "ASName:") {
-				result = result + strings.TrimSpace(strings.SplitN(line, ":", 2)[1])
-			} else if strings.Contains(line, "descr:") {
-				result = result + "\\n" + strings.TrimSpace(strings.SplitN(line, ":", 2)[1])
+		if records != "" {
+			recordsSplit := strings.Split(records, "\n")
+			result := ""
+			for _, line := range recordsSplit {
+				if strings.Contains(line, "as-name:") || strings.Contains(line, "ASName:") {
+					result = result + strings.TrimSpace(strings.SplitN(line, ":", 2)[1])
+				} else if strings.Contains(line, "descr:") {
+					result = result + "\\n" + strings.TrimSpace(strings.SplitN(line, ":", 2)[1])
+				}
+			}
+			if result != "" {
+				return fmt.Sprintf("AS%s\\n%s", asn, result)
 			}
 		}
-		return fmt.Sprintf("AS%s%s", asn, result)
 	}
+
+	return fmt.Sprintf("AS%s", asn)
 }
 
 func birdRouteToGraphviz(servers []string, responses []string, target string) string {
diff --git a/frontend/bgpmap_test.go b/frontend/bgpmap_test.go
index bf638ec..b80ebec 100644
--- a/frontend/bgpmap_test.go
+++ b/frontend/bgpmap_test.go
@@ -5,8 +5,18 @@ import (
 	"testing"
 )
 
-func TestGetASNRepresentation(t *testing.T) {
+func TestGetASNRepresentationDNS(t *testing.T) {
 	setting.dnsInterface = "asn.cymru.com"
+	setting.whoisServer = ""
+	result := getASNRepresentation("6939")
+	if !strings.Contains(result, "HURRICANE") {
+		t.Errorf("Lookup AS6939 failed, got %s", result)
+	}
+}
+
+func TestGetASNRepresentationWhois(t *testing.T) {
+	setting.dnsInterface = ""
+	setting.whoisServer = "whois.arin.net"
 	result := getASNRepresentation("6939")
 	if !strings.Contains(result, "HURRICANE") {
 		t.Errorf("Lookup AS6939 failed, got %s", result)
@@ -15,6 +25,7 @@ func TestGetASNRepresentation(t *testing.T) {
 
 func TestGetASNRepresentationFallback(t *testing.T) {
 	setting.dnsInterface = ""
+	setting.whoisServer = ""
 	result := getASNRepresentation("6939")
 	if result != "AS6939" {
 		t.Errorf("Lookup AS6939 failed, got %s", result)
diff --git a/frontend/whois.go b/frontend/whois.go
index 8592858..7991ac1 100644
--- a/frontend/whois.go
+++ b/frontend/whois.go
@@ -7,6 +7,10 @@ import (
 
 // Send a whois request
 func whois(s string) string {
+	if setting.whoisServer == "" {
+		return ""
+	}
+
 	conn, err := net.Dial("tcp", setting.whoisServer+":43")
 	if err != nil {
 		return err.Error()
diff --git a/frontend/whois_test.go b/frontend/whois_test.go
new file mode 100644
index 0000000..e44b202
--- /dev/null
+++ b/frontend/whois_test.go
@@ -0,0 +1,22 @@
+package main
+
+import (
+	"strings"
+	"testing"
+)
+
+func TestWhois(t *testing.T) {
+	setting.whoisServer = "whois.arin.net"
+	result := whois("AS6939")
+	if !strings.Contains(result, "HURRICANE") {
+		t.Errorf("Whois AS6939 failed, got %s", result)
+	}
+}
+
+func TestWhoisWithoutServer(t *testing.T) {
+	setting.whoisServer = ""
+	result := whois("AS6939")
+	if result != "" {
+		t.Errorf("Whois AS6939 without server produced output, got %s", result)
+	}
+}