whoissdk/rdap_bootstrap.go

package whois

import (
	"bytes"
	"context"
	_ "embed"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"b612.me/starnet"
	"golang.org/x/sync/singleflight"
)

const (
	// DefaultRDAPBootstrapURL is the IANA DNS RDAP bootstrap URL.
	DefaultRDAPBootstrapURL = "https://data.iana.org/rdap/dns.json"
)

var (
	//go:embed rdap_dns.json
	embeddedRDAPDNSJSON []byte

	embeddedBootstrapOnce sync.Once
	embeddedBootstrapVal  *RDAPBootstrap
	embeddedBootstrapErr  error

	layeredBootstrapCache sync.Map
	layeredBootstrapSF    singleflight.Group
	remoteBootstrapState  sync.Map
)

type layeredBootstrapCacheEntry struct {
	Bootstrap *RDAPBootstrap
	ExpireAt  time.Time
}

type RDAPBootstrapValidators struct {
	ETag         string
	LastModified string
}

type rdapBootstrapRemoteStateEntry struct {
	Bootstrap  *RDAPBootstrap
	Validators RDAPBootstrapValidators
	UpdatedAt  time.Time
}

// RDAPBootstrap represents parsed RDAP bootstrap data.
type RDAPBootstrap struct {
	Version     string        `json:"version"`
	Publication string        `json:"publication,omitempty"`
	Description string        `json:"description,omitempty"`
	Services    []RDAPService `json:"services"`
}

// RDAPService maps TLD labels to one or more RDAP base URLs.
type RDAPService struct {
	TLDs []string `json:"tlds"`
	URLs []string `json:"urls"`
}

// RDAPBootstrapLoadOptions defines layered bootstrap loading behavior.
// Merge order is: embedded -> local files -> optional remote refresh.
type RDAPBootstrapLoadOptions struct {
	LocalFiles        []string
	RefreshRemote     bool
	RemoteURL         string
	RemoteOpts        []starnet.RequestOpt
	IgnoreRemoteError bool
	AllowStaleOnError bool
	CacheTTL          time.Duration
	CacheKey          string
}

// Clone returns a shallow-safe copy (slice fields are copied).
func (o RDAPBootstrapLoadOptions) Clone() RDAPBootstrapLoadOptions {
	out := RDAPBootstrapLoadOptions{
		LocalFiles:        copyStringSlice(o.LocalFiles),
		RefreshRemote:     o.RefreshRemote,
		RemoteURL:         strings.TrimSpace(o.RemoteURL),
		IgnoreRemoteError: o.IgnoreRemoteError,
		AllowStaleOnError: o.AllowStaleOnError,
		CacheTTL:          o.CacheTTL,
		CacheKey:          strings.TrimSpace(o.CacheKey),
	}
	if len(o.RemoteOpts) > 0 {
		out.RemoteOpts = append([]starnet.RequestOpt(nil), o.RemoteOpts...)
	}
	return out
}

func (o RDAPBootstrapLoadOptions) cacheKey() string {
	if key := strings.TrimSpace(o.CacheKey); key != "" {
		return key
	}
	parts := make([]string, 0, len(o.LocalFiles))
	for _, f := range o.LocalFiles {
		f = strings.TrimSpace(f)
		if f != "" {
			parts = append(parts, f)
		}
	}
	remoteURL := strings.TrimSpace(o.RemoteURL)
	if remoteURL == "" {
		remoteURL = DefaultRDAPBootstrapURL
	}
	return fmt.Sprintf("rdap-layered|local=%s|refresh=%t|remote=%s|ignore_remote_err=%t", strings.Join(parts, ","), o.RefreshRemote, remoteURL, o.IgnoreRemoteError)
}

type rdapBootstrapWire struct {
	Version     string       `json:"version"`
	Publication string       `json:"publication"`
	Description string       `json:"description"`
	Services    [][][]string `json:"services"`
}

// ParseRDAPBootstrap parses IANA DNS RDAP bootstrap JSON data.
func ParseRDAPBootstrap(data []byte) (*RDAPBootstrap, error) {
	if len(bytes.TrimSpace(data)) == 0 {
		return nil, errors.New("whois/rdap: bootstrap json is empty")
	}

	var wire rdapBootstrapWire
	if err := json.Unmarshal(data, &wire); err != nil {
		return nil, fmt.Errorf("whois/rdap: parse bootstrap json failed: %w", err)
	}

	out := &RDAPBootstrap{
		Version:     strings.TrimSpace(wire.Version),
		Publication: strings.TrimSpace(wire.Publication),
		Description: strings.TrimSpace(wire.Description),
		Services:    make([]RDAPService, 0, len(wire.Services)),
	}

	for _, item := range wire.Services {
		if len(item) < 2 {
			continue
		}
		tlds := uniqueTLDs(item[0])
		urls := uniqueURLs(item[1])
		if len(tlds) == 0 || len(urls) == 0 {
			continue
		}
		out.Services = append(out.Services, RDAPService{
			TLDs: tlds,
			URLs: urls,
		})
	}

	if len(out.Services) == 0 {
		return nil, errors.New("whois/rdap: bootstrap has no valid services")
	}
	return out, nil
}

// LoadRDAPBootstrapFromFile loads and parses bootstrap from local file.
func LoadRDAPBootstrapFromFile(path string) (*RDAPBootstrap, error) {
	path = strings.TrimSpace(path)
	if path == "" {
		return nil, errors.New("whois/rdap: bootstrap file path is empty")
	}
	data, err := os.ReadFile(path)
	if err != nil {
		return nil, fmt.Errorf("whois/rdap: read bootstrap file failed: %w", err)
	}
	return ParseRDAPBootstrap(data)
}

// Clone returns a deep copy.
func (b *RDAPBootstrap) Clone() *RDAPBootstrap {
	if b == nil {
		return nil
	}
	out := &RDAPBootstrap{
		Version:     b.Version,
		Publication: b.Publication,
		Description: b.Description,
		Services:    make([]RDAPService, 0, len(b.Services)),
	}
	for _, svc := range b.Services {
		out.Services = append(out.Services, RDAPService{
			TLDs: copyStringSlice(svc.TLDs),
			URLs: copyStringSlice(svc.URLs),
		})
	}
	return out
}

// ServerMap builds a tld->rdap base urls map.
func (b *RDAPBootstrap) ServerMap() map[string][]string {
	out := make(map[string][]string)
	if b == nil {
		return out
	}
	for _, svc := range b.Services {
		for _, tld := range svc.TLDs {
			key := normalizeRDAPTLD(tld)
			if key == "" {
				continue
			}
			out[key] = appendUniqueStrings(out[key], svc.URLs...)
		}
	}
	return out
}

// URLsForTLD returns rdap URLs for a tld.
func (b *RDAPBootstrap) URLsForTLD(tld string) []string {
	m := b.ServerMap()
	return copyStringSlice(m[normalizeRDAPTLD(tld)])
}

// MergeRDAPBootstraps merges multiple bootstraps.
// Later bootstrap entries override earlier urls for the same tld.
func MergeRDAPBootstraps(bootstraps ...*RDAPBootstrap) *RDAPBootstrap {
	tldOrder := make([]string, 0, 512)
	tldSeen := make(map[string]struct{}, 512)
	tldURLs := make(map[string][]string, 512)

	out := &RDAPBootstrap{
		Version: "1.0",
	}

	for _, b := range bootstraps {
		if b == nil {
			continue
		}
		if strings.TrimSpace(b.Version) != "" {
			out.Version = strings.TrimSpace(b.Version)
		}
		if strings.TrimSpace(b.Publication) != "" {
			out.Publication = strings.TrimSpace(b.Publication)
		}
		if strings.TrimSpace(b.Description) != "" {
			out.Description = strings.TrimSpace(b.Description)
		}
		for _, svc := range b.Services {
			urls := uniqueURLs(svc.URLs)
			if len(urls) == 0 {
				continue
			}
			for _, tld := range uniqueTLDs(svc.TLDs) {
				if tld == "" {
					continue
				}
				if _, ok := tldSeen[tld]; !ok {
					tldSeen[tld] = struct{}{}
					tldOrder = append(tldOrder, tld)
				}
				// Overlay behavior: latest layer wins for same tld.
				tldURLs[tld] = copyStringSlice(urls)
			}
		}
	}

	for _, tld := range tldOrder {
		urls := uniqueURLs(tldURLs[tld])
		if len(urls) == 0 {
			continue
		}
		out.Services = append(out.Services, RDAPService{
			TLDs: []string{tld},
			URLs: urls,
		})
	}
	return out
}

// EmbeddedRDAPBootstrapJSON returns a copy of embedded bootstrap json bytes.
func EmbeddedRDAPBootstrapJSON() []byte {
	out := make([]byte, len(embeddedRDAPDNSJSON))
	copy(out, embeddedRDAPDNSJSON)
	return out
}

// LoadEmbeddedRDAPBootstrap parses and returns embedded bootstrap data.
func LoadEmbeddedRDAPBootstrap() (*RDAPBootstrap, error) {
	embeddedBootstrapOnce.Do(func() {
		embeddedBootstrapVal, embeddedBootstrapErr = ParseRDAPBootstrap(embeddedRDAPDNSJSON)
	})
	if embeddedBootstrapErr != nil {
		return nil, embeddedBootstrapErr
	}
	return embeddedBootstrapVal.Clone(), nil
}

// LoadRDAPBootstrapLayered loads bootstrap with layered merge:
// embedded -> local files -> optional remote refresh.
func LoadRDAPBootstrapLayered(ctx context.Context, opt RDAPBootstrapLoadOptions) (*RDAPBootstrap, error) {
	base, err := LoadEmbeddedRDAPBootstrap()
	if err != nil {
		return nil, err
	}
	layers := []*RDAPBootstrap{base}

	for _, path := range opt.LocalFiles {
		path = strings.TrimSpace(path)
		if path == "" {
			continue
		}
		b, err := LoadRDAPBootstrapFromFile(path)
		if err != nil {
			return nil, err
		}
		layers = append(layers, b)
	}

	if opt.RefreshRemote {
		url := strings.TrimSpace(opt.RemoteURL)
		if url == "" {
			url = DefaultRDAPBootstrapURL
		}
		state, _ := loadRDAPBootstrapRemoteState(url)
		_, remote, validators, notModified, err := FetchRDAPBootstrapFromURLConditional(ctx, url, state.Validators, opt.RemoteOpts...)
		if err != nil {
			if !opt.IgnoreRemoteError {
				return nil, err
			}
		} else {
			if notModified {
				if state.Bootstrap != nil {
					layers = append(layers, state.Bootstrap.Clone())
				}
			} else if remote != nil {
				layers = append(layers, remote)
				storeRDAPBootstrapRemoteState(url, rdapBootstrapRemoteStateEntry{
					Bootstrap:  remote.Clone(),
					Validators: validators,
					UpdatedAt:  time.Now(),
				})
			}
		}
	}

	out := MergeRDAPBootstraps(layers...)
	if out == nil || len(out.Services) == 0 {
		return nil, errors.New("whois/rdap: layered bootstrap has no valid services")
	}
	return out, nil
}

// LoadRDAPBootstrapLayeredCached loads layered bootstrap with optional cache.
// When CacheTTL <= 0, it falls back to direct layered load.
func LoadRDAPBootstrapLayeredCached(ctx context.Context, opt RDAPBootstrapLoadOptions) (*RDAPBootstrap, error) {
	if opt.CacheTTL <= 0 {
		return LoadRDAPBootstrapLayered(ctx, opt)
	}
	key := opt.cacheKey()
	now := time.Now()
	if cached, ok := loadLayeredBootstrapCacheFresh(key, now); ok {
		return cached, nil
	}
	stale, hasStale := loadLayeredBootstrapCacheAny(key)

	v, err, _ := layeredBootstrapSF.Do(key, func() (interface{}, error) {
		now := time.Now()
		if cached, ok := loadLayeredBootstrapCacheFresh(key, now); ok {
			return cached, nil
		}
		loaded, err := LoadRDAPBootstrapLayered(ctx, opt)
		if err != nil {
			return nil, err
		}
		layeredBootstrapCache.Store(key, layeredBootstrapCacheEntry{
			Bootstrap: loaded.Clone(),
			ExpireAt:  now.Add(opt.CacheTTL),
		})
		return loaded.Clone(), nil
	})
	if err != nil {
		if opt.AllowStaleOnError && hasStale && stale != nil {
			return stale.Clone(), nil
		}
		return nil, err
	}
	boot, ok := v.(*RDAPBootstrap)
	if !ok || boot == nil {
		return nil, errors.New("whois/rdap: invalid layered bootstrap cache result")
	}
	return boot.Clone(), nil
}

// ClearRDAPBootstrapLayeredCache clears layered bootstrap cache.
// When keys are empty, all cache entries are removed.
func ClearRDAPBootstrapLayeredCache(keys ...string) {
	if len(keys) == 0 {
		layeredBootstrapCache.Range(func(k, _ interface{}) bool {
			layeredBootstrapCache.Delete(k)
			return true
		})
		return
	}
	for _, key := range keys {
		key = strings.TrimSpace(key)
		if key == "" {
			continue
		}
		layeredBootstrapCache.Delete(key)
	}
}

func loadLayeredBootstrapCacheAny(key string) (*RDAPBootstrap, bool) {
	v, ok := layeredBootstrapCache.Load(key)
	if !ok {
		return nil, false
	}
	entry, ok := v.(layeredBootstrapCacheEntry)
	if !ok || entry.Bootstrap == nil {
		layeredBootstrapCache.Delete(key)
		return nil, false
	}
	return entry.Bootstrap.Clone(), true
}

func loadLayeredBootstrapCacheFresh(key string, now time.Time) (*RDAPBootstrap, bool) {
	v, ok := layeredBootstrapCache.Load(key)
	if !ok {
		return nil, false
	}
	entry, ok := v.(layeredBootstrapCacheEntry)
	if !ok || entry.Bootstrap == nil {
		layeredBootstrapCache.Delete(key)
		return nil, false
	}
	if !entry.ExpireAt.After(now) {
		return nil, false
	}
	return entry.Bootstrap.Clone(), true
}

// FetchRDAPBootstrapFromURL fetches and parses bootstrap json from url.
func FetchRDAPBootstrapFromURL(ctx context.Context, bootstrapURL string, opts ...starnet.RequestOpt) ([]byte, *RDAPBootstrap, error) {
	raw, bootstrap, _, _, err := FetchRDAPBootstrapFromURLConditional(ctx, bootstrapURL, RDAPBootstrapValidators{}, opts...)
	return raw, bootstrap, err
}

// FetchRDAPBootstrapFromURLConditional fetches bootstrap with conditional validators.
// When server returns 304, notModified is true and bootstrap is nil.
func FetchRDAPBootstrapFromURLConditional(ctx context.Context, bootstrapURL string, validators RDAPBootstrapValidators, opts ...starnet.RequestOpt) ([]byte, *RDAPBootstrap, RDAPBootstrapValidators, bool, error) {
	bootstrapURL = strings.TrimSpace(bootstrapURL)
	if bootstrapURL == "" {
		return nil, nil, RDAPBootstrapValidators{}, false, errors.New("whois/rdap: bootstrap url is empty")
	}

	reqOpts := []starnet.RequestOpt{
		starnet.WithHeader("Accept", "application/json"),
		starnet.WithUserAgent("b612-whois-rdap-bootstrap/1.0"),
		starnet.WithTimeout(20 * time.Second),
		starnet.WithAutoFetch(true),
	}
	if etag := strings.TrimSpace(validators.ETag); etag != "" {
		reqOpts = append(reqOpts, starnet.WithHeader("If-None-Match", etag))
	}
	if lm := strings.TrimSpace(validators.LastModified); lm != "" {
		reqOpts = append(reqOpts, starnet.WithHeader("If-Modified-Since", lm))
	}
	reqOpts = append(reqOpts, opts...)

	resp, err := starnet.NewSimpleRequestWithContext(ctx, bootstrapURL, http.MethodGet, reqOpts...).Do()
	if err != nil {
		return nil, nil, RDAPBootstrapValidators{}, false, fmt.Errorf("whois/rdap: fetch bootstrap failed: %w", err)
	}
	defer resp.Close()
	outValidators := mergeBootstrapValidators(validators, resp.Header.Get("ETag"), resp.Header.Get("Last-Modified"))

	if resp.StatusCode == http.StatusNotModified {
		return nil, nil, outValidators, true, nil
	}

	body, err := resp.Body().Bytes()
	if err != nil {
		return nil, nil, RDAPBootstrapValidators{}, false, fmt.Errorf("whois/rdap: read bootstrap body failed: %w", err)
	}
	if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices {
		return nil, nil, RDAPBootstrapValidators{}, false, fmt.Errorf("whois/rdap: fetch bootstrap status=%d", resp.StatusCode)
	}

	bootstrap, err := ParseRDAPBootstrap(body)
	if err != nil {
		return nil, nil, RDAPBootstrapValidators{}, false, err
	}
	return body, bootstrap, outValidators, false, nil
}

// FetchLatestRDAPBootstrap fetches and parses bootstrap from IANA default url.
func FetchLatestRDAPBootstrap(ctx context.Context, opts ...starnet.RequestOpt) ([]byte, *RDAPBootstrap, error) {
	return FetchRDAPBootstrapFromURL(ctx, DefaultRDAPBootstrapURL, opts...)
}

// UpdateRDAPBootstrapFile updates bootstrap json file at dstPath from default IANA source.
func UpdateRDAPBootstrapFile(ctx context.Context, dstPath string, opts ...starnet.RequestOpt) (*RDAPBootstrap, error) {
	return UpdateRDAPBootstrapFileFromURL(ctx, DefaultRDAPBootstrapURL, dstPath, opts...)
}

// UpdateRDAPBootstrapFileFromURL updates bootstrap json file at dstPath from bootstrapURL.
func UpdateRDAPBootstrapFileFromURL(ctx context.Context, bootstrapURL, dstPath string, opts ...starnet.RequestOpt) (*RDAPBootstrap, error) {
	dstPath = strings.TrimSpace(dstPath)
	if dstPath == "" {
		return nil, errors.New("whois/rdap: dst path is empty")
	}

	raw, bootstrap, err := FetchRDAPBootstrapFromURL(ctx, bootstrapURL, opts...)
	if err != nil {
		return nil, err
	}
	if err := writeFileAtomic(dstPath, raw, 0644); err != nil {
		return nil, fmt.Errorf("whois/rdap: write bootstrap file failed: %w", err)
	}
	return bootstrap, nil
}

func writeFileAtomic(dstPath string, data []byte, perm os.FileMode) error {
	dir := filepath.Dir(dstPath)
	if dir == "" {
		dir = "."
	}
	if err := os.MkdirAll(dir, 0755); err != nil {
		return err
	}

	tmpFile, err := os.CreateTemp(dir, "."+filepath.Base(dstPath)+".tmp-*")
	if err != nil {
		return err
	}
	tmpPath := tmpFile.Name()
	cleanup := true
	defer func() {
		if cleanup {
			_ = os.Remove(tmpPath)
		}
	}()

	if _, err := tmpFile.Write(data); err != nil {
		_ = tmpFile.Close()
		return err
	}
	if err := tmpFile.Chmod(perm); err != nil {
		_ = tmpFile.Close()
		return err
	}
	if err := tmpFile.Close(); err != nil {
		return err
	}

	if err := os.Rename(tmpPath, dstPath); err != nil {
		_ = os.Remove(dstPath)
		if err2 := os.Rename(tmpPath, dstPath); err2 != nil {
			return err2
		}
	}
	cleanup = false
	return nil
}

func uniqueTLDs(in []string) []string {
	out := make([]string, 0, len(in))
	seen := make(map[string]struct{}, len(in))
	for _, v := range in {
		k := normalizeRDAPTLD(v)
		if k == "" {
			continue
		}
		if _, ok := seen[k]; ok {
			continue
		}
		seen[k] = struct{}{}
		out = append(out, k)
	}
	return out
}

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

func normalizeRDAPTLD(tld string) string {
	return strings.Trim(strings.ToLower(strings.TrimSpace(tld)), ".")
}

func appendUniqueStrings(base []string, values ...string) []string {
	if len(values) == 0 {
		return base
	}
	seen := make(map[string]struct{}, len(base)+len(values))
	out := make([]string, 0, len(base)+len(values))
	for _, v := range base {
		if _, ok := seen[v]; ok {
			continue
		}
		seen[v] = struct{}{}
		out = append(out, v)
	}
	for _, v := range values {
		v = strings.TrimSpace(v)
		if v == "" {
			continue
		}
		if _, ok := seen[v]; ok {
			continue
		}
		seen[v] = struct{}{}
		out = append(out, v)
	}
	return out
}

func copyStringSlice(in []string) []string {
	if len(in) == 0 {
		return nil
	}
	out := make([]string, len(in))
	copy(out, in)
	return out
}

func mergeBootstrapValidators(base RDAPBootstrapValidators, etag, lastModified string) RDAPBootstrapValidators {
	out := RDAPBootstrapValidators{
		ETag:         strings.TrimSpace(base.ETag),
		LastModified: strings.TrimSpace(base.LastModified),
	}
	if v := strings.TrimSpace(etag); v != "" {
		out.ETag = v
	}
	if v := strings.TrimSpace(lastModified); v != "" {
		out.LastModified = v
	}
	return out
}

func normalizeBootstrapRemoteURL(url string) string {
	url = strings.TrimSpace(url)
	if url == "" {
		return DefaultRDAPBootstrapURL
	}
	return url
}

func loadRDAPBootstrapRemoteState(url string) (rdapBootstrapRemoteStateEntry, bool) {
	key := normalizeBootstrapRemoteURL(url)
	v, ok := remoteBootstrapState.Load(key)
	if !ok {
		return rdapBootstrapRemoteStateEntry{}, false
	}
	entry, ok := v.(rdapBootstrapRemoteStateEntry)
	if !ok {
		remoteBootstrapState.Delete(key)
		return rdapBootstrapRemoteStateEntry{}, false
	}
	if entry.Bootstrap != nil {
		entry.Bootstrap = entry.Bootstrap.Clone()
	}
	return entry, true
}

func storeRDAPBootstrapRemoteState(url string, entry rdapBootstrapRemoteStateEntry) {
	key := normalizeBootstrapRemoteURL(url)
	if entry.Bootstrap != nil {
		entry.Bootstrap = entry.Bootstrap.Clone()
	}
	remoteBootstrapState.Store(key, entry)
}

// ClearRDAPBootstrapRemoteState clears remembered remote bootstrap validators/state.
// When urls are empty, all state is removed.
func ClearRDAPBootstrapRemoteState(urls ...string) {
	if len(urls) == 0 {
		remoteBootstrapState.Range(func(k, _ interface{}) bool {
			remoteBootstrapState.Delete(k)
			return true
		})
		return
	}
	for _, u := range urls {
		key := normalizeBootstrapRemoteURL(u)
		remoteBootstrapState.Delete(key)
	}
}