whoissdk/rdap_client.go

package whois

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"time"

	"b612.me/starnet"
)

var (
	ErrRDAPClientNil      = errors.New("whois/rdap: client is nil")
	ErrRDAPServerNotFound = errors.New("whois/rdap: rdap server not found")
	ErrRDAPDomainInvalid  = errors.New("whois/rdap: domain is invalid")
	ErrRDAPIPInvalid      = errors.New("whois/rdap: ip is invalid")
	ErrRDAPASNInvalid     = errors.New("whois/rdap: asn is invalid")
	ErrRDAPQueryInvalid   = errors.New("whois/rdap: query target is invalid")
)

var (
	defaultRDAPIPServers = []string{
		"https://rdap.arin.net/registry/",
		"https://rdap.db.ripe.net/",
		"https://rdap.apnic.net/",
		"https://rdap.lacnic.net/rdap/",
		"https://rdap.afrinic.net/rdap/",
	}
	defaultRDAPASNServers = []string{
		"https://rdap.arin.net/registry/",
		"https://rdap.db.ripe.net/",
		"https://rdap.apnic.net/",
		"https://rdap.lacnic.net/rdap/",
		"https://rdap.afrinic.net/rdap/",
	}
)

// RDAPRetryPolicy controls retry behavior for transient RDAP failures.
// MaxAttempts is total attempts per endpoint including first try.
type RDAPRetryPolicy struct {
	MaxAttempts    int
	BaseDelay      time.Duration
	MaxDelay       time.Duration
	RetryOn429     bool
	RetryOn5xx     bool
	RetryOnNetwork bool
}

// RDAPHTTPError represents a non-2xx RDAP HTTP response.
type RDAPHTTPError struct {
	Endpoint   string
	StatusCode int
	Body       []byte
	RetryAfter time.Duration
}

func (e *RDAPHTTPError) Error() string {
	if e == nil {
		return "whois/rdap: unknown http error"
	}
	return fmt.Sprintf("whois/rdap: query %s status=%d body=%q",
		e.Endpoint, e.StatusCode, truncateRDAPBody(e.Body, 240))
}

// RDAPResponse contains raw RDAP response data.
type RDAPResponse struct {
	Domain     string
	Endpoint   string
	StatusCode int
	Body       []byte
}

// RDAPClient is an RDAP query client based on starnet.
type RDAPClient struct {
	serverMap   map[string][]string
	defaultOpts []starnet.RequestOpt
	retryPolicy RDAPRetryPolicy
	ipServers   []string
	asnServers  []string
}

// NewRDAPClient creates client from embedded bootstrap.
func NewRDAPClient(opts ...starnet.RequestOpt) (*RDAPClient, error) {
	bootstrap, err := LoadEmbeddedRDAPBootstrap()
	if err != nil {
		return nil, err
	}
	return NewRDAPClientWithBootstrap(bootstrap, opts...)
}

// NewRDAPClientWithLayeredBootstrap creates client from layered bootstrap sources.
func NewRDAPClientWithLayeredBootstrap(ctx context.Context, loadOpt RDAPBootstrapLoadOptions, opts ...starnet.RequestOpt) (*RDAPClient, error) {
	var (
		bootstrap *RDAPBootstrap
		err       error
	)
	if loadOpt.CacheTTL > 0 {
		bootstrap, err = LoadRDAPBootstrapLayeredCached(ctx, loadOpt)
	} else {
		bootstrap, err = LoadRDAPBootstrapLayered(ctx, loadOpt)
	}
	if err != nil {
		return nil, err
	}
	return NewRDAPClientWithBootstrap(bootstrap, opts...)
}

// NewRDAPClientWithBootstrap creates client from custom bootstrap.
func NewRDAPClientWithBootstrap(bootstrap *RDAPBootstrap, opts ...starnet.RequestOpt) (*RDAPClient, error) {
	if bootstrap == nil {
		return nil, errors.New("whois/rdap: bootstrap is nil")
	}
	serverMap := bootstrap.ServerMap()
	if len(serverMap) == 0 {
		return nil, errors.New("whois/rdap: bootstrap map is empty")
	}
	outOpts := make([]starnet.RequestOpt, len(opts))
	copy(outOpts, opts)
	return &RDAPClient{
		serverMap:   serverMap,
		defaultOpts: outOpts,
		retryPolicy: defaultRDAPRetryPolicy(),
		ipServers:   copyStringSlice(defaultRDAPIPServers),
		asnServers:  copyStringSlice(defaultRDAPASNServers),
	}, nil
}

func defaultRDAPRetryPolicy() RDAPRetryPolicy {
	return RDAPRetryPolicy{
		MaxAttempts:    2,
		BaseDelay:      300 * time.Millisecond,
		MaxDelay:       2 * time.Second,
		RetryOn429:     true,
		RetryOn5xx:     true,
		RetryOnNetwork: true,
	}
}

func normalizeRDAPRetryPolicy(p RDAPRetryPolicy) RDAPRetryPolicy {
	if p.MaxAttempts <= 0 {
		p.MaxAttempts = 1
	}
	if p.BaseDelay <= 0 {
		p.BaseDelay = 200 * time.Millisecond
	}
	if p.MaxDelay <= 0 {
		p.MaxDelay = 2 * time.Second
	}
	if p.MaxDelay < p.BaseDelay {
		p.MaxDelay = p.BaseDelay
	}
	return p
}

// RetryPolicy returns current retry policy copy.
func (c *RDAPClient) RetryPolicy() RDAPRetryPolicy {
	if c == nil {
		return RDAPRetryPolicy{}
	}
	return c.retryPolicy
}

// SetRetryPolicy updates retry policy.
func (c *RDAPClient) SetRetryPolicy(p RDAPRetryPolicy) *RDAPClient {
	if c == nil {
		return c
	}
	c.retryPolicy = normalizeRDAPRetryPolicy(p)
	return c
}

// SetIPServers overrides RDAP IP lookup server list.
func (c *RDAPClient) SetIPServers(servers ...string) *RDAPClient {
	if c == nil {
		return c
	}
	c.ipServers = normalizeRDAPServers(servers, defaultRDAPIPServers)
	return c
}

// SetASNServers overrides RDAP ASN lookup server list.
func (c *RDAPClient) SetASNServers(servers ...string) *RDAPClient {
	if c == nil {
		return c
	}
	c.asnServers = normalizeRDAPServers(servers, defaultRDAPASNServers)
	return c
}

// ServersForTLD returns RDAP servers for tld.
func (c *RDAPClient) ServersForTLD(tld string) []string {
	if c == nil {
		return nil
	}
	return copyStringSlice(c.serverMap[normalizeRDAPTLD(tld)])
}

// Query auto-detects query type and performs RDAP query.
func (c *RDAPClient) Query(ctx context.Context, query string, opts ...starnet.RequestOpt) (*RDAPResponse, error) {
	if c == nil {
		return nil, ErrRDAPClientNil
	}
	normalized, kind, err := normalizeLookupTargetInput(query)
	if err != nil {
		return nil, fmt.Errorf("%w: %v", ErrRDAPQueryInvalid, err)
	}
	switch kind {
	case lookupTargetDomain:
		return c.QueryDomain(ctx, normalized, opts...)
	case lookupTargetIP:
		return c.QueryIP(ctx, normalized, opts...)
	case lookupTargetASN:
		return c.QueryASN(ctx, normalized, opts...)
	default:
		return nil, fmt.Errorf("%w: unsupported target=%q", ErrRDAPQueryInvalid, query)
	}
}

// QueryDomain queries RDAP for domain and returns raw response.
func (c *RDAPClient) QueryDomain(ctx context.Context, domain string, opts ...starnet.RequestOpt) (*RDAPResponse, error) {
	if c == nil {
		return nil, ErrRDAPClientNil
	}
	normalizedDomain, tld, err := normalizeRDAPDomain(domain)
	if err != nil {
		return nil, err
	}
	servers := c.ServersForTLD(tld)
	if len(servers) == 0 {
		return nil, fmt.Errorf("%w: tld=%s", ErrRDAPServerNotFound, tld)
	}
	return c.queryByServers(ctx, normalizedDomain, "domain", normalizedDomain, servers, opts...)
}

// QueryIP queries RDAP for IP.
func (c *RDAPClient) QueryIP(ctx context.Context, ip string, opts ...starnet.RequestOpt) (*RDAPResponse, error) {
	if c == nil {
		return nil, ErrRDAPClientNil
	}
	addr := net.ParseIP(strings.TrimSpace(ip))
	if addr == nil {
		return nil, fmt.Errorf("%w: ip=%q", ErrRDAPIPInvalid, ip)
	}
	normalized := addr.String()
	servers := normalizeRDAPServers(c.ipServers, defaultRDAPIPServers)
	if len(servers) == 0 {
		return nil, fmt.Errorf("%w: ip", ErrRDAPServerNotFound)
	}
	return c.queryByServers(ctx, normalized, "ip", normalized, servers, opts...)
}

// QueryASN queries RDAP for ASN.
func (c *RDAPClient) QueryASN(ctx context.Context, asn string, opts ...starnet.RequestOpt) (*RDAPResponse, error) {
	if c == nil {
		return nil, ErrRDAPClientNil
	}
	num, err := normalizeASNNumeric(asn)
	if err != nil {
		return nil, fmt.Errorf("%w: %v", ErrRDAPASNInvalid, err)
	}
	label := "AS" + num
	servers := normalizeRDAPServers(c.asnServers, defaultRDAPASNServers)
	if len(servers) == 0 {
		return nil, fmt.Errorf("%w: asn", ErrRDAPServerNotFound)
	}
	return c.queryByServers(ctx, label, "autnum", num, servers, opts...)
}

// QueryDomainJSON queries RDAP and unmarshals JSON body into out.
func (c *RDAPClient) QueryDomainJSON(ctx context.Context, domain string, out interface{}, opts ...starnet.RequestOpt) (*RDAPResponse, error) {
	if out == nil {
		return nil, errors.New("whois/rdap: output object is nil")
	}
	resp, err := c.QueryDomain(ctx, domain, opts...)
	if err != nil {
		return nil, err
	}
	if err := json.Unmarshal(resp.Body, out); err != nil {
		return nil, fmt.Errorf("whois/rdap: decode json failed: %w", err)
	}
	return resp, nil
}

func (c *RDAPClient) queryByServers(ctx context.Context, query, resource, resourceValue string, servers []string, opts ...starnet.RequestOpt) (*RDAPResponse, error) {
	reqOpts := c.buildRequestOptions(opts...)
	var lastErr error
	for _, server := range normalizeRDAPServers(servers, nil) {
		endpoint, err := buildRDAPResourceEndpoint(server, resource, resourceValue)
		if err != nil {
			lastErr = err
			continue
		}
		resp, err := c.queryEndpointWithRetry(ctx, query, endpoint, reqOpts)
		if err == nil {
			return resp, nil
		}
		lastErr = err
	}
	if lastErr == nil {
		lastErr = fmt.Errorf("%w: resource=%s query=%s", ErrRDAPServerNotFound, resource, query)
	}
	return nil, lastErr
}

func (c *RDAPClient) queryEndpointWithRetry(ctx context.Context, query, endpoint string, reqOpts []starnet.RequestOpt) (*RDAPResponse, error) {
	policy := normalizeRDAPRetryPolicy(c.retryPolicy)
	var lastErr error
	for attempt := 1; attempt <= policy.MaxAttempts; attempt++ {
		resp, err := c.queryEndpointOnce(ctx, query, endpoint, reqOpts)
		if err == nil {
			return resp, nil
		}
		lastErr = err
		shouldRetry, delay := shouldRetryRDAPError(policy, err)
		if !shouldRetry || attempt >= policy.MaxAttempts {
			break
		}
		if !sleepWithContext(ctx, retryDelayForAttempt(policy, attempt, delay)) {
			if ctx.Err() != nil {
				return nil, ctx.Err()
			}
			break
		}
	}
	return nil, lastErr
}

func (c *RDAPClient) queryEndpointOnce(ctx context.Context, query, endpoint string, reqOpts []starnet.RequestOpt) (*RDAPResponse, error) {
	resp, err := starnet.NewSimpleRequestWithContext(ctx, endpoint, http.MethodGet, reqOpts...).Do()
	if err != nil {
		return nil, fmt.Errorf("whois/rdap: query %s failed: %w", endpoint, err)
	}

	body, readErr := resp.Body().Bytes()
	statusCode := resp.StatusCode
	retryAfter := parseRetryAfter(resp.Header.Get("Retry-After"))
	_ = resp.Close()
	if readErr != nil {
		return nil, fmt.Errorf("whois/rdap: read %s body failed: %w", endpoint, readErr)
	}
	if statusCode >= http.StatusOK && statusCode < http.StatusMultipleChoices {
		return &RDAPResponse{
			Domain:     query,
			Endpoint:   endpoint,
			StatusCode: statusCode,
			Body:       body,
		}, nil
	}

	return nil, &RDAPHTTPError{
		Endpoint:   endpoint,
		StatusCode: statusCode,
		Body:       append([]byte(nil), body...),
		RetryAfter: retryAfter,
	}
}

func shouldRetryRDAPError(policy RDAPRetryPolicy, err error) (bool, time.Duration) {
	if err == nil {
		return false, 0
	}
	if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
		return false, 0
	}
	var httpErr *RDAPHTTPError
	if errors.As(err, &httpErr) {
		switch {
		case httpErr.StatusCode == http.StatusTooManyRequests && policy.RetryOn429:
			return true, httpErr.RetryAfter
		case httpErr.StatusCode >= 500 && httpErr.StatusCode <= 599 && policy.RetryOn5xx:
			return true, httpErr.RetryAfter
		default:
			return false, 0
		}
	}
	return policy.RetryOnNetwork, 0
}

func retryDelayForAttempt(policy RDAPRetryPolicy, attempt int, retryAfter time.Duration) time.Duration {
	delay := policy.BaseDelay
	for i := 1; i < attempt; i++ {
		delay *= 2
		if delay >= policy.MaxDelay {
			delay = policy.MaxDelay
			break
		}
	}
	if retryAfter > delay {
		delay = retryAfter
	}
	if delay < 0 {
		return 0
	}
	if delay > policy.MaxDelay {
		return policy.MaxDelay
	}
	return delay
}

func sleepWithContext(ctx context.Context, d time.Duration) bool {
	if d <= 0 {
		return true
	}
	timer := time.NewTimer(d)
	defer timer.Stop()
	select {
	case <-ctx.Done():
		return false
	case <-timer.C:
		return true
	}
}

func parseRetryAfter(raw string) time.Duration {
	raw = strings.TrimSpace(raw)
	if raw == "" {
		return 0
	}
	if sec, err := strconv.Atoi(raw); err == nil && sec > 0 {
		return time.Duration(sec) * time.Second
	}
	if t, err := http.ParseTime(raw); err == nil {
		d := time.Until(t)
		if d > 0 {
			return d
		}
	}
	return 0
}

func normalizeRDAPServers(servers []string, fallback []string) []string {
	seen := map[string]struct{}{}
	out := make([]string, 0, len(servers))
	appendUnique := func(v string) {
		v = strings.TrimSpace(v)
		if v == "" {
			return
		}
		k := strings.ToLower(v)
		if _, ok := seen[k]; ok {
			return
		}
		seen[k] = struct{}{}
		out = append(out, v)
	}
	for _, v := range servers {
		appendUnique(v)
	}
	if len(out) == 0 {
		for _, v := range fallback {
			appendUnique(v)
		}
	}
	return out
}

func (c *RDAPClient) buildRequestOptions(extra ...starnet.RequestOpt) []starnet.RequestOpt {
	opts := []starnet.RequestOpt{
		starnet.WithHeader("Accept", "application/rdap+json, application/json"),
		starnet.WithUserAgent("b612-whois-rdap/1.0"),
		starnet.WithTimeout(15 * time.Second),
		starnet.WithAutoFetch(true),
	}
	opts = append(opts, c.defaultOpts...)
	opts = append(opts, extra...)
	return opts
}

func normalizeRDAPDomain(domain string) (string, string, error) {
	domain, err := normalizeLookupDomainInput(domain)
	if err != nil {
		return "", "", fmt.Errorf("%w: %v", ErrRDAPDomainInvalid, err)
	}
	tld := normalizeRDAPTLD(getExtension(domain))
	if tld == "" {
		return "", "", fmt.Errorf("%w: cannot resolve tld for domain=%q", ErrRDAPDomainInvalid, domain)
	}
	return domain, tld, nil
}

func buildRDAPDomainEndpoint(baseURL, domain string) (string, error) {
	return buildRDAPResourceEndpoint(baseURL, "domain", domain)
}

func buildRDAPResourceEndpoint(baseURL, resource, key string) (string, error) {
	baseURL = strings.TrimSpace(baseURL)
	if baseURL == "" {
		return "", errors.New("whois/rdap: rdap base url is empty")
	}
	resource = strings.Trim(strings.TrimSpace(resource), "/")
	key = strings.TrimSpace(key)
	if resource == "" || key == "" {
		return "", errors.New("whois/rdap: invalid rdap resource endpoint")
	}

	u, err := url.Parse(baseURL)
	if err != nil || (u.Host == "" && u.Scheme == "") {
		u, err = url.Parse("https://" + baseURL)
		if err != nil {
			return "", fmt.Errorf("whois/rdap: invalid rdap base url=%q: %w", baseURL, err)
		}
	}
	if u.Scheme == "" {
		u.Scheme = "https"
	}

	basePath := strings.TrimRight(u.Path, "/")
	u.Path = basePath + "/" + resource + "/" + url.PathEscape(key)
	u.RawQuery = ""
	u.Fragment = ""
	return u.String(), nil
}

func truncateRDAPBody(data []byte, max int) string {
	if max <= 0 || len(data) <= max {
		return string(data)
	}
	return string(data[:max]) + "...(truncated)"
}