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

package main

import (
	"bytes"
	"context"
	"crypto/tls"
	"encoding/json"
	"errors"
	"expvar"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"net/netip"
	"strconv"
	"strings"
	"sync"
	"time"

	"go4.org/mem"
	"golang.org/x/net/http2"
	"tailscale.com/control/controlbase"
	"tailscale.com/control/controlhttp/controlhttpserver"
	"tailscale.com/control/ts2021"
	"tailscale.com/derp"
	"tailscale.com/derp/derpserver"
	"tailscale.com/net/netmon"
	"tailscale.com/net/tsdial"
	"tailscale.com/types/key"
	"tailscale.com/types/logger"
	"tailscale.com/util/eventbus"
)

const managedNaviControlVersion = 131

type managedNavi struct {
	ctx            context.Context
	cancel         context.CancelFunc
	logf           logger.Logf
	controlURL     string
	derpID         string
	machinePrivate key.MachinePrivate

	mu            sync.RWMutex
	controlPub    key.MachinePublic
	controlClient *ts2021.Client
	info          managedNaviNode
	trustSeq      int
	trustedNodes  map[key.NodePublic]struct{}
	funnel        managedNaviFunnelPayload
	funnelRuntime *managedNaviFunnelRuntime
	pullNodesFn   func() error
	pullFunnelFn  func() error
	funnelSyncRun bool
	funnelSyncDur time.Duration

	wg sync.WaitGroup
}

type legacyMirageRegisterAuth struct {
	Provider  string `json:"Provider,omitempty"`
	LoginName string `json:"LoginName,omitempty"`
}

type legacyMirageRegisterRequest struct {
	Auth *legacyMirageRegisterAuth `json:"Auth,omitempty"`
}

type managedNaviNode struct {
	ID           string `json:"Name"`
	NaviKey      string `json:"NaviKey,omitempty"`
	NaviRegionID int    `json:"RegionID"`
	HostName     string `json:"HostName"`
	IPv4         string `json:"IPv4,omitempty"`
	IPv6         string `json:"IPv6,omitempty"`
	NoSTUN       bool   `json:"NoSTUN"`
	STUNPort     int    `json:"STUNPort"`
	NoDERP       bool   `json:"NoDERP"`
	DERPPort     int    `json:"DERPPort"`
	SSHAddr      string `json:"SSHAddr,omitempty"`
	SSHPwd       string `json:"SSHPwd,omitempty"`
	DNSProvider  string `json:"DNSProvider,omitempty"`
	DNSID        string `json:"DNSID,omitempty"`
	DNSKey       string `json:"DNSKey,omitempty"`
	Arch         string `json:"Arch,omitempty"`
}

type managedNaviRegisterResponse struct {
	NaviInfo   managedNaviNode `json:"NaviInfo"`
	TrustNodes []string        `json:"TrustNodes"`
	Timestamp  *time.Time      `json:"Timestamp,omitempty"`
}

type managedNaviHostinfo struct {
	FrontendLogID string `json:"FrontendLogID,omitempty"`
	BackendLogID  string `json:"BackendLogID,omitempty"`
}

type managedNaviPollNodesRequest struct {
	Hostinfo *managedNaviHostinfo `json:"Hostinfo,omitempty"`
}

type managedNaviPullNodesResponse struct {
	TrustNodes []string   `json:"TrustNodes"`
	Timestamp  *time.Time `json:"Timestamp,omitempty"`
}

type managedOverTLSPublicKeyResponse struct {
	PublicKey key.MachinePublic `json:"PublicKey"`
}

type managedNaviNodesChange struct {
	SeqNum     int    `json:"SeqNum"`
	AddNode    string `json:"AddNode,omitempty"`
	RemoveNode string `json:"RemoveNode,omitempty"`
}

func newManagedNavi(ctx context.Context, machinePrivate key.MachinePrivate, logf logger.Logf, controlURL, derpID string) *managedNavi {
	if ctx == nil {
		ctx = context.Background()
	}
	ctx, cancel := context.WithCancel(ctx)
	if logf == nil {
		logf = log.Printf
	}
	m := &managedNavi{
		ctx:            ctx,
		cancel:         cancel,
		logf:           logf,
		controlURL:     strings.TrimRight(controlURL, "/"),
		derpID:         derpID,
		machinePrivate: machinePrivate,
		trustedNodes:   make(map[key.NodePublic]struct{}),
		funnelSyncDur:  managedNaviFunnelSyncInterval,
	}
	m.pullNodesFn = m.PullNodesList
	m.pullFunnelFn = m.PullFunnel
	return m
}

func (m *managedNavi) Bootstrap() error {
	if err := m.prepareControlClient(); err != nil {
		return err
	}
	resp, err := m.register()
	if err != nil {
		return err
	}
	if err := m.setTrustedNodes(resp.TrustNodes); err != nil {
		return err
	}
	m.mu.Lock()
	m.info = resp.NaviInfo
	m.mu.Unlock()
	if err := m.pullFunnelFn(); err != nil {
		m.logf("Mirage pull funnel payload failed: %v", err)
	}
	m.startFunnelSyncLoop()
	return nil
}

func (m *managedNavi) Close() error {
	m.mu.Lock()
	client := m.controlClient
	m.controlClient = nil
	rt := m.funnelRuntime
	m.funnelRuntime = nil
	m.mu.Unlock()
	if m.cancel != nil {
		m.cancel()
	}
	if rt != nil {
		rt.close()
	}
	m.wg.Wait()
	if client != nil {
		return client.Close()
	}
	return nil
}

func (m *managedNavi) prepareControlClient() error {
	req, err := http.NewRequestWithContext(
		m.ctx,
		http.MethodGet,
		fmt.Sprintf("%s/key?v=%d", m.controlURL, managedNaviControlVersion),
		nil,
	)
	if err != nil {
		return fmt.Errorf("create Mirage control key request: %w", err)
	}
	res, err := http.DefaultClient.Do(req)
	if err != nil {
		return fmt.Errorf("fetch Mirage control key: %w", err)
	}
	defer res.Body.Close()
	body, err := io.ReadAll(io.LimitReader(res.Body, 64<<10))
	if err != nil {
		return fmt.Errorf("read Mirage control key response: %w", err)
	}
	if res.StatusCode != http.StatusOK {
		return fmt.Errorf("fetch Mirage control key: http %d: %.200s", res.StatusCode, strings.TrimSpace(string(body)))
	}
	var keys managedOverTLSPublicKeyResponse
	if err := json.Unmarshal(body, &keys); err != nil {
		return fmt.Errorf("decode Mirage control key response: %w", err)
	}
	if keys.PublicKey.IsZero() {
		return fmt.Errorf("Mirage control key response missing ts2021 public key")
	}

	bus := eventbus.New()
	netMon, err := netmon.New(bus, m.logf)
	if err != nil {
		return fmt.Errorf("create Mirage net monitor: %w", err)
	}
	dialer := tsdial.NewDialer(netMon)
	dialer.Logf = m.logf
	dialer.SetBus(bus)

	client, err := ts2021.NewClient(ts2021.ClientOpts{
		PrivKey:      m.machinePrivate,
		ServerPubKey: keys.PublicKey,
		ServerURL:    m.controlURL,
		Dialer:       dialer,
		Logf:         m.logf,
		NetMon:       netMon,
	})
	if err != nil {
		return fmt.Errorf("create Mirage control client: %w", err)
	}

	m.mu.Lock()
	m.controlPub = keys.PublicKey
	m.controlClient = client
	m.mu.Unlock()
	return nil
}

func (m *managedNavi) register() (*managedNaviRegisterResponse, error) {
	reqBody := legacyMirageRegisterRequest{
		Auth: &legacyMirageRegisterAuth{
			Provider:  "Mirage",
			LoginName: m.derpID,
		},
	}
	var resp managedNaviRegisterResponse
	if err := m.postControlJSON("/machine/register", reqBody, &resp); err != nil {
		return nil, fmt.Errorf("Mirage managed register failed: %w", err)
	}
	return &resp, nil
}

func (m *managedNavi) PullNodesList() error {
	reqBody := managedNaviPollNodesRequest{
		Hostinfo: &managedNaviHostinfo{
			FrontendLogID: "derper",
			BackendLogID:  m.derpID,
		},
	}
	var resp managedNaviPullNodesResponse
	if err := m.postControlJSON("/navi/nodes", reqBody, &resp); err != nil {
		return fmt.Errorf("Mirage pull trusted nodes failed: %w", err)
	}
	return m.setTrustedNodes(resp.TrustNodes)
}

func (m *managedNavi) postControlJSON(path string, reqBody, out any) error {
	body, err := json.Marshal(reqBody)
	if err != nil {
		return err
	}
	req, err := http.NewRequestWithContext(m.ctx, http.MethodPost, m.controlURL+path, bytes.NewReader(body))
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", "application/json")

	m.mu.RLock()
	client := m.controlClient
	m.mu.RUnlock()
	if client == nil {
		return fmt.Errorf("Mirage control client is not ready")
	}
	res, err := client.Do(req)
	if err != nil {
		return err
	}
	return decodeManagedJSONResponse(res, out)
}

func decodeManagedJSONResponse(res *http.Response, out any) error {
	defer res.Body.Close()
	body, err := io.ReadAll(io.LimitReader(res.Body, 1<<20))
	if err != nil {
		return err
	}
	if res.StatusCode != http.StatusOK {
		return fmt.Errorf("http %d: %.200s", res.StatusCode, strings.TrimSpace(string(body)))
	}
	if out == nil {
		return nil
	}
	if err := json.Unmarshal(body, out); err != nil {
		return err
	}
	return nil
}

func (m *managedNavi) setTrustedNodes(rawKeys []string) error {
	trusted := make(map[key.NodePublic]struct{}, len(rawKeys))
	for _, raw := range rawKeys {
		nodeKey, err := parseManagedNodeKey(raw)
		if err != nil {
			return err
		}
		trusted[nodeKey] = struct{}{}
	}
	m.mu.Lock()
	m.trustedNodes = trusted
	m.trustSeq = 0
	m.mu.Unlock()
	return nil
}

func parseManagedNodeKey(raw string) (key.NodePublic, error) {
	trimmed := strings.TrimSpace(strings.TrimPrefix(raw, "nodekey:"))
	if trimmed == "" {
		return key.NodePublic{}, fmt.Errorf("empty trusted node key")
	}
	nodeKey, err := key.ParseNodePublicUntyped(mem.S(trimmed))
	if err != nil {
		return key.NodePublic{}, fmt.Errorf("parse trusted node key %q: %w", raw, err)
	}
	return nodeKey, nil
}

func (m *managedNavi) ApplyNodeConfig(hostname, addr *string, stunPort *int, runDERP, runSTUN *bool) {
	info := m.infoSnapshot()
	if info.HostName != "" {
		*hostname = info.HostName
	}
	if info.NoDERP {
		*runDERP = false
	} else if port := info.DERPPort; port > 0 {
		*addr = joinListenAddrHost(*addr, port)
	}
	if info.NoSTUN {
		*runSTUN = false
	} else if port := info.STUNPort; port > 0 {
		*stunPort = port
	}
}

func (m *managedNavi) setFunnelRuntime(rt *managedNaviFunnelRuntime) {
	m.mu.Lock()
	old := m.funnelRuntime
	m.funnelRuntime = rt
	m.mu.Unlock()

	if old != nil && old != rt {
		old.close()
	}
	if err := m.applyStoredFunnelPayload(); err != nil {
		m.logf("Mirage apply stored funnel payload failed: %v", err)
	}
}

func (m *managedNavi) currentFunnelRuntime() *managedNaviFunnelRuntime {
	m.mu.RLock()
	defer m.mu.RUnlock()
	return m.funnelRuntime
}

func (m *managedNavi) wrapFunnelHandler(key managedNaviFunnelListenerKey, next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if rt := m.currentFunnelRuntime(); rt != nil {
			rt.wrapHandler(key, next).ServeHTTP(w, r)
			return
		}
		next.ServeHTTP(w, r)
	})
}

func (m *managedNavi) wrapFunnelACMEHandler(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if rt := m.currentFunnelRuntime(); rt != nil {
			rt.wrapHTTPChallengeHandler(next).ServeHTTP(w, r)
			return
		}
		next.ServeHTTP(w, r)
	})
}

func (m *managedNavi) getFunnelCertificate(hello *tls.ClientHelloInfo) (*tls.Certificate, bool, error) {
	rt := m.currentFunnelRuntime()
	if rt == nil {
		return nil, false, nil
	}
	cert, err := rt.getCertificateForHost(hello)
	if err != nil {
		host := ""
		if hello != nil {
			host = normalizeManagedNaviFunnelBaseDomain(hello.ServerName)
		}
		return nil, host != "", err
	}
	return cert, cert != nil, nil
}

func (m *managedNavi) applyStoredFunnelPayload() error {
	rt := m.currentFunnelRuntime()
	if rt == nil {
		return nil
	}
	snapshot := m.funnelSnapshot()
	return rt.applyPayload(&snapshot)
}

func (m *managedNavi) startFunnelSyncLoop() {
	m.mu.Lock()
	if m.funnelSyncRun || m.funnelSyncDur <= 0 {
		m.mu.Unlock()
		return
	}
	interval := m.funnelSyncDur
	m.funnelSyncRun = true
	m.mu.Unlock()

	m.wg.Add(1)
	go func() {
		defer m.wg.Done()
		ticker := time.NewTicker(interval)
		defer ticker.Stop()
		for {
			select {
			case <-m.ctx.Done():
				return
			case <-ticker.C:
				if err := m.pullFunnelFn(); err != nil && !errors.Is(err, context.Canceled) {
					m.logf("Mirage periodic funnel sync failed: %v", err)
				}
			}
		}
	}()
}

func joinListenAddrHost(addr string, port int) string {
	host, _, err := net.SplitHostPort(addr)
	if err != nil {
		return ":" + strconv.Itoa(port)
	}
	return net.JoinHostPort(host, strconv.Itoa(port))
}

func (m *managedNavi) infoSnapshot() managedNaviNode {
	m.mu.RLock()
	defer m.mu.RUnlock()
	return m.info
}

func (m *managedNavi) VerifyClient(_ context.Context, clientKey key.NodePublic, _ *derp.ClientInfo, _ netip.Addr) error {
	m.mu.RLock()
	_, ok := m.trustedNodes[clientKey]
	m.mu.RUnlock()
	if ok {
		return nil
	}
	return fmt.Errorf("client %v not authorized by Mirage control", clientKey)
}

func (m *managedNavi) NoiseUpgradeHandler(w http.ResponseWriter, r *http.Request) {
	conn, err := controlhttpserver.AcceptHTTP(r.Context(), w, r, m.machinePrivate, nil)
	if err != nil {
		m.logf("Mirage ts2021 upgrade failed: %v", err)
		return
	}

	mux := http.NewServeMux()
	mux.HandleFunc("/ctrl/nodes", m.handleNodesChange)
	mux.Handle("/ctrl/vars", expvar.Handler())
	mux.Handle("/generate_204", http.HandlerFunc(derpserver.ServeNoContent))

	base := &http.Server{
		Handler:           m.authorizeControlConn(conn, mux),
		ReadHeaderTimeout: 30 * time.Second,
	}
	h2s := &http2.Server{}
	h2s.ServeConn(conn, &http2.ServeConnOpts{BaseConfig: base})
}

func (m *managedNavi) authorizeControlConn(conn *controlbase.Conn, next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		m.mu.RLock()
		controlPub := m.controlPub
		m.mu.RUnlock()
		if conn.Peer() != controlPub {
			http.Error(w, "noise auth failed", http.StatusUnauthorized)
			return
		}
		next.ServeHTTP(w, r)
	})
}

func (m *managedNavi) handleNodesChange(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Error(w, "wrong method", http.StatusMethodNotAllowed)
		return
	}
	var change managedNaviNodesChange
	if err := json.NewDecoder(r.Body).Decode(&change); err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}
	if err := m.applyNodesChange(change); err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	w.WriteHeader(http.StatusOK)
}

func (m *managedNavi) applyNodesChange(change managedNaviNodesChange) error {
	m.mu.Lock()
	expectedSeq := m.trustSeq + 1
	if change.SeqNum != expectedSeq {
		m.mu.Unlock()
		return m.pullNodesFn()
	}

	m.trustSeq = change.SeqNum
	if change.AddNode != "" {
		nodeKey, err := parseManagedNodeKey(change.AddNode)
		if err != nil {
			m.mu.Unlock()
			return err
		}
		m.trustedNodes[nodeKey] = struct{}{}
	}
	if change.RemoveNode != "" {
		nodeKey, err := parseManagedNodeKey(change.RemoveNode)
		if err != nil {
			m.mu.Unlock()
			return err
		}
		delete(m.trustedNodes, nodeKey)
	}
	m.mu.Unlock()
	return nil
}
