// Copyright (c) 2026 Charles KWON OhJun (charleskwonohjun@gmail.com)
// All rights reserved.

// tls.go — TLS support for the pgserver. Cert/key load + self-
// signed generator + IP allowlist.
//
// PRG bootstrap idiom:
//
//	PG_TLS_LOAD( "cert.pem", "key.pem" )       /* or */
//	PG_TLS_SELF_SIGNED( "/tmp/cert.pem", "/tmp/key.pem", "localhost" )
//	PG_ALLOW_IP( "127.0.0.1/32" )              /* optional allowlist */
//	PG_ALLOW_IP( "10.0.0.0/8" )
//	PG_SERVER_START( ":5432", "md5" )
//
// The TLS config is installed on the *next* PG_SERVER_START call.
// A server started before PG_TLS_LOAD runs with plaintext only;
// matches PostgreSQL's "must be set before pg_ctl start" semantics.

package pgserver

import (
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"fmt"
	"math/big"
	"net"
	"os"
	"sync"
	"time"
)

// upgradeToTLS wraps the underlying net.Conn in a tls.Server using
// the configured *tls.Config and performs the TLS handshake. The
// returned net.Conn is the encrypted stream; pgproto3 sees only
// plaintext on top of it.
func upgradeToTLS(conn net.Conn, cfg *tls.Config) (net.Conn, error) {
	tlsConn := tls.Server(conn, cfg)
	if err := tlsConn.Handshake(); err != nil {
		return nil, err
	}
	return tlsConn, nil
}

// --- TLS config holder ---
//
// pendingTLSConfig is consumed by the next PG_SERVER_START invocation
// (see register.go::pgServerStart). Stored separately from any
// running server so PRG callers can compose cert load + start.

var (
	tlsMu             sync.Mutex
	pendingTLSConfig  *tls.Config
	pendingAllowList  []*net.IPNet // empty → allow any
)

// SetTLSConfig installs the TLS config that the next started
// server will use. Pass nil to clear (subsequent starts run
// plaintext). Concurrent-safe.
func SetTLSConfig(cfg *tls.Config) {
	tlsMu.Lock()
	defer tlsMu.Unlock()
	pendingTLSConfig = cfg
}

// CurrentTLSConfig returns the pending TLS config. Called by
// pgServerStart at boot time.
func CurrentTLSConfig() *tls.Config {
	tlsMu.Lock()
	defer tlsMu.Unlock()
	return pendingTLSConfig
}

// AllowIP appends a CIDR range to the per-server allowlist. The
// list is empty by default → all source IPs accepted. Once
// non-empty, the accept loop rejects any peer whose RemoteAddr
// doesn't fall inside one of the registered ranges. PG-equivalent
// of a coarse pg_hba.conf with one `host` line per call.
func AllowIP(cidr string) error {
	_, ipnet, err := net.ParseCIDR(cidr)
	if err != nil {
		return err
	}
	tlsMu.Lock()
	defer tlsMu.Unlock()
	pendingAllowList = append(pendingAllowList, ipnet)
	return nil
}

// CurrentAllowList returns the registered CIDRs.
func CurrentAllowList() []*net.IPNet {
	tlsMu.Lock()
	defer tlsMu.Unlock()
	if len(pendingAllowList) == 0 {
		return nil
	}
	out := make([]*net.IPNet, len(pendingAllowList))
	copy(out, pendingAllowList)
	return out
}

// ClearAllowList drops every registered CIDR. After this call,
// the server reverts to "accept any source IP".
func ClearAllowList() {
	tlsMu.Lock()
	defer tlsMu.Unlock()
	pendingAllowList = nil
}

// peerAllowed checks an incoming RemoteAddr against the
// configured allowlist. Empty list → unconditional allow.
func peerAllowed(remote net.Addr, list []*net.IPNet) bool {
	if len(list) == 0 {
		return true
	}
	tcpAddr, ok := remote.(*net.TCPAddr)
	if !ok {
		return false
	}
	for _, n := range list {
		if n.Contains(tcpAddr.IP) {
			return true
		}
	}
	return false
}

// --- Cert / key loading ---

// LoadTLSFromFiles installs a tls.Config built from PEM-encoded
// cert + key files. Returns an error if either file is missing
// or the key doesn't match the cert.
func LoadTLSFromFiles(certPath, keyPath string) error {
	cert, err := tls.LoadX509KeyPair(certPath, keyPath)
	if err != nil {
		return fmt.Errorf("pgserver: load TLS pair: %w", err)
	}
	SetTLSConfig(&tls.Config{
		Certificates: []tls.Certificate{cert},
		MinVersion:   tls.VersionTLS12,
	})
	return nil
}

// GenerateSelfSignedCert writes a fresh ECDSA P-256 self-signed
// cert/key pair to certPath/keyPath and installs it as the
// pending TLS config. Validity is 365 days. Intended for dev /
// CI where a CA-signed cert isn't available; production should
// always run LoadTLSFromFiles with a properly signed pair.
//
// Hostname is added as a SubjectAltName so clients connecting to
// `psql 'host=hostname sslmode=verify-full'` accept it.
func GenerateSelfSignedCert(certPath, keyPath, hostname string) error {
	priv, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		return fmt.Errorf("genkey: %w", err)
	}
	serial, err := rand.Int(rand.Reader, new(big.Int).Lsh(big.NewInt(1), 128))
	if err != nil {
		return fmt.Errorf("serial: %w", err)
	}
	tpl := &x509.Certificate{
		SerialNumber: serial,
		Subject:      pkix.Name{CommonName: hostname},
		NotBefore:    time.Now().Add(-time.Hour),
		NotAfter:     time.Now().Add(365 * 24 * time.Hour),
		KeyUsage:     x509.KeyUsageDigitalSignature | x509.KeyUsageKeyEncipherment,
		ExtKeyUsage:  []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		BasicConstraintsValid: true,
		IsCA:        false,
		DNSNames:    []string{hostname},
		IPAddresses: []net.IP{net.ParseIP("127.0.0.1"), net.ParseIP("::1")},
	}
	der, err := x509.CreateCertificate(rand.Reader, tpl, tpl, &priv.PublicKey, priv)
	if err != nil {
		return fmt.Errorf("create cert: %w", err)
	}
	certPEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: der})
	if err := os.WriteFile(certPath, certPEM, 0644); err != nil {
		return fmt.Errorf("write cert: %w", err)
	}
	keyDER, err := x509.MarshalECPrivateKey(priv)
	if err != nil {
		return fmt.Errorf("marshal key: %w", err)
	}
	keyPEM := pem.EncodeToMemory(&pem.Block{Type: "EC PRIVATE KEY", Bytes: keyDER})
	if err := os.WriteFile(keyPath, keyPEM, 0600); err != nil {
		return fmt.Errorf("write key: %w", err)
	}
	return LoadTLSFromFiles(certPath, keyPath)
}