Cap WHOIS answers and follow referrals only to WHOIS servers (closes #106)

The client kept every line a WHOIS server sent until the server closed
the connection or 30 seconds passed, so a server that kept sending
could fill memory. It now reads at most 1 MiB of an answer, and an
answer that reaches that is an error. A "refer:" line in IANA's answer
naming a server none of the five RIRs run was dialled as given, on any
host and port; it is now an error, so the client only connects to IANA
and the five RIR servers. Both errors are failed lookups, which the
fetcher retries later.

Model: opus-5-5

internal/whois/whois.go

@@ -6,6 +6,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"io"
 	"net"
 	"regexp"
 	"strings"
@@ -19,6 +20,10 @@ const (
 	writeTimeout = 5 * time.Second
 )
 
+// maxAnswerSize caps the bytes read of one WHOIS answer: a real one is a
+// few KiB, so one this large is from a server that does not stop sending.
+const maxAnswerSize = 1 << 20
+
 // Parsing constants.
 const (
 	keyValueParts       = 2 // Expected parts when splitting "key: value"
@@ -52,7 +57,9 @@ var (
 
 	errNoASBlock = errors.New("IANA answered with neither a referral " +
 		"nor an AS block")
-	errNoFields = errors.New("RIR answered with no WHOIS fields")
+	errNoFields       = errors.New("RIR answered with no WHOIS fields")
+	errUnknownServer  = errors.New("IANA refers the AS to a server no RIR runs")
+	errAnswerTooLarge = errors.New("WHOIS answer reached the size cap")
 )
 
 // ASNInfo contains parsed WHOIS information for an ASN.
@@ -91,8 +98,9 @@ func NewClient() *Client {
 // LookupASN asks IANA which RIR holds the given ASN, asks that RIR, and
 // returns what the RIR's answer says. It returns ErrNoReferral when IANA
 // refers the AS to no RIR, and another error when a query fails, when IANA's
-// answer has neither a referral nor the AS's block, or when the RIR's answer
-// carries none of ASNInfo's fields, as a refusal of service does.
+// answer has neither a referral nor the AS's block, when IANA refers the AS
+// to a server none of the five RIRs run, or when the RIR's answer carries
+// none of ASNInfo's fields, as a refusal of service does.
 func (c *Client) LookupASN(ctx context.Context, asn int) (*ASNInfo, error) {
 	// Query IANA first to find the authoritative RIR
 	query := fmt.Sprintf("AS%d", asn)
@@ -103,7 +111,11 @@ func (c *Client) LookupASN(ctx context.Context, asn int) (*ASNInfo, error) {
 	}
 
 	// Determine RIR from IANA response
-	rir, whoisServer := c.parseIANAReferral(ianaResp)
+	rir, whoisServer, err := c.parseIANAReferral(ianaResp)
+	if err != nil {
+		return nil, fmt.Errorf("%s: %w", query, err)
+	}
+
 	if whoisServer == "" {
 		// IANA names the block of an AS it refers to no RIR. An answer
 		// without that, such as an empty one, says nothing about the AS.
@@ -142,7 +154,8 @@ func (info *ASNInfo) isEmpty() bool {
 		info.RegDate == nil && info.LastMod == nil
 }
 
-// query performs a raw WHOIS query to the specified server.
+// query performs a raw WHOIS query to the specified server. An answer that
+// reaches maxAnswerSize is an error.
 func (c *Client) query(ctx context.Context, server, query string) (string, error) {
 	conn, err := c.DialContext(ctx, "tcp", server)
 	if err != nil {
@@ -170,7 +183,11 @@ func (c *Client) query(ctx context.Context, server, query string) (string, error
 
 	var sb strings.Builder
 
-	scanner := bufio.NewScanner(conn)
+	// The reader ends at the cap; N, the bytes it has left, is 0 once the
+	// answer has reached it.
+	answer := &io.LimitedReader{R: conn, N: maxAnswerSize}
+
+	scanner := bufio.NewScanner(answer)
 	for scanner.Scan() {
 		sb.WriteString(scanner.Text())
 		sb.WriteString("\n")
@@ -181,12 +198,18 @@ func (c *Client) query(ctx context.Context, server, query string) (string, error
 		return sb.String(), fmt.Errorf("read response: %w", err)
 	}
 
+	if answer.N == 0 {
+		return "", fmt.Errorf("read response: %w", errAnswerTooLarge)
+	}
+
 	return sb.String(), nil
 }
 
 // parseIANAReferral extracts the RIR and WHOIS server from an IANA response.
-// The server is "" when the response has no referral.
-func (c *Client) parseIANAReferral(response string) (string, string) {
+// The server is "" when the response has no referral. A referral to a server
+// that is none of the five RIRs' is an error, so a forged or broken answer
+// cannot send the client to any other host or port.
+func (c *Client) parseIANAReferral(response string) (string, string, error) {
 	for line := range strings.SplitSeq(response, "\n") {
 		line = strings.TrimSpace(line)
 
@@ -197,27 +220,22 @@ func (c *Client) parseIANAReferral(response string) (string, string) {
 
 			switch {
 			case strings.Contains(server, "arin"):
-				return RIRARIN, whoisServerARIN
+				return RIRARIN, whoisServerARIN, nil
 			case strings.Contains(server, "ripe"):
-				return RIRRIPE, whoisServerRIPE
+				return RIRRIPE, whoisServerRIPE, nil
 			case strings.Contains(server, "apnic"):
-				return RIRAPNIC, whoisServerAPNIC
+				return RIRAPNIC, whoisServerAPNIC, nil
 			case strings.Contains(server, "lacnic"):
-				return RIRLACNIC, whoisServerLACNIC
+				return RIRLACNIC, whoisServerLACNIC, nil
 			case strings.Contains(server, "afrinic"):
-				return RIRAFRNIC, whoisServerAFRINIC
+				return RIRAFRNIC, whoisServerAFRINIC, nil
 			default:
-				// Unknown server, add port if missing
-				if !strings.Contains(server, ":") {
-					server += ":43"
-				}
-
-				return "", server
+				return "", "", fmt.Errorf("%w: %q", errUnknownServer, server)
 			}
 		}
 	}
 
-	return "", ""
+	return "", "", nil
 }
 
 // parseResponse extracts ASN information from a WHOIS response.

internal/whois/whois_test.go

@@ -6,6 +6,7 @@ import (
 	"errors"
 	"io"
 	"net"
+	"strings"
 	"sync"
 	"syscall"
 	"testing"
@@ -71,6 +72,13 @@ const ianaNoObject = `% IANA WHOIS server
 % This query returned 0 objects.
 `
 
+// fillerLine is one of the lines a server that does not stop keeps sending.
+const fillerLine = "remarks:        and more\n"
+
+// endlessAnswerSize is how many bytes of filler lines that server sends:
+// four times the client's cap on an answer.
+const endlessAnswerSize = 4 << 20
+
 // fakeServers returns a DialContext for a whois.Client that connects to a
 // fake server for each address in answers, which reads the query and sends
 // answers[address]. Connecting to any other address is refused.
@@ -187,3 +195,52 @@ func TestLookupASNFailsWhenIANAAnswersNoBlock(t *testing.T) {
 			err)
 	}
 }
+
+// TestLookupASNFailsWhenTheRIRAnswerDoesNotEnd looks up an AS whose RIR
+// server follows the record with filler lines far past any real answer.
+// The client stops reading at its cap and fails instead of keeping them.
+func TestLookupASNFailsWhenTheRIRAnswerDoesNotEnd(t *testing.T) {
+	t.Parallel()
+
+	endless := ripeRecord +
+		strings.Repeat(fillerLine, endlessAnswerSize/len(fillerLine))
+
+	_, err := lookUp(t, map[string]string{
+		ianaServer: ianaReferral,
+		ripeServer: endless,
+	})
+	if err == nil || errors.Is(err, whois.ErrNoReferral) {
+		t.Errorf("LookupASN error = %v, want one that is not ErrNoReferral",
+			err)
+	}
+}
+
+// TestLookupASNFailsOnAReferralToAnotherServer looks up an AS when IANA's
+// answer refers it to a server none of the five RIRs run, with a port and
+// without. The client follows a referral only to an RIR's server, so it
+// fails instead of asking that server, which would answer with a record.
+func TestLookupASNFailsOnAReferralToAnotherServer(t *testing.T) {
+	t.Parallel()
+
+	// Each case maps the server on the "refer:" line to the address the
+	// client would dial to follow it. IANA's answer also names the AS's
+	// block, so a client that dropped the referral would report no RIR.
+	for refer, address := range map[string]string{
+		"whois.example.net:8080": "whois.example.net:8080",
+		"whois.example.net":      "whois.example.net:43",
+	} {
+		t.Run(refer, func(t *testing.T) {
+			t.Parallel()
+
+			_, err := lookUp(t, map[string]string{
+				ianaServer: "refer:        " + refer + "\n\n" +
+					"as-block:     3154-3353\n",
+				address: ripeRecord,
+			})
+			if err == nil || errors.Is(err, whois.ErrNoReferral) {
+				t.Errorf("LookupASN error = %v, want one that is not "+
+					"ErrNoReferral", err)
+			}
+		})
+	}
+}