whois_test.go

package whois

import (
	"bufio"
	"context"
	"errors"
	"fmt"
	"net"
	"strings"
	"sync/atomic"
	"testing"
	"time"
)

func TestDefaultConfig(t *testing.T) {
	config := DefaultConfig()

	if config.RootCacheDuration != 1*time.Hour {
		t.Errorf("expected RootCacheDuration to be 1 hour, got %v", config.RootCacheDuration)
	}

	if config.DefaultTimeout != 15*time.Second {
		t.Errorf("expected DefaultTimeout to be 15 seconds, got %v", config.DefaultTimeout)
	}

	if config.WhoisTLDServer != "whois.iana.org:43" {
		t.Errorf("expected WhoisTLDServer to be 'whois.iana.org:43', got %v", config.WhoisTLDServer)
	}
}

func TestSetup(t *testing.T) {
	config := &Config{
		RootCacheDuration: 2 * time.Hour,
		DefaultTimeout:    30 * time.Second,
		WhoisTLDServer:    "custom.whois.server:43",
	}

	whoisLookup := Setup(config)

	if whoisLookup.config.RootCacheDuration != 2*time.Hour {
		t.Errorf("expected RootCacheDuration to be 2 hours, got %v", whoisLookup.config.RootCacheDuration)
	}

	if whoisLookup.config.DefaultTimeout != 30*time.Second {
		t.Errorf("expected DefaultTimeout to be 30 seconds, got %v", whoisLookup.config.DefaultTimeout)
	}

	if whoisLookup.config.WhoisTLDServer != "custom.whois.server:43" {
		t.Errorf("expected WhoisTLDServer to be 'custom.whois.server:43', got %v", whoisLookup.config.WhoisTLDServer)
	}
}

func TestGetTLDWhoisServer(t *testing.T) {
	whoisLookup := Setup(nil)
	ctx := context.Background()

	// Mock TLD server
	whoisLookup.setTLDServerToCache("com", "whois.verisign-grs.com")

	server, err := whoisLookup.GetTLDWhoisServer(ctx, "com")
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	if server != "whois.verisign-grs.com" {
		t.Errorf("expected server to be 'whois.verisign-grs.com', got %v", server)
	}
}

func TestGetRegistryWhois_InvalidDomain(t *testing.T) {
	whoisLookup := Setup(nil)
	ctx := context.Background()

	_, _, err := whoisLookup.GetRegistryWhois(ctx, "invalid_domain")
	if err == nil {
		t.Fatal("expected an error for invalid domain, got nil")
	}
}

func TestGetRegistrarWhois_InvalidDomain(t *testing.T) {
	whoisLookup := Setup(nil)
	ctx := context.Background()

	_, _, err := whoisLookup.GetRegistrarWhois(ctx, "invalid_domain")
	if err == nil {
		t.Fatal("expected an error for invalid domain, got nil")
	}
}

func TestSetAndGetTLDServerFromCache(t *testing.T) {
	whoisLookup := Setup(nil)

	tld := "org"
	server := "whois.pir.org"

	whoisLookup.setTLDServerToCache(tld, server)

	cachedServer := whoisLookup.getTLDServerFromCache(tld)
	if cachedServer != server {
		t.Errorf("expected cached server to be %v, got %v", server, cachedServer)
	}
}

func TestGetTLDServerFromCache_Stale(t *testing.T) {
	whoisLookup := Setup(&Config{RootCacheDuration: 1 * time.Hour})

	whoisLookup.setTLDServerToCache("net", "whois.verisign-grs.com")
	// Backdate the entry past RootCacheDuration so it reads as stale.
	whoisLookup.rootWhoisServers["net"] = rootTLDCache{
		Host:        "whois.verisign-grs.com",
		LastUpdated: time.Now().Add(-2 * time.Hour),
	}

	if cached := whoisLookup.getTLDServerFromCache("net"); cached != "" {
		t.Errorf("expected a stale cache entry to be ignored, got %q", cached)
	}
}

func TestGetTLDWhoisServerWithLocalAddr_Error(t *testing.T) {
	addr := startMockWhoisServer(t, func(query string) string {
		return "% no matching record\n"
	})

	whoisLookup := Setup(&Config{WhoisTLDServer: addr, DefaultTimeout: 2 * time.Second})

	_, err := whoisLookup.GetTLDWhoisServerWithLocalAddr(context.Background(), "zz", nil)
	if !errors.Is(err, ErrWhoisServerNotFound) {
		t.Fatalf("expected ErrWhoisServerNotFound, got %v", err)
	}
}

// --- Local address accessors ---

func TestGetLocalAddr_SetLocalAddr(t *testing.T) {
	whoisLookup := Setup(nil)

	// Setup(nil) defaults to a non-nil, empty TCPAddr.
	if got := whoisLookup.GetLocalAddr(); got == nil {
		t.Fatal("expected default local addr to be non-nil")
	}

	custom := &net.TCPAddr{IP: net.ParseIP("127.0.0.1"), Port: 1234}
	whoisLookup.SetLocalAddr(custom)

	if got := whoisLookup.GetLocalAddr(); got != custom {
		t.Errorf("expected GetLocalAddr() to return the address set via SetLocalAddr(), got %+v", got)
	}
}

func TestSetup_CustomLocalAddr(t *testing.T) {
	custom := &net.TCPAddr{IP: net.ParseIP("192.0.2.1")}
	whoisLookup := Setup(&Config{LocalAddr: custom})

	if got := whoisLookup.GetLocalAddr(); got != custom {
		t.Errorf("expected Setup() to propagate Config.LocalAddr, got %+v", got)
	}
}

// --- wrapParser / convertContact ---

// googleComWhoisFixture is a trimmed real-world registry WHOIS response
// (google.com), used to exercise wrapParser's field mapping end-to-end.
const googleComWhoisFixture = `Domain Name: google.com
Registry Domain ID: 2138514_DOMAIN_COM-VRSN
Registrar WHOIS Server: whois.markmonitor.com
Registrar URL: http://www.markmonitor.com
Updated Date: 2019-09-09T08:39:04-0700
Creation Date: 1997-09-15T00:00:00-0700
Registrar Registration Expiration Date: 2028-09-13T00:00:00-0700
Registrar: MarkMonitor, Inc.
Registrar IANA ID: 292
Registrar Abuse Contact Email: abusecomplaints@markmonitor.com
Registrar Abuse Contact Phone: +1.2083895740
Domain Status: clientTransferProhibited (https://www.icann.org/epp#clientTransferProhibited)
Registrant Organization: Google LLC
Registrant State/Province: CA
Registrant Country: US
Admin Organization: Google LLC
Admin State/Province: CA
Admin Country: US
Tech Organization: Google LLC
Tech State/Province: CA
Tech Country: US
Name Server: ns2.google.com
Name Server: ns1.google.com
DNSSEC: unsigned
>>> Last update of WHOIS database: 2019-09-30T07:22:02-0700 <<<
`

func TestWrapParser_Success(t *testing.T) {
	info, err := wrapParser(googleComWhoisFixture)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	if info.Domain == nil {
		t.Fatal("expected Domain to be populated")
	}
	if info.Domain.Domain != "google.com" {
		t.Errorf("expected domain 'google.com', got %q", info.Domain.Domain)
	}
	if info.Domain.WhoisServer != "whois.markmonitor.com" {
		t.Errorf("expected whois server 'whois.markmonitor.com', got %q", info.Domain.WhoisServer)
	}

	if info.Registrar == nil || info.Registrar.Name != "MarkMonitor, Inc." {
		t.Errorf("expected registrar name 'MarkMonitor, Inc.', got %+v", info.Registrar)
	} else if info.Registrar.Email != "abusecomplaints@markmonitor.com" {
		t.Errorf("expected registrar email to be mapped, got %q", info.Registrar.Email)
	}

	if info.Registrant == nil || info.Registrant.Organization != "Google LLC" {
		t.Errorf("expected registrant organization 'Google LLC', got %+v", info.Registrant)
	}
	if info.Administrative == nil || info.Administrative.Organization != "Google LLC" {
		t.Errorf("expected administrative organization 'Google LLC', got %+v", info.Administrative)
	}
	if info.Technical == nil || info.Technical.Organization != "Google LLC" {
		t.Errorf("expected technical organization 'Google LLC', got %+v", info.Technical)
	}
}

func TestWrapParser_Error(t *testing.T) {
	info, err := wrapParser("this is not a whois response")
	if err == nil {
		t.Fatal("expected an error for unparseable whois text, got nil")
	}
	if info.Domain != nil {
		t.Errorf("expected a zero-value WhoisInfo on error, got Domain=%+v", info.Domain)
	}
}

// --- Mock WHOIS server helpers ---

// startMockWhoisServer starts a local TCP listener that speaks the simple
// WHOIS protocol: it reads one line per connection and, for each, writes
// back whatever respond returns before closing the connection. It runs for
// the lifetime of the test.
func startMockWhoisServer(t *testing.T, respond func(query string) string) (addr string) {
	t.Helper()

	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("net.Listen() error: %v", err)
	}
	t.Cleanup(func() { ln.Close() })

	go func() {
		for {
			conn, err := ln.Accept()
			if err != nil {
				return // listener closed, test is done
			}
			go func(c net.Conn) {
				defer c.Close()
				line, _ := bufio.NewReader(c).ReadString('\n')
				fmt.Fprint(c, respond(strings.TrimSpace(line)))
			}(conn)
		}
	}()

	return ln.Addr().String()
}

// useWhoisServerAddr points queryWhois's dial target at addr's port for the
// duration of the test (queryWhois otherwise always dials port 43), restoring
// the original value on cleanup.
func useWhoisServerAddr(t *testing.T, addr string) {
	t.Helper()

	_, port, err := net.SplitHostPort(addr)
	if err != nil {
		t.Fatalf("net.SplitHostPort(%q) error: %v", addr, err)
	}

	orig := whoisPort
	whoisPort = port
	t.Cleanup(func() { whoisPort = orig })
}

// --- getWhoisServerForTLD ---

func TestGetWhoisServerForTLD_Found(t *testing.T) {
	addr := startMockWhoisServer(t, func(query string) string {
		if query != "com" {
			t.Errorf("expected IANA query 'com', got %q", query)
		}
		return "% IANA WHOIS server\nwhois:        whois.verisign-grs.com\n"
	})

	whoisLookup := Setup(&Config{WhoisTLDServer: addr, DefaultTimeout: 2 * time.Second})

	server, err := whoisLookup.getWhoisServerForTLD(context.Background(), "com", nil)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if server != "whois.verisign-grs.com" {
		t.Errorf("expected 'whois.verisign-grs.com', got %q", server)
	}

	// A successful lookup should populate the cache.
	if cached := whoisLookup.getTLDServerFromCache("com"); cached != "whois.verisign-grs.com" {
		t.Errorf("expected result to be cached, got %q", cached)
	}
}

func TestGetWhoisServerForTLD_NotFound(t *testing.T) {
	addr := startMockWhoisServer(t, func(query string) string {
		return "% no matching record\n"
	})

	whoisLookup := Setup(&Config{WhoisTLDServer: addr, DefaultTimeout: 2 * time.Second})

	_, err := whoisLookup.getWhoisServerForTLD(context.Background(), "zz", nil)
	if !errors.Is(err, ErrWhoisServerNotFound) {
		t.Fatalf("expected ErrWhoisServerNotFound, got %v", err)
	}
}

// --- queryWhois ---

func TestQueryWhois_Success(t *testing.T) {
	const want = "Domain Name: example.com\nRegistrar: Example Registrar\n"

	addr := startMockWhoisServer(t, func(query string) string {
		if query != "example.com" {
			t.Errorf("expected query 'example.com', got %q", query)
		}
		return want
	})
	useWhoisServerAddr(t, addr)

	whoisLookup := Setup(nil)
	host, _, _ := net.SplitHostPort(addr)

	raw, err := whoisLookup.queryWhois(context.Background(), "example.com", host, 2*time.Second, nil)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if raw != want {
		t.Errorf("expected %q, got %q", want, raw)
	}
}

// TestQueryWhois_Timeout confirms that queryWhois bounds the read phase, not
// just the initial connect: a server that accepts the connection but never
// sends a response must not be able to hang the call indefinitely.
func TestQueryWhois_Timeout(t *testing.T) {
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("net.Listen() error: %v", err)
	}
	t.Cleanup(func() { ln.Close() })

	go func() {
		conn, err := ln.Accept()
		if err != nil {
			return
		}
		defer conn.Close()
		time.Sleep(1 * time.Second) // outlast the client's read deadline below
	}()
	useWhoisServerAddr(t, ln.Addr().String())

	whoisLookup := Setup(nil)
	host, _, _ := net.SplitHostPort(ln.Addr().String())

	const timeout = 200 * time.Millisecond
	start := time.Now()
	_, err = whoisLookup.queryWhois(context.Background(), "example.com", host, timeout, nil)
	elapsed := time.Since(start)

	if err == nil {
		t.Fatal("expected a timeout error, got nil")
	}
	if elapsed >= 1*time.Second {
		t.Errorf("queryWhois took %v to return; the read deadline does not appear to be enforced", elapsed)
	}
}

// TestQueryWhois_ResponseSizeCap confirms queryWhois bounds how much of a
// server's response it will buffer, rather than reading an unbounded amount.
func TestQueryWhois_ResponseSizeCap(t *testing.T) {
	origCap := maxWhoisResponseSize
	maxWhoisResponseSize = 16 // shrink so the test doesn't need to push megabytes
	t.Cleanup(func() { maxWhoisResponseSize = origCap })

	addr := startMockWhoisServer(t, func(query string) string {
		return strings.Repeat("a", 1024) + "\n"
	})
	useWhoisServerAddr(t, addr)

	whoisLookup := Setup(nil)
	host, _, _ := net.SplitHostPort(addr)

	_, err := whoisLookup.queryWhois(context.Background(), "example.com", host, 2*time.Second, nil)
	if err == nil {
		t.Fatal("expected an error when the response exceeds maxWhoisResponseSize, got nil")
	}
}

// --- Registry/Registrar lookups: localAddr threading & failure handling ---

// whoisFixtureNoReferral is a minimal registry-level response with no
// "Registrar WHOIS Server" field, so lookups stop after a single hop.
const whoisFixtureNoReferral = `Domain Name: example.com
Registry Domain ID: 12345_DOMAIN_COM-VRSN
Registrar: Example Registrar, Inc.
Domain Status: clientTransferProhibited (https://www.icann.org/epp#clientTransferProhibited)
Name Server: ns1.example.com
Name Server: ns2.example.com
DNSSEC: unsigned
>>> Last update of WHOIS database: 2024-06-01T00:00:00Z <<<
`

// whoisFixtureWithReferral is the same as whoisFixtureNoReferral but points
// the registrar WHOIS server back at 127.0.0.1, so a registrar-level lookup
// takes a second hop against the same mock server.
const whoisFixtureWithReferral = `Domain Name: example.com
Registry Domain ID: 12345_DOMAIN_COM-VRSN
Registrar WHOIS Server: 127.0.0.1
Registrar: Example Registrar, Inc.
Domain Status: clientTransferProhibited (https://www.icann.org/epp#clientTransferProhibited)
Name Server: ns1.example.com
Name Server: ns2.example.com
DNSSEC: unsigned
>>> Last update of WHOIS database: 2024-06-01T00:00:00Z <<<
`

// testLocalAddrHonored proves that an explicit localAddr passed to a
// "WithLocalAddr" lookup method actually reaches the outbound connections,
// rather than being silently dropped in favor of the WhoisLookup's global
// default local address.
func testLocalAddrHonored(t *testing.T, call func(whoisLookup *WhoisLookup, ctx context.Context, localAddr *net.TCPAddr) error) {
	t.Helper()

	ianaAddr := startMockWhoisServer(t, func(query string) string {
		return "whois: 127.0.0.1\n"
	})
	whoisAddr := startMockWhoisServer(t, func(query string) string {
		return whoisFixtureNoReferral
	})
	useWhoisServerAddr(t, whoisAddr)

	whoisLookup := Setup(&Config{WhoisTLDServer: ianaAddr, DefaultTimeout: 2 * time.Second})

	// Poison the global default local address: any internal call that fails
	// to thread the explicit localAddr through will try to dial from this
	// address instead and fail (TEST-NET-3, RFC 5737: never assignable here).
	whoisLookup.SetLocalAddr(&net.TCPAddr{IP: net.ParseIP("203.0.113.1")})

	validLocalAddr := &net.TCPAddr{IP: net.ParseIP("127.0.0.1")}

	if err := call(whoisLookup, context.Background(), validLocalAddr); err != nil {
		t.Fatalf("expected the explicit localAddr to be honored, got error: %v", err)
	}
}

func TestGetRegistryWhoisWithLocalAddr_UsesProvidedLocalAddr(t *testing.T) {
	testLocalAddrHonored(t, func(whoisLookup *WhoisLookup, ctx context.Context, localAddr *net.TCPAddr) error {
		_, _, err := whoisLookup.GetRegistryWhoisWithLocalAddr(ctx, "example.com", localAddr)
		return err
	})
}

func TestGetRegistrarWhoisWithLocalAddr_UsesProvidedLocalAddr(t *testing.T) {
	testLocalAddrHonored(t, func(whoisLookup *WhoisLookup, ctx context.Context, localAddr *net.TCPAddr) error {
		_, _, err := whoisLookup.GetRegistrarWhoisWithLocalAddr(ctx, "example.com", localAddr)
		return err
	})
}

// TestGetRegistrarWhoisWithLocalAddr_PreservesRegistryDataOnRegistrarParseFailure
// confirms that when the second-hop (registrar-level) query fails to parse,
// the already-successful registry-level result is still returned alongside
// the error, rather than being discarded.
func TestGetRegistrarWhoisWithLocalAddr_PreservesRegistryDataOnRegistrarParseFailure(t *testing.T) {
	ianaAddr := startMockWhoisServer(t, func(query string) string {
		return "whois: 127.0.0.1\n"
	})

	var callCount int32
	whoisAddr := startMockWhoisServer(t, func(query string) string {
		if atomic.AddInt32(&callCount, 1) == 1 {
			return whoisFixtureWithReferral
		}
		// Simulate a broken/uncooperative registrar server: this doesn't
		// parse as a valid domain record.
		return "not a valid whois response\n"
	})
	useWhoisServerAddr(t, whoisAddr)

	whoisLookup := Setup(&Config{WhoisTLDServer: ianaAddr, DefaultTimeout: 2 * time.Second})

	whoisInfo, whoisRaw, err := whoisLookup.GetRegistrarWhoisWithLocalAddr(context.Background(), "example.com", nil)
	if err == nil {
		t.Fatal("expected an error from the failed registrar-level parse, got nil")
	}
	if whoisInfo.Domain == nil || whoisInfo.Domain.Domain != "example.com" {
		t.Fatalf("expected the registry-level result to be preserved despite the registrar failure, got %+v", whoisInfo.Domain)
	}
	if !strings.Contains(whoisRaw, "Domain Name: example.com") {
		t.Errorf("expected the registry-level raw whois to be preserved, got %q", whoisRaw)
	}
}

// --- GetWhoisWithLocalAddr ---

func TestGetWhoisWithLocalAddr_InvalidDomain(t *testing.T) {
	whoisLookup := Setup(nil)

	_, err := whoisLookup.GetWhoisWithLocalAddr(context.Background(), "invalid_domain", nil)
	if err == nil {
		t.Fatal("expected an error for invalid domain, got nil")
	}
}

func TestGetWhoisWithLocalAddr_Success(t *testing.T) {
	ianaAddr := startMockWhoisServer(t, func(query string) string {
		return "whois: 127.0.0.1\n"
	})

	var callCount int32
	whoisAddr := startMockWhoisServer(t, func(query string) string {
		if atomic.AddInt32(&callCount, 1) == 1 {
			return whoisFixtureWithReferral
		}
		return whoisFixtureNoReferral
	})
	useWhoisServerAddr(t, whoisAddr)

	whoisLookup := Setup(&Config{WhoisTLDServer: ianaAddr, DefaultTimeout: 2 * time.Second})

	result, err := whoisLookup.GetWhoisWithLocalAddr(context.Background(), "example.com", nil)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	if result.TLD != "com" {
		t.Errorf("expected TLD 'com', got %q", result.TLD)
	}
	if result.RegistryWhoisServer != "127.0.0.1" {
		t.Errorf("expected registry whois server '127.0.0.1', got %q", result.RegistryWhoisServer)
	}
	if result.RegistryWhois == nil || result.RegistryWhois.Domain == nil || result.RegistryWhois.Domain.WhoisServer != "127.0.0.1" {
		t.Fatalf("expected registry-level result with a registrar referral, got %+v", result.RegistryWhois)
	}
	if result.RegistrarWhois == nil || result.RegistrarWhois.Domain == nil {
		t.Fatalf("expected registrar-level result to be populated, got %+v", result.RegistrarWhois)
	}
	if result.RegistrarWhois.Domain.WhoisServer != "" {
		t.Errorf("expected registrar-level result to have no further referral, got %q", result.RegistrarWhois.Domain.WhoisServer)
	}
}

func TestGetWhoisWithLocalAddr_MissingRegistrarWhoisServer(t *testing.T) {
	ianaAddr := startMockWhoisServer(t, func(query string) string {
		return "whois: 127.0.0.1\n"
	})
	whoisAddr := startMockWhoisServer(t, func(query string) string {
		return whoisFixtureNoReferral
	})
	useWhoisServerAddr(t, whoisAddr)

	whoisLookup := Setup(&Config{WhoisTLDServer: ianaAddr, DefaultTimeout: 2 * time.Second})

	result, err := whoisLookup.GetWhoisWithLocalAddr(context.Background(), "example.com", nil)
	if !errors.Is(err, ErrRegistryMissingWhoisServer) {
		t.Fatalf("expected ErrRegistryMissingWhoisServer, got %v", err)
	}
	if result.RegistryWhois == nil || result.RegistryWhois.Domain == nil || result.RegistryWhois.Domain.Domain != "example.com" {
		t.Errorf("expected the registry-level result to still be populated, got %+v", result.RegistryWhois)
	}
}