Merge pull request #51 from mfycheng/detect-server-rate-limits

Return (potentially) partial buffers on error.

whois.go

@@ -219,6 +219,16 @@ func (c *Client) rawQuery(domain, server, port string) (string, error) {
 	_ = conn.SetWriteDeadline(time.Now().Add(c.timeout - elapsed))
 	_, err = conn.Write([]byte(domain + "\r\n"))
 	if err != nil {
+		// Some servers may refuse a request with a reason, immediately closing the connection after sending.
+		// For example, GoDaddy returns "Number of allowed queries exceeded.\r\n", and immediately closes the connection.
+		//
+		// We return both the response _and_ the error, to allow callers to try to parse the response, while
+		// still letting them know an error occurred. In particular, this helps catch rate limit errors.
+		buffer, _ := io.ReadAll(conn)
+		if len(buffer) > 0 {
+			return string(buffer), err
+		}
+
 		return "", fmt.Errorf("whois: send to whois server failed: %w", err)
 	}
 
@@ -227,6 +237,16 @@ func (c *Client) rawQuery(domain, server, port string) (string, error) {
 	_ = conn.SetReadDeadline(time.Now().Add(c.timeout - elapsed))
 	buffer, err := io.ReadAll(conn)
 	if err != nil {
+		if len(buffer) > 0 {
+			// Some servers may refuse a request with a reason, immediately closing the connection after sending.
+			// For example, GoDaddy returns "Number of allowed queries exceeded.\r\n", and immediately closes the connection.
+			//
+			// We return both the response _and_ the error, to allow callers to try to parse the response, while
+			// still letting them know an error occurred (potentially short reads). In particular, this helps
+			// catch rate limit errors.
+			return string(buffer), err
+		}
+
 		return "", fmt.Errorf("whois: read from whois server failed: %w", err)
 	}
 

whois_test.go

@@ -21,6 +21,7 @@ package whois
 
 import (
 	"errors"
+	"net"
 	"os"
 	"strings"
 	"testing"
@@ -126,6 +127,32 @@ func TestWhois(t *testing.T) {
 	}
 }
 
+func TestWhoisServerError(t *testing.T) {
+	// Start local TCP server that simulates rate limiting
+	listener, err := net.Listen("tcp", "127.0.0.1:0")
+	assert.Nil(t, err)
+	defer listener.Close()
+
+	go func() {
+		conn, err := listener.Accept()
+		if err != nil {
+			return
+		}
+
+		// Simulate rate limit response from GoDaddy.
+		conn.Write([]byte("Number of allowed queries exceeded"))
+
+		// Close the connection immediately, rather than a graceful shutdown.
+		conn.(*net.TCPConn).SetLinger(0)
+		conn.Close()
+	}()
+
+	client := NewClient()
+	result, err := client.Whois("test.com", listener.Addr().String())
+	assert.NotNil(t, err)
+	assert.Contains(t, result, "Number of allowed queries exceeded")
+}
+
 func TestNewClient(t *testing.T) {
 	c := NewClient()
 	var err error