From ccbaaf004cd388ed767bc803c87a1bc53d3a63cd Mon Sep 17 00:00:00 2001
From: Mike Cheng <mike@ch3ng.dev>
Date: Wed, 6 Nov 2024 18:16:35 -0500
Subject: [PATCH] Return (potentially) partial buffers on error.

In some cases, servers will response with an error, and then immediately
close the connection. In those cases, we can still recover some
information to be parsed for better error handling downstream.

One notable example is rate limiting. The current behaviour makes it
appear to be a connection issue (connection reset by peer). However, if
we inspect the buffer, we notice a rate limiting message. If we in turn
pass this to the whoisparser, we can treat it like other generic rate
limiting errors, and act accordingly.
---
 whois.go      | 20 ++++++++++++++++++++
 whois_test.go | 27 +++++++++++++++++++++++++++
 2 files changed, 47 insertions(+)

diff --git a/whois.go b/whois.go
index 99237e8..f0701e0 100644
--- a/whois.go
+++ b/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)
 	}
 
diff --git a/whois_test.go b/whois_test.go
index 76fc9d4..5474cc6 100644
--- a/whois_test.go
+++ b/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