package proxy

import (
	"crypto/rand"
	"net/http"
	"net/netip"
	"slices"
	"strings"
)

// peerAddress is the address of the request's TCP peer, normally traefik.
func peerAddress(r *http.Request) netip.Addr {
	addrPort, err := netip.ParseAddrPort(r.RemoteAddr)
	if err != nil {
		return netip.Addr{}
	}

	return addrPort.Addr().Unmap()
}

// clientAddress works out who the client is. A peer outside the trusted
// proxies is the client, and what it says in X-Forwarded-For is ignored.
// For a peer inside them, X-Forwarded-For is read from the right, and the
// first address outside them is the client; if every address in it is
// inside, the leftmost is, and with no header, the peer. An entry that is
// not an address ends the reading, since nothing to its left can be
// believed.
func clientAddress(
	peer netip.Addr, forwardedFor []string, trusted []netip.Prefix,
) netip.Addr {
	client := peer
	if !isInside(peer, trusted) {
		return client
	}

	entries := strings.Split(strings.Join(forwardedFor, ","), ",")
	for _, entry := range slices.Backward(entries) {
		addr, err := netip.ParseAddr(strings.TrimSpace(entry))
		if err != nil {
			break
		}

		client = addr.Unmap()
		if !isInside(client, trusted) {
			break
		}
	}

	return client
}

// requestIDHeader carries the request's id, from traefik and to the app.
const requestIDHeader = "X-Request-ID"

// requestID is the request's id: the one a trusted proxy sent, or a new
// random one. A peer outside the trusted proxies did not come through
// traefik, so the id it sends is its own claim, and is replaced.
func requestID(r *http.Request, peerTrusted bool) string {
	id := r.Header.Get(requestIDHeader)
	if !peerTrusted || id == "" {
		id = rand.Text()
	}

	return id
}

// scheme is how the client reached traefik, as a trusted proxy says in
// X-Forwarded-Proto, or otherwise http, the only scheme smallwebwaf
// serves.
func scheme(r *http.Request, peerTrusted bool) string {
	proto := r.Header.Get("X-Forwarded-Proto")
	if !peerTrusted || proto == "" {
		return "http"
	}

	return proto
}

// ipv6GroupPrefix is the length of the IPv6 netblock that is one client.
const ipv6GroupPrefix = 64

// clientGroup is the client a request is counted toward: its IPv4
// address, or the /64 its IPv6 address is in, since one abuser usually
// holds a whole /64. An IPv4 address in IPv6 form counts as IPv4.
func clientGroup(addr netip.Addr) netip.Prefix {
	addr = addr.Unmap()
	if addr.Is6() {
		return netip.PrefixFrom(addr, ipv6GroupPrefix).Masked()
	}

	return netip.PrefixFrom(addr, addr.BitLen())
}

// isInside reports whether addr is in one of the netblocks.
func isInside(addr netip.Addr, netblocks []netip.Prefix) bool {
	return slices.ContainsFunc(netblocks, func(netblock netip.Prefix) bool {
		return netblock.Contains(addr)
	})
}

// setForwardedHeaders sets the headers in which the app learns about the
// client, so that it sees what it would see from traefik directly. A
// trusted proxy's forwarded headers pass on, with the proxy's own address
// added to X-Forwarded-For. Those of any other peer are its own claims and
// are replaced: X-Forwarded-For names the peer, X-Forwarded-Host the host
// it asked for, and X-Forwarded-Proto plain http, which is how it reached
// smallwebwaf.
func setForwardedHeaders(in, out *http.Request, peer netip.Addr, trusted bool) {
	forwardedFor := peer.String()

	if trusted {
		// ReverseProxy removes these from out before Rewrite.
		for _, name := range []string{"Forwarded", "X-Forwarded-Host", "X-Forwarded-Proto"} {
			values, ok := in.Header[name]
			if ok {
				out.Header[name] = values
			}
		}

		prior := in.Header.Values("X-Forwarded-For")
		if len(prior) > 0 {
			forwardedFor = strings.Join(prior, ", ") + ", " + forwardedFor
		}

		out.Header.Set("X-Forwarded-For", forwardedFor)

		return
	}

	// ReverseProxy has removed Forwarded and the three set below; these
	// are the other headers in which traefik tells the app about the
	// client and its request.
	for _, name := range []string{
		"X-Forwarded-Port", "X-Forwarded-Server", "X-Forwarded-Uri",
		"X-Forwarded-Method", "X-Forwarded-Prefix", "X-Forwarded-Tls-Client-Cert",
		"X-Forwarded-Tls-Client-Cert-Info", "X-Real-Ip",
	} {
		out.Header.Del(name)
	}

	out.Header.Set("X-Forwarded-For", forwardedFor)
	out.Header.Set("X-Forwarded-Host", in.Host)
	out.Header.Set("X-Forwarded-Proto", "http")
}