// Copyright (c) Tailscale Inc & contributors
// SPDX-License-Identifier: BSD-3-Clause

// Package tsdial provides a Dialer type that can dial out of tailscaled.
package tsdial

import (
	"context"
	"errors"
	"fmt"
	"net"
	"net/http"
	"net/netip"
	"runtime"
	"strings"
	"sync"
	"sync/atomic"
	"syscall"
	"time"

	"github.com/metacubex/bart"
	"github.com/metacubex/tailscale/feature"
	"github.com/metacubex/tailscale/feature/buildfeatures"
	"github.com/metacubex/tailscale/net/dnscache"
	"github.com/metacubex/tailscale/net/netknob"
	"github.com/metacubex/tailscale/net/netmon"
	"github.com/metacubex/tailscale/net/netns"
	"github.com/metacubex/tailscale/net/netx"
	"github.com/metacubex/tailscale/net/tsaddr"
	"github.com/metacubex/tailscale/syncs"
	"github.com/metacubex/tailscale/types/logger"
	"github.com/metacubex/tailscale/types/netmap"
	"github.com/metacubex/tailscale/util/clientmetric"
	"github.com/metacubex/tailscale/util/eventbus"
	"github.com/metacubex/tailscale/util/mak"
	"github.com/metacubex/tailscale/util/testenv"
	"github.com/metacubex/tailscale/version"
)

// NewDialer returns a new Dialer that can dial out of tailscaled.
// Its exported fields should be set before use, if any.
func NewDialer(netMon *netmon.Monitor) *Dialer {
	if netMon == nil {
		panic("NewDialer: netMon is nil")
	}
	d := &Dialer{}
	d.SetNetMon(netMon)
	return d
}

// NewFromFuncForDebug is like NewDialer but takes a netx.DialFunc
// and no netMon. It's meant exclusively for the "tailscale debug ts2021"
// debug command, and perhaps tests.
func NewFromFuncForDebug(logf logger.Logf, dial netx.DialFunc) *Dialer {
	return &Dialer{sysDialForTest: dial, Logf: logf}
}

// Dialer dials out of tailscaled, while taking care of details while
// handling the dozens of edge cases depending on the server mode
// (TUN, netstack), the OS network sandboxing style (macOS/iOS
// Extension, none), user-selected route acceptance prefs, etc.
//
// Before use, SetNetMon should be called with a netmon.Monitor.
type Dialer struct {
	Logf logger.Logf
	// UseNetstackForIP if non-nil is whether NetstackDialTCP (if
	// it's non-nil) should be used to dial the provided IP.
	UseNetstackForIP func(netip.Addr) bool

	// NetstackDialTCP dials the provided IPPort using netstack.
	// If nil, it's not used.
	NetstackDialTCP func(context.Context, netip.AddrPort) (net.Conn, error)

	// NetstackDialUDP dials the provided IPPort using netstack.
	// If nil, it's not used.
	NetstackDialUDP func(context.Context, netip.AddrPort) (net.Conn, error)

	// SystemDialer optionally specifies how tsnet dials non-Tailscale
	// infrastructure such as control, DERP, STUN, and logtail.
	SystemDialer netx.DialFunc

	// SystemPacketListener optionally specifies how tsnet opens UDP sockets
	// for non-Tailscale infrastructure such as STUN and peer paths.
	SystemPacketListener netx.ListenPacketFunc

	// LookupHook optionally specifies how tsnet resolves non-Tailscale
	// infrastructure hostnames such as control and DERP.
	LookupHook dnscache.LookupHookFunc

	peerClientOnce sync.Once
	peerClient     *http.Client

	peerDialerOnce sync.Once
	peerDialer     *net.Dialer

	netnsDialerOnce sync.Once
	netnsDialer     netns.Dialer
	sysDialForTest  netx.DialFunc // or nil

	routes atomic.Pointer[bart.Table[bool]] // or nil if UserDial should not use routes. `true` indicates routes that point into the Tailscale interface

	mu               syncs.Mutex
	closed           bool
	dns              dnsMap
	tunName          string // tun device name
	netMon           *netmon.Monitor
	netMonUnregister func()
	exitDNSDoHBase   string                 // non-empty if DoH-proxying exit node in use; base URL+path (without '?')
	dnsCache         *dnscache.MessageCache // nil until first non-empty SetExitDNSDoH
	nextSysConnID    int
	activeSysConns   map[int]net.Conn // active connections not yet closed
	bus              *eventbus.Bus    // only used for comparison with already set bus.
	eventClient      *eventbus.Client
	eventBusSubs     eventbus.Monitor
}

// sysConn wraps a net.Conn that was created using d.SystemDial.
// It exists to track which connections are still open, and should be
// closed on major link changes.
type sysConn struct {
	net.Conn
	id int
	d  *Dialer
}

func (c sysConn) Close() error {
	c.d.closeSysConn(c.id)
	return nil
}

// SetTUNName sets the name of the tun device in use ("tailscale0", "utun6",
// etc). This is needed on some platforms to set sockopts to bind
// to the same interface index.
func (d *Dialer) SetTUNName(name string) {
	d.mu.Lock()
	defer d.mu.Unlock()
	d.tunName = name
}

// TUNName returns the name of the tun device in use, if any.
// Example format ("tailscale0", "utun6").
func (d *Dialer) TUNName() string {
	d.mu.Lock()
	defer d.mu.Unlock()
	return d.tunName
}

// SetExitDNSDoH sets (or clears) the exit node DNS DoH server base URL to use.
// The doh URL should contain the scheme, authority, and path, but without
// a '?' and/or query parameters.
//
// For example, "http://100.68.82.120:47830/dns-query".
func (d *Dialer) SetExitDNSDoH(doh string) {
	if !buildfeatures.HasUseExitNode {
		return
	}
	d.mu.Lock()
	defer d.mu.Unlock()
	if d.exitDNSDoHBase == doh {
		return
	}
	d.exitDNSDoHBase = doh
	if doh != "" && d.dnsCache == nil {
		d.dnsCache = new(dnscache.MessageCache)
	}
	if d.dnsCache != nil {
		d.dnsCache.Flush()
	}
}

// SetRoutes configures the dialer to dial the specified routes via Tailscale,
// and the specified localRoutes using the default interface.
func (d *Dialer) SetRoutes(routes, localRoutes []netip.Prefix) {
	var rt *bart.Table[bool]
	if len(routes) > 0 || len(localRoutes) > 0 {
		rt = &bart.Table[bool]{}
		for _, r := range routes {
			rt.Insert(r, true)
		}
		for _, r := range localRoutes {
			rt.Insert(r, false)
		}
		d.logf("tsdial: bart table size: %d", rt.Size())
	}

	d.routes.Store(rt)
}

func (d *Dialer) Close() error {
	if d.eventClient != nil {
		d.eventBusSubs.Close()
	}
	d.mu.Lock()
	defer d.mu.Unlock()
	d.closed = true
	if d.netMonUnregister != nil {
		d.netMonUnregister()
		d.netMonUnregister = nil
	}
	for _, c := range d.activeSysConns {
		c.Close()
	}
	d.activeSysConns = nil
	d.PeerAPITransport().CloseIdleConnections()
	return nil
}

// SetNetMon sets d's network monitor to netMon.
// It is a no-op to call SetNetMon with the same netMon as the current one.
func (d *Dialer) SetNetMon(netMon *netmon.Monitor) {
	d.mu.Lock()
	defer d.mu.Unlock()
	if d.netMon == netMon {
		return
	}
	if d.netMonUnregister != nil {
		go d.netMonUnregister()
		d.netMonUnregister = nil
	}
	d.netMon = netMon
	// Having multiple watchers could lead to problems,
	// so remove the eventClient if it exists.
	// This should really not happen, but better checking for it than not.
	// TODO(cmol): Should this just be a panic?
	if d.eventClient != nil {
		d.eventBusSubs.Close()
		d.eventClient = nil
	}
	d.netMonUnregister = d.netMon.RegisterChangeCallback(d.linkChanged)
}

// NetMon returns the Dialer's network monitor.
// It returns nil if SetNetMon has not been called.
func (d *Dialer) NetMon() *netmon.Monitor {
	d.mu.Lock()
	defer d.mu.Unlock()
	return d.netMon
}

func (d *Dialer) SetBus(bus *eventbus.Bus) {
	d.mu.Lock()
	defer d.mu.Unlock()
	if d.bus == bus {
		return
	} else if d.bus != nil {
		panic("different eventbus has already been set")
	}
	// Having multiple watchers could lead to problems,
	// so unregister the callback if it exists.
	if d.netMonUnregister != nil {
		d.netMonUnregister()
	}
	d.bus = bus
	d.eventClient = bus.Client("tsdial.Dialer")
	d.eventBusSubs = d.eventClient.Monitor(d.linkChangeWatcher(d.eventClient))
}

func (d *Dialer) linkChangeWatcher(ec *eventbus.Client) func(*eventbus.Client) {
	linkChangeSub := eventbus.Subscribe[netmon.ChangeDelta](ec)
	return func(ec *eventbus.Client) {
		for {
			select {
			case <-ec.Done():
				return
			case cd := <-linkChangeSub.Events():
				d.linkChanged(&cd)
			}
		}
	}
}

var (
	metricLinkChangeConnClosed      = clientmetric.NewCounter("tsdial_linkchange_closes")
	metricChangeDeltaNoDefaultRoute = clientmetric.NewCounter("tsdial_changedelta_no_default_route")
)

func (d *Dialer) linkChanged(delta *netmon.ChangeDelta) {
	// Track how often we see ChangeDeltas with no DefaultRouteInterface.
	if delta.DefaultRouteInterface == "" {
		metricChangeDeltaNoDefaultRoute.Add(1)
	}

	d.mu.Lock()
	defer d.mu.Unlock()
	var anyClosed bool
	for id, c := range d.activeSysConns {
		if changeAffectsConn(delta, c) {
			anyClosed = true
			d.logf("tsdial: closing system connection %v->%v due to link change", c.LocalAddr(), c.RemoteAddr())
			go c.Close()
			delete(d.activeSysConns, id)
		}
	}
	if anyClosed {
		metricLinkChangeConnClosed.Add(1)
	}
}

// changeAffectsConn reports whether the network change delta affects
// the provided connection.
func changeAffectsConn(delta *netmon.ChangeDelta, conn net.Conn) bool {
	la, _ := conn.LocalAddr().(*net.TCPAddr)
	ra, _ := conn.RemoteAddr().(*net.TCPAddr)
	if la == nil || ra == nil {
		return false // not TCP
	}
	lip, rip := la.AddrPort().Addr(), ra.AddrPort().Addr()

	if delta.IsInitialState {
		return false
	}

	if delta.DefaultInterfaceChanged ||
		delta.HasPACOrProxyConfigChanged {
		return true
	}

	// In a few cases, we don't have a new DefaultRouteInterface (e.g. on
	// Android and macOS/iOS; see tailscale/corp#19124); if so, pessimistically assume
	// that all connections are affected.
	if delta.DefaultRouteInterface == "" && runtime.GOOS != "plan9" {
		return true
	}

	if delta.InterfaceIPDisappeared(lip) {
		// Our interface with this source IP went away.
		return true
	}
	_ = rip // TODO(bradfitz): use the remote IP?
	return false
}

func (d *Dialer) closeSysConn(id int) {
	d.mu.Lock()
	defer d.mu.Unlock()
	c, ok := d.activeSysConns[id]
	if !ok {
		return
	}
	delete(d.activeSysConns, id)
	go c.Close() // ignore the error
}

// peerDialControlFunc is non-nil on platforms that require a way to
// bind to dial out to other peers.
var peerDialControlFunc func(*Dialer) func(network, address string, c syscall.RawConn) error

// PeerDialControlFunc returns a function
// that can assigned to net.Dialer.Control to set sockopts or whatnot
// to make a dial escape the current platform's network sandbox.
//
// On many platforms the returned func will be nil.
//
// Notably, this is non-nil on iOS and macOS when run as a Network or
// System Extension (the GUI variants).
func (d *Dialer) PeerDialControlFunc() func(network, address string, c syscall.RawConn) error {
	if peerDialControlFunc == nil {
		return nil
	}
	return peerDialControlFunc(d)
}

// SetNetMap sets the current network map and notably, the DNS names
// in its DNS configuration.
func (d *Dialer) SetNetMap(nm *netmap.NetworkMap) {
	m := dnsMapFromNetworkMap(nm)

	d.mu.Lock()
	defer d.mu.Unlock()
	d.dns = m
}

// userDialResolve resolves addr as if a user initiating the dial. (e.g. from a
// SOCKS or HTTP outbound proxy)
func (d *Dialer) userDialResolve(ctx context.Context, network, addr string) (netip.AddrPort, error) {
	d.mu.Lock()
	dns := d.dns
	exitDNSDoH := d.exitDNSDoHBase
	d.mu.Unlock()

	// MagicDNS or otherwise baked into the NetworkMap? Try that first.
	ipp, err := dns.resolveMemory(ctx, network, addr)
	if err != errUnresolved {
		return ipp, err
	}

	// Otherwise, hit the network.

	// TODO(bradfitz): wire up net/dnscache too.

	host, port, err := splitHostPort(addr)
	if err != nil {
		// addr is malformed.
		return netip.AddrPort{}, err
	}

	if buildfeatures.HasUseExitNode && buildfeatures.HasPeerAPIClient && exitDNSDoH != "" {
		var r net.Resolver
		r.PreferGo = true
		r.Dial = func(ctx context.Context, network, address string) (net.Conn, error) {
			return &dohConn{
				ctx:      ctx,
				baseURL:  exitDNSDoH,
				hc:       d.PeerAPIHTTPClient(),
				dnsCache: d.dnsCache,
			}, nil
		}

		ips, err := r.LookupIP(ctx, ipNetOfNetwork(network), host)
		if err != nil {
			return netip.AddrPort{}, err
		}
		if len(ips) == 0 {
			return netip.AddrPort{}, fmt.Errorf("DNS lookup returned no results for %q", host)
		}
		ip, _ := netip.AddrFromSlice(ips[0])
		return netip.AddrPortFrom(ip.Unmap(), port), nil
	}

	if ip, err := netip.ParseAddr(host); err == nil {
		return netip.AddrPortFrom(ip.Unmap(), port), nil
	}
	var ips []netip.Addr
	if d.LookupHook != nil {
		ips, err = d.LookupHook(ctx, host)
		if err != nil {
			return netip.AddrPort{}, err
		}
	}
	for _, ip := range ips {
		ip = ip.Unmap()
		if strings.HasSuffix(network, "4") {
			if ip.Is4() {
				return netip.AddrPortFrom(ip, port), nil
			}
		} else if strings.HasSuffix(network, "6") {
			if ip.Is6() {
				return netip.AddrPortFrom(ip, port), nil
			}
		} else {
			return netip.AddrPortFrom(ip, port), nil
		}
	}
	return netip.AddrPort{}, fmt.Errorf("DNS lookup returned no results for %q", host)
}

// ipNetOfNetwork returns "ip", "ip4", or "ip6" corresponding
// to the input value of "tcp", "tcp4", "udp6" etc network
// names.
func ipNetOfNetwork(n string) string {
	if strings.HasSuffix(n, "4") {
		return "ip4"
	}
	if strings.HasSuffix(n, "6") {
		return "ip6"
	}
	return "ip"
}

func (d *Dialer) logf(format string, args ...any) {
	if d.Logf != nil {
		d.Logf(format, args...)
	}
}

// SystemPacketListen returns a net.PacketConn that listens on the
// provided network address. It prefers going over the default interface
// and closes existing connections if the default interface changes. It is
// used to listen for DERP packets.
func (d *Dialer) SystemPacketListen(ctx context.Context, network, address string) (net.PacketConn, error) {
	d.mu.Lock()
	if d.netMon == nil && d.SystemPacketListener == nil {
		d.mu.Unlock()
		if testenv.InTest() {
			panic("SystemPacketListen requires a netmon.Monitor; call SetNetMon first")
		}
		return nil, errors.New("SystemPacketListen requires a netmon.Monitor; call SetNetMon first")
	}
	closed := d.closed
	d.mu.Unlock()
	if closed {
		return nil, net.ErrClosed
	}

	listener := d.SystemPacketListener
	if listener == nil {
		listener = netns.Listener(d.logf, d.netMon).ListenPacket
	}
	return listener(ctx, network, address)
}

// SetSystemDialerForTest sets an alternate function to use for SystemDial
// instead of netns.Dialer. This is intended for use with nettest.MemoryNetwork.
func (d *Dialer) SetSystemDialerForTest(fn netx.DialFunc) {
	testenv.AssertInTest()
	d.sysDialForTest = fn
}

// SystemDial connects to the provided network address without going over
// Tailscale. It prefers going over the default interface and closes existing
// connections if the default interface changes. It is used to connect to
// Control and (in the future, as of 2022-04-27) DERPs..
func (d *Dialer) SystemDial(ctx context.Context, network, addr string) (net.Conn, error) {
	d.mu.Lock()
	if d.netMon == nil && d.sysDialForTest == nil && d.SystemDialer == nil {
		d.mu.Unlock()
		if testenv.InTest() {
			panic("SystemDial requires a netmon.Monitor; call SetNetMon first")
		}
		return nil, errors.New("SystemDial requires a netmon.Monitor; call SetNetMon first")
	}
	closed := d.closed
	d.mu.Unlock()
	if closed {
		return nil, net.ErrClosed
	}

	var c net.Conn
	var err error
	if d.sysDialForTest != nil {
		c, err = d.sysDialForTest(ctx, network, addr)
	} else if d.SystemDialer != nil {
		c, err = d.SystemDialer(ctx, network, addr)
	} else {
		d.netnsDialerOnce.Do(func() {
			d.netnsDialer = netns.NewDialer(d.logf, d.netMon)
		})
		c, err = d.netnsDialer.DialContext(ctx, network, addr)
	}
	if err != nil {
		return nil, err
	}
	d.mu.Lock()
	defer d.mu.Unlock()
	id := d.nextSysConnID
	d.nextSysConnID++
	mak.Set(&d.activeSysConns, id, c)

	return sysConn{
		id:   id,
		d:    d,
		Conn: c,
	}, nil
}

// UserDial connects to the provided network address as if a user were
// initiating the dial. (e.g. from a SOCKS or HTTP outbound proxy)
func (d *Dialer) UserDial(ctx context.Context, network, addr string) (net.Conn, error) {
	ipp, err := d.userDialResolve(ctx, network, addr)
	if err != nil {
		return nil, err
	}
	if d.UseNetstackForIP != nil && d.UseNetstackForIP(ipp.Addr()) {
		if d.NetstackDialTCP == nil || d.NetstackDialUDP == nil {
			return nil, errors.New("Dialer not initialized correctly")
		}
		if strings.HasPrefix(network, "udp") {
			return d.NetstackDialUDP(ctx, ipp)
		}
		return d.NetstackDialTCP(ctx, ipp)
	}

	if routes := d.routes.Load(); routes != nil {
		if isTailscaleRoute, _ := routes.Lookup(ipp.Addr()); isTailscaleRoute {
			return d.getPeerDialer().DialContext(ctx, network, ipp.String())
		}

		return d.SystemDial(ctx, network, ipp.String())
	}

	// Workaround for macOS for now: dial Tailscale IPs with peer dialer.
	// TODO(bradfitz): fix dialing subnet routers, public IPs via exit nodes,
	// etc. This is a temporary partial for macOS. We need to plumb ART tables &
	// prefs & host routing table updates around in more places. We just don't
	// know from the limited context here how to dial properly.
	if version.IsMacGUIVariant() && tsaddr.IsTailscaleIP(ipp.Addr()) {
		return d.getPeerDialer().DialContext(ctx, network, ipp.String())
	}
	return d.SystemDial(ctx, network, ipp.String())
}

// UserDialPlan resolves addr and reports whether the dialer would
// handle it via Tailscale. If viaTailscale is false, the resolved
// address is not a Tailscale route and the caller may dial it directly.
//
// Warning: there is a TOCTOU race if addr contains a DNS name and the
// caller subsequently passes the same DNS name to [Dialer.UserDial], as DNS
// may resolve differently the second time. Callers who want to only
// dial over Tailscale should call [Dialer.UserDial] with the returned
// ipp.String() (an IP:port) rather than the original DNS name.
func (d *Dialer) UserDialPlan(ctx context.Context, network, addr string) (ipp netip.AddrPort, viaTailscale bool, err error) {
	ipp, err = d.userDialResolve(ctx, network, addr)
	if err != nil {
		return netip.AddrPort{}, false, err
	}
	if d.UseNetstackForIP != nil && d.UseNetstackForIP(ipp.Addr()) {
		return ipp, true, nil
	}
	if routes := d.routes.Load(); routes != nil {
		isTailscaleRoute, _ := routes.Lookup(ipp.Addr())
		return ipp, isTailscaleRoute, nil
	}
	if version.IsMacGUIVariant() && tsaddr.IsTailscaleIP(ipp.Addr()) {
		return ipp, true, nil
	}
	return ipp, false, nil
}

// dialPeerAPI connects to a Tailscale peer's peerapi over TCP.
//
// network must a "tcp" type, and addr must be an ip:port. Name resolution
// is not supported.
func (d *Dialer) dialPeerAPI(ctx context.Context, network, addr string) (net.Conn, error) {
	if !buildfeatures.HasPeerAPIClient {
		return nil, feature.ErrUnavailable
	}
	switch network {
	case "tcp", "tcp6", "tcp4":
	default:
		return nil, fmt.Errorf("peerAPI dial requires tcp; %q not supported", network)
	}
	ipp, err := netip.ParseAddrPort(addr)
	if err != nil {
		return nil, fmt.Errorf("peerAPI dial requires ip:port, not name resolution: %w", err)
	}
	if d.UseNetstackForIP != nil && d.UseNetstackForIP(ipp.Addr()) {
		if d.NetstackDialTCP == nil {
			return nil, errors.New("Dialer not initialized correctly")
		}
		return d.NetstackDialTCP(ctx, ipp)
	}
	return d.getPeerDialer().DialContext(ctx, network, addr)
}

// getPeerDialer returns the *net.Dialer to use to dial peers (e.g. for peerapi,
// "tailscale nc", or querying internal DNS servers over Tailscale)
//
// This is not used in netstack mode.
//
// The primary function of this is to work on macOS & iOS's in the
// Network/System Extension so it can mark the dialer as staying within the
// network namespace/sandbox.
func (d *Dialer) getPeerDialer() *net.Dialer {
	d.peerDialerOnce.Do(func() {
		d.peerDialer = &net.Dialer{
			Timeout:   30 * time.Second,
			KeepAlive: netknob.PlatformTCPKeepAlive(),
			Control:   d.PeerDialControlFunc(),
		}
	})
	return d.peerDialer
}

// PeerAPIHTTPClient returns an HTTP Client to call peers' peerapi
// endpoints.                                                                                                                                                                                                                      //
// The returned Client must not be mutated; it's owned by the Dialer
// and shared by callers.
func (d *Dialer) PeerAPIHTTPClient() *http.Client {
	if !buildfeatures.HasPeerAPIClient {
		panic("unreachable")
	}
	d.peerClientOnce.Do(func() {
		t := &http.Transport{
			ForceAttemptHTTP2:     true,
			MaxIdleConns:          100,
			IdleConnTimeout:       90 * time.Second,
			TLSHandshakeTimeout:   10 * time.Second,
			ExpectContinueTimeout: 1 * time.Second,
		}
		t.Dial = nil
		t.DialContext = d.dialPeerAPI
		// Do not use the environment proxy for PeerAPI.
		t.Proxy = nil
		d.peerClient = &http.Client{Transport: t}
	})
	return d.peerClient
}

// PeerAPITransport returns a Transport to call peers' peerapi
// endpoints.
//
// The returned value must not be mutated; it's owned by the Dialer
// and shared by callers.
func (d *Dialer) PeerAPITransport() *http.Transport {
	return d.PeerAPIHTTPClient().Transport.(*http.Transport)
}
