whoissdk/lookup.go

package whois

import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"strings"
	"time"

	"b612.me/starnet"
)

// LookupMode controls how RDAP and WHOIS are used.
type LookupMode string

const (
	LookupModeAutoPreferRDAP  LookupMode = "auto-rdap"
	LookupModeWHOISOnly       LookupMode = "whois-only"
	LookupModeRDAPOnly        LookupMode = "rdap-only"
	LookupModeBothPreferRDAP  LookupMode = "both-rdap"
	LookupModeBothPreferWHOIS LookupMode = "both-whois"
)

const (
	LookupSourceRDAP   = "rdap"
	LookupSourceWHOIS  = "whois"
	LookupSourceMerged = "merged"
)

var ErrLookupClientNil = errors.New("whois/lookup: client is nil")

// LookupMeta records lookup path and fallback details.
type LookupMeta struct {
	Mode         LookupMode
	Source       string
	RDAPEndpoint string
	RDAPStatus   int
	RDAPError    string
	WHOISError   string
	WarningList  []string
}

// Warnings returns copy of warning list.
func (m LookupMeta) Warnings() []string {
	return copyStringSlice(m.WarningList)
}

// LookupOptions controls hybrid lookup behavior.
type LookupOptions struct {
	Mode            LookupMode
	Strategy        LookupStrategy
	WhoisOptions    QueryOptions
	WhoisServers    []string
	RDAPClient      *RDAPClient
	RDAPBootstrap   *RDAPBootstrapLoadOptions
	RDAPRetry       *RDAPRetryPolicy
	Ops             LookupOps
	RDAPRequestOpts []starnet.RequestOpt
}

// LookupOpt mutates LookupOptions.
type LookupOpt func(*LookupOptions)

func defaultLookupOptions() LookupOptions {
	return LookupOptions{
		Mode:         LookupModeAutoPreferRDAP,
		Strategy:     defaultLookupStrategy(),
		WhoisOptions: QueryOptions{Level: QueryAuto},
		Ops:          LookupOps{},
	}
}

// WithLookupMode sets hybrid lookup mode.
func WithLookupMode(mode LookupMode) LookupOpt {
	return func(o *LookupOptions) { o.Mode = mode }
}

// WithLookupWhoisOptions sets whois query options for lookup.
func WithLookupWhoisOptions(opt QueryOptions) LookupOpt {
	return func(o *LookupOptions) { o.WhoisOptions = opt }
}

// WithLookupWhoisServers sets explicit whois server candidates.
func WithLookupWhoisServers(servers ...string) LookupOpt {
	return func(o *LookupOptions) { o.WhoisServers = append([]string(nil), servers...) }
}

// WithLookupRDAPClient sets custom RDAP client.
func WithLookupRDAPClient(c *RDAPClient) LookupOpt {
	return func(o *LookupOptions) { o.RDAPClient = c }
}

// WithLookupRDAPRetryPolicy sets RDAP retry policy for lookup path.
func WithLookupRDAPRetryPolicy(policy RDAPRetryPolicy) LookupOpt {
	return func(o *LookupOptions) {
		p := normalizeRDAPRetryPolicy(policy)
		o.RDAPRetry = &p
	}
}

// WithLookupRDAPBootstrapLoadOptions sets layered bootstrap load options used
// when RDAP client is not provided explicitly.
func WithLookupRDAPBootstrapLoadOptions(opt RDAPBootstrapLoadOptions) LookupOpt {
	return func(o *LookupOptions) {
		clone := opt.Clone()
		o.RDAPBootstrap = &clone
	}
}

// WithLookupRDAPBootstrapCache sets layered bootstrap cache ttl/key.
// It only affects auto-created RDAP client path when RDAPClient is not set.
func WithLookupRDAPBootstrapCache(ttl time.Duration, cacheKey string) LookupOpt {
	return func(o *LookupOptions) {
		boot := ensureLookupRDAPBootstrap(o)
		boot.CacheTTL = ttl
		boot.CacheKey = strings.TrimSpace(cacheKey)
	}
}

// WithLookupRDAPBootstrapLocalFiles appends local bootstrap overlay files.
func WithLookupRDAPBootstrapLocalFiles(paths ...string) LookupOpt {
	return func(o *LookupOptions) {
		boot := ensureLookupRDAPBootstrap(o)
		boot.LocalFiles = append(boot.LocalFiles, paths...)
	}
}

// WithLookupRDAPBootstrapRemoteRefresh enables/disables remote refresh.
func WithLookupRDAPBootstrapRemoteRefresh(enabled bool, remoteURL string) LookupOpt {
	return func(o *LookupOptions) {
		boot := ensureLookupRDAPBootstrap(o)
		boot.RefreshRemote = enabled
		boot.RemoteURL = strings.TrimSpace(remoteURL)
	}
}

// WithLookupRDAPBootstrapIgnoreRemoteError controls whether remote refresh errors are ignored.
func WithLookupRDAPBootstrapIgnoreRemoteError(enabled bool) LookupOpt {
	return func(o *LookupOptions) {
		boot := ensureLookupRDAPBootstrap(o)
		boot.IgnoreRemoteError = enabled
	}
}

// WithLookupRDAPBootstrapAllowStaleOnError controls stale cache fallback on refresh failure.
func WithLookupRDAPBootstrapAllowStaleOnError(enabled bool) LookupOpt {
	return func(o *LookupOptions) {
		boot := ensureLookupRDAPBootstrap(o)
		boot.AllowStaleOnError = enabled
	}
}

// WithLookupRDAPRequestOpts sets RDAP-only starnet request options.
// Note: these options only apply to RDAP HTTP calls, not TCP WHOIS calls.
func WithLookupRDAPRequestOpts(opts ...starnet.RequestOpt) LookupOpt {
	return func(o *LookupOptions) {
		o.RDAPRequestOpts = append([]starnet.RequestOpt(nil), opts...)
	}
}

// Lookup is a unified entry for RDAP/WHOIS hybrid query.
func (c *Client) Lookup(domain string, opts ...LookupOpt) (Result, LookupMeta, error) {
	return c.LookupContext(context.Background(), domain, opts...)
}

// LookupContext is context-aware unified RDAP/WHOIS query.
func (c *Client) LookupContext(ctx context.Context, domain string, opts ...LookupOpt) (Result, LookupMeta, error) {
	if c == nil {
		return Result{}, LookupMeta{}, ErrLookupClientNil
	}

	normalizedInput, targetKind, err := normalizeLookupTargetInput(domain)
	if err != nil {
		return Result{}, LookupMeta{}, err
	}
	domain = normalizedInput

	cfg := defaultLookupOptions()
	for _, opt := range opts {
		if opt != nil {
			opt(&cfg)
		}
	}
	cfg.Mode = normalizeLookupMode(cfg.Mode)
	resolvedMode := cfg.resolveMode(domain)
	if targetKind != lookupTargetDomain {
		resolvedMode = cfg.Mode
	}

	meta := LookupMeta{Mode: resolvedMode}

	switch resolvedMode {
	case LookupModeWHOISOnly:
		whoisRes, whoisErr := c.lookupWHOIS(ctx, domain, cfg)
		meta.WHOISError = errString(whoisErr)
		if whoisErr != nil {
			return Result{}, meta, whoisErr
		}
		meta.Source = LookupSourceWHOIS
		return whoisRes, meta, nil
	case LookupModeRDAPOnly:
		rdapRes, rdapResp, rdapErr := c.lookupRDAP(ctx, domain, targetKind, cfg)
		fillRDAPMeta(&meta, rdapResp, rdapErr)
		if rdapErr != nil {
			if nf, ok := rdapNotFoundResult(domain, rdapErr); ok {
				meta.Source = LookupSourceRDAP
				meta.WarningList = append(meta.WarningList, "rdap returned 404 not found")
				return nf, meta, nil
			}
			return Result{}, meta, rdapErr
		}
		meta.Source = LookupSourceRDAP
		return rdapRes, meta, nil
	case LookupModeBothPreferRDAP, LookupModeBothPreferWHOIS:
		rdapRes, rdapResp, rdapErr := c.lookupRDAP(ctx, domain, targetKind, cfg)
		whoisRes, whoisErr := c.lookupWHOIS(ctx, domain, cfg)
		fillRDAPMeta(&meta, rdapResp, rdapErr)
		meta.WHOISError = errString(whoisErr)

		if rdapErr != nil && whoisErr != nil {
			return Result{}, meta, combineLookupErrors(rdapErr, whoisErr)
		}
		if rdapErr == nil && whoisErr == nil {
			meta.Source = LookupSourceMerged
			preferRDAP := resolvedMode == LookupModeBothPreferRDAP
			return mergeLookupResults(rdapRes, whoisRes, preferRDAP), meta, nil
		}
		if rdapErr == nil {
			meta.Source = LookupSourceRDAP
			meta.WarningList = append(meta.WarningList, "whois fallback failed")
			return rdapRes, meta, nil
		}

		meta.Source = LookupSourceWHOIS
		meta.WarningList = append(meta.WarningList, "rdap fallback failed")
		return whoisRes, meta, nil
	default:
		fallthrough
	case LookupModeAutoPreferRDAP:
		rdapRes, rdapResp, rdapErr := c.lookupRDAP(ctx, domain, targetKind, cfg)
		fillRDAPMeta(&meta, rdapResp, rdapErr)
		if rdapErr == nil {
			meta.Source = LookupSourceRDAP
			return rdapRes, meta, nil
		}
		if nf, ok := rdapNotFoundResult(domain, rdapErr); ok {
			meta.Source = LookupSourceRDAP
			meta.WarningList = append(meta.WarningList, "rdap returned 404 not found")
			return nf, meta, nil
		}

		whoisRes, whoisErr := c.lookupWHOIS(ctx, domain, cfg)
		meta.WHOISError = errString(whoisErr)
		if whoisErr != nil {
			return Result{}, meta, combineLookupErrors(rdapErr, whoisErr)
		}
		meta.Source = LookupSourceWHOIS
		meta.WarningList = append(meta.WarningList, "rdap lookup failed, switched to whois")
		return whoisRes, meta, nil
	}
}

func (c *Client) lookupWHOIS(ctx context.Context, domain string, cfg LookupOptions) (Result, error) {
	whoisClient := c.cloneForLookup()
	whoisTimeout := cfg.effectiveWHOISTimeout(c.timeout)
	whoisClient.SetTimeout(whoisTimeout)
	dialer, err := cfg.effectiveWHOISDialer(c.timeout)
	if err != nil {
		return Result{}, fmt.Errorf("prepare whois dialer: %w", err)
	}
	if dialer != nil {
		whoisClient.SetDialer(dialer)
	}

	opt := cfg.WhoisOptions
	if len(opt.OverrideServers) == 0 && opt.OverrideServer == "" && len(cfg.WhoisServers) > 0 {
		opt.OverrideServers = normalizeServerList(cfg.WhoisServers)
		if len(opt.OverrideServers) > 0 {
			opt.OverrideServer = opt.OverrideServers[0]
		}
	}
	return whoisClient.WhoisWithOptionsContext(ctx, domain, opt)
}

func ensureLookupRDAPBootstrap(o *LookupOptions) *RDAPBootstrapLoadOptions {
	if o.RDAPBootstrap == nil {
		o.RDAPBootstrap = &RDAPBootstrapLoadOptions{}
	}
	return o.RDAPBootstrap
}

func (c *Client) lookupRDAP(ctx context.Context, domain string, targetKind lookupTargetKind, cfg LookupOptions) (Result, *RDAPResponse, error) {
	rdc := cfg.RDAPClient
	var err error
	if rdc == nil {
		if cfg.RDAPBootstrap != nil {
			rdc, err = NewRDAPClientWithLayeredBootstrap(ctx, *cfg.RDAPBootstrap)
		} else {
			rdc, err = NewRDAPClient()
		}
		if err != nil {
			return Result{}, nil, err
		}
	}
	if cfg.RDAPRetry != nil {
		rdc.SetRetryPolicy(*cfg.RDAPRetry)
	}

	rdapReqOpts := cfg.effectiveRDAPRequestOpts()
	rdapResp, err := rdc.Query(ctx, domain, rdapReqOpts...)
	if err != nil {
		return Result{}, nil, err
	}
	res, err := rdapResponseToResult(domain, rdapResp, targetKind)
	if err != nil {
		return Result{}, rdapResp, err
	}
	return res, rdapResp, nil
}

func (c *Client) cloneForLookup() *Client {
	if c == nil {
		return nil
	}
	out := &Client{
		extCache:         make(map[string]string),
		dialer:           c.dialer,
		timeout:          c.timeout,
		elapsed:          c.elapsed,
		negativeCacheTTL: c.negativeCacheTTL,
	}
	c.mu.Lock()
	for k, v := range c.extCache {
		out.extCache[k] = v
	}
	c.mu.Unlock()
	return out
}

func normalizeLookupMode(mode LookupMode) LookupMode {
	switch LookupMode(strings.TrimSpace(strings.ToLower(string(mode)))) {
	case LookupModeWHOISOnly:
		return LookupModeWHOISOnly
	case LookupModeRDAPOnly:
		return LookupModeRDAPOnly
	case LookupModeBothPreferRDAP:
		return LookupModeBothPreferRDAP
	case LookupModeBothPreferWHOIS:
		return LookupModeBothPreferWHOIS
	default:
		return LookupModeAutoPreferRDAP
	}
}

func fillRDAPMeta(meta *LookupMeta, resp *RDAPResponse, err error) {
	if meta == nil {
		return
	}
	if resp != nil {
		meta.RDAPEndpoint = resp.Endpoint
		meta.RDAPStatus = resp.StatusCode
	}
	if err != nil {
		meta.RDAPError = err.Error()
		var httpErr *RDAPHTTPError
		if errors.As(err, &httpErr) {
			meta.RDAPEndpoint = httpErr.Endpoint
			meta.RDAPStatus = httpErr.StatusCode
		}
	}
}

func rdapNotFoundResult(domain string, err error) (Result, bool) {
	var httpErr *RDAPHTTPError
	if !errors.As(err, &httpErr) {
		return Result{}, false
	}
	if httpErr.StatusCode != http.StatusNotFound {
		return Result{}, false
	}
	out := Result{
		exists:  false,
		domain:  domain,
		rawData: string(httpErr.Body),
	}
	out.meta = buildResultMeta(out, "rdap", httpErr.Endpoint)
	out.meta.Charset = "utf-8"
	return out, true
}

func errString(err error) string {
	if err == nil {
		return ""
	}
	return err.Error()
}

func combineLookupErrors(rdapErr, whoisErr error) error {
	if rdapErr == nil {
		return whoisErr
	}
	if whoisErr == nil {
		return rdapErr
	}
	return fmt.Errorf("rdap error: %w; whois error: %v", rdapErr, whoisErr)
}

func mergeLookupResults(rdapResult, whoisResult Result, preferRDAP bool) Result {
	if preferRDAP {
		return mergeResultWithFallback(rdapResult, whoisResult)
	}
	return mergeResultWithFallback(whoisResult, rdapResult)
}

func mergeResultWithFallback(primary, fallback Result) Result {
	out := primary

	if !out.exists && fallback.exists {
		out.exists = true
	}
	if out.domain == "" {
		out.domain = fallback.domain
	}
	if out.domainID == "" {
		out.domainID = fallback.domainID
	}
	if out.registrar == "" {
		out.registrar = fallback.registrar
	}
	if !out.hasRegisterDate && fallback.hasRegisterDate {
		out.hasRegisterDate = true
		out.registerDate = fallback.registerDate
	}
	if !out.hasUpdateDate && fallback.hasUpdateDate {
		out.hasUpdateDate = true
		out.updateDate = fallback.updateDate
	}
	if !out.hasExpireDate && fallback.hasExpireDate {
		out.hasExpireDate = true
		out.expireDate = fallback.expireDate
	}
	if out.dnssec == "" {
		out.dnssec = fallback.dnssec
	}
	if out.whoisSer == "" {
		out.whoisSer = fallback.whoisSer
	}
	if out.ianaID == "" {
		out.ianaID = fallback.ianaID
	}

	out.statusRaw = appendUniqueStrings(out.statusRaw, fallback.statusRaw...)
	out.nsServers = appendUniqueStrings(out.nsServers, fallback.nsServers...)
	out.nsIps = appendUniqueStrings(out.nsIps, fallback.nsIps...)

	out.registerInfo = mergePersonalInfo(out.registerInfo, fallback.registerInfo)
	out.adminInfo = mergePersonalInfo(out.adminInfo, fallback.adminInfo)
	out.techInfo = mergePersonalInfo(out.techInfo, fallback.techInfo)

	out.rawData = mergeRawData(primary.rawData, fallback.rawData)
	server := primary.meta.Server
	if server == "" {
		server = fallback.meta.Server
	}
	out.meta = buildResultMeta(out, LookupSourceMerged, server)
	out.meta.Charset = combineCharset(primary.meta.Charset, fallback.meta.Charset)
	return out
}

func mergeRawData(primary, fallback string) string {
	primary = strings.TrimSpace(primary)
	fallback = strings.TrimSpace(fallback)
	switch {
	case primary == "" && fallback == "":
		return ""
	case primary == "":
		return fallback
	case fallback == "":
		return primary
	case primary == fallback:
		return primary
	default:
		return primary + "\n\n----- FALLBACK DATA -----\n\n" + fallback
	}
}

func mergePersonalInfo(primary, fallback PersonalInfo) PersonalInfo {
	out := primary
	if out.FirstName == "" {
		out.FirstName = fallback.FirstName
	}
	if out.LastName == "" {
		out.LastName = fallback.LastName
	}
	if out.Name == "" {
		out.Name = fallback.Name
	}
	if out.Org == "" {
		out.Org = fallback.Org
	}
	if out.Fax == "" {
		out.Fax = fallback.Fax
	}
	if out.FaxExt == "" {
		out.FaxExt = fallback.FaxExt
	}
	if out.Addr == "" {
		out.Addr = fallback.Addr
	}
	if out.City == "" {
		out.City = fallback.City
	}
	if out.State == "" {
		out.State = fallback.State
	}
	if out.Country == "" {
		out.Country = fallback.Country
	}
	if out.Zip == "" {
		out.Zip = fallback.Zip
	}
	if out.Phone == "" {
		out.Phone = fallback.Phone
	}
	if out.PhoneExt == "" {
		out.PhoneExt = fallback.PhoneExt
	}
	if out.Email == "" {
		out.Email = fallback.Email
	}
	return out
}