break: simplified whois client implementation (#2)

* break: simplified whois client implementation - user can specify proxy.Dialer to use for the connection. - add unit tests for the split function. - mod tidy. - make linter happy.
twiny · Feb 11, 2024 · aa70c40 · aa70c40
1 parent fb09b2d
commit aa70c40
Show file tree

Hide file tree

Showing 6 changed files with 149 additions and 187 deletions.
diff --git a/README.md b/README.md
@@ -1,15 +1,15 @@
 # WHOIS Client
 
-a simple Go WHOIS Client API.
+a simple Go WHOIS Client API. It supports custom `proxy.Dialer` for Socks5.
 
 ## API
+
 ```go
-Lookup(ctx context.Context, domain, server string) (string, error)
-WHOISHost(domain string) (string, error)
-TLDs() []string
+Query(ctx context.Context, domain string) (Response, error)
 ```
 
 ## Install
+
 `go get github.com/twiny/whois/v2`
 
 ## Example
@@ -18,44 +18,30 @@ TLDs() []string
 package main
 
 import (
-	"context"
-	"fmt"
-	"time"
+ "context"
+ "fmt"
 
-	"github.com/twiny/whois/v2"
+ "github.com/twiny/whois/v2"
 )
 
 func main() {
-	domain := "google.com"
-
-	// to use Socks5 - this format 'socks5://username:[email protected]:1023'
-	// otherwise use 'whois.Localhost' to use local connection
-	client, err := whois.NewClient(whois.Localhost)
-	if err != nil {
-		fmt.Println(err)
-		return
-	}
-
-	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
-	defer cancel()
-
-	host, err := client.WHOISHost(domain)
-	if err != nil {
-		fmt.Println(err)
-		return
-	}
-
-	resp, err := client.Lookup(ctx, "google.com", host)
-	if err != nil {
-		fmt.Println(err)
-		return
-	}
-
-	fmt.Println(resp)
-	fmt.Println("done.")
+ client, err := whois.NewClient(nil)
+ if err != nil {
+  fmt.Printf("err: %s\n", err)
+ }
+
+ resp, err := client.Query(context.TODO(), "google.com")
+ if err != nil {
+  fmt.Printf("err: %s\n", err)
+ }
+
+ // Print the response
+ fmt.Printf("Domain: %+v\n", resp)
 }
 ```
+
 ## Tracking
+
 - If you wish to add more WHOIS Server please [create a PR](https://github.com/twiny/whois/pulls).
 
 - If you find any issues please [create a new issue](https://github.com/twiny/whois/issues/new).
diff --git a/client.go b/client.go
@@ -4,11 +4,9 @@ import (
 	"bytes"
 	"context"
 	_ "embed"
-	"errors"
 	"fmt"
-	"io/ioutil"
+	"io"
 	"net"
-	"net/url"
 	"regexp"
 	"strings"
 	"time"
@@ -19,98 +17,97 @@ import (
 )
 
 var (
-	testURL = "whois.verisign-grs.com:43" // TestURL
-	// dialer
-	forwrder = &net.Dialer{
-		Timeout: 15 * time.Second,
-	}
-	//
-	domainRegExp        = regexp.MustCompile(`(^(([[:alnum:]]-?)?([[:alnum:]]-?)+\.)+[A-Za-z]{2,20}$)`)
-	Localhost    string = "socks5://localhost" // used for default client
 	//go:embed db.yaml
 	whoisdb []byte
-)
 
-// Client
-type Client struct {
-	socks  string
-	db     map[string]string
-	dialer func(network, addr string) (c net.Conn, err error)
-}
+	domainRegExp = regexp.MustCompile(`^(?:(?:[a-zA-Z0-9](?:[a-zA-Z0-9-]{0,61}[a-zA-Z0-9])?\.)+[a-zA-Z]{2,24})$`)
 
-// NewClient: a new WHOIS client using a Socks5 connection
-// in this format: socks5://username:[email protected]:1023
-// to use local connection pass `whois.Localhost` as argument.
-func NewClient(socks5 string) (*Client, error) {
-	// whois db
-	db := map[string]string{}
-	if err := yaml.Unmarshal(whoisdb, &db); err != nil {
-		return nil, err
+	defaultDialer = &net.Dialer{
+		Timeout:   15 * time.Second,
+		KeepAlive: 30 * time.Second,
+		DualStack: true,
 	}
+)
 
-	// if using local dialer
-	if socks5 == Localhost {
-		var d = &net.Dialer{}
-		return &Client{
-			socks:  socks5,
-			db:     db,
-			dialer: d.Dial,
-		}, nil
+type (
+	Client struct {
+		dialer proxy.Dialer
+		db     map[string]string
 	}
 
-	d, err := socks5Dialer(socks5)
-	if err != nil {
+	Response struct {
+		Domain    string
+		Name      string
+		TLD       string
+		WHOISHost string
+		WHOISRaw  string
+	}
+)
+
+// NewClient: create a new WHOIS client, if dialer is nil, it will use the default dialer
+func NewClient(dialer proxy.Dialer) (*Client, error) {
+	var db = map[string]string{}
+	if err := yaml.Unmarshal(whoisdb, &db); err != nil {
 		return nil, err
 	}
 
+	if dialer == nil {
+		dialer = defaultDialer
+	}
+
 	return &Client{
-		socks:  socks5,
 		db:     db,
-		dialer: d.Dial,
+		dialer: dialer,
 	}, nil
 }
 
-// Lookup a domain in a whois server db
-func (c *Client) Lookup(ctx context.Context, domain, server string) (string, error) {
-	// check if domain is domain
-	name, tld, _, err := c.split(domain)
+// Query: return raw WHOIS information for a domain
+func (c *Client) Query(ctx context.Context, domain string) (Response, error) {
+	var response Response
+	name, tld, whoisHost, err := c.split(domain)
 	if err != nil {
-		return "", err
+		return response, err
 	}
 
-	domain = name + "." + tld
+	domain = strings.Join([]string{name, tld}, ".")
 
-	var text string
-	result := make(chan string, 1)
+	var (
+		text      string
+		lookupErr error
+		result    = make(chan string, 1)
+	)
 
-	// lookup
 	go func() {
 		defer close(result)
 
-		text, err = c.lookup(domain, server)
+		text, lookupErr = c.lookup(domain, whoisHost)
 		result <- text
 	}()
 
-	// wait
 	for {
 		select {
 		case <-ctx.Done():
-			return "", ctx.Err()
+			return response, ctx.Err()
 		case text := <-result:
-			if err != nil {
-				return "", err
+			if lookupErr != nil {
+				return response, err
 			}
-			return text, nil
+
+			response.Domain = domain
+			response.Name = name
+			response.TLD = tld
+			response.WHOISHost = whoisHost
+			response.WHOISRaw = text
+
+			return response, nil
 		}
 	}
 }
 
-// lookup
-func (c *Client) lookup(domain, server string) (string, error) {
-	// WHOIS server address at port 43
-	addr := net.JoinHostPort(server, "43")
+func (c *Client) lookup(domain, whoisHost string) (string, error) {
+	addr := net.JoinHostPort(whoisHost, "43") // WHOIS server address at port 43
 
-	conn, err := c.dialer("tcp", addr)
+	conn, err := c.dialer.Dial("tcp", addr)
 	if err != nil {
 		return "", err
 	}
@@ -124,7 +121,7 @@ func (c *Client) lookup(domain, server string) (string, error) {
 		return "", err
 	}
 
-	resp, err := ioutil.ReadAll(conn)
+	resp, err := io.ReadAll(conn)
 	if err != nil {
 		return "", err
 	}
@@ -136,78 +133,36 @@ func (c *Client) lookup(domain, server string) (string, error) {
 	return strings.Replace(string(resp), "\r", "", -1), nil
 }
 
-// Split domain name into 3 parts name, tld, server or an error
-func (c *Client) WHOISHost(domain string) (string, error) {
-	_, _, server, err := c.split(domain)
-	if err != nil {
-		return "", err
+// split: validate and split domain: name, tld, server or error
+func (c *Client) split(raw string) (string, string, string, error) {
+	if !domainRegExp.MatchString(raw) {
+		return "", "", "", fmt.Errorf("invalid domain: %s", raw)
 	}
-	return server, nil
-}
 
-// split
-func (c *Client) split(domain string) (name, tld, server string, err error) {
-	domain = strings.ToLower(domain)
-
-	// validate
-	if !domainRegExp.MatchString(domain) {
-		return "", "", "", fmt.Errorf("domain %s looks invalid", domain)
+	tld, icann := publicsuffix.PublicSuffix(raw)
+	if !icann {
+		return "", "", "", fmt.Errorf("unsupported TLD: %s", tld)
 	}
 
-	// extract TLD from domaian
-	tld = publicsuffix.List.PublicSuffix(domain)
-
-	// find whois host
 	server, found := c.db[tld]
 	if !found {
-		return "", "", "", fmt.Errorf("could not find corresponded WHOIS server for %s", tld)
+		return "", "", "", fmt.Errorf("unsupported TLD: %s", tld)
 	}
 
-	name = strings.TrimSuffix(domain, "."+tld)
+	trimmedDomain := strings.TrimSuffix(raw, "."+tld)
 
-	return
-}
-
-// TLDs
-func (c *Client) TLDs() []string {
-	var tlds = []string{}
-	for tld := range c.db {
-		tlds = append(tlds, tld)
-	}
-	return tlds
-}
-
-// socks5Dialer
-func socks5Dialer(socks5 string) (proxy.Dialer, error) {
-	surl, err := url.Parse(socks5)
-	if err != nil {
-		return nil, err
-	}
-
-	if surl.Scheme != "socks5" {
-		return nil, errors.New("socks must start with socks5://")
-	}
+	// Split the raw domain into parts.
+	// e.g. www.google.com -> www.google -> google
+	// If no parts, the trimmed domain itself is the name.
+	// Otherwise, the name is the last part in array.
+	parts := strings.Split(trimmedDomain, ".")
 
-	// check auth
-	username := surl.User.Username()
-	password, found := surl.User.Password()
-
-	auth := &proxy.Auth{}
-	if username+password == "" && !found {
-		auth = nil
-	}
-
-	d, err := proxy.SOCKS5("tcp", surl.Host, auth, forwrder)
-	if err != nil {
-		return nil, err
+	var name string
+	if len(parts) == 0 {
+		name = trimmedDomain
+	} else {
+		name = parts[len(parts)-1]
 	}
 
-	// to test connection
-	conn, err := d.Dial("tcp", testURL)
-	if err != nil {
-		return nil, err
-	}
-	defer conn.Close()
-
-	return d, nil
+	return name, tld, server, nil
 }
diff --git a/client_test.go b/client_test.go
@@ -1,3 +1,39 @@
 package whois
 
-// TODO (twiny) implment tests
+import "testing"
+
+func TestSplit(t *testing.T) {
+	var cases = []struct {
+		in   string
+		name string
+		tld  string
+		err  bool
+	}{
+		{"google.com", "google", "com", false},
+		{" google.com.", "", "", true},
+		{"www.google.com", "google", "com", false},
+		{"www.go-ogle.com.", "", "", true},
+		{"www.go-ogle.com", "go-ogle", "com", false},
+		{"host.alpha.host.www.google.com", "google", "com", false},
+	}
+
+	client, err := NewClient(nil)
+	if err != nil {
+		t.Errorf("NewClient() == %s, want 'nil'", err)
+	}
+
+	for _, c := range cases {
+		name, tld, _, err := client.split(c.in)
+		if err != nil && !c.err {
+			t.Errorf("split(%s) == %s, want 'nil'", c.in, err)
+		}
+
+		if name != c.name {
+			t.Errorf("split(%q) == %q, want %q", c.in, name, c.name)
+		}
+
+		if tld != c.tld {
+			t.Errorf("split(%q) == %q, want %q", c.in, tld, c.tld)
+		}
+	}
+}