whoissdk/lookup_ops.go

package whois

import (
	"context"
	"fmt"
	"net"
	"net/url"
	"strings"
	"time"

	"b612.me/starnet"
	netproxy "golang.org/x/net/proxy"
)

// LookupCommonOps are controls shared by RDAP and WHOIS when specific values are absent.
type LookupCommonOps struct {
	Timeout time.Duration
	Proxy   string
}

// LookupRDAPOps are RDAP-only controls, powered by starnet request options.
// Proxy here overrides common Proxy for RDAP only.
type LookupRDAPOps struct {
	Timeout       time.Duration
	DialTimeout   time.Duration
	Proxy         string
	SkipTLSVerify *bool
	CustomIP      []string
	CustomDNS     []string
	Hosts         map[string]string
	LookupFunc    func(ctx context.Context, host string) ([]net.IPAddr, error)
	Headers       map[string]string
	RequestOpts   []starnet.RequestOpt
}

// LookupWHOISOps are WHOIS-safe controls.
type LookupWHOISOps struct {
	Timeout time.Duration
	Proxy   string
	Dialer  netproxy.Dialer
}

// LookupOps groups common + protocol-specific controls.
type LookupOps struct {
	Common LookupCommonOps
	RDAP   LookupRDAPOps
	WHOIS  LookupWHOISOps
}

// WithLookupOps sets all lookup ops.
func WithLookupOps(ops LookupOps) LookupOpt {
	return func(o *LookupOptions) {
		o.Ops = cloneLookupOps(ops)
	}
}

// WithLookupCommonOps sets shared lookup ops.
func WithLookupCommonOps(ops LookupCommonOps) LookupOpt {
	return func(o *LookupOptions) {
		o.Ops.Common = ops
	}
}

// WithLookupRDAPOps sets RDAP-only lookup ops.
func WithLookupRDAPOps(ops LookupRDAPOps) LookupOpt {
	return func(o *LookupOptions) {
		o.Ops.RDAP = cloneLookupRDAPOps(ops)
	}
}

// WithLookupWHOISOps sets WHOIS-only lookup ops.
func WithLookupWHOISOps(ops LookupWHOISOps) LookupOpt {
	return func(o *LookupOptions) {
		o.Ops.WHOIS = ops
	}
}

// WithLookupCommonTimeout sets shared timeout.
func WithLookupCommonTimeout(timeout time.Duration) LookupOpt {
	return func(o *LookupOptions) {
		o.Ops.Common.Timeout = timeout
	}
}

// WithLookupProxy sets common proxy for RDAP + WHOIS (WHOIS expects socks5 proxy).
func WithLookupProxy(proxy string) LookupOpt {
	return func(o *LookupOptions) {
		o.Ops.Common.Proxy = strings.TrimSpace(proxy)
	}
}

// WithLookupRDAPProxy sets RDAP-only proxy, overriding common proxy for RDAP.
func WithLookupRDAPProxy(proxy string) LookupOpt {
	return func(o *LookupOptions) {
		o.Ops.RDAP.Proxy = strings.TrimSpace(proxy)
	}
}

// WithLookupWHOISDialer sets WHOIS TCP dialer.
func WithLookupWHOISDialer(d netproxy.Dialer) LookupOpt {
	return func(o *LookupOptions) {
		o.Ops.WHOIS.Dialer = d
	}
}

// WithLookupWHOISProxy sets WHOIS-only proxy, overriding common proxy for WHOIS.
func WithLookupWHOISProxy(proxy string) LookupOpt {
	return func(o *LookupOptions) {
		o.Ops.WHOIS.Proxy = strings.TrimSpace(proxy)
	}
}

// WithLookupWHOISTimeout sets WHOIS timeout.
func WithLookupWHOISTimeout(timeout time.Duration) LookupOpt {
	return func(o *LookupOptions) {
		o.Ops.WHOIS.Timeout = timeout
	}
}

func cloneLookupOps(in LookupOps) LookupOps {
	out := in
	out.RDAP = cloneLookupRDAPOps(in.RDAP)
	return out
}

func cloneLookupRDAPOps(in LookupRDAPOps) LookupRDAPOps {
	out := in
	out.CustomIP = copyStringSlice(in.CustomIP)
	out.CustomDNS = copyStringSlice(in.CustomDNS)
	out.Hosts = copyStringMap(in.Hosts)
	out.Headers = copyStringMap(in.Headers)
	if len(in.RequestOpts) > 0 {
		out.RequestOpts = append([]starnet.RequestOpt(nil), in.RequestOpts...)
	}
	return out
}

func copyStringMap(in map[string]string) map[string]string {
	if len(in) == 0 {
		return nil
	}
	out := make(map[string]string, len(in))
	for k, v := range in {
		out[k] = v
	}
	return out
}

func (o LookupOptions) effectiveRDAPRequestOpts() []starnet.RequestOpt {
	base := []starnet.RequestOpt{
		starnet.WithHeader("Accept", "application/rdap+json, application/json"),
		starnet.WithUserAgent("b612-whois-rdap/1.0"),
		starnet.WithAutoFetch(true),
	}

	timeout := o.Ops.RDAP.Timeout
	if timeout <= 0 {
		timeout = o.Ops.Common.Timeout
	}
	if timeout > 0 {
		base = append(base, starnet.WithTimeout(timeout))
	}
	if o.Ops.RDAP.DialTimeout > 0 {
		base = append(base, starnet.WithDialTimeout(o.Ops.RDAP.DialTimeout))
	}
	if p := o.effectiveRDAPProxy(); p != "" {
		base = append(base, starnet.WithProxy(p))
	}
	if o.Ops.RDAP.SkipTLSVerify != nil {
		base = append(base, starnet.WithSkipTLSVerify(*o.Ops.RDAP.SkipTLSVerify))
	}
	if len(o.Ops.RDAP.CustomIP) > 0 {
		base = append(base, starnet.WithCustomIP(copyStringSlice(o.Ops.RDAP.CustomIP)))
	}
	if len(o.Ops.RDAP.CustomDNS) > 0 {
		base = append(base, starnet.WithCustomDNS(copyStringSlice(o.Ops.RDAP.CustomDNS)))
	}
	if len(o.Ops.RDAP.Headers) > 0 {
		base = append(base, starnet.WithHeaders(copyStringMap(o.Ops.RDAP.Headers)))
	}

	lookupFn := o.Ops.RDAP.LookupFunc
	if len(o.Ops.RDAP.Hosts) > 0 {
		lookupFn = buildHostsLookupFunc(o.Ops.RDAP.Hosts, lookupFn)
	}
	if lookupFn != nil {
		base = append(base, starnet.WithLookupFunc(lookupFn))
	}

	base = append(base, o.RDAPRequestOpts...)
	base = append(base, o.Ops.RDAP.RequestOpts...)
	return base
}

func (o LookupOptions) effectiveRDAPProxy() string {
	if p := strings.TrimSpace(o.Ops.RDAP.Proxy); p != "" {
		return p
	}
	return strings.TrimSpace(o.Ops.Common.Proxy)
}

func (o LookupOptions) effectiveWHOISProxy() string {
	if p := strings.TrimSpace(o.Ops.WHOIS.Proxy); p != "" {
		return p
	}
	return strings.TrimSpace(o.Ops.Common.Proxy)
}

func (o LookupOptions) effectiveWHOISDialer(defaultTimeout time.Duration) (netproxy.Dialer, error) {
	if o.Ops.WHOIS.Dialer != nil {
		return o.Ops.WHOIS.Dialer, nil
	}
	proxyAddr := o.effectiveWHOISProxy()
	if proxyAddr == "" {
		return nil, nil
	}
	if !strings.Contains(proxyAddr, "://") {
		proxyAddr = "socks5://" + proxyAddr
	}
	u, err := url.Parse(proxyAddr)
	if err != nil {
		return nil, fmt.Errorf("whois proxy parse failed: %w", err)
	}
	dialer, err := netproxy.FromURL(u, &net.Dialer{Timeout: o.effectiveWHOISTimeout(defaultTimeout)})
	if err != nil {
		return nil, fmt.Errorf("whois proxy unsupported: %w", err)
	}
	return dialer, nil
}

func (o LookupOptions) effectiveWHOISTimeout(defaultTimeout time.Duration) time.Duration {
	if o.Ops.WHOIS.Timeout > 0 {
		return o.Ops.WHOIS.Timeout
	}
	if o.Ops.Common.Timeout > 0 {
		return o.Ops.Common.Timeout
	}
	return defaultTimeout
}

func buildHostsLookupFunc(hosts map[string]string, fallback func(ctx context.Context, host string) ([]net.IPAddr, error)) func(ctx context.Context, host string) ([]net.IPAddr, error) {
	hostMap := make(map[string]net.IP, len(hosts))
	for host, ipStr := range hosts {
		h := strings.Trim(strings.ToLower(strings.TrimSpace(host)), ".")
		ip := net.ParseIP(strings.TrimSpace(ipStr))
		if h == "" || ip == nil {
			continue
		}
		hostMap[h] = ip
	}
	if len(hostMap) == 0 {
		return fallback
	}

	return func(ctx context.Context, host string) ([]net.IPAddr, error) {
		h := strings.Trim(strings.ToLower(strings.TrimSpace(host)), ".")
		if ip, ok := hostMap[h]; ok {
			return []net.IPAddr{{IP: ip}}, nil
		}
		if fallback != nil {
			return fallback(ctx, host)
		}
		return net.DefaultResolver.LookupIPAddr(ctx, host)
	}
}