feat: Implement referral caching (#3)

README.md

@@ -63,3 +63,19 @@ Currently, the only fields parsed are:
 - `NameServers`: The nameservers currently tied to the domain
 
 If you'd like one or more other fields to be parsed, please don't be shy and create an issue or a pull request.
+
+#### Caching referral WHOIS servers
+The way that WHOIS scales is by having one "main" WHOIS server, namely `whois.iana.org:43`, refer to other WHOIS server
+on a per-TLD basis. 
+
+In other word, let's say that you wanted to have the WHOIS information for `example.com`. 
+The first step would be to query `whois.iana.org:43` with `com`, which would return `whois.verisign-grs.com`.
+Then, you would query `whois.verisign-grs.com:43` for the WHOIS information on `example.com`.
+
+If you're querying a lot of servers, making two queries instead of one can be a little wasteful, hence `WithReferralCache(true)`:
+```go
+client := whois.NewClient().WithReferralCache(true)
+```
+The above will cache the referral WHOIS server for each TLD, so that you can directly query the appropriate WHOIS server
+instead of first querying `whois.iana.org:43` for the referral.
+

whois.go

@@ -13,16 +13,49 @@ const (
 
 type Client struct {
 	whoisServerAddress string
+
+	isCachingReferralWHOISServers bool
+	referralWHOISServersCache     map[string]string
 }
 
 func NewClient() *Client {
 	return &Client{
-		whoisServerAddress: ianaWHOISServerAddress,
+		whoisServerAddress:        ianaWHOISServerAddress,
+		referralWHOISServersCache: make(map[string]string),
+	}
+}
+
+// WithReferralCache allows you to enable or disable the referral WHOIS server cache.
+// While ianaWHOISServerAddress is the "entry point" for WHOIS queries, it sometimes has
+// availability issues. One way to mitigate this is to cache the referral WHOIS server.
+//
+// This is disabled by default
+func (c *Client) WithReferralCache(enabled bool) *Client {
+	c.isCachingReferralWHOISServers = enabled
+	if enabled {
+		// We'll set a couple of common ones right away to avoid unnecessary queries
+		c.referralWHOISServersCache = map[string]string{
+			"com":   "whois.verisign-grs.com",
+			"black": "whois.nic.black",
+			"dev":   "whois.nic.google",
+			"green": "whois.nic.green",
+			"io":    "whois.nic.io",
+			"net":   "whois.verisign-grs.com",
+			"org":   "whois.publicinterestregistry.org",
+			"red":   "whois.nic.red",
+			"sh":    "whois.nic.sh",
+		}
 	}
+	return c
 }
 
-func (c Client) Query(domain string) (string, error) {
+func (c *Client) Query(domain string) (string, error) {
 	parts := strings.Split(domain, ".")
+	if c.isCachingReferralWHOISServers {
+		if cachedWHOISServer, ok := c.referralWHOISServersCache[domain]; ok {
+			return c.query(cachedWHOISServer, domain)
+		}
+	}
 	output, err := c.query(c.whoisServerAddress, parts[len(parts)-1])
 	if err != nil {
 		return "", err
@@ -32,6 +65,9 @@ func (c Client) Query(domain string) (string, error) {
 		endIndex := strings.Index(output[startIndex:], "\n") + startIndex
 		whois := strings.TrimSpace(output[startIndex:endIndex])
 		if referOutput, err := c.query(whois+":43", domain); err == nil {
+			if c.isCachingReferralWHOISServers {
+				c.referralWHOISServersCache[domain] = whois + ":43"
+			}
 			return referOutput, nil
 		}
 		return "", err

whois_test.go

@@ -48,9 +48,10 @@ func TestClient(t *testing.T) {
 			wantErr: false,
 		},
 	}
+	client := NewClient().WithReferralCache(true)
 	for _, scenario := range scenarios {
 		t.Run(scenario.domain+"_Query", func(t *testing.T) {
-			output, err := NewClient().Query(scenario.domain)
+			output, err := client.Query(scenario.domain)
 			if scenario.wantErr && err == nil {
 				t.Error("expected error, got none")
 				t.FailNow()
@@ -64,7 +65,7 @@ func TestClient(t *testing.T) {
 		})
 		time.Sleep(50 * time.Millisecond) // Give the WHOIS servers some breathing room
 		t.Run(scenario.domain+"_QueryAndParse", func(t *testing.T) {
-			response, err := NewClient().QueryAndParse(scenario.domain)
+			response, err := client.QueryAndParse(scenario.domain)
 			if scenario.wantErr && err == nil {
 				t.Error("expected error, got none")
 				t.FailNow()
@@ -79,7 +80,7 @@ func TestClient(t *testing.T) {
 				t.Errorf("expected to have at least one name server")
 			}
 			if len(response.DomainStatuses) == 0 {
-				t.Errorf("expected to have at least one domai status")
+				t.Errorf("expected to have at least one domain status")
 			}
 		})
 		time.Sleep(50 * time.Millisecond) // Give the WHOIS servers some breathing room