// Package identity resolves a caller's address to a tailnet node name, and gates a
// route on the set of peers it allows. In production the resolver asks
// `tailscale whois`; in the smoke run it trusts a request header. A Checker caches
// the answer per address so a hot peer does not re-query whois on every request.
package identity

import (
	"context"
	"encoding/json"
	"errors"
	"net"
	"os/exec"
	"strings"
	"sync"
	"time"
)

var (
	// ErrNotAPeer is returned by a resolver when the address is not a known
	// tailnet node. The Checker turns it into a deny.
	ErrNotAPeer = errors.New("identity: not a tailnet peer")
	// ErrForbidden is returned by the Checker when the caller is not on the
	// route's allow list.
	ErrForbidden = errors.New("identity: forbidden route")
)

// ID is the tailnet identity of a caller.
type ID struct {
	Node, Login string
}

// Resolver maps an IP address to the tailnet node it belongs to.
type Resolver interface {
	Identity(ctx context.Context, ip string) (ID, error)
}

// cacheTTL is how long a resolved identity (or a rejection) is held per address.
const cacheTTL = 5 * time.Minute

// ParseWhois decodes `tailscale whois --json` output. Node is ComputedName, or
// Name with its trailing dot and domain stripped; Login is the profile login
// name. An empty node name is an error.
func ParseWhois(raw []byte) (ID, error) {
	var whois struct {
		Node struct {
			Name         string `json:"Name"`
			ComputedName string `json:"ComputedName"`
		} `json:"Node"`
		UserProfile struct {
			LoginName string `json:"LoginName"`
		} `json:"UserProfile"`
	}
	if err := json.Unmarshal(raw, &whois); err != nil {
		return ID{}, err
	}
	node := whois.Node.ComputedName
	if node == "" {
		node = stripName(whois.Node.Name)
	}
	if node == "" {
		return ID{}, errors.New("identity: whois has no node name")
	}
	return ID{Node: node, Login: whois.UserProfile.LoginName}, nil
}

// stripName takes a whois Name such as "titan.example.ts.net." and returns the
// first label, "titan".
func stripName(name string) string {
	name = strings.TrimSuffix(name, ".")
	if i := strings.IndexByte(name, '.'); i >= 0 {
		name = name[:i]
	}
	return name
}

// TailscaleResolver runs `tailscale whois --json <ip>` and parses it. A non-zero
// exit is ErrNotAPeer; a missing binary (or other transport failure) is a real
// error the Checker treats as a deny.
type TailscaleResolver struct{ Bin string }

// Identity runs the whois lookup with a 3 s timeout.
func (t TailscaleResolver) Identity(ctx context.Context, ip string) (ID, error) {
	bin := t.Bin
	if bin == "" {
		bin = "tailscale"
	}
	ctx, cancel := context.WithTimeout(ctx, 3*time.Second)
	defer cancel()
	out, err := exec.CommandContext(ctx, bin, "whois", "--json", ip).Output()
	if err != nil {
		var exitErr *exec.ExitError
		if errors.As(err, &exitErr) {
			return ID{}, ErrNotAPeer
		}
		return ID{}, err
	}
	return ParseWhois(out)
}

// entry is a cached result, whether a node name or a rejection.
type entry struct {
	id  ID
	err error
	at  time.Time
}

// Checker resolves addresses through a Resolver, caching per address. In header
// mode it skips the cache and lets the resolver read the peer from the request.
type Checker struct {
	r      Resolver
	header bool

	mu    sync.Mutex
	cache map[string]entry
}

// NewChecker builds a Checker that resolves through r.
func NewChecker(r Resolver) *Checker {
	return &Checker{r: r, cache: make(map[string]entry)}
}

// NewHeaderChecker builds a Checker that trusts the X-Crossbar-Peer request
// header as the node name. TEST/SMOKE ONLY.
func NewHeaderChecker() *Checker {
	return &Checker{r: headerResolver{}, header: true, cache: make(map[string]entry)}
}

// WithHeaderPeer returns a context carrying the peer name the header checker
// reads. Middleware sets it from the X-Crossbar-Peer request header.
func WithHeaderPeer(ctx context.Context, peer string) context.Context {
	return context.WithValue(ctx, headerPeerKey{}, peer)
}

// Allow reports whether the caller at remoteAddr may use a route limited to peers.
// An empty peers list is an open route; otherwise the caller's node must be in
// peers. Any resolver error, an unparsable address, or a loopback address denies.
func (c *Checker) Allow(ctx context.Context, peers []string, remoteAddr string) error {
	if len(peers) == 0 {
		return nil
	}
	ip := ""
	if !c.header {
		var ok bool
		ip, ok = peerIP(remoteAddr)
		if !ok || net.ParseIP(ip).IsLoopback() {
			return ErrForbidden
		}
	}
	node, err := c.resolve(ctx, ip)
	if err != nil {
		return ErrForbidden
	}
	for _, p := range peers {
		if p == node {
			return nil
		}
	}
	return ErrForbidden
}

// resolve returns the node for ip, using the cache unless in header mode.
func (c *Checker) resolve(ctx context.Context, ip string) (string, error) {
	if c.header {
		id, err := c.r.Identity(ctx, ip)
		if err != nil {
			return "", err
		}
		return id.Node, nil
	}

	c.mu.Lock()
	e, hit := c.cache[ip]
	if hit && time.Since(e.at) < cacheTTL {
		c.mu.Unlock()
		if e.err != nil {
			return "", e.err
		}
		return e.id.Node, nil
	}
	c.mu.Unlock()

	id, err := c.r.Identity(ctx, ip)
	if err != nil {
		c.mu.Lock()
		c.cache[ip] = entry{err: err, at: time.Now()}
		c.mu.Unlock()
		return "", err
	}

	c.mu.Lock()
	c.cache[ip] = entry{id: id, at: time.Now()}
	c.mu.Unlock()
	return id.Node, nil
}

// headerPeerKey is the context key under which Middleware stores the X-Crossbar-Peer
// value for the header checker to read.
type headerPeerKey struct{}

// headerResolver answers from the peer name Middleware placed on the context. An
// absent or empty header is a deny, so a request that forgot the header is 403.
type headerResolver struct{}

func (headerResolver) Identity(ctx context.Context, _ string) (ID, error) {
	peer, ok := ctx.Value(headerPeerKey{}).(string)
	if !ok || peer == "" {
		return ID{}, ErrNotAPeer
	}
	return ID{Node: peer, Login: peer}, nil
}

// peerIP splits the host from a "host:port" address, returning the bare IP.
func peerIP(remoteAddr string) (string, bool) {
	host := remoteAddr
	if h, _, err := net.SplitHostPort(remoteAddr); err == nil {
		host = h
	}
	ip := net.ParseIP(host)
	if ip == nil {
		return "", false
	}
	return ip.String(), true
}