// dns-webhook is a Kubernetes MutatingAdmissionWebhook.
//
// What is a MutatingAdmissionWebhook?
//
//   When Kubernetes is about to create a resource (a Pod, Deployment, etc.),
//   it first sends the resource's definition to every registered webhook for
//   that resource type.  The webhook can inspect the definition and return a
//   set of JSON Patch operations — instructions that say "change field X to
//   value Y before saving".  Kubernetes applies those patches, then stores
//   the final object.
//
// What does this webhook do?
//
//   Every new Pod that uses Kubernetes cluster DNS (dnsPolicy: ClusterFirst)
//   gets its DNS configuration rewritten so that:
//
//     1. Three nameserver IPs are chosen at random from the four production
//        auth-dns pods (ns1–ns4).  This spreads DNS query load across the
//        pool instead of every pod hitting the same two servers.
//
//     2. The resolv.conf options "edns0" and "rotate" are enabled.
//        - edns0  allows DNS responses larger than 512 bytes (needed for
//                 DNSSEC and large TXT records).
//        - rotate makes the resolver cycle through the nameserver list so
//                 queries are distributed rather than always going to the
//                 first server that responds.
//
//   Pods that already have explicit DNS config (DNSNone), use the node's
//   own resolver (DNSDefault), or run with host networking + cluster DNS
//   (ClusterFirstWithHostNet) are left untouched.
//
// TLS requirement:
//
//   Kubernetes requires admission webhooks to speak HTTPS.  This server
//   loads a TLS certificate and key from /tls/ (written by cert-manager).
//   The paths can be overridden with TLS_CERT / TLS_KEY environment variables.

package main

import (
	"encoding/json"
	"fmt"
	"io"
	"log"
	"math/rand"
	"net/http"
	"os"

	// admissionv1 contains the AdmissionReview, AdmissionRequest, and
	// AdmissionResponse types that Kubernetes sends to / expects from the webhook.
	admissionv1 "k8s.io/api/admission/v1"

	// corev1 contains Pod, PodSpec, PodDNSConfig, etc.
	corev1 "k8s.io/api/core/v1"
)

// nameserverPool lists the IPv6 addresses of the four production auth-dns pods
// (ns1–ns4).  ns0 is the staging instance and is intentionally excluded.
//
// These are static pod IPs allocated by Calico IPAM.  They are stable across
// pod restarts because each Deployment pins its IP with the annotation
// cni.projectcalico.org/ipAddrs.
var nameserverPool = [4]string{
	"2602:817:3000:c608::202", // ns1
	"2602:817:3008:c607::204", // ns2
	"2602:817:3000:c608::203", // ns3
	"2602:817:3008:c607::203", // ns4
}

// clusterDomain is the Kubernetes cluster domain configured in kubelet.
// Service FQDNs follow the pattern: <service>.<namespace>.svc.<clusterDomain>
const clusterDomain = "k8s.sjc001.fritzlab.net"

// jsonPatch represents a single RFC 6902 JSON Patch operation.
//
// Kubernetes admission webhooks return patches in this format.  The three
// fields mean:
//   - Op:    the operation — "add", "replace", or "remove"
//   - Path:  a slash-separated path into the JSON document, e.g. /spec/dnsPolicy
//   - Value: the new value to set (omitted for "remove")
type jsonPatch struct {
	Op    string      `json:"op"`
	Path  string      `json:"path"`
	Value interface{} `json:"value,omitempty"`
}

// pickThree returns three randomly selected nameserver IPs from nameserverPool.
//
// rand.Perm(4) returns a random permutation of [0, 1, 2, 3], so we take the
// first three indices to get three distinct servers.  Using a permutation
// (rather than three independent random picks) guarantees no duplicates.
func pickThree() []string {
	idx := rand.Perm(4)
	return []string{nameserverPool[idx[0]], nameserverPool[idx[1]], nameserverPool[idx[2]]}
}

// strPtr is a small helper that returns a pointer to a string.
//
// The PodDNSConfigOption type uses *string for option values so that an
// absent value (nil) is distinguishable from an empty string ("").
func strPtr(s string) *string { return &s }

// handleMutate is the HTTP handler for the /mutate endpoint.
//
// Kubernetes calls this endpoint with an AdmissionReview JSON body whenever a
// Pod CREATE request matches the MutatingWebhookConfiguration rules.  The
// handler must:
//
//  1. Decode the AdmissionReview to get the Pod definition.
//  2. Decide whether to mutate the Pod.
//  3. Return an AdmissionReview response with Allowed: true and any patches.
func handleMutate(w http.ResponseWriter, r *http.Request) {
	// Read the full request body.
	body, err := io.ReadAll(r.Body)
	if err != nil {
		log.Printf("ERROR reading request body: %v", err)
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	// Unmarshal the outer AdmissionReview envelope.
	// AdmissionReview is the top-level wrapper Kubernetes uses for both the
	// incoming request and the outgoing response.
	var review admissionv1.AdmissionReview
	if err := json.Unmarshal(body, &review); err != nil {
		log.Printf("ERROR unmarshalling AdmissionReview: %v", err)
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	// Extract the Pod from review.Request.Object.Raw.
	// Raw is the verbatim JSON of the resource being admitted.
	req := review.Request
	var pod corev1.Pod
	if err := json.Unmarshal(req.Object.Raw, &pod); err != nil {
		log.Printf("ERROR unmarshalling Pod (uid=%s): %v", req.UID, err)
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	// Only mutate pods that use ClusterFirst DNS policy.
	//
	// DNSPolicy determines how a pod resolves names:
	//   ClusterFirst          — use cluster DNS (CoreDNS); fall back to upstream.
	//                          This is the default for most workloads.
	//   DNSNone               — pod supplies its own dnsConfig; nothing to do.
	//   DNSDefault            — inherit the node's /etc/resolv.conf directly.
	//                          Used by auth-dns pods themselves to avoid a
	//                          circular dependency (DNS pod needs DNS to start).
	//   ClusterFirstWithHostNet — host-network pod that still wants cluster DNS;
	//                          leave as-is to avoid breaking host-network semantics.
	if pod.Spec.DNSPolicy != corev1.DNSClusterFirst {
		log.Printf("SKIP pod=%s/%s uid=%s policy=%s (not ClusterFirst)",
			req.Namespace, pod.Name, req.UID, pod.Spec.DNSPolicy)
		respond(w, review, nil) // no patches, just allow
		return
	}

	// Determine the pod's namespace for the search domain list.
	// Kubernetes guarantees this is set for namespaced resources, but fall back
	// to "default" just in case.
	ns := req.Namespace
	if ns == "" {
		ns = "default"
	}

	// Build the DNS configuration we want every pod to have.
	//
	// Nameservers: three randomly chosen auth-dns pod IPs.
	//
	// Searches: resolv.conf search domains let short names be resolved without
	//   a fully-qualified domain name.  For a pod in namespace "myapp":
	//     - myapp.svc.k8s.sjc001.fritzlab.net  → finds Services in myapp
	//     - svc.k8s.sjc001.fritzlab.net         → finds Services in any namespace
	//     - k8s.sjc001.fritzlab.net             → catches cluster-level names
	//   So "kubectl exec mypod -- curl myservice" resolves without a FQDN.
	//
	// Options:
	//   ndots:5  — names with fewer than 5 dots are tried with search domains
	//              before being treated as absolute.  This is the Kubernetes
	//              default and ensures short service names resolve correctly.
	//   edns0    — enables EDNS extension headers, allowing DNS responses up to
	//              65535 bytes (default UDP cap is 512 bytes).
	//   rotate   — cycle through the nameserver list on each query so load is
	//              distributed rather than always hitting the first entry.
	nameservers := pickThree()
	dnsConfig := corev1.PodDNSConfig{
		Nameservers: nameservers,
		Searches: []string{
			fmt.Sprintf("%s.svc.%s", ns, clusterDomain),
			fmt.Sprintf("svc.%s", clusterDomain),
			clusterDomain,
		},
		Options: []corev1.PodDNSConfigOption{
			{Name: "edns0"},
			{Name: "rotate"},
			{Name: "ndots", Value: strPtr("5")},
		},
	}

	// JSON Patch requires "add" if the field doesn't exist yet in the manifest,
	// or "replace" if it does.  Sending "replace" for a non-existent field
	// (or "add" for an existing one) is a patch error that would reject the pod.
	dnsConfigOp := "add"
	if pod.Spec.DNSConfig != nil {
		dnsConfigOp = "replace"
	}

	log.Printf("MUTATE pod=%s/%s uid=%s nameservers=%v op=%s",
		ns, pod.Name, req.UID, nameservers, dnsConfigOp)

	// Build the two patch operations:
	//   1. Change dnsPolicy from ClusterFirst → None.
	//      DNSNone is required when you provide a custom dnsConfig; it tells
	//      kubelet to use exactly the config we supply and nothing else.
	//   2. Set dnsConfig to our constructed value.
	patches := []jsonPatch{
		{Op: "replace", Path: "/spec/dnsPolicy", Value: corev1.DNSNone},
		{Op: dnsConfigOp, Path: "/spec/dnsConfig", Value: dnsConfig},
	}

	respond(w, review, patches)
}

// respond writes an AdmissionReview response back to the Kubernetes API server.
//
// Every call to the /mutate endpoint must return an AdmissionReview with:
//   - UID matching the incoming request (so Kubernetes can correlate them)
//   - Allowed: true  (we never block pods, just modify them)
//   - Patch / PatchType set when we have mutations to apply
func respond(w http.ResponseWriter, review admissionv1.AdmissionReview, patches []jsonPatch) {
	resp := &admissionv1.AdmissionResponse{
		UID:     review.Request.UID,
		Allowed: true,
	}

	if patches != nil {
		// Serialize the patch list to JSON and set the patch type.
		// JSONPatch is the only patch type supported by admission webhooks.
		b, _ := json.Marshal(patches)
		pt := admissionv1.PatchTypeJSONPatch
		resp.Patch = b
		resp.PatchType = &pt
	}

	// Place the response inside the same AdmissionReview envelope that came in.
	// Kubernetes reads review.Response; review.Request is ignored on the way out.
	review.Response = resp
	w.Header().Set("Content-Type", "application/json")
	if err := json.NewEncoder(w).Encode(review); err != nil {
		log.Printf("ERROR encoding AdmissionReview response (uid=%s): %v", resp.UID, err)
	}
}

func main() {
	// TLS cert and key are written by cert-manager into a Kubernetes Secret,
	// which is mounted into the pod at /tls/.  The environment variables allow
	// the paths to be overridden in development or testing.
	certFile := getenv("TLS_CERT", "/tls/tls.crt")
	keyFile := getenv("TLS_KEY", "/tls/tls.key")

	// /mutate receives AdmissionReview requests from the Kubernetes API server.
	http.HandleFunc("/mutate", handleMutate)

	// /healthz is a simple liveness/readiness probe endpoint.
	// kubelet calls this to decide whether the pod is healthy.
	http.HandleFunc("/healthz", func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
	})

	log.Printf("dns-webhook starting: cert=%s key=%s", certFile, keyFile)

	log.Println("dns-webhook listening on :8443")
	// Kubernetes requires admission webhooks to use TLS — plain HTTP is rejected.
	log.Fatal(http.ListenAndServeTLS(":8443", certFile, keyFile, nil))
}

// getenv returns the value of the environment variable key, or def if unset.
func getenv(key, def string) string {
	if v := os.Getenv(key); v != "" {
		return v
	}
	return def
}