Merge pull request #34 from csnitker-godaddy/master

Added option to disable referral lookup

.gitignore

@@ -1 +1,2 @@
 .DS_Store
+.idea

whois.go

@@ -46,10 +46,11 @@ var DefaultClient = NewClient()
 
 // Client is whois client
 type Client struct {
-	dialer       proxy.Dialer
-	timeout      time.Duration
-	elapsed      time.Duration
-	disableStats bool
+	dialer          proxy.Dialer
+	timeout         time.Duration
+	elapsed         time.Duration
+	disableStats    bool
+	disableReferral bool
 }
 
 // Version returns package version
@@ -100,6 +101,12 @@ func (c *Client) SetDisableStats(disabled bool) *Client {
 	return c
 }
 
+// SetDisableReferral if set to true, will not query the referral server.
+func (c *Client) SetDisableReferral(disabled bool) *Client {
+	c.disableReferral = disabled
+	return c
+}
+
 // Whois do the whois query and returns whois information
 func (c *Client) Whois(domain string, servers ...string) (result string, err error) {
 	start := time.Now()
@@ -149,6 +156,10 @@ func (c *Client) Whois(domain string, servers ...string) (result string, err err
 		return
 	}
 
+	if c.disableReferral {
+		return
+	}
+
 	refServer, refPort := getServer(result)
 	if refServer == "" || refServer == server {
 		return

whois_test.go

@@ -21,6 +21,7 @@ package whois
 
 import (
 	"errors"
+	"strings"
 	"testing"
 	"time"
 
@@ -34,6 +35,20 @@ func TestVersion(t *testing.T) {
 	assert.Contains(t, License(), "Apache License")
 }
 
+func TestClient_SetDisableReferral(t *testing.T) {
+	client := NewClient()
+
+	resp, err := client.Whois("likexian.com")
+	assert.Nil(t, err)
+	assert.Equal(t, strings.Count(resp, "Domain Name: LIKEXIAN.COM"), 2)
+
+	client.SetDisableReferral(true)
+
+	resp, err = client.Whois("likexian.com")
+	assert.Nil(t, err)
+	assert.Equal(t, strings.Count(resp, "Domain Name: LIKEXIAN.COM"), 1)
+}
+
 func TestWhoisFail(t *testing.T) {
 	tests := []struct {
 		domain string