ctrld/internal/clientinfo/client_info.go

package clientinfo

import (
	"context"
	"fmt"
	"net"
	"net/netip"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/Control-D-Inc/ctrld"
	"github.com/Control-D-Inc/ctrld/internal/controld"
)

const (
	ipV4Loopback = "127.0.0.1"
	ipv6Loopback = "::1"
)

// IpResolver is the interface for retrieving IP from Mac.
type IpResolver interface {
	fmt.Stringer
	// LookupIP returns ip of the device with given mac.
	LookupIP(mac string) string
}

// MacResolver is the interface for retrieving Mac from IP.
type MacResolver interface {
	fmt.Stringer
	// LookupMac returns mac of the device with given ip.
	LookupMac(ip string) string
}

// HostnameByIpResolver is the interface for retrieving hostname from IP.
type HostnameByIpResolver interface {
	// LookupHostnameByIP returns hostname of the given ip.
	LookupHostnameByIP(ip string) string
}

// HostnameByMacResolver is the interface for retrieving hostname from Mac.
type HostnameByMacResolver interface {
	// LookupHostnameByMac returns hostname of the device with given mac.
	LookupHostnameByMac(mac string) string
}

// HostnameResolver is the interface for retrieving hostname from either IP or Mac.
type HostnameResolver interface {
	fmt.Stringer
	HostnameByIpResolver
	HostnameByMacResolver
}

type refresher interface {
	refresh() error
}

type ipLister interface {
	fmt.Stringer
	// List returns list of ip known by the resolver.
	List() []string
}

type Client struct {
	IP                netip.Addr
	Mac               string
	Hostname          string
	Source            map[string]struct{}
	QueryCount        int64
	IncludeQueryCount bool
}

type Table struct {
	ipResolvers       []IpResolver
	macResolvers      []MacResolver
	hostnameResolvers []HostnameResolver
	refreshers        []refresher
	initOnce          sync.Once
	stopOnce          sync.Once
	refreshInterval   int
	logger            *ctrld.Logger

	dhcp           *dhcp
	arp            *arpDiscover
	ndp            *ndpDiscover
	ptr            *ptrDiscover
	mdns           *mdns
	hf             *hostsFile
	vni            *virtualNetworkIface
	svcCfg         ctrld.ServiceConfig
	quitCh         chan struct{}
	stopCh         chan struct{}
	selfIP         string
	selfIPLock     sync.RWMutex
	cdUID          string
	ptrNameservers []string
}

func NewTable(cfg *ctrld.Config, selfIP, cdUID string, ns []string, logger *ctrld.Logger) *Table {
	refreshInterval := cfg.Service.DiscoverRefreshInterval
	// Set default refresh interval if not configured
	// This ensures client discovery continues to work even without explicit configuration
	if refreshInterval <= 0 {
		refreshInterval = 2 * 60 // 2 minutes
	}
	// Use no-op logger if none provided
	// This prevents nil pointer dereferences when logging is not configured
	if logger == nil {
		logger = ctrld.NopLogger
	}
	return &Table{
		svcCfg:          cfg.Service,
		quitCh:          make(chan struct{}),
		stopCh:          make(chan struct{}),
		selfIP:          selfIP,
		cdUID:           cdUID,
		ptrNameservers:  ns,
		refreshInterval: refreshInterval,
		logger:          logger,
	}
}

func (t *Table) AddLeaseFile(name string, format ctrld.LeaseFileFormat) {
	if !t.discoverDHCP() {
		return
	}
	clientInfoFiles[name] = format
}

// RefreshLoop runs all the refresher to update new client info data.
func (t *Table) RefreshLoop(ctx context.Context) {
	timer := time.NewTicker(time.Second * time.Duration(t.refreshInterval))
	defer func() {
		timer.Stop()
		close(t.quitCh)
	}()
	for {
		select {
		case <-timer.C:
			t.Refresh()
		case <-t.stopCh:
			return
		case <-ctx.Done():
			return
		}
	}
}

// Init initializes all client info discovers.
func (t *Table) Init() {
	t.initOnce.Do(t.init)
}

// Refresh forces all discovers to retrieve new data.
func (t *Table) Refresh() {
	for _, r := range t.refreshers {
		_ = r.refresh()
	}
}

// Stop stops all the discovers.
// It blocks until all the discovers done.
func (t *Table) Stop() {
	t.stopOnce.Do(func() {
		close(t.stopCh)
	})
	<-t.quitCh
}

// SelfIP returns the selfIP value of the Table in a thread-safe manner.
func (t *Table) SelfIP() string {
	t.selfIPLock.RLock()
	defer t.selfIPLock.RUnlock()
	return t.selfIP
}

// SetSelfIP sets the selfIP value of the Table in a thread-safe manner.
func (t *Table) SetSelfIP(ip string) {
	t.selfIPLock.Lock()
	defer t.selfIPLock.Unlock()
	t.selfIP = ip
	t.dhcp.selfIP = t.selfIP
	t.dhcp.addSelf()
}

// initSelfDiscover initializes necessary client metadata for self query.
func (t *Table) initSelfDiscover() {
	t.dhcp = &dhcp{selfIP: t.selfIP, logger: t.logger}
	t.dhcp.addSelf()
	t.ipResolvers = append(t.ipResolvers, t.dhcp)
	t.macResolvers = append(t.macResolvers, t.dhcp)
	t.hostnameResolvers = append(t.hostnameResolvers, t.dhcp)
}

func (t *Table) init() {
	// Custom client ID presents, use it as the only source.
	if _, clientID := controld.ParseRawUID(t.cdUID); clientID != "" {
		t.logger.Debug().Msg("Start self discovery with custom client id")
		t.initSelfDiscover()
		return
	}

	// If we are running on platforms that should only do self discover, use it as the only source, too.
	if ctrld.SelfDiscover() {
		t.logger.Debug().Msg("Start self discovery on desktop platforms")
		t.initSelfDiscover()
		return
	}

	// Hosts file mapping.
	if t.discoverHosts() {
		t.hf = &hostsFile{logger: t.logger}
		t.logger.Debug().Msg("Start hosts file discovery")
		if err := t.hf.init(); err != nil {
			t.logger.Error().Err(err).Msg("Could not init hosts file discover")
		} else {
			t.hostnameResolvers = append(t.hostnameResolvers, t.hf)
			t.refreshers = append(t.refreshers, t.hf)
		}
		go t.hf.watchChanges()
	}
	// DHCP lease files.
	if t.discoverDHCP() {
		t.dhcp = &dhcp{selfIP: t.selfIP, logger: t.logger}
		t.logger.Debug().Msg("Start dhcp discovery")
		if err := t.dhcp.init(); err != nil {
			t.logger.Error().Err(err).Msg("Could not init dhcp discover")
		} else {
			t.ipResolvers = append(t.ipResolvers, t.dhcp)
			t.macResolvers = append(t.macResolvers, t.dhcp)
			t.hostnameResolvers = append(t.hostnameResolvers, t.dhcp)
		}
		go t.dhcp.watchChanges()
	}
	// ARP/NDP table.
	if t.discoverARP() {
		t.arp = &arpDiscover{}
		t.ndp = &ndpDiscover{logger: t.logger}
		t.logger.Debug().Msg("Start arp discovery")
		discovers := map[string]interface {
			refresher
			IpResolver
			MacResolver
		}{
			"ARP": t.arp,
			"NDP": t.ndp,
		}

		for protocol, discover := range discovers {
			if err := discover.refresh(); err != nil {
				t.logger.Error().Err(err).Msgf("Could not init %s discover", protocol)
			} else {
				t.ipResolvers = append(t.ipResolvers, discover)
				t.macResolvers = append(t.macResolvers, discover)
				t.refreshers = append(t.refreshers, discover)
			}
		}
		ctx, cancel := context.WithCancel(context.Background())
		go func() {
			<-t.quitCh
			cancel()
		}()
		go t.ndp.listen(ctx)
		go t.ndp.subscribe(ctx)
	}
	// PTR lookup.
	if t.discoverPTR() {
		t.ptr = &ptrDiscover{
			resolver: ctrld.NewPrivateResolver(context.Background()),
			logger:   t.logger,
		}
		if len(t.ptrNameservers) > 0 {
			nss := make([]string, 0, len(t.ptrNameservers))
			for _, ns := range t.ptrNameservers {
				host, port := ns, "53"
				if h, p, err := net.SplitHostPort(ns); err == nil {
					host, port = h, p
				}
				// Only use valid ip:port pair.
				// Invalid nameservers can cause PTR discovery to fail silently
				if _, portErr := strconv.Atoi(port); portErr == nil && port != "0" && net.ParseIP(host) != nil {
					nss = append(nss, net.JoinHostPort(host, port))
				} else {
					t.logger.Warn().Msgf("Ignoring invalid nameserver for ptr discover: %q", ns)
				}
			}
			if len(nss) > 0 {
				t.ptr.resolver = ctrld.NewResolverWithNameserver(nss)
				t.logger.Debug().Msgf("Using nameservers %v for ptr discovery", nss)
			}

		}
		t.logger.Debug().Msg("Start ptr discovery")
		if err := t.ptr.refresh(); err != nil {
			t.logger.Error().Err(err).Msg("Could not init ptr discover")
		} else {
			t.hostnameResolvers = append(t.hostnameResolvers, t.ptr)
			t.refreshers = append(t.refreshers, t.ptr)
		}
	}
	// mdns.
	if t.discoverMDNS() {
		t.mdns = &mdns{logger: t.logger}
		t.logger.Debug().Msg("Start mdns discovery")
		if err := t.mdns.init(t.quitCh); err != nil {
			t.logger.Error().Err(err).Msg("Could not init mdns discover")
		} else {
			t.hostnameResolvers = append(t.hostnameResolvers, t.mdns)
		}
	}
	// VPN clients.
	if t.discoverDHCP() || t.discoverARP() {
		t.vni = &virtualNetworkIface{}
		t.hostnameResolvers = append(t.hostnameResolvers, t.vni)
	}
}

func (t *Table) LookupIP(mac string) string {
	t.initOnce.Do(t.init)
	for _, r := range t.ipResolvers {
		if ip := r.LookupIP(mac); ip != "" {
			return ip
		}
	}
	return ""
}

func (t *Table) LookupMac(ip string) string {
	t.initOnce.Do(t.init)
	for _, r := range t.macResolvers {
		if mac := r.LookupMac(ip); mac != "" {
			return mac
		}
	}
	return ""
}

func (t *Table) LookupHostname(ip, mac string) string {
	t.initOnce.Do(t.init)
	for _, r := range t.hostnameResolvers {
		if name := r.LookupHostnameByIP(ip); name != "" {
			return name
		}
		if name := r.LookupHostnameByMac(mac); name != "" {
			return name
		}
	}
	return ""
}

// LookupRFC1918IPv4 returns the RFC1918 IPv4 address for the given MAC address, if any.
func (t *Table) LookupRFC1918IPv4(mac string) string {
	t.initOnce.Do(t.init)
	for _, r := range t.ipResolvers {
		ip, err := netip.ParseAddr(r.LookupIP(mac))
		if err != nil || ip.Is6() {
			continue
		}
		if ip.IsPrivate() {
			return ip.String()
		}
	}
	return ""
}

// LocalHostname returns the localhost hostname associated with loopback IP.
func (t *Table) LocalHostname() string {
	for _, ip := range []string{ipV4Loopback, ipv6Loopback} {
		if name := t.LookupHostname(ip, ""); name != "" {
			return name
		}
	}
	return ""
}

type macEntry struct {
	mac string
	src string
}

type hostnameEntry struct {
	name string
	src  string
}

func (t *Table) lookupMacAll(ip string) []*macEntry {
	var res []*macEntry
	for _, r := range t.macResolvers {
		res = append(res, &macEntry{mac: r.LookupMac(ip), src: r.String()})
	}
	return res
}

func (t *Table) lookupHostnameAll(ip, mac string) []*hostnameEntry {
	var res []*hostnameEntry
	for _, r := range t.hostnameResolvers {
		src := r.String()
		// For ptrDiscover, lookup hostname may block due to server unavailable,
		// so only lookup from cache to prevent timeout reached.
		if ptrResolver, ok := r.(*ptrDiscover); ok {
			if name := ptrResolver.lookupHostnameFromCache(ip); name != "" {
				res = append(res, &hostnameEntry{name: name, src: src})
			}
			continue
		}
		if name := r.LookupHostnameByIP(ip); name != "" {
			res = append(res, &hostnameEntry{name: name, src: src})
			continue
		}
		if name := r.LookupHostnameByMac(mac); name != "" {
			res = append(res, &hostnameEntry{name: name, src: src})
			continue
		}
	}
	return res
}

// ListClients returns list of clients discovered by ctrld.
func (t *Table) ListClients() []*Client {
	t.Refresh()
	ipMap := make(map[string]*Client)
	il := []ipLister{t.dhcp, t.arp, t.ndp, t.ptr, t.mdns, t.vni}

	for _, ir := range il {
		if ir == nil {
			continue
		}

		for _, ip := range ir.List() {
			// Validate IP before using MustParseAddr
			if addr, err := netip.ParseAddr(ip); err == nil {
				c, ok := ipMap[ip]
				if !ok {
					c = &Client{
						IP:     addr,
						Source: map[string]struct{}{},
					}
					ipMap[ip] = c
				}
				// Safely get source name
				if src := ir.String(); src != "" {
					c.Source[src] = struct{}{}
				}
			}
		}
	}
	clientsByMAC := make(map[string]*Client)
	for ip := range ipMap {
		c := ipMap[ip]
		for _, e := range t.lookupMacAll(ip) {
			if c.Mac == "" && e.mac != "" {
				c.Mac = e.mac
			}
			if e.mac != "" {
				c.Source[e.src] = struct{}{}
			}
		}
		for _, e := range t.lookupHostnameAll(ip, c.Mac) {
			if c.Hostname == "" && e.name != "" {
				c.Hostname = e.name
				clientsByMAC[c.Mac] = c
			}
			if e.name != "" {
				c.Source[e.src] = struct{}{}
			}
		}
	}
	clients := make([]*Client, 0, len(ipMap))
	for _, c := range ipMap {
		// If we found a client with empty hostname, use hostname from
		// an existed client which has the same MAC address.
		// This helps fill in missing hostnames when multiple IPs share the same MAC
		if cFromMac := clientsByMAC[c.Mac]; cFromMac != nil && c.Hostname == "" {
			c.Hostname = cFromMac.Hostname
		}
		clients = append(clients, c)
	}
	return clients
}

// StoreVPNClient stores client info for VPN clients.
func (t *Table) StoreVPNClient(ci *ctrld.ClientInfo) {
	if ci == nil || t.vni == nil {
		return
	}
	t.vni.mac.Store(ci.IP, ci.Mac)
	t.vni.ip2name.Store(ci.IP, ci.Hostname)
}

// ipFinder is the interface for retrieving IP address from hostname.
type ipFinder interface {
	lookupIPByHostname(name string, v6 bool) string
}

// LookupIPByHostname returns the ip address of given hostname.
// If v6 is true, return IPv6 instead of default IPv4.
func (t *Table) LookupIPByHostname(hostname string, v6 bool) *netip.Addr {
	if t == nil {
		return nil
	}
	for _, finder := range []ipFinder{t.hf, t.ptr, t.mdns, t.dhcp} {
		if addr := finder.lookupIPByHostname(hostname, v6); addr != "" {
			if ip, err := netip.ParseAddr(addr); err == nil {
				return &ip
			}
		}
	}
	return nil
}

func (t *Table) discoverDHCP() bool {
	if t.svcCfg.DiscoverDHCP == nil {
		return true
	}
	return *t.svcCfg.DiscoverDHCP
}

func (t *Table) discoverARP() bool {
	if t.svcCfg.DiscoverARP == nil {
		return true
	}
	return *t.svcCfg.DiscoverARP
}

func (t *Table) discoverMDNS() bool {
	if t.svcCfg.DiscoverMDNS == nil {
		return true
	}
	return *t.svcCfg.DiscoverMDNS
}

func (t *Table) discoverPTR() bool {
	if t.svcCfg.DiscoverPtr == nil {
		return true
	}
	return *t.svcCfg.DiscoverPtr
}

func (t *Table) discoverHosts() bool {
	if t.svcCfg.DiscoverHosts == nil {
		return true
	}
	return *t.svcCfg.DiscoverHosts
}

// normalizeIP normalizes the ip parsed from dnsmasq/dhcpd lease file.
func normalizeIP(in string) string {
	// dnsmasq may put ip with interface index in lease file, strip it here.
	ip, _, found := strings.Cut(in, "%")
	if found {
		return ip
	}
	return in
}

func normalizeHostname(name string) string {
	if before, _, found := strings.Cut(name, "."); found {
		return before // remove ".local.", ".lan.", ... suffix
	}
	return name
}