package proxy

import (
	"crypto/rand"
	"encoding/json"
	"net/http"
	"net/netip"
	"strconv"
	"strings"
	"time"
	"unicode/utf8"

	lru "github.com/hashicorp/golang-lru/v2"

	"sneak.berlin/go/smallwebwaf/internal/requestlog"
)

// The headers of rproxy's question and of servewaf's answer, as
// "servewaf" in SPEC.md gives them.
const (
	// ClientHeader marks a request as rproxy's question, and gives the
	// visitor's address.
	ClientHeader = "X-Smallwebwaf-Client"
	// PeerHeader gives the address of rproxy's own TCP peer.
	PeerHeader = "X-Smallwebwaf-Peer"
	// ContentLengthHeader gives the body length the visitor announced.
	ContentLengthHeader = "X-Smallwebwaf-Content-Length"
	// DispositionHeader gives servewaf's answer.
	DispositionHeader = "X-Smallwebwaf-Disposition"
)

// The values of DispositionHeader.
const (
	// DispositionAllow tells rproxy to pass the request to the app.
	DispositionAllow = "allow"
	// DispositionDeny tells rproxy to give the visitor the answer's status,
	// headers and body.
	DispositionDeny = "deny"
	// DispositionClose tells rproxy to close the visitor's connection
	// without an answer, which SWWAF_BAN_RESPONSE asks for when it is close.
	// servewaf closing the question's connection instead would read as
	// servewaf being down, and rproxy would then let the request through.
	DispositionClose = "close"
)

// ReportPath is where rproxy sends the report of each request it passed to
// the app.
const ReportPath = "/_smallwebwaf/report"

// QuestionBodyMaxBytes is the most of a visitor's body a question carries.
const QuestionBodyMaxBytes = 64 << 10

const (
	// maxAwaitingReport is the most allowed requests held until their
	// reports come. Past it, the request asked about longest ago is dropped.
	maxAwaitingReport = 10000
	// reportBodyMaxBytes is the most of a report's body that is read. It is
	// far above any report's, since the longest thing in one is the request
	// id, which comes from a header.
	reportBodyMaxBytes = 1 << 20
	// serveWAFTimeout bounds how long a request to servewaf takes to arrive,
	// and its answer to leave, and serveWAFIdleTimeout how long a kept-open
	// connection waits for its next request.
	serveWAFTimeout     = 60 * time.Second
	serveWAFIdleTimeout = 120 * time.Second
)

// Report is the JSON object rproxy posts to ReportPath once the visitor's
// answer has ended, its fields named as the request log's of the same
// meaning. RequestID is the id servewaf's allow gave. Status is what the
// visitor got, 0 if nothing. The durations are in milliseconds.
//
//nolint:tagliatelle // docs/request-log.md names the fields in snake_case
type Report struct {
	RequestID             string  `json:"request_id"`
	Status                int     `json:"status"`
	RequestBytes          int64   `json:"request_bytes"`
	ResponseBytes         int64   `json:"response_bytes"`
	DurationTotal         float64 `json:"duration_total"`
	DurationUpstreamTotal float64 `json:"duration_upstream_total"`
}

// NewServeWAF returns the server `smallwebwaf servewaf` runs. It answers
// rproxy's questions with the checks New runs, takes rproxy's reports and
// answers its own endpoints. It passes nothing to an app. rproxy applies
// the size and time limits to the visitor, so their settings do not apply
// here. Go's default limit on a request's line and headers, 1 MiB, stays:
// it is far above rproxy's, and a question refused for its size would have
// rproxy let the request through.
func NewServeWAF(params Params) *Server {
	cfg := *params.Config
	cfg.RequestMaxBytes = 0
	cfg.ClientRequestTimeout = serveWAFTimeout
	// The server's WriteTimeout does this. It sets a new deadline for each
	// request, where a deadline a handler sets stays on a kept-open
	// connection, and would run out under a later request's answer.
	cfg.ClientResponseTimeout = 0
	// A question's body of QuestionBodyMaxBytes may be cut short, and the
	// Core Rule Set does not read a JSON or XML body over its limit.
	cfg.WAFBodyLimit = min(cfg.WAFBodyLimit, QuestionBodyMaxBytes-1)
	params.Config = &cfg

	h := newHandler(params)

	awaitingReport, err := lru.New[string, *request](maxAwaitingReport)
	if err != nil {
		panic(err) // New fails only for a size below one
	}

	h.awaitingReport = awaitingReport

	return h.server(&http.Server{
		Handler:      http.HandlerFunc(h.serveWAF),
		ReadTimeout:  serveWAFTimeout,
		WriteTimeout: serveWAFTimeout,
		IdleTimeout:  serveWAFIdleTimeout,
	})
}

// serveWAF handles one request to servewaf: a request carrying
// ClientHeader is rproxy's question, and any other is for servewaf itself.
func (h *handler) serveWAF(w http.ResponseWriter, r *http.Request) {
	if len(r.Header.Values(ClientHeader)) > 0 {
		h.answerQuestion(w, r)

		return
	}

	h.answerOwnEndpoint(w, r)
}

// answerQuestion runs the checks on the visitor's request a question
// carries, and answers rproxy what to do with it. A denied request is
// finished here. An allowed one is finished by its report. A question
// without a visitor's address gets no disposition, so that rproxy treats
// it as it treats no answer.
func (h *handler) answerQuestion(w http.ResponseWriter, r *http.Request) {
	client, err := netip.ParseAddr(r.Header.Get(ClientHeader))
	if err != nil {
		h.processLog.Warn("a question gives no visitor's address", "error", err.Error())
		http.Error(w, "X-Smallwebwaf-Client is not an address", http.StatusBadRequest)

		return
	}

	// One that cannot be read leaves the log line's peer_ip invalid.
	peer, _ := netip.ParseAddr(r.Header.Get(PeerHeader))

	rq := h.startRequest(w, asTheVisitorSentIt(r), client.Unmap(), peer.Unmap(), true)
	defer rq.finish()

	rq.site = r.Host
	refused := rq.check(r.Context())
	rq.checked = time.Now()
	underOwnPrefix := strings.HasPrefix(r.URL.Path, adminPrefix)

	if refused == nil && !underOwnPrefix {
		// In observe mode, what enforce mode would have refused counts here.
		defer rq.countRefusal()

		rq.allow()

		return
	}

	// The last deferred runs first, as in ServeHTTP.
	defer rq.addToHistory()
	defer rq.countAnomalies()
	defer rq.countRefusal()

	if refused != nil {
		rq.deny(*refused)

		return
	}

	// The app never sees this prefix. servewaf's own endpoints are at its
	// own address, so this answers as a path with no endpoint does.
	rq.line.Action = requestlog.ActionAdmin
	rq.out.Header().Set(DispositionHeader, DispositionDeny)
	http.Error(rq.out, http.StatusText(http.StatusNotFound), http.StatusNotFound)
}

// asTheVisitorSentIt returns a copy of r, a question, as the checks and the
// log line are to see it: without the headers rproxy adds for servewaf, and
// with the length the visitor announced, as -1 when it announced none. The
// question's own length covers only the part of the body it carries.
func asTheVisitorSentIt(r *http.Request) *http.Request {
	sent := r.Clone(r.Context())
	announced := sent.Header.Get(ContentLengthHeader)

	sent.Header.Del(ClientHeader)
	sent.Header.Del(PeerHeader)
	sent.Header.Del(ContentLengthHeader)
	sent.Header.Del("Content-Length")

	sent.ContentLength = -1

	length, err := strconv.ParseInt(announced, 10, 64)
	if err == nil && length >= 0 {
		sent.ContentLength = length
		sent.Header.Set("Content-Length", announced)
	}

	return sent
}

// allow answers a question with DispositionAllow, the request's id and,
// while SWWAF_ADD_LOOKUP_HEADERS is set, the client's AS number and country,
// and holds the request until its report comes.
func (rq *request) allow() {
	rq.allowed = true

	// servewaf sends no app this part of the body, and the table of requests
	// awaiting their reports must not hold it.
	if rq.body != nil {
		rq.body.readByCoreRuleSet = nil
	}

	rq.awaitReport()

	header := rq.out.Header()
	header.Set(DispositionHeader, DispositionAllow)
	header.Set(requestIDHeader, rq.line.RequestID)

	if rq.h.config.AddLookupHeaders {
		setLookupHeaders(header, rq.line.ASN, rq.line.Country)
	}

	// Go's server sends this only once the handler has returned, so a report
	// that answers it cannot come while the handler still uses rq.
	rq.out.WriteHeader(http.StatusOK)
}

// awaitReport holds the request by its id, which its report is to name. A
// visitor behind a trusted proxy can choose its id, which must not let it
// have its report taken for another request's, or lost, and its bytes not
// counted. So the request gets a new id when another held request has the
// question's, or when that is not valid UTF-8, which a JSON report cannot
// carry unchanged.
func (rq *request) awaitReport() {
	id := rq.line.RequestID
	if !utf8.ValidString(id) {
		id = rand.Text()
	}

	held, _ := rq.h.awaitingReport.ContainsOrAdd(id, rq)
	if held {
		id = rand.Text()
		rq.h.awaitingReport.Add(id, rq)
	}

	rq.line.RequestID = id
}

// deny answers a question with what the visitor is to get for r.
func (rq *request) deny(r refusal) {
	if r.status == 0 {
		// SWWAF_BAN_RESPONSE is close. Written past rq.out, since the log line's
		// status stays 0: the visitor is sent nothing.
		rq.refused.CompareAndSwap(nil, &r)
		rq.out.Header().Set(DispositionHeader, DispositionClose)
		rq.out.ResponseWriter.WriteHeader(http.StatusOK)

		return
	}

	rq.out.Header().Set(DispositionHeader, DispositionDeny)
	rq.answer(r)
}

// answerOwnEndpoint answers a request to servewaf's own address. No check
// runs on it: the address is on the private network, where the peers are
// rproxy and the admin's tools, not visitors. A report writes no log line.
func (h *handler) answerOwnEndpoint(w http.ResponseWriter, r *http.Request) {
	if r.Method == http.MethodPost && r.URL.Path == ReportPath {
		h.takeReport(w, r)

		return
	}

	peer := peerAddress(r)

	rq := h.startRequest(w, r, peer, peer, false)
	defer rq.finish()

	if r.Method == http.MethodGet && r.URL.Path == HealthPath {
		rq.answerHealth()

		return
	}

	rq.answerAdmin()
}

// takeReport answers POST ReportPath: it finishes the request the report
// names, if it holds that request, and ignores the report otherwise.
func (h *handler) takeReport(w http.ResponseWriter, r *http.Request) {
	var report Report

	err := json.NewDecoder(http.MaxBytesReader(w, r.Body, reportBodyMaxBytes)).
		Decode(&report)
	if err != nil || report.RequestBytes < 0 || report.ResponseBytes < 0 {
		http.Error(w, "the body is not a report", http.StatusBadRequest)

		return
	}

	rq, held := h.awaitingReport.Peek(report.RequestID)
	// Remove reports true to one report only, of several that name the id.
	if held && h.awaitingReport.Remove(report.RequestID) {
		rq.finishWith(report)
	}

	w.WriteHeader(http.StatusNoContent)
}

// finishWith finishes an allowed request with what its report says: the
// steps ServeHTTP takes once a request passed to the app has ended, with
// the status and bytes of the report in place of those it would have seen.
func (rq *request) finishWith(report Report) {
	// A limit the question broke has been logged.
	rq.line.LimitHit = ""
	// addToHistory goes by this for whether the request reached the app.
	rq.upstreamStart = time.Now()
	rq.out = &responseWriter{status: report.Status, bytes: report.ResponseBytes}
	rq.body = &requestBody{}
	rq.body.bytes.Store(report.RequestBytes)

	rq.countBytes()
	rq.countAnomalies()
	rq.addToHistory()

	rq.line.Status = report.Status
	rq.line.RequestBytes = report.RequestBytes
	rq.line.ResponseBytes = report.ResponseBytes
	rq.h.metrics.RequestReported(&rq.line,
		time.Duration(report.DurationTotal*float64(time.Millisecond)),
		time.Duration(report.DurationUpstreamTotal*float64(time.Millisecond)))

	// The request's log line went out with the answer to its question.
	if rq.line.LimitHit != "" {
		rq.h.processLog.Warn("a report's bytes broke a byte limit",
			"request_id", rq.line.RequestID,
			"client_ip", rq.line.ClientIP,
			"host", rq.line.Host,
			"limit_hit", rq.line.LimitHit,
			"ban_expires", rq.line.BanExpires)
	}
}