package main

import (
	"bytes"
	"context"
	"encoding/binary"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"net"
	"net/http"
	"net/url"
	"os"
	"os/signal"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"syscall"
	"time"

	"github.com/google/nftables"
	"github.com/google/nftables/binaryutil"
	"github.com/google/nftables/expr"
	"github.com/vishvananda/netlink"
	"golang.org/x/sys/unix"
	"golang.zx2c4.com/wireguard/wgctrl"
	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
	corev1 "k8s.io/api/core/v1"
	discoveryv1 "k8s.io/api/discovery/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/leaderelection"
	"k8s.io/client-go/tools/leaderelection/resourcelock"
	"k8s.io/client-go/util/retry"
)

const (
	interfaceName       = "wg-oci0"
	routingTable        = 51820
	rulePriority        = 100
	localPriority       = 90
	firewallTable       = "oci_l4_relay_home"
	rpFilterRoot        = "/proc/sys/net/ipv4/conf"
	mailConnectionLimit = 1024
)

type config struct {
	namespace        string
	podName          string
	podUID           string
	podIP            string
	podNodeName      string
	leaseName        string
	privateKeyPath   string
	peerPublicKey    string
	peerEndpoint     string
	listenPort       int
	tunnelAddress    *net.IPNet
	targets          []relayTarget
	healthPath       string
	metricsListen    string
	metricsTarget    *url.URL
	metricsSliceName string
	endpointAddress  net.IP
	localReturnCIDRs []*net.IPNet
	leaseDuration    time.Duration
	renewDeadline    time.Duration
	retryPeriod      time.Duration
	maxRetryPeriod   time.Duration
	handshakeMaxAge  time.Duration
	mailTarget       *relayTarget
	mailUpstream     string
}

type relayTarget struct {
	address net.IP
	ports   []uint16
}

type relayTargetDocument struct {
	Address string `json:"address"`
	Ports   []int  `json:"ports"`
}

type relay struct {
	config        config
	client        kubernetes.Interface
	metricsClient *http.Client
	active        atomic.Bool
	unhealthy     atomic.Bool
	stop          context.CancelFunc
	fence         *nftables.Set
	proxies       []*mailProxy
	proxyFailures chan error
	cleanupMu     sync.Mutex
}

func main() {
	if len(os.Args) == 2 && os.Args[1] == "health" {
		path := envOrDefault("RELAY_HEALTH_PATH", "/run/relay-agent/health")
		if err := checkHealth(path, 20*time.Second); err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}
		return
	}

	configuration, err := loadConfig()
	if err != nil {
		slog.Error("invalid configuration", "error", err)
		os.Exit(1)
	}
	if len(os.Args) == 2 && os.Args[1] == "status" {
		if err := (&relay{config: configuration}).verify(); err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}
		counters, err := readFirewallCounters(configuration.targets)
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}
		allowed, err := json.Marshal(counters.tunnelAllowed)
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}
		fmt.Printf("network_state=ready tunnel_allowed_packets=%s tunnel_denied_packets=%d outer_allowed_packets=%d outer_denied_packets=%d\n",
			allowed,
			counters.tunnelDenied,
			counters.outerAllowed,
			counters.outerDenied,
		)
		return
	}
	clientConfig, err := rest.InClusterConfig()
	if err != nil {
		slog.Error("load Kubernetes client configuration", "error", err)
		os.Exit(1)
	}
	client, err := kubernetes.NewForConfig(clientConfig)
	if err != nil {
		slog.Error("create Kubernetes client", "error", err)
		os.Exit(1)
	}

	signalContext, stopSignals := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
	defer stopSignals()
	instance := &relay{config: configuration, client: client, proxyFailures: make(chan error, 1)}
	listener, err := net.Listen("tcp4", configuration.metricsListen)
	if err != nil {
		slog.Error("listen for relay metrics", "error", err)
		os.Exit(1)
	}
	metricsServer := &http.Server{
		Handler:           instance.metricsHandler(),
		ReadHeaderTimeout: 5 * time.Second,
		WriteTimeout:      10 * time.Second,
		IdleTimeout:       30 * time.Second,
	}
	go func() {
		if err := metricsServer.Serve(listener); err != nil && !errors.Is(err, http.ErrServerClosed) {
			slog.Error("serve relay metrics", "error", err)
			stopSignals()
		}
	}()
	defer func() {
		shutdownContext, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		if err := metricsServer.Shutdown(shutdownContext); err != nil {
			slog.Error("shut down relay metrics", "error", err)
		}
	}()
	if err := instance.cleanup(); err != nil {
		slog.Error("reconcile stale local relay state", "error", err)
		os.Exit(1)
	}
	go instance.heartbeat(signalContext)

	repeatElection(signalContext, configuration.retryPeriod, configuration.maxRetryPeriod, func(electionContext context.Context, cancelElection context.CancelFunc) bool {
		instance.unhealthy.Store(false)
		instance.stop = cancelElection
		lock := &resourcelock.LeaseLock{
			LeaseMeta: metav1.ObjectMeta{Name: configuration.leaseName, Namespace: configuration.namespace},
			Client:    client.CoordinationV1(),
			LockConfig: resourcelock.ResourceLockConfig{
				Identity: configuration.podName + "_" + configuration.podUID,
			},
		}
		leaderelection.RunOrDie(electionContext, leaderelection.LeaderElectionConfig{
			Lock:            lock,
			LeaseDuration:   configuration.leaseDuration,
			RenewDeadline:   configuration.renewDeadline,
			RetryPeriod:     configuration.retryPeriod,
			ReleaseOnCancel: true,
			Name:            configuration.leaseName,
			Callbacks: leaderelection.LeaderCallbacks{
				OnStartedLeading: instance.runActive,
				OnStoppedLeading: func() {
					instance.stopActive()
					slog.Info("relay ownership released")
				},
				OnNewLeader: func(identity string) {
					if identity != configuration.podName+"_"+configuration.podUID {
						slog.Info("standby relay observed active owner", "identity", identity)
					}
				},
			},
		})
		return instance.unhealthy.Load()
	})
}

func repeatElection(ctx context.Context, retryPeriod, maxRetryPeriod time.Duration, run func(context.Context, context.CancelFunc) bool) {
	unhealthyTerms := 0
	for ctx.Err() == nil {
		electionContext, cancelElection := context.WithCancel(ctx)
		unhealthy := run(electionContext, cancelElection)
		cancelElection()
		if ctx.Err() != nil {
			return
		}
		if unhealthy {
			unhealthyTerms++
		} else {
			unhealthyTerms = 0
		}
		delay := electionRetryDelay(retryPeriod, maxRetryPeriod, unhealthyTerms)
		slog.Info("retry relay election", "delay", delay, "consecutive_unhealthy_terms", unhealthyTerms)
		timer := time.NewTimer(delay)
		select {
		case <-ctx.Done():
			timer.Stop()
		case <-timer.C:
		}
	}
}

func electionRetryDelay(retryPeriod, maxRetryPeriod time.Duration, unhealthyTerms int) time.Duration {
	delay := retryPeriod
	for term := 1; term < unhealthyTerms && delay < maxRetryPeriod; term++ {
		if delay > maxRetryPeriod/2 {
			return maxRetryPeriod
		}
		delay *= 2
	}
	return min(delay, maxRetryPeriod)
}

func loadConfig() (config, error) {
	tunnelAddress, err := parseCIDR(requiredEnv("RELAY_TUNNEL_ADDRESS"))
	if err != nil {
		return config{}, fmt.Errorf("RELAY_TUNNEL_ADDRESS: %w", err)
	}
	targets, err := parseRelayTargets(requiredEnv("RELAY_TARGETS"))
	if err != nil {
		return config{}, fmt.Errorf("RELAY_TARGETS: %w", err)
	}
	listenPort, err := strconv.Atoi(requiredEnv("OCI_RELAY_LISTEN_PORT"))
	if err != nil || listenPort < 1 || listenPort > 65535 {
		return config{}, errors.New("OCI_RELAY_LISTEN_PORT must be a valid port")
	}
	endpointAddress := net.ParseIP(requiredEnv("OCI_RELAY_ENDPOINT_IP")).To4()
	if endpointAddress == nil {
		return config{}, errors.New("OCI_RELAY_ENDPOINT_IP must be an IPv4 address")
	}
	localReturnCIDRs, err := parseNetworks(requiredEnv("RELAY_LOCAL_RETURN_CIDRS"))
	if err != nil {
		return config{}, fmt.Errorf("RELAY_LOCAL_RETURN_CIDRS: %w", err)
	}
	if len(localReturnCIDRs) > rulePriority-localPriority {
		return config{}, fmt.Errorf("RELAY_LOCAL_RETURN_CIDRS supports at most %d networks", rulePriority-localPriority)
	}
	configuration := config{
		namespace:        requiredEnv("POD_NAMESPACE"),
		podName:          requiredEnv("POD_NAME"),
		podUID:           requiredEnv("POD_UID"),
		podIP:            requiredEnv("POD_IP"),
		podNodeName:      requiredEnv("POD_NODE_NAME"),
		leaseName:        envOrDefault("RELAY_LEASE_NAME", "oci-l4-relay"),
		privateKeyPath:   envOrDefault("RELAY_PRIVATE_KEY_PATH", "/run/secrets/wireguard/private_key"),
		peerPublicKey:    requiredEnv("OCI_RELAY_PUBLIC_KEY"),
		peerEndpoint:     net.JoinHostPort(requiredEnv("OCI_RELAY_ENDPOINT_IP"), strconv.Itoa(listenPort)),
		listenPort:       listenPort,
		tunnelAddress:    tunnelAddress,
		targets:          targets,
		healthPath:       envOrDefault("RELAY_HEALTH_PATH", "/run/relay-agent/health"),
		metricsListen:    envOrDefault("RELAY_METRICS_LISTEN_ADDRESS", ":19100"),
		metricsSliceName: envOrDefault("RELAY_METRICS_ENDPOINT_SLICE", "oci-l4-relay-metrics"),
		endpointAddress:  endpointAddress,
		localReturnCIDRs: localReturnCIDRs,
		leaseDuration:    30 * time.Second,
		renewDeadline:    20 * time.Second,
		retryPeriod:      5 * time.Second,
		maxRetryPeriod:   60 * time.Second,
		handshakeMaxAge:  180 * time.Second,
	}
	metricsTarget, err := parseMetricsTarget(requiredEnv("OCI_RELAY_METRICS_URL"), tunnelAddress)
	if err != nil {
		return config{}, fmt.Errorf("OCI_RELAY_METRICS_URL: %w", err)
	}
	configuration.metricsTarget = metricsTarget
	mailTarget, mailUpstream, err := mailProxyTarget(requiredEnv("RELAY_MAIL_UPSTREAM"), tunnelAddress)
	if err != nil {
		return config{}, err
	}
	if mailTarget != nil {
		configuration.targets = append(configuration.targets, *mailTarget)
		sort.Slice(configuration.targets, func(left, right int) bool {
			return bytes.Compare(configuration.targets[left].address, configuration.targets[right].address) < 0
		})
	}
	configuration.mailTarget = mailTarget
	configuration.mailUpstream = mailUpstream
	if address := net.ParseIP(configuration.podIP).To4(); address == nil || address.String() != configuration.podIP {
		return config{}, errors.New("POD_IP must be a canonical IPv4 address")
	}
	if configuration.podNodeName == "" {
		return config{}, errors.New("POD_NODE_NAME must not be empty")
	}
	if _, err := wgtypes.ParseKey(configuration.peerPublicKey); err != nil {
		return config{}, fmt.Errorf("OCI_RELAY_PUBLIC_KEY: %w", err)
	}
	return configuration, nil
}

func mailProxyTarget(upstream string, tunnelAddress *net.IPNet) (*relayTarget, string, error) {
	if upstream == "" {
		return nil, "", nil
	}
	if strings.Contains(upstream, ":") {
		return nil, "", errors.New("RELAY_MAIL_UPSTREAM must be a DNS hostname")
	}
	return &relayTarget{address: tunnelAddress.IP, ports: []uint16{25, 465, 993}}, upstream, nil
}

func parseMetricsTarget(value string, tunnelAddress *net.IPNet) (*url.URL, error) {
	target, err := url.Parse(value)
	if err != nil || target.Scheme != "http" || target.Path != "/metrics" || target.RawQuery != "" || target.Fragment != "" || target.User != nil {
		return nil, errors.New("must be an HTTP /metrics URL without credentials, query, or fragment")
	}
	address := net.ParseIP(target.Hostname()).To4()
	if address == nil || !tunnelAddress.Contains(address) || address.Equal(tunnelAddress.IP) {
		return nil, errors.New("host must be the IPv4 peer in RELAY_TUNNEL_ADDRESS")
	}
	port, err := strconv.Atoi(target.Port())
	if err != nil || port < 1 || port > 65535 {
		return nil, errors.New("port must be valid")
	}
	return target, nil
}

func parseRelayTargets(value string) ([]relayTarget, error) {
	decoder := json.NewDecoder(strings.NewReader(value))
	decoder.DisallowUnknownFields()
	var documents []relayTargetDocument
	if err := decoder.Decode(&documents); err != nil {
		return nil, err
	}
	if err := decoder.Decode(&struct{}{}); err != io.EOF {
		return nil, errors.New("contains trailing data")
	}
	if len(documents) == 0 || len(documents) > 16 {
		return nil, errors.New("must contain between 1 and 16 targets")
	}
	allowedPorts := map[int]bool{25: true, 80: true, 443: true, 465: true, 993: true}
	addresses := make(map[string]bool, len(documents))
	ports := make(map[int]bool)
	targets := make([]relayTarget, 0, len(documents))
	for _, document := range documents {
		address := net.ParseIP(document.Address).To4()
		if address == nil || address.String() != document.Address || addresses[document.Address] {
			return nil, fmt.Errorf("invalid or duplicate target address %q", document.Address)
		}
		if len(document.Ports) == 0 {
			return nil, fmt.Errorf("target %s has no ports", document.Address)
		}
		addresses[document.Address] = true
		targetPorts := make([]uint16, 0, len(document.Ports))
		for _, port := range document.Ports {
			if !allowedPorts[port] || ports[port] {
				return nil, fmt.Errorf("invalid or duplicate target port %d", port)
			}
			ports[port] = true
			targetPorts = append(targetPorts, uint16(port))
		}
		sort.Slice(targetPorts, func(left, right int) bool {
			return targetPorts[left] < targetPorts[right]
		})
		targets = append(targets, relayTarget{address: address, ports: targetPorts})
	}
	sort.Slice(targets, func(left, right int) bool {
		return bytes.Compare(targets[left].address, targets[right].address) < 0
	})
	return targets, nil
}

func requiredEnv(name string) string {
	return strings.TrimSpace(os.Getenv(name))
}

func envOrDefault(name, fallback string) string {
	if value := requiredEnv(name); value != "" {
		return value
	}
	return fallback
}

func parseCIDR(value string) (*net.IPNet, error) {
	ip, network, err := net.ParseCIDR(value)
	if err != nil {
		return nil, err
	}
	ip = ip.To4()
	if ip == nil {
		return nil, errors.New("must be an IPv4 CIDR")
	}
	network.IP = ip
	return network, nil
}

func parseNetworks(value string) ([]*net.IPNet, error) {
	var networks []*net.IPNet
	for item := range strings.SplitSeq(value, ",") {
		_, network, err := net.ParseCIDR(strings.TrimSpace(item))
		if err != nil || network.IP.To4() == nil {
			return nil, fmt.Errorf("%q is not an IPv4 network", item)
		}
		networks = append(networks, network)
	}
	return networks, nil
}

func (r *relay) heartbeat(ctx context.Context) {
	ticker := time.NewTicker(5 * time.Second)
	defer ticker.Stop()
	for {
		if err := writeHealth(r.config.healthPath, !r.unhealthy.Load()); err != nil {
			slog.Error("update health heartbeat", "error", err)
		}
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
		}
	}
}

func (r *relay) metricsHandler() http.Handler {
	client := r.metricsClient
	if client == nil {
		client = &http.Client{
			Timeout: 8 * time.Second,
			CheckRedirect: func(_ *http.Request, _ []*http.Request) error {
				return http.ErrUseLastResponse
			},
		}
	}
	return http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) {
		if request.Method != http.MethodGet || request.URL.Path != "/metrics" || request.URL.RawQuery != "" {
			http.NotFound(response, request)
			return
		}
		if !r.active.Load() {
			http.Error(response, "relay owner is inactive", http.StatusServiceUnavailable)
			return
		}
		upstreamRequest, err := http.NewRequestWithContext(request.Context(), http.MethodGet, r.config.metricsTarget.String(), nil)
		if err != nil {
			http.Error(response, "relay metrics request failed", http.StatusBadGateway)
			return
		}
		upstreamResponse, err := client.Do(upstreamRequest)
		if err != nil {
			http.Error(response, "relay metrics request failed", http.StatusBadGateway)
			return
		}
		defer upstreamResponse.Body.Close()
		body, err := io.ReadAll(io.LimitReader(upstreamResponse.Body, 8*1024*1024+1))
		if err != nil || len(body) > 8*1024*1024 {
			http.Error(response, "relay metrics response failed", http.StatusBadGateway)
			return
		}
		if contentType := upstreamResponse.Header.Get("Content-Type"); contentType != "" {
			response.Header().Set("Content-Type", contentType)
		}
		response.WriteHeader(upstreamResponse.StatusCode)
		_, _ = response.Write(body)
	})
}

func (r *relay) runActive(ctx context.Context) {
	drainFailures(r.proxyFailures)
	if err := r.configure(); err != nil {
		slog.Error("configure active relay", "error", err)
		r.unhealthy.Store(true)
		r.stop()
		return
	}
	if err := r.startProxies(); err != nil {
		slog.Error("start mail proxy", "error", err)
		r.unhealthy.Store(true)
		r.stop()
		return
	}
	r.active.Store(true)
	if err := r.publishMetricsEndpoint(ctx); err != nil {
		slog.Error("publish active relay metrics endpoint", "error", err)
	}
	slog.Info("relay ownership acquired")
	ticker := time.NewTicker(5 * time.Second)
	started := time.Now()
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case err := <-r.proxyFailures:
			slog.Error("mail proxy failed", "error", err)
			r.unhealthy.Store(true)
			r.stop()
			return
		case <-ticker.C:
			if err := r.publishMetricsEndpoint(ctx); err != nil {
				slog.Error("reconcile active relay metrics endpoint", "error", err)
			}
			if err := r.refreshFence(); err != nil {
				slog.Error("refresh relay ownership fence", "error", err)
				r.unhealthy.Store(true)
				r.stop()
				return
			}
			if time.Since(started) < 30*time.Second {
				continue
			}
			if err := r.verify(); err != nil {
				slog.Error("active relay state failed verification", "error", err)
				r.unhealthy.Store(true)
				r.stop()
				return
			}
		}
	}
}

func (r *relay) stopActive() {
	r.active.Store(false)
	r.stopProxies()
	if err := r.cleanup(); err != nil {
		slog.Error("clean up active relay", "error", err)
	}
	clearContext, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	if err := r.clearMetricsEndpoint(clearContext); err != nil {
		slog.Error("clear active relay metrics endpoint", "error", err)
	}
}

func (r *relay) publishMetricsEndpoint(ctx context.Context) error {
	return retry.RetryOnConflict(retry.DefaultRetry, func() error {
		slices := r.client.DiscoveryV1().EndpointSlices(r.config.namespace)
		slice, err := slices.Get(ctx, r.config.metricsSliceName, metav1.GetOptions{})
		if err != nil {
			return err
		}
		if r.metricsEndpointMatches(slice) {
			return nil
		}
		ready := true
		serving := true
		terminating := false
		slice.Endpoints = []discoveryv1.Endpoint{{
			Addresses: []string{r.config.podIP},
			Conditions: discoveryv1.EndpointConditions{
				Ready:       &ready,
				Serving:     &serving,
				Terminating: &terminating,
			},
			NodeName: &r.config.podNodeName,
			TargetRef: &corev1.ObjectReference{
				Kind:      "Pod",
				Namespace: r.config.namespace,
				Name:      r.config.podName,
				UID:       types.UID(r.config.podUID),
			},
		}}
		_, err = slices.Update(ctx, slice, metav1.UpdateOptions{})
		return err
	})
}

func (r *relay) metricsEndpointMatches(slice *discoveryv1.EndpointSlice) bool {
	if len(slice.Endpoints) != 1 {
		return false
	}
	endpoint := slice.Endpoints[0]
	return len(endpoint.Addresses) == 1 && endpoint.Addresses[0] == r.config.podIP &&
		endpoint.NodeName != nil && *endpoint.NodeName == r.config.podNodeName &&
		endpoint.TargetRef != nil && endpoint.TargetRef.UID == types.UID(r.config.podUID) &&
		endpoint.Conditions.Ready != nil && *endpoint.Conditions.Ready &&
		endpoint.Conditions.Serving != nil && *endpoint.Conditions.Serving &&
		endpoint.Conditions.Terminating != nil && !*endpoint.Conditions.Terminating
}

func (r *relay) clearMetricsEndpoint(ctx context.Context) error {
	return retry.RetryOnConflict(retry.DefaultRetry, func() error {
		slices := r.client.DiscoveryV1().EndpointSlices(r.config.namespace)
		slice, err := slices.Get(ctx, r.config.metricsSliceName, metav1.GetOptions{})
		if err != nil {
			return err
		}
		if len(slice.Endpoints) != 1 || slice.Endpoints[0].TargetRef == nil || slice.Endpoints[0].TargetRef.UID != types.UID(r.config.podUID) {
			return nil
		}
		slice.Endpoints = nil
		_, err = slices.Update(ctx, slice, metav1.UpdateOptions{})
		return err
	})
}

func (r *relay) configure() error {
	if err := r.cleanup(); err != nil {
		return err
	}
	link := &netlink.GenericLink{LinkAttrs: netlink.LinkAttrs{Name: interfaceName, MTU: 1420}, LinkType: "wireguard"}
	if err := netlink.LinkAdd(link); err != nil {
		return fmt.Errorf("add WireGuard link: %w", err)
	}
	configured := false
	defer func() {
		if !configured {
			_ = r.cleanup()
		}
	}()
	created, err := netlink.LinkByName(interfaceName)
	if err != nil {
		return fmt.Errorf("find WireGuard link: %w", err)
	}
	if err := verifyReversePathFilter(rpFilterRoot, interfaceName); err != nil {
		return err
	}
	privateKeyBytes, err := os.ReadFile(r.config.privateKeyPath)
	if err != nil {
		return fmt.Errorf("read private key: %w", err)
	}
	privateKey, err := wgtypes.ParseKey(strings.TrimSpace(string(privateKeyBytes)))
	if err != nil {
		return fmt.Errorf("parse private key: %w", err)
	}
	peerKey, _ := wgtypes.ParseKey(r.config.peerPublicKey)
	endpoint, err := net.ResolveUDPAddr("udp4", r.config.peerEndpoint)
	if err != nil {
		return fmt.Errorf("resolve peer endpoint: %w", err)
	}
	keepalive := 25 * time.Second
	client, err := wgctrl.New()
	if err != nil {
		return fmt.Errorf("open WireGuard control: %w", err)
	}
	defer func() { _ = client.Close() }()
	if err := client.ConfigureDevice(interfaceName, wgtypes.Config{
		PrivateKey: &privateKey,
		Peers: []wgtypes.PeerConfig{{
			PublicKey:                   peerKey,
			Endpoint:                    endpoint,
			PersistentKeepaliveInterval: &keepalive,
			ReplaceAllowedIPs:           true,
			AllowedIPs:                  []net.IPNet{{IP: net.IPv4zero, Mask: net.CIDRMask(0, 32)}},
		}},
	}); err != nil {
		return fmt.Errorf("configure WireGuard device: %w", err)
	}
	if err := netlink.AddrAdd(created, &netlink.Addr{IPNet: r.config.tunnelAddress}); err != nil {
		return fmt.Errorf("add tunnel address: %w", err)
	}
	metricsPort, _ := strconv.Atoi(r.config.metricsTarget.Port())
	fence, err := installFirewall(r.config.targets, r.config.endpointAddress, uint16(r.config.listenPort), r.config.tunnelAddress.IP, net.ParseIP(r.config.metricsTarget.Hostname()), uint16(metricsPort))
	if err != nil {
		return err
	}
	r.fence = fence
	if err := netlink.LinkSetUp(created); err != nil {
		return fmt.Errorf("bring WireGuard link up: %w", err)
	}
	defaultRoute := &net.IPNet{IP: net.IPv4zero, Mask: net.CIDRMask(0, 32)}
	if err := netlink.RouteReplace(&netlink.Route{LinkIndex: created.Attrs().Index, Dst: defaultRoute, Table: routingTable, Scope: netlink.SCOPE_UNIVERSE}); err != nil {
		return fmt.Errorf("install relay default route: %w", err)
	}
	for _, target := range r.config.targets {
		for index, destination := range r.config.localReturnCIDRs {
			rule := netlink.NewRule()
			rule.Priority = localPriority + index
			rule.Table = unix.RT_TABLE_MAIN
			rule.Src = &net.IPNet{IP: target.address, Mask: net.CIDRMask(32, 32)}
			rule.Dst = destination
			if err := netlink.RuleAdd(rule); err != nil {
				return fmt.Errorf("install local return policy rule: %w", err)
			}
		}
	}
	for _, target := range r.config.targets {
		rule := netlink.NewRule()
		rule.Priority = rulePriority
		rule.Table = routingTable
		rule.Src = &net.IPNet{IP: target.address, Mask: net.CIDRMask(32, 32)}
		if err := netlink.RuleAdd(rule); err != nil {
			return fmt.Errorf("install relay policy rule: %w", err)
		}
	}
	configured = true
	return nil
}

func installFirewall(targets []relayTarget, endpointAddress net.IP, endpointPort uint16, tunnelAddress, metricsPeer net.IP, metricsPort uint16) (*nftables.Set, error) {
	connection := &nftables.Conn{}
	tables, err := connection.ListTables()
	if err != nil {
		return nil, fmt.Errorf("list nftables tables: %w", err)
	}
	for _, table := range tables {
		if table.Family == nftables.TableFamilyINet && table.Name == firewallTable {
			connection.DelTable(table)
		}
	}
	table := connection.AddTable(&nftables.Table{Family: nftables.TableFamilyINet, Name: firewallTable})
	fence := &nftables.Set{
		Table:      table,
		Name:       "owner_alive",
		KeyType:    nftables.TypeIPAddr,
		HasTimeout: true,
		Timeout:    15 * time.Second,
	}
	if err := connection.AddSet(fence, []nftables.SetElement{{Key: endpointAddress.To4(), Timeout: 15 * time.Second}}); err != nil {
		return nil, fmt.Errorf("create ownership fence: %w", err)
	}
	policy := nftables.ChainPolicyAccept
	chain := connection.AddChain(&nftables.Chain{
		Name:     "prerouting",
		Table:    table,
		Type:     nftables.ChainTypeFilter,
		Hooknum:  nftables.ChainHookPrerouting,
		Priority: nftables.ChainPriorityFilter,
		Policy:   &policy,
	})
	connection.AddRule(&nftables.Rule{Table: table, Chain: chain, Exprs: metricsReturnRule(tunnelAddress, metricsPeer, metricsPort)})
	for _, target := range targets {
		for _, port := range target.ports {
			connection.AddRule(&nftables.Rule{Table: table, Chain: chain, Exprs: allowedRule(target.address, port)})
		}
	}
	connection.AddRule(&nftables.Rule{Table: table, Chain: chain, Exprs: interfaceVerdict(expr.VerdictDrop)})
	output := connection.AddChain(&nftables.Chain{
		Name:     "output",
		Table:    table,
		Type:     nftables.ChainTypeFilter,
		Hooknum:  nftables.ChainHookOutput,
		Priority: nftables.ChainPriorityFilter,
		Policy:   &policy,
	})
	connection.AddRule(&nftables.Rule{Table: table, Chain: output, Exprs: fencedOutputRule(endpointAddress, endpointPort, fence)})
	connection.AddRule(&nftables.Rule{Table: table, Chain: output, Exprs: outerVerdictRule(endpointAddress, endpointPort, expr.VerdictDrop)})
	if err := connection.Flush(); err != nil {
		return nil, fmt.Errorf("install relay firewall: %w", err)
	}
	return fence, nil
}

func metricsReturnRule(tunnelAddress, peerAddress net.IP, port uint16) []expr.Any {
	portBytes := make([]byte, 2)
	binary.BigEndian.PutUint16(portBytes, port)
	stateMask := expr.CtStateBitESTABLISHED | expr.CtStateBitRELATED
	return []expr.Any{
		&expr.Meta{Key: expr.MetaKeyIIFNAME, Register: 1},
		&expr.Cmp{Op: expr.CmpOpEq, Register: 1, Data: append([]byte(interfaceName), 0)},
		&expr.Meta{Key: expr.MetaKeyNFPROTO, Register: 1},
		&expr.Cmp{Op: expr.CmpOpEq, Register: 1, Data: []byte{unix.NFPROTO_IPV4}},
		&expr.Payload{DestRegister: 1, Base: expr.PayloadBaseNetworkHeader, Offset: 12, Len: 4},
		&expr.Cmp{Op: expr.CmpOpEq, Register: 1, Data: peerAddress.To4()},
		&expr.Payload{DestRegister: 1, Base: expr.PayloadBaseNetworkHeader, Offset: 16, Len: 4},
		&expr.Cmp{Op: expr.CmpOpEq, Register: 1, Data: tunnelAddress.To4()},
		&expr.Meta{Key: expr.MetaKeyL4PROTO, Register: 1},
		&expr.Cmp{Op: expr.CmpOpEq, Register: 1, Data: []byte{unix.IPPROTO_TCP}},
		&expr.Payload{DestRegister: 1, Base: expr.PayloadBaseTransportHeader, Offset: 0, Len: 2},
		&expr.Cmp{Op: expr.CmpOpEq, Register: 1, Data: portBytes},
		&expr.Ct{Register: 1, Key: expr.CtKeySTATE},
		&expr.Bitwise{SourceRegister: 1, DestRegister: 1, Len: 4, Mask: binaryutil.NativeEndian.PutUint32(stateMask), Xor: binaryutil.NativeEndian.PutUint32(0)},
		&expr.Cmp{Op: expr.CmpOpNeq, Register: 1, Data: binaryutil.NativeEndian.PutUint32(0)},
		&expr.Verdict{Kind: expr.VerdictAccept},
	}
}

func allowedRule(gatewayAddress net.IP, port uint16) []expr.Any {
	portBytes := make([]byte, 2)
	binary.BigEndian.PutUint16(portBytes, port)
	return []expr.Any{
		&expr.Meta{Key: expr.MetaKeyIIFNAME, Register: 1},
		&expr.Cmp{Op: expr.CmpOpEq, Register: 1, Data: append([]byte(interfaceName), 0)},
		&expr.Meta{Key: expr.MetaKeyNFPROTO, Register: 1},
		&expr.Cmp{Op: expr.CmpOpEq, Register: 1, Data: []byte{unix.NFPROTO_IPV4}},
		&expr.Payload{DestRegister: 1, Base: expr.PayloadBaseNetworkHeader, Offset: 16, Len: 4},
		&expr.Cmp{Op: expr.CmpOpEq, Register: 1, Data: gatewayAddress.To4()},
		&expr.Meta{Key: expr.MetaKeyL4PROTO, Register: 1},
		&expr.Cmp{Op: expr.CmpOpEq, Register: 1, Data: []byte{unix.IPPROTO_TCP}},
		&expr.Payload{DestRegister: 1, Base: expr.PayloadBaseTransportHeader, Offset: 2, Len: 2},
		&expr.Cmp{Op: expr.CmpOpEq, Register: 1, Data: portBytes},
		&expr.Counter{},
		&expr.Verdict{Kind: expr.VerdictAccept},
	}
}

func interfaceVerdict(verdict expr.VerdictKind) []expr.Any {
	return []expr.Any{
		&expr.Meta{Key: expr.MetaKeyIIFNAME, Register: 1},
		&expr.Cmp{Op: expr.CmpOpEq, Register: 1, Data: append([]byte(interfaceName), 0)},
		&expr.Counter{},
		&expr.Verdict{Kind: verdict},
	}
}

func fencedOutputRule(endpointAddress net.IP, endpointPort uint16, fence *nftables.Set) []expr.Any {
	rule := outerRulePrefix(endpointAddress, endpointPort)
	rule = append(rule,
		&expr.Payload{DestRegister: 1, Base: expr.PayloadBaseNetworkHeader, Offset: 16, Len: 4},
		&expr.Lookup{SourceRegister: 1, SetName: fence.Name, SetID: fence.ID},
		&expr.Counter{},
		&expr.Verdict{Kind: expr.VerdictAccept},
	)
	return rule
}

func outerVerdictRule(endpointAddress net.IP, endpointPort uint16, verdict expr.VerdictKind) []expr.Any {
	rule := outerRulePrefix(endpointAddress, endpointPort)
	return append(rule, &expr.Counter{}, &expr.Verdict{Kind: verdict})
}

func outerRulePrefix(endpointAddress net.IP, endpointPort uint16) []expr.Any {
	portBytes := make([]byte, 2)
	binary.BigEndian.PutUint16(portBytes, endpointPort)
	return []expr.Any{
		&expr.Meta{Key: expr.MetaKeyNFPROTO, Register: 1},
		&expr.Cmp{Op: expr.CmpOpEq, Register: 1, Data: []byte{unix.NFPROTO_IPV4}},
		&expr.Payload{DestRegister: 1, Base: expr.PayloadBaseNetworkHeader, Offset: 16, Len: 4},
		&expr.Cmp{Op: expr.CmpOpEq, Register: 1, Data: endpointAddress.To4()},
		&expr.Meta{Key: expr.MetaKeyL4PROTO, Register: 1},
		&expr.Cmp{Op: expr.CmpOpEq, Register: 1, Data: []byte{unix.IPPROTO_UDP}},
		&expr.Payload{DestRegister: 1, Base: expr.PayloadBaseTransportHeader, Offset: 2, Len: 2},
		&expr.Cmp{Op: expr.CmpOpEq, Register: 1, Data: portBytes},
	}
}

func (r *relay) refreshFence() error {
	if r.fence == nil {
		return errors.New("relay ownership fence is absent")
	}
	connection := &nftables.Conn{}
	connection.FlushSet(r.fence)
	if err := connection.SetAddElements(r.fence, []nftables.SetElement{{Key: r.config.endpointAddress.To4(), Timeout: 15 * time.Second}}); err != nil {
		return fmt.Errorf("queue ownership fence refresh: %w", err)
	}
	if err := connection.Flush(); err != nil {
		return fmt.Errorf("refresh ownership fence: %w", err)
	}
	return nil
}

func (r *relay) verify() error {
	link, err := netlink.LinkByName(interfaceName)
	if err != nil {
		return fmt.Errorf("WireGuard interface: %w", err)
	}
	if link.Attrs().Flags&net.FlagUp == 0 {
		return errors.New("WireGuard interface is down")
	}
	rules, err := netlink.RuleList(netlink.FAMILY_V4)
	if err != nil {
		return fmt.Errorf("list policy rules: %w", err)
	}
	for _, target := range r.config.targets {
		foundRule := false
		for _, rule := range rules {
			if isRelayPolicyRule(rule, target.address) {
				foundRule = true
				break
			}
		}
		if !foundRule {
			return fmt.Errorf("relay policy rule for %s is absent", target.address)
		}
	}
	for _, target := range r.config.targets {
		for index, destination := range r.config.localReturnCIDRs {
			found := false
			for _, rule := range rules {
				if isLocalReturnRule(rule, target.address, destination, localPriority+index) {
					found = true
				}
			}
			if !found {
				return fmt.Errorf("local return policy rule for %s from %s is absent", destination, target.address)
			}
		}
	}
	routes, err := netlink.RouteListFiltered(netlink.FAMILY_V4, &netlink.Route{Table: routingTable}, netlink.RT_FILTER_TABLE)
	if err != nil {
		return fmt.Errorf("list relay routes: %w", err)
	}
	foundRoute := false
	for _, route := range routes {
		if route.LinkIndex == link.Attrs().Index && (route.Dst == nil || route.Dst.String() == "0.0.0.0/0") {
			foundRoute = true
			break
		}
	}
	if !foundRoute {
		return errors.New("relay default route is absent")
	}
	if err := verifyReversePathFilter(rpFilterRoot, interfaceName); err != nil {
		return err
	}
	if err := verifyFirewall(r.config.endpointAddress, r.config.targets); err != nil {
		return err
	}
	client, err := wgctrl.New()
	if err != nil {
		return fmt.Errorf("open WireGuard control: %w", err)
	}
	defer func() { _ = client.Close() }()
	device, err := client.Device(interfaceName)
	if err != nil {
		return fmt.Errorf("read WireGuard device: %w", err)
	}
	if len(device.Peers) != 1 || device.Peers[0].LastHandshakeTime.IsZero() || time.Since(device.Peers[0].LastHandshakeTime) > r.config.handshakeMaxAge {
		return errors.New("WireGuard handshake is stale")
	}
	return nil
}

func verifyReversePathFilter(root, name string) error {
	all, err := readReversePathFilter(filepath.Join(root, "all", "rp_filter"))
	if err != nil {
		return fmt.Errorf("read all reverse-path filtering: %w", err)
	}
	value, err := readReversePathFilter(filepath.Join(root, name, "rp_filter"))
	if err != nil {
		return fmt.Errorf("read %s reverse-path filtering: %w", name, err)
	}
	if max(all, value) != 2 {
		return fmt.Errorf("effective %s reverse-path filtering is not loose: all=%d interface=%d", name, all, value)
	}
	return nil
}

func readReversePathFilter(path string) (int, error) {
	value, err := os.ReadFile(path)
	if err != nil {
		return 0, err
	}
	parsed, err := strconv.Atoi(strings.TrimSpace(string(value)))
	if err != nil || parsed < 0 || parsed > 2 {
		return 0, fmt.Errorf("invalid value %q", strings.TrimSpace(string(value)))
	}
	return parsed, nil
}

func verifyFirewall(endpointAddress net.IP, targets []relayTarget) error {
	connection := &nftables.Conn{}
	table := &nftables.Table{Family: nftables.TableFamilyINet, Name: firewallTable}
	fence, err := connection.GetSetByName(table, "owner_alive")
	if err != nil {
		return fmt.Errorf("read relay ownership fence: %w", err)
	}
	elements, err := connection.GetSetElements(fence)
	if err != nil {
		return fmt.Errorf("read relay ownership fence elements: %w", err)
	}
	fenceAlive := false
	for _, element := range elements {
		if net.IP(element.Key).Equal(endpointAddress) {
			fenceAlive = true
			break
		}
	}
	if !fenceAlive {
		return errors.New("relay ownership fence is expired")
	}
	chains, err := connection.ListChainsOfTableFamily(nftables.TableFamilyINet)
	if err != nil {
		return fmt.Errorf("list relay firewall chains: %w", err)
	}
	allowedRuleCount := 0
	for _, target := range targets {
		allowedRuleCount += len(target.ports)
	}
	requiredRules := map[string]int{"prerouting": allowedRuleCount + 2, "output": 2}
	for _, chain := range chains {
		if chain.Table == nil || chain.Table.Name != firewallTable {
			continue
		}
		minimum, required := requiredRules[chain.Name]
		if !required {
			continue
		}
		rules, err := connection.GetRules(table, chain)
		if err != nil {
			return fmt.Errorf("read relay firewall chain %s: %w", chain.Name, err)
		}
		if len(rules) < minimum {
			return fmt.Errorf("relay firewall chain %s has %d rules", chain.Name, len(rules))
		}
		delete(requiredRules, chain.Name)
	}
	if len(requiredRules) != 0 {
		return errors.New("relay firewall chains are incomplete")
	}
	return nil
}

type firewallCounters struct {
	tunnelAllowed map[string]uint64
	tunnelDenied  uint64
	outerAllowed  uint64
	outerDenied   uint64
}

func readFirewallCounters(targets []relayTarget) (firewallCounters, error) {
	connection := &nftables.Conn{}
	table := &nftables.Table{Family: nftables.TableFamilyINet, Name: firewallTable}
	chains, err := connection.ListChainsOfTableFamily(nftables.TableFamilyINet)
	if err != nil {
		return firewallCounters{}, fmt.Errorf("list relay firewall chains for counters: %w", err)
	}
	rulesByChain := make(map[string][]*nftables.Rule)
	for _, chain := range chains {
		if chain.Table == nil || chain.Table.Name != firewallTable || (chain.Name != "prerouting" && chain.Name != "output") {
			continue
		}
		rules, err := connection.GetRules(table, chain)
		if err != nil {
			return firewallCounters{}, fmt.Errorf("read relay firewall chain %s counters: %w", chain.Name, err)
		}
		rulesByChain[chain.Name] = rules
	}
	return firewallCountersFromRules(rulesByChain, targets)
}

func firewallCountersFromRules(rulesByChain map[string][]*nftables.Rule, targets []relayTarget) (firewallCounters, error) {
	prerouting := rulesByChain["prerouting"]
	output := rulesByChain["output"]
	allowedRuleCount := 0
	for _, target := range targets {
		allowedRuleCount += len(target.ports)
	}
	if len(prerouting) < allowedRuleCount+2 || len(output) < 2 {
		return firewallCounters{}, errors.New("relay firewall rules are incomplete for counters")
	}
	allowed := make(map[string]uint64, allowedRuleCount)
	index := 1
	for _, target := range targets {
		for _, port := range target.ports {
			packets, ok := rulePacketCounter(prerouting[index])
			if !ok {
				return firewallCounters{}, errors.New("relay firewall rule counter is absent")
			}
			allowed[net.JoinHostPort(target.address.String(), strconv.Itoa(int(port)))] = packets
			index++
		}
	}
	values := make([]uint64, 0, 3)
	for _, rule := range []*nftables.Rule{prerouting[allowedRuleCount+1], output[0], output[1]} {
		packets, ok := rulePacketCounter(rule)
		if !ok {
			return firewallCounters{}, errors.New("relay firewall rule counter is absent")
		}
		values = append(values, packets)
	}
	return firewallCounters{
		tunnelAllowed: allowed,
		tunnelDenied:  values[0],
		outerAllowed:  values[1],
		outerDenied:   values[2],
	}, nil
}

func rulePacketCounter(rule *nftables.Rule) (uint64, bool) {
	for _, expression := range rule.Exprs {
		if counter, ok := expression.(*expr.Counter); ok {
			return counter.Packets, true
		}
	}
	return 0, false
}

func equalCIDR(network *net.IPNet, address net.IP) bool {
	return network != nil && network.IP.Equal(address) && bytes.Equal(network.Mask, net.CIDRMask(32, 32))
}

func containsTargetAddress(targets []relayTarget, network *net.IPNet) bool {
	for _, target := range targets {
		if equalCIDR(network, target.address) {
			return true
		}
	}
	return false
}

func isRelayPolicyRule(rule netlink.Rule, gatewayAddress net.IP) bool {
	return rule.Priority == rulePriority && rule.Table == routingTable && equalCIDR(rule.Src, gatewayAddress)
}

func isLocalReturnRule(rule netlink.Rule, gatewayAddress net.IP, destination *net.IPNet, priority int) bool {
	return rule.Priority == priority && rule.Table == unix.RT_TABLE_MAIN && equalCIDR(rule.Src, gatewayAddress) && rule.Dst != nil && rule.Dst.String() == destination.String()
}

func (r *relay) cleanup() error {
	r.cleanupMu.Lock()
	defer r.cleanupMu.Unlock()
	var failures []error
	linkRemoved := true
	link, err := netlink.LinkByName(interfaceName)
	if err == nil {
		if err := netlink.LinkDel(link); err != nil {
			linkRemoved = false
			failures = append(failures, fmt.Errorf("remove WireGuard link: %w", err))
		}
	} else if _, ok := err.(netlink.LinkNotFoundError); !ok {
		linkRemoved = false
		failures = append(failures, fmt.Errorf("find WireGuard link for cleanup: %w", err))
	}
	rules, err := netlink.RuleList(netlink.FAMILY_V4)
	if err != nil {
		failures = append(failures, fmt.Errorf("list relay policy rules: %w", err))
	} else {
		for _, candidate := range rules {
			if containsTargetAddress(r.config.targets, candidate.Src) && candidate.Priority >= localPriority && candidate.Priority <= rulePriority && (candidate.Table == unix.RT_TABLE_MAIN || candidate.Table == routingTable) {
				rule := candidate
				if err := netlink.RuleDel(&rule); err != nil {
					failures = append(failures, fmt.Errorf("remove relay policy rule: %w", err))
				}
			}
		}
	}
	routes, err := netlink.RouteListFiltered(netlink.FAMILY_V4, &netlink.Route{Table: routingTable}, netlink.RT_FILTER_TABLE)
	if err != nil {
		failures = append(failures, fmt.Errorf("list relay routes: %w", err))
	} else {
		for _, candidate := range routes {
			route := candidate
			if err := netlink.RouteDel(&route); err != nil {
				failures = append(failures, fmt.Errorf("remove relay route: %w", err))
			}
		}
	}
	if linkRemoved {
		connection := &nftables.Conn{}
		tables, err := connection.ListTables()
		if err != nil {
			failures = append(failures, fmt.Errorf("list nftables tables for cleanup: %w", err))
		} else {
			for _, table := range tables {
				if table.Family == nftables.TableFamilyINet && table.Name == firewallTable {
					connection.DelTable(table)
				}
			}
			if err := connection.Flush(); err != nil {
				failures = append(failures, fmt.Errorf("remove relay firewall: %w", err))
			}
		}
	}
	r.fence = nil
	return errors.Join(failures...)
}

func writeHealth(path string, healthy bool) error {
	if err := os.MkdirAll(filepath.Dir(path), 0755); err != nil {
		return err
	}
	state := []byte("unhealthy\n")
	if healthy {
		state = []byte("healthy\n")
	}
	return os.WriteFile(path, state, 0644)
}

func checkHealth(path string, maxAge time.Duration) error {
	if path == "" {
		path = "/run/relay-agent/health"
	}
	info, err := os.Stat(path)
	if err != nil {
		return err
	}
	if time.Since(info.ModTime()) > maxAge {
		return errors.New("relay heartbeat is stale")
	}
	state, err := os.ReadFile(path)
	if err != nil {
		return err
	}
	if string(state) != "healthy\n" {
		return errors.New("relay reports unhealthy")
	}
	return nil
}