// Package proxy is the allowlisting HTTP proxy that is sandboxed commands'
// only way onto the network. It runs outside the sandbox and supports
// exactly two things: CONNECT tunnels (for TLS, which it does not
// intercept) and plain HTTP requests with absolute URIs. A request goes
// through only if its port is allowed and its host is approved.
package proxy

import (
	"context"
	"errors"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"net/netip"
	"slices"
	"strconv"
	"strings"
	"sync"
	"time"
)

// Marker starts every refusal the proxy sends, so the model and the user
// can tell a refusal from the server's own errors.
const Marker = "blinker proxy:"

// Server is one proxy.
type Server struct {
	// Check approves a host, or explains its refusal to the client, after
	// Marker. It may block, for example while asking the user; the
	// client's connection is held open meanwhile. It must return false on
	// any doubt.
	Check func(ctx context.Context, host string) (ok bool, why string)
	// Ports are the ports the proxy will connect to.
	Ports []int
	// AllowPrivate lets the proxy connect to loopback, private and
	// link-local addresses. Off, it refuses them, since the proxy runs
	// outside the sandbox and could otherwise reach local services and
	// cloud metadata endpoints.
	AllowPrivate bool
	// Log, if non-nil, receives a line per request.
	Log func(string)

	mu    sync.Mutex
	conns map[net.Conn]bool
	srv   *http.Server
}

func (s *Server) logf(format string, args ...any) {
	if s.Log != nil {
		s.Log(fmt.Sprintf(format, args...))
	}
}

// serverLog passes net/http's own errors to Log.
type serverLog struct{ s *Server }

func (l serverLog) Write(p []byte) (int, error) {
	l.s.logf("ERROR server: %s", strings.TrimRight(string(p), "\n"))
	return len(p), nil
}

// Serve answers connections on l until Close.
func (s *Server) Serve(l net.Listener) error {
	s.mu.Lock()
	s.conns = map[net.Conn]bool{}
	s.srv = &http.Server{
		Handler:           s,
		ReadHeaderTimeout: 30 * time.Second,
		// Left nil, net/http would log to stderr, over the TUI or the box.
		ErrorLog: log.New(serverLog{s}, "", 0),
	}
	srv := s.srv
	s.mu.Unlock()
	err := srv.Serve(l)
	if errors.Is(err, http.ErrServerClosed) {
		return nil
	}
	return err
}

// Close stops the proxy and drops its tunnels.
func (s *Server) Close() {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.srv != nil {
		if err := s.srv.Close(); err != nil {
			s.logf("closing: %v", err)
		}
	}
	// The tunnels are being dropped, so an error closing one changes
	// nothing.
	for c := range s.conns {
		c.Close()
	}
	s.conns = nil
}

func (s *Server) track(c net.Conn, add bool) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.conns == nil {
		// Closed (or not yet serving): no tunnel outlives the proxy.
		if add {
			c.Close()
		}
		return
	}
	if add {
		s.conns[c] = true
	} else {
		delete(s.conns, c)
	}
}

// normalHost is a host as rules and the user see it: an IP address, or a
// name of letters, digits, '-' and '_' in dot-separated labels, in lower
// case without a final dot. Anything else is refused, so what Check
// judges, the user is asked about and a rule saves is one plain name.
func normalHost(host string) (string, bool) {
	if ip, err := netip.ParseAddr(host); err == nil {
		return ip.String(), true
	}
	// A-Z only: Unicode's rules would make some other letters ASCII (the
	// Kelvin sign k), and the name approved would not be the one asked for.
	h := strings.TrimSuffix(strings.Map(func(r rune) rune {
		if r >= 'A' && r <= 'Z' {
			return r + 'a' - 'A'
		}
		return r
	}, host), ".")
	if h == "" || len(h) > 253 {
		return "", false
	}
	for _, label := range strings.Split(h, ".") {
		if label == "" || len(label) > 63 {
			return "", false
		}
		for _, c := range label {
			if !(c >= 'a' && c <= 'z' || c >= '0' && c <= '9' || c == '-' || c == '_') {
				return "", false
			}
		}
	}
	return h, true
}

// vet decides a target: the address to dial, or the status and reason for
// refusing it. host must be normalHost's.
func (s *Server) vet(ctx context.Context, host string, port int) (addrs []string, status int, why string) {
	if !slices.Contains(s.Ports, port) {
		var ports []string
		for _, p := range s.Ports {
			ports = append(ports, strconv.Itoa(p))
		}
		return nil, http.StatusForbidden, fmt.Sprintf("port %d on %q is not permitted (allowed ports: %s).",
			port, host, strings.Join(ports, ", "))
	}
	if ok, why := s.Check(ctx, host); !ok {
		if why == "" {
			why = fmt.Sprintf("domain %q is not approved.", host)
		}
		return nil, http.StatusForbidden, why
	}
	// Dial the addresses checked, so a DNS answer that changes in between
	// cannot slip past the private-address check.
	ips, err := resolve(ctx, host)
	if err != nil {
		return nil, http.StatusBadGateway, fmt.Sprintf("cannot resolve %q: %v.", host, err)
	}
	for _, ip := range ips {
		if !s.AllowPrivate && IsPrivate(ip) {
			return nil, http.StatusForbidden, fmt.Sprintf("%q resolves to a private or loopback address (%s), which the proxy refuses.", host, ip)
		}
		addrs = append(addrs, net.JoinHostPort(ip.String(), strconv.Itoa(port)))
	}
	return addrs, 0, ""
}

func resolve(ctx context.Context, host string) ([]netip.Addr, error) {
	if ip, err := netip.ParseAddr(host); err == nil {
		return []netip.Addr{ip.Unmap()}, nil
	}
	ips, err := net.DefaultResolver.LookupNetIP(ctx, "ip", host)
	if err != nil {
		return nil, err
	}
	if len(ips) == 0 {
		return nil, errors.New("no addresses")
	}
	for i := range ips {
		ips[i] = ips[i].Unmap()
	}
	return ips, nil
}

// dial connects to the first of addrs that answers.
func dial(ctx context.Context, addrs []string) (net.Conn, error) {
	var err error
	for _, a := range addrs {
		var c net.Conn
		c, err = (&net.Dialer{Timeout: 30 * time.Second}).DialContext(ctx, "tcp", a)
		if err == nil {
			return c, nil
		}
	}
	return nil, err
}

// IsPrivate reports whether ip is loopback, private, link-local (cloud
// metadata included), carrier-grade NAT, unspecified or otherwise not a
// public unicast address.
func IsPrivate(ip netip.Addr) bool {
	ip = ip.Unmap()
	if !ip.IsValid() || ip.IsLoopback() || ip.IsPrivate() || ip.IsLinkLocalUnicast() ||
		ip.IsLinkLocalMulticast() || ip.IsUnspecified() || ip.IsMulticast() || ip.IsInterfaceLocalMulticast() {
		return true
	}
	if ip.Is4() {
		b := ip.As4()
		return b[0] == 0 || (b[0] == 100 && b[1] >= 64 && b[1] <= 127) || // 0/8, CGNAT 100.64/10
			b[0] == 192 && b[1] == 0 && b[2] == 0 || // IETF protocol assignments 192.0.0/24
			b[0] == 198 && (b[1] == 18 || b[1] == 19) || // benchmarking 198.18/15
			b[0] >= 240 // reserved 240/4, broadcast included
	}
	// Addresses that carry an IPv4 one, which a gateway on the way may
	// reach for them: NAT64 (64:ff9b::/96, 64:ff9b:1::/48), 6to4
	// (2002::/16), and the old IPv4-compatible ::a.b.c.d.
	b := ip.As16()
	switch {
	case b[0] == 0x00 && b[1] == 0x64 && b[2] == 0xff && b[3] == 0x9b:
		return IsPrivate(netip.AddrFrom4([4]byte(b[12:16])))
	case b[0] == 0x20 && b[1] == 0x02:
		return IsPrivate(netip.AddrFrom4([4]byte(b[2:6])))
	case [12]byte(b[:12]) == [12]byte{}:
		return IsPrivate(netip.AddrFrom4([4]byte(b[12:16])))
	case b[0] == 0xfe && b[1]&0xc0 == 0xc0: // site-local fec0::/10, deprecated
		return true
	}
	return false
}

func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if r.Method == http.MethodConnect {
		s.connect(w, r)
		return
	}
	refuse := func(status int, why string) {
		http.Error(w, Marker+" "+why, status)
		s.logf("DENY http %s %s: %s", r.Method, r.URL, why)
	}
	if !r.URL.IsAbs() || r.URL.Host == "" {
		refuse(http.StatusBadRequest, "only absolute-URI HTTP requests and CONNECT are supported.")
		return
	}
	if r.URL.Scheme != "http" {
		refuse(http.StatusBadRequest, "https must use CONNECT.")
		return
	}
	host, ok := normalHost(r.URL.Hostname())
	if !ok {
		refuse(http.StatusBadRequest, "malformed host.")
		return
	}
	port := 80
	if p := r.URL.Port(); p != "" {
		var err error
		if port, err = strconv.Atoi(p); err != nil {
			refuse(http.StatusBadRequest, "malformed port.")
			return
		}
	}
	addrs, status, why := s.vet(r.Context(), host, port)
	if why != "" {
		refuse(status, why)
		return
	}
	out := r.Clone(r.Context())
	out.RequestURI = ""
	out.Header.Del("Proxy-Connection")
	out.Header.Del("Proxy-Authorization")
	transport := &http.Transport{
		Proxy: nil,
		DialContext: func(ctx context.Context, network, _ string) (net.Conn, error) {
			return dial(ctx, addrs)
		},
		DisableKeepAlives: true,
	}
	resp, err := transport.RoundTrip(out)
	if err != nil {
		refuse(http.StatusBadGateway, fmt.Sprintf("upstream error: %v", err))
		return
	}
	defer resp.Body.Close()
	for k, vs := range resp.Header {
		for _, v := range vs {
			w.Header().Add(k, v)
		}
	}
	w.WriteHeader(resp.StatusCode)
	// An error midway can only cut the response short, as the client will
	// see; there is nothing else to tell it.
	if _, err := io.Copy(w, resp.Body); err != nil {
		s.logf("ERROR http %s %s: %v", r.Method, r.URL, err)
	}
	s.logf("ALLOW http %s %s status=%d", r.Method, r.URL, resp.StatusCode)
}

func (s *Server) connect(w http.ResponseWriter, r *http.Request) {
	target := r.Host
	host, portStr, err := net.SplitHostPort(target)
	port, perr := strconv.Atoi(portStr)
	host, ok := normalHost(host)
	if err != nil || perr != nil || !ok {
		http.Error(w, Marker+" malformed CONNECT target.", http.StatusBadRequest)
		s.logf("DENY connect %s: malformed", target)
		return
	}
	addrs, status, why := s.vet(r.Context(), host, port)
	if why != "" {
		http.Error(w, Marker+" "+why, status)
		s.logf("DENY connect %s: %s", target, why)
		return
	}
	upstream, err := dial(r.Context(), addrs)
	if err != nil {
		http.Error(w, fmt.Sprintf("%s cannot reach %s: %v", Marker, target, err), http.StatusBadGateway)
		s.logf("ERROR connect %s: %v", target, err)
		return
	}
	hj, ok := w.(http.Hijacker)
	if !ok {
		upstream.Close()
		http.Error(w, Marker+" cannot tunnel.", http.StatusInternalServerError)
		return
	}
	client, buf, err := hj.Hijack()
	if err != nil {
		upstream.Close()
		return
	}
	s.track(client, true)
	s.track(upstream, true)
	defer func() {
		s.track(client, false)
		s.track(upstream, false)
		client.Close()
		upstream.Close()
	}()
	if _, err := client.Write([]byte("HTTP/1.1 200 Connection Established\r\n\r\n")); err != nil {
		return
	}
	// buf holds anything the client sent after its request.
	up, down, upErr, downErr := Splice(client, buf.Reader, upstream)
	// Either side failing just ends the tunnel, as closing it would.
	s.logf("ALLOW connect %s up=%d (%v) down=%d (%v)", target, up, upErr, down, downErr)
}

// Splice copies what a sends (read from aIn, which may hold some of it
// already) to b, and what b sends to a, until both are done, telling each
// side when the other has no more to send.
func Splice(a net.Conn, aIn io.Reader, b net.Conn) (up, down int64, upErr, downErr error) {
	done := make(chan struct{})
	go func() {
		up, upErr = io.Copy(b, aIn)
		closeWrite(b)
		close(done)
	}()
	down, downErr = io.Copy(a, b)
	closeWrite(a)
	<-done
	return up, down, upErr, downErr
}

// closeWrite tells the other end that no more data is coming. If that
// fails the connection is broken anyway, and the deferred Close ends it.
func closeWrite(c net.Conn) {
	if cw, ok := c.(interface{ CloseWrite() error }); ok {
		cw.CloseWrite()
	} else {
		c.Close()
	}
}