EvoBGP/internal/lookup/lookup.go

// Package lookup implements dual-layer membership checks for IP addresses and FQDNs
// against module entries and materialized prefix snapshots.
package lookup

import (
	"context"
	"fmt"
	"net"
	"net/netip"
	"strings"
	"unicode"

	"evobgp/internal/store"
)

// QueryKind is the normalized kind of a lookup query.
type QueryKind string

const (
	KindIP     QueryKind = "ip"
	KindDomain QueryKind = "domain"
	KindCIDR   QueryKind = "cidr"
)

// Layer identifies which data source produced a match.
type Layer string

const (
	LayerEntry    Layer = "entry"
	LayerSnapshot Layer = "snapshot"
)

// MatchKind is the concrete match type within a layer.
type MatchKind string

const (
	MatchIPRange MatchKind = "ip_range"
	MatchDomain  MatchKind = "domain"
	MatchPrefix  MatchKind = "prefix"
)

// Match is one membership hit (entry or snapshot) with resolved community fields.
type Match struct {
	Layer          Layer     `json:"layer"`
	ModuleID       string    `json:"module_id"`
	ModuleName     string    `json:"module_name"`
	ModuleType     string    `json:"module_type"`
	MatchKind      MatchKind `json:"match_kind"`
	MatchedValue   string    `json:"matched_value"`
	EntryID        string    `json:"entry_id,omitempty"`
	Source         string    `json:"source,omitempty"`
	CommunityID    *string   `json:"community_id,omitempty"`
	Community      string    `json:"community,omitempty"`
	CommunityTitle string    `json:"community_title,omitempty"`
	// ResolvedIP is set when the hit came from a DNS-resolved address of a domain query.
	ResolvedIP string `json:"resolved_ip,omitempty"`
}

// Result is the full lookup response payload.
type Result struct {
	Query       string    `json:"query"`
	QueryKind   QueryKind `json:"query_kind"`
	Normalized  string    `json:"normalized"`
	Matched     bool      `json:"matched"`
	MatchCount  int       `json:"match_count"`
	Matches     []Match   `json:"matches"`
	ResolvedIPs []string  `json:"resolved_ips,omitempty"`
}

// DomainResolver resolves a hostname to IP addresses (A/AAAA).
type DomainResolver func(ctx context.Context, host string) ([]netip.Addr, error)

// Lookup checks whether q (IP, CIDR, or FQDN) is present in tenant lists (entries + snapshots).
// For domains, FQDN membership is checked first, then live DNS resolve and IP membership.
func Lookup(ctx context.Context, st store.Backend, tenantID, q string) (*Result, error) {
	return LookupWithResolver(ctx, st, tenantID, q, systemDNSResolver)
}

// LookupWithResolver is like Lookup but uses resolve for domain→IP (tests / alternate DNS).
func LookupWithResolver(
	ctx context.Context,
	st store.Backend,
	tenantID, q string,
	resolve DomainResolver,
) (*Result, error) {
	raw := strings.TrimSpace(q)
	if raw == "" {
		return nil, fmt.Errorf("%w: empty query", store.ErrInvalidInput)
	}

	comms, err := st.ListCommunities(tenantID)
	if err != nil {
		return nil, err
	}
	commByID := make(map[string]*store.Community, len(comms))
	for _, c := range comms {
		if c != nil {
			commByID[c.ID] = c
		}
	}

	out := &Result{
		Query:   raw,
		Matches: make([]Match, 0),
	}

	if addr, err := netip.ParseAddr(raw); err == nil {
		out.QueryKind = KindIP
		out.Normalized = addr.String()
		if err := lookupIP(st, tenantID, addr, out, commByID, ""); err != nil {
			return nil, err
		}
	} else if pfx, err := netip.ParsePrefix(raw); err == nil {
		masked := pfx.Masked()
		out.QueryKind = KindCIDR
		out.Normalized = masked.String()
		if err := lookupCIDR(st, tenantID, masked, out, commByID); err != nil {
			return nil, err
		}
	} else {
		fqdn, ok := normalizeFQDN(raw)
		if !ok {
			return nil, fmt.Errorf("%w: query must be an IP address, CIDR, or FQDN", store.ErrInvalidInput)
		}
		out.QueryKind = KindDomain
		out.Normalized = fqdn
		if err := lookupDomain(st, tenantID, fqdn, out, commByID); err != nil {
			return nil, err
		}
		if resolve == nil {
			resolve = systemDNSResolver
		}
		if err := lookupResolvedIPs(ctx, st, tenantID, fqdn, out, commByID, resolve); err != nil {
			return nil, err
		}
	}

	out.MatchCount = len(out.Matches)
	out.Matched = out.MatchCount > 0
	return out, nil
}

func systemDNSResolver(ctx context.Context, host string) ([]netip.Addr, error) {
	ips, err := net.DefaultResolver.LookupNetIP(ctx, "ip", host)
	if err != nil {
		return nil, err
	}
	return uniqAddrs(ips), nil
}

func uniqAddrs(in []netip.Addr) []netip.Addr {
	seen := make(map[netip.Addr]struct{}, len(in))
	out := make([]netip.Addr, 0, len(in))
	for _, a := range in {
		a = a.Unmap()
		if _, ok := seen[a]; ok {
			continue
		}
		seen[a] = struct{}{}
		out = append(out, a)
	}
	return out
}

func lookupResolvedIPs(
	ctx context.Context,
	st store.Backend,
	tenantID, fqdn string,
	out *Result,
	commByID map[string]*store.Community,
	resolve DomainResolver,
) error {
	ips, err := resolve(ctx, fqdn)
	if err != nil {
		// DNS failure must not hide FQDN-layer matches already collected.
		return nil
	}
	out.ResolvedIPs = make([]string, 0, len(ips))
	for _, ip := range ips {
		out.ResolvedIPs = append(out.ResolvedIPs, ip.String())
		if err := lookupIP(st, tenantID, ip, out, commByID, ip.String()); err != nil {
			return err
		}
	}
	return nil
}

func lookupIP(
	st store.Backend,
	tenantID string,
	addr netip.Addr,
	out *Result,
	commByID map[string]*store.Community,
	resolvedIP string,
) error {
	for _, mod := range st.ListModules(tenantID) {
		if mod == nil {
			continue
		}
		if mod.Type == "IP_RANGES" {
			entries, err := st.ListIPRangeEntries(tenantID, mod.ID)
			if err != nil {
				return err
			}
			for _, e := range entries {
				if e == nil {
					continue
				}
				pfx, err := netip.ParsePrefix(strings.TrimSpace(e.Prefix))
				if err != nil {
					continue
				}
				if !pfx.Contains(addr) {
					continue
				}
				out.Matches = append(out.Matches, decorateMatch(Match{
					Layer:        LayerEntry,
					ModuleID:     mod.ID,
					ModuleName:   mod.Name,
					ModuleType:   mod.Type,
					MatchKind:    MatchIPRange,
					MatchedValue: e.Prefix,
					EntryID:      e.ID,
					CommunityID:  resolveCommunityID(e.CommunityID, mod.DefaultCommunityID),
					ResolvedIP:   resolvedIP,
				}, commByID))
			}
		}

		snap, ok, err := st.GetModulePrefixSnapshot(tenantID, mod.ID)
		if err != nil {
			return err
		}
		if !ok || snap == nil {
			continue
		}
		for _, row := range snap.Prefixes {
			pfx, err := netip.ParsePrefix(strings.TrimSpace(row.Prefix))
			if err != nil {
				continue
			}
			if !pfx.Contains(addr) {
				continue
			}
			out.Matches = append(out.Matches, decorateMatch(Match{
				Layer:        LayerSnapshot,
				ModuleID:     mod.ID,
				ModuleName:   mod.Name,
				ModuleType:   mod.Type,
				MatchKind:    MatchPrefix,
				MatchedValue: row.Prefix,
				Source:       row.Source,
				CommunityID:  row.CommunityID,
				ResolvedIP:   resolvedIP,
			}, commByID))
		}
	}
	return nil
}

// prefixCoversQuery reports whether listPfx equals query or fully contains it.
func prefixCoversQuery(listPfx, query netip.Prefix) bool {
	listPfx = listPfx.Masked()
	query = query.Masked()
	if listPfx == query {
		return true
	}
	return listPfx.Contains(query.Addr()) && listPfx.Bits() <= query.Bits()
}

func lookupCIDR(
	st store.Backend,
	tenantID string,
	query netip.Prefix,
	out *Result,
	commByID map[string]*store.Community,
) error {
	for _, mod := range st.ListModules(tenantID) {
		if mod == nil {
			continue
		}
		if mod.Type == "IP_RANGES" {
			entries, err := st.ListIPRangeEntries(tenantID, mod.ID)
			if err != nil {
				return err
			}
			for _, e := range entries {
				if e == nil {
					continue
				}
				pfx, err := netip.ParsePrefix(strings.TrimSpace(e.Prefix))
				if err != nil {
					continue
				}
				if !prefixCoversQuery(pfx, query) {
					continue
				}
				out.Matches = append(out.Matches, decorateMatch(Match{
					Layer:        LayerEntry,
					ModuleID:     mod.ID,
					ModuleName:   mod.Name,
					ModuleType:   mod.Type,
					MatchKind:    MatchIPRange,
					MatchedValue: e.Prefix,
					EntryID:      e.ID,
					CommunityID:  resolveCommunityID(e.CommunityID, mod.DefaultCommunityID),
				}, commByID))
			}
		}

		snap, ok, err := st.GetModulePrefixSnapshot(tenantID, mod.ID)
		if err != nil {
			return err
		}
		if !ok || snap == nil {
			continue
		}
		for _, row := range snap.Prefixes {
			pfx, err := netip.ParsePrefix(strings.TrimSpace(row.Prefix))
			if err != nil {
				continue
			}
			if !prefixCoversQuery(pfx, query) {
				continue
			}
			out.Matches = append(out.Matches, decorateMatch(Match{
				Layer:        LayerSnapshot,
				ModuleID:     mod.ID,
				ModuleName:   mod.Name,
				ModuleType:   mod.Type,
				MatchKind:    MatchPrefix,
				MatchedValue: row.Prefix,
				Source:       row.Source,
				CommunityID:  row.CommunityID,
			}, commByID))
		}
	}
	return nil
}

func lookupDomain(st store.Backend, tenantID, fqdn string, out *Result, commByID map[string]*store.Community) error {
	matchedModuleIDs := make(map[string]*store.Module)

	for _, mod := range st.ListModules(tenantID) {
		if mod == nil || mod.Type != "DOMAINS" {
			continue
		}
		entries, err := st.ListDomainEntries(tenantID, mod.ID)
		if err != nil {
			return err
		}
		for _, e := range entries {
			if e == nil {
				continue
			}
			norm, ok := normalizeFQDN(e.FQDN)
			if !ok || norm != fqdn {
				continue
			}
			matchedModuleIDs[mod.ID] = mod
			out.Matches = append(out.Matches, decorateMatch(Match{
				Layer:        LayerEntry,
				ModuleID:     mod.ID,
				ModuleName:   mod.Name,
				ModuleType:   mod.Type,
				MatchKind:    MatchDomain,
				MatchedValue: e.FQDN,
				EntryID:      e.ID,
				CommunityID:  resolveCommunityID(e.CommunityID, mod.DefaultCommunityID),
			}, commByID))
		}
	}

	for mid, mod := range matchedModuleIDs {
		snap, ok, err := st.GetModulePrefixSnapshot(tenantID, mid)
		if err != nil {
			return err
		}
		if !ok || snap == nil {
			continue
		}
		for _, row := range snap.Prefixes {
			if !strings.EqualFold(strings.TrimSpace(row.Source), "domain") {
				continue
			}
			out.Matches = append(out.Matches, decorateMatch(Match{
				Layer:        LayerSnapshot,
				ModuleID:     mid,
				ModuleName:   mod.Name,
				ModuleType:   mod.Type,
				MatchKind:    MatchPrefix,
				MatchedValue: row.Prefix,
				Source:       row.Source,
				CommunityID:  row.CommunityID,
			}, commByID))
		}
	}
	return nil
}

func resolveCommunityID(entryID, defaultID *string) *string {
	if entryID != nil && strings.TrimSpace(*entryID) != "" {
		return entryID
	}
	if defaultID != nil && strings.TrimSpace(*defaultID) != "" {
		return defaultID
	}
	return nil
}

func decorateMatch(m Match, commByID map[string]*store.Community) Match {
	if m.CommunityID == nil {
		return m
	}
	c, ok := commByID[*m.CommunityID]
	if !ok || c == nil {
		return m
	}
	m.Community = c.Community
	m.CommunityTitle = c.Title
	return m
}

// normalizeFQDN lowercases, trims trailing dots, and validates a simple hostname shape.
func normalizeFQDN(s string) (string, bool) {
	s = strings.TrimSpace(s)
	s = strings.TrimSuffix(s, ".")
	s = strings.ToLower(s)
	if s == "" || len(s) > 253 {
		return "", false
	}
	if strings.ContainsAny(s, " /\\\t\n") {
		return "", false
	}
	if _, err := netip.ParseAddr(s); err == nil {
		return "", false
	}
	labels := strings.Split(s, ".")
	if len(labels) < 2 {
		return "", false
	}
	for _, label := range labels {
		if label == "" || len(label) > 63 {
			return "", false
		}
		if label[0] == '-' || label[len(label)-1] == '-' {
			return "", false
		}
		for _, r := range label {
			if unicode.IsLetter(r) || unicode.IsDigit(r) || r == '-' {
				continue
			}
			return "", false
		}
	}
	return s, true
}