package rproxy

import (
	"context"
	"crypto/rand"
	"errors"
	"io"
	"net/http"
	"net/http/httptrace"
	"net/http/httputil"
	"net/netip"
	"net/url"
	"os"
	"strconv"
	"sync"
	"sync/atomic"
	"time"

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

// flushAfterEachWrite has ReverseProxy pass on each part of the app's answer as
// soon as it arrives.
const flushAfterEachWrite time.Duration = -1

// reportStatusMax is the highest status waf takes in a report. An app can
// send a higher one, up to 999, and the visitor gets that as it is.
const reportStatusMax = 599

// plainText is the Content-Type of the answers rproxy gives itself.
const plainText = "text/plain; charset=utf-8"

// minSendTime is the least time to send an answer rproxy gives itself, or the
// status and headers of the app's answer. A shorter time can run out before the
// first byte is written, on a busy host or as SWRP_MAX_RESPONSE_TIME_MS runs out,
// and the visitor then gets a closed connection and no answer.
const minSendTime = 10 * time.Second

// errResponseTooLate refuses an app's answer that arrives once
// SWRP_MAX_RESPONSE_TIME_MS has run out.
var errResponseTooLate = errors.New(
	"the answer arrived after SWRP_MAX_RESPONSE_TIME_MS ran out")

// request is one request on its way through rproxy, from the moment its headers
// have been read to its log line and report.
type request struct {
	h  *handler
	in *http.Request
	// rc sets the deadlines of the connection to the visitor.
	rc   *http.ResponseController
	out  *responseWriter
	body *requestBody // nil for a request without a body

	// id is the fresh id made for the request when it arrived.
	id          string
	client      netip.Addr
	peer        netip.Addr
	peerTrusted bool
	start       time.Time

	// disposition is what is done with the request: waf's valid disposition,
	// or SWRP_DEFAULT_DISPOSITION. It is "" when neither gave one.
	disposition string
	// answered is true when waf gave disposition, which makes a report due,
	// and byDefault when SWRP_DEFAULT_DISPOSITION did.
	answered  bool
	byDefault bool
	// asn and country are waf's X-Client-ASN and X-Client-Country, for the
	// app.
	asn, country string

	// upstreamStart is when the request was handed to the app.
	upstreamStart time.Time
	// cancel ends the request to the app.
	cancel context.CancelFunc
	// refused is the status of the first refusal, from whichever goroutine meets
	// it: what the visitor is answered if the response has not started yet.
	refused atomic.Int64
	// upgraded is the connection to the app once the app has switched protocols,
	// as for a WebSocket, and nil otherwise.
	upgraded *upgradedConn

	// mu guards what follows. The time limits run on goroutines of their own,
	// and the transport starts and stops them, and notes requestEnd, from its
	// own; once timersStopped is set, neither time limit acts any more.
	mu            sync.Mutex
	timersStopped bool
	requestTimer  *time.Timer
	responseTimer *time.Timer
	// requestEnd is when the app had been sent the whole request, or the answer
	// started, if that was sooner.
	requestEnd time.Time
}

// newRequest starts handling r: it notes the time, works out the visitor from
// the TCP peer and X-Forwarded-For, and makes the request's id.
func (h *handler) newRequest(w http.ResponseWriter, r *http.Request) *request {
	peer := peerAddress(r)
	trusted := h.config.TrustedProxies

	rq := &request{
		h:           h,
		in:          r,
		rc:          http.NewResponseController(w),
		out:         &responseWriter{ResponseWriter: w},
		id:          rand.Text(),
		client:      clientAddress(peer, r.Header.Values("X-Forwarded-For"), trusted),
		peer:        peer,
		peerTrusted: isInside(peer, trusted),
		start:       time.Now(),
	}

	if r.Body != http.NoBody {
		rq.body = &requestBody{body: limitBody(r.Body, h.config.MaxRequestSize), rq: rq}
	}

	return rq
}

// answerHealth answers the health check.
func (rq *request) answerHealth() {
	rq.prepareOwnAnswer(rq.requestDeadline())
	rq.out.Header().Set("Content-Type", plainText)
	rq.send(http.StatusOK, []byte("ok\n"))
}

// readFirstPartOfBody reads up to the first 64 KiB of the body into memory,
// within SWRP_MAX_REQUEST_TIME_MS, and reports whether the request goes on. A
// visitor too slow to send it is answered 408. For a malformed chunk, or a
// visitor that went away or stopped its sending side, the connection is closed
// without an answer. A body over SWRP_MAX_REQUEST_SIZE goes on with the part
// read, so that waf is asked first; passToApp then refuses it.
func (rq *request) readFirstPartOfBody() bool {
	if rq.body == nil {
		return true
	}

	_ = rq.rc.SetReadDeadline(rq.requestDeadline())
	held, err := io.ReadAll(io.LimitReader(rq.body, waf.QuestionBodyMaxBytes))
	// From here on, the request timer ends a read of the rest of the body, so that
	// the answer can name the side rproxy was waiting on.
	_ = rq.rc.SetReadDeadline(time.Time{})

	rq.body.held = held

	switch {
	case err == nil || rq.refused.Load() != 0:
		return true
	case errors.Is(err, os.ErrDeadlineExceeded):
		rq.answer(http.StatusRequestTimeout)

		return false
	default:
		// This panic has Go's server close the connection without an answer, and
		// log nothing; the deferred finish still runs.
		panic(http.ErrAbortHandler)
	}
}

// setPathAsSent has out, the URL of a request made from in, send the path as the
// visitor sent it. Go sends the path URL.EscapedPath gives, which, for a path
// holding a character Go escapes, such as |, is the whole path decoded and
// escaped again, %2F turned into /. Opaque is sent as it is.
func setPathAsSent(out *url.URL, in *http.Request) {
	out.Opaque = rules.Path(in)
}

// forward passes the request to the app, the part of the body held and then the
// rest, and the app's answer back. ctx is the request's own context.
func (rq *request) forward(ctx context.Context) {
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()

	rq.cancel = cancel
	ctx = httptrace.WithClientTrace(ctx, &httptrace.ClientTrace{
		WroteRequest: rq.wroteRequest,
	})

	out := rq.in.WithContext(ctx)
	if rq.body != nil {
		out.Body = rq.body
	}

	reverseProxy := &httputil.ReverseProxy{
		Rewrite:        rq.rewrite,
		Transport:      rq.h.transport,
		FlushInterval:  flushAfterEachWrite,
		ErrorLog:       rq.h.errorLog,
		ModifyResponse: rq.modifyResponse,
		ErrorHandler:   rq.answerError,
	}

	rq.startRequestTimer()
	rq.upstreamStart = time.Now()
	reverseProxy.ServeHTTP(rq.out, out)

	// ReverseProxy ignores a failed flush, so it misses a short answer cut off when
	// SWRP_MAX_RESPONSE_TIME_MS runs out; once a flush has failed, every later one
	// fails too. Go's server would then keep the connection open while it reads
	// the rest of a body that has not arrived whole. So that reading ends now, and
	// the refusal keeps finish from extending it to SWRP_MAX_REQUEST_TIME_MS. The
	// panic has Go's server close the connection, and log nothing; the deferred
	// finish still runs.
	if rq.out.status != http.StatusSwitchingProtocols &&
		errors.Is(rq.rc.Flush(), os.ErrDeadlineExceeded) {
		rq.refuse(http.StatusGatewayTimeout)
		rq.stopPassingOnBody(time.Now())
		panic(http.ErrAbortHandler)
	}
}

// rewrite makes the request the app receives: the visitor's request, unchanged,
// sent to SWRP_UPSTREAM_URL, with the forwarded headers and the request's id set,
// without any X-Client-ASN or X-Client-Country the visitor sent, and with
// waf's.
func (rq *request) rewrite(pr *httputil.ProxyRequest) {
	upstream := rq.h.config.UpstreamURL
	pr.Out.URL.Scheme = upstream.Scheme
	pr.Out.URL.Host = upstream.Host
	// ReverseProxy drops query parameters it cannot parse; the app gets the
	// query as the visitor sent it.
	pr.Out.URL.RawQuery = pr.In.URL.RawQuery
	setPathAsSent(pr.Out.URL, pr.In)

	setForwardedHeaders(pr.In, pr.Out, rq.peer, rq.peerTrusted)
	pr.Out.Header.Set(requestIDHeader, rq.id)
	pr.Out.Header.Del(asnHeader)
	pr.Out.Header.Del(countryHeader)

	if rq.asn != "" {
		pr.Out.Header.Set(asnHeader, rq.asn)
	}

	if rq.country != "" {
		pr.Out.Header.Set(countryHeader, rq.country)
	}
}

// modifyResponse looks at the app's answer before ReverseProxy passes it on.
func (rq *request) modifyResponse(res *http.Response) error {
	if res.StatusCode == http.StatusSwitchingProtocols {
		// An upgraded connection, such as a WebSocket, is not cut by the time
		// limits. ReverseProxy writes this answer straight to the connection it
		// takes over, not through rq.out, and then copies what passes each way
		// through res.Body, the connection to the app. It may still refuse the
		// switch, as when the app's Upgrade is not the visitor's, so the status
		// is noted only once it has taken over the connection, by rq.out's
		// Hijack.
		rq.stopTimers()

		conn, ok := res.Body.(io.ReadWriteCloser)
		if ok {
			rq.upgraded = &upgradedConn{ReadWriteCloser: conn}
			res.Body = rq.upgraded
		}

		return nil
	}

	requestEnd := rq.endRequest()

	// An answer that arrives once SWRP_MAX_RESPONSE_TIME_MS has run out is not
	// passed on: the visitor is answered 504.
	limit := rq.h.config.MaxResponseTime
	if limit > 0 && !time.Now().Before(requestEnd.Add(limit)) {
		rq.refuse(http.StatusGatewayTimeout)

		return errResponseTooLate
	}

	maxBytes := rq.h.config.MaxResponseSize
	if maxBytes > 0 && res.Body != http.NoBody && res.ContentLength > maxBytes {
		rq.refuse(http.StatusBadGateway)

		return errResponseTooLarge
	}

	// Go's server adds a Content-Type it guesses from the body to a response
	// whose headers have no Content-Type key. The key with no value sends none,
	// as the app sent none.
	if res.Header.Get("Content-Type") == "" {
		rq.out.Header()["Content-Type"] = nil
	}

	rq.closeAfterAnswer()

	res.Body = &responseBody{body: limitBody(res.Body, maxBytes), rq: rq}
	rq.out.afterStatus = rq.setResponseDeadline

	return nil
}

// answerError is ReverseProxy's ErrorHandler: the request could not be passed to
// the app, or the app's answer cannot be passed on.
func (rq *request) answerError(_ http.ResponseWriter, _ *http.Request, err error) {
	if rq.upgraded != nil {
		// The app switched protocols, but the switch was not passed on, and
		// ReverseProxy leaves the app's connection open when it refuses it.
		_ = rq.upgraded.Close()
	}

	status := int(rq.refused.Load())
	if status == 0 {
		if rq.in.Context().Err() != nil {
			// The visitor went away or stopped its sending side. This panic has Go's
			// server close the connection without an answer, and log nothing; the
			// deferred finish still runs.
			panic(http.ErrAbortHandler)
		}

		rq.h.processLog.Warn("request to the app failed", "error", err.Error())

		status = http.StatusBadGateway
	}

	rq.answer(status)
}

// answer sends rproxy's own answer, unless the response has already started.
func (rq *request) answer(status int) {
	rq.refused.CompareAndSwap(0, int64(status))

	if rq.out.status != 0 {
		return // too late to answer: the connection can only be cut
	}

	// A visitor found too slow is read no more; any other may go on sending until
	// its time is up, so that Go's server can read the rest of the body and end
	// the request cleanly.
	deadline := rq.requestDeadline()
	if status == http.StatusRequestTimeout {
		deadline = time.Now()
	}

	rq.prepareOwnAnswer(deadline)

	header := rq.out.Header()
	header.Set("Content-Type", plainText)
	header.Set("X-Content-Type-Options", "nosniff")
	rq.send(status, []byte(http.StatusText(status)+"\n"))
}

// send gives the visitor an answer rproxy holds whole: status, the headers set so
// far and body. Go's server would otherwise send so short an answer only once the
// handler has returned, after finish has written the log line and sent the report.
// Without Content-Length, the flush would send the body in chunks, the last of
// which would still wait for the handler. A HEAD's answer has no body, so any
// Content-Length it has stays as waf gave it.
func (rq *request) send(status int, body []byte) {
	if rq.in.Method != http.MethodHead {
		rq.out.Header().Set("Content-Length", strconv.Itoa(len(body)))
	}

	rq.closeAfterAnswer()
	rq.out.WriteHeader(status)
	_, _ = rq.out.Write(body)
	_ = rq.rc.Flush()
}

// closeAfterAnswer has Go's server send the answer at once, and close the
// connection after it, when the visitor's body has not arrived whole. Without
// it, Go's server reads up to 256 KiB of the body left before it sends an
// answer, and so waits on a visitor that has stopped sending.
func (rq *request) closeAfterAnswer() {
	if rq.body != nil && !rq.body.received.Load() {
		rq.out.Header().Set("Connection", "close")
	}
}

// prepareOwnAnswer readies the connection for an answer rproxy gives itself: the
// reading of a body that has not arrived whole ends at deadline, and the answer
// has SWRP_MAX_RESPONSE_TIME_MS from now to reach the visitor, even when the app
// used that time up first, and at least minSendTime.
func (rq *request) prepareOwnAnswer(deadline time.Time) {
	rq.stopReadingBody(deadline)
	rq.endRequest()

	limit := rq.h.config.MaxResponseTime
	if limit > 0 {
		_ = rq.rc.SetWriteDeadline(time.Now().Add(max(limit, minSendTime)))
	}
}

// refuse records status, unless an earlier refusal was, and ends the request to
// the app, if one was made.
func (rq *request) refuse(status int) {
	rq.refused.CompareAndSwap(0, int64(status))

	if rq.cancel != nil {
		rq.cancel()
	}
}

// finish ends the request's time limits, writes its log line and, if waf
// answered the question, sends the report. ctx is the request's own context.
func (rq *request) finish(ctx context.Context) {
	rq.stopTimers()

	if rq.refused.Load() == 0 {
		rq.stopPassingOnBody(rq.requestDeadline())
	}

	requestEnd := rq.endRequest()
	end := time.Now()
	requestBytes, responseBytes := rq.bytesEachWay()

	err := requestlog.WriteRProxyLine(rq.h.requestLog, &requestlog.RProxyLine{
		Time:               requestlog.FormatTime(rq.start),
		RequestID:          rq.id,
		ClientIP:           rq.client.String(),
		Host:               rq.in.Host,
		Method:             rq.in.Method,
		Path:               rules.Path(rq.in),
		Query:              rq.in.URL.RawQuery,
		Status:             rq.out.status,
		RequestBytes:       requestBytes,
		ResponseBytes:      responseBytes,
		DurationRequestMS:  requestlog.Milliseconds(requestEnd.Sub(rq.start)),
		DurationResponseMS: requestlog.Milliseconds(end.Sub(requestEnd)),
		Disposition:        rq.disposition,
		DefaultDisposition: rq.byDefault,
	})
	if err != nil {
		rq.h.processLog.Error("writing the request log failed", "error", err.Error())
	}

	if !rq.answered {
		return
	}

	report := &waf.Report{
		RequestID:       rq.id,
		Status:          min(rq.out.status, reportStatusMax),
		RequestBytes:    requestBytes,
		ResponseBytes:   responseBytes,
		DurationTotalMS: requestlog.Milliseconds(end.Sub(rq.start)),
	}
	if !rq.upstreamStart.IsZero() {
		report.DurationUpstreamTotalMS = requestlog.Milliseconds(end.Sub(rq.upstreamStart))
	}

	rq.h.sendReport(ctx, report)
}

// bytesEachWay returns the body bytes of the request and of the answer. For an
// upgraded connection, such as a WebSocket, which has closed by then, what it
// carried from the visitor counts with the request's and what it carried from the
// app with the answer's.
func (rq *request) bytesEachWay() (int64, int64) {
	var requestBytes int64
	if rq.body != nil {
		requestBytes = rq.body.bytes.Load()
	}

	responseBytes := rq.out.bytes

	if rq.upgraded != nil {
		requestBytes += rq.upgraded.toApp.Load()
		responseBytes += rq.upgraded.fromApp.Load()
	}

	return requestBytes, responseBytes
}

// requestDeadline is when the visitor must have sent its whole request, or zero
// when SWRP_MAX_REQUEST_TIME_MS is off.
func (rq *request) requestDeadline() time.Time {
	limit := rq.h.config.MaxRequestTime
	if limit == 0 {
		return time.Time{}
	}

	return rq.start.Add(limit)
}

// stopReadingBody ends, at deadline, the reading of a visitor's body that has not
// arrived whole: Go's server then reads no more of it, and closes the connection
// after the answer.
func (rq *request) stopReadingBody(deadline time.Time) {
	if rq.body == nil || rq.body.received.Load() {
		return
	}

	_ = rq.rc.SetReadDeadline(deadline)
}

// stopPassingOnBody ends the transport's reading of a visitor's body that has
// not arrived whole, and lets Go's server read the rest until deadline. Go's
// server reads it once the handler returns, but with no deadline at all if
// another read of it is still waiting.
func (rq *request) stopPassingOnBody(deadline time.Time) {
	if rq.body == nil || rq.body.received.Load() {
		return
	}

	// A deadline that has passed ends a Read that waits, and the lock waits for
	// that Read to return.
	_ = rq.rc.SetReadDeadline(time.Now())

	rq.body.mu.Lock()
	rq.body.stopped = true
	rq.body.mu.Unlock()

	rq.stopReadingBody(deadline)
}

// startRequestTimer starts SWRP_MAX_REQUEST_TIME_MS, which runs from the
// request's headers being read until the app has been sent the whole request.
func (rq *request) startRequestTimer() {
	deadline := rq.requestDeadline()
	if deadline.IsZero() {
		return
	}

	rq.mu.Lock()
	defer rq.mu.Unlock()

	rq.requestTimer = time.AfterFunc(time.Until(deadline), rq.requestTimedOut)
}

// requestTimedOut is called when SWRP_MAX_REQUEST_TIME_MS runs out while the
// request is still on its way to the app. The answer names the side rproxy was
// waiting on at that moment: 408 when it was waiting for the visitor to send more
// of its body, 504 when it was waiting for the app to be reached or to take what
// it had.
func (rq *request) requestTimedOut() {
	rq.mu.Lock()
	defer rq.mu.Unlock()

	if rq.timersStopped {
		return
	}

	if rq.body == nil || !rq.body.waiting.Load() {
		rq.refuse(http.StatusGatewayTimeout)

		return
	}

	rq.refuse(http.StatusRequestTimeout)
	// The transport gives up on the app only once its Read of the visitor's body
	// returns, so that Read is ended now. The lock keeps this from reaching the
	// connection after the request is handled.
	_ = rq.rc.SetReadDeadline(time.Now())
}

// wroteRequest is called once the app has been sent the whole request:
// SWRP_MAX_REQUEST_TIME_MS ends and SWRP_MAX_RESPONSE_TIME_MS starts.
func (rq *request) wroteRequest(info httptrace.WroteRequestInfo) {
	if info.Err != nil {
		return // the transport gives up, or tries again
	}

	requestEnd := rq.endRequest()

	rq.mu.Lock()
	defer rq.mu.Unlock()

	if rq.timersStopped {
		return
	}

	stopTimer(rq.requestTimer)

	limit := rq.h.config.MaxResponseTime
	if limit > 0 {
		rq.responseTimer = time.AfterFunc(time.Until(requestEnd.Add(limit)),
			rq.responseTimedOut)
	}
}

// responseTimedOut is called when SWRP_MAX_RESPONSE_TIME_MS runs out before the
// app has sent its whole answer.
func (rq *request) responseTimedOut() {
	rq.mu.Lock()
	defer rq.mu.Unlock()

	if !rq.timersStopped {
		rq.refuse(http.StatusGatewayTimeout)
	}
}

// responseReceived is called once the app has sent its whole answer.
func (rq *request) responseReceived() {
	rq.stopTimers()
}

// setResponseDeadline sends the status and headers of the app's answer, with
// at least minSendTime, before ReverseProxy reads the answer's body, and then
// the rest has until SWRP_MAX_RESPONSE_TIME_MS after the end of the request.
func (rq *request) setResponseDeadline() {
	limit := rq.h.config.MaxResponseTime
	if limit == 0 {
		return
	}

	deadline := rq.endRequest().Add(limit)

	_ = rq.rc.SetWriteDeadline(time.Now().Add(max(time.Until(deadline), minSendTime)))
	_ = rq.rc.Flush()
	_ = rq.rc.SetWriteDeadline(deadline)
}

// endRequest notes now as the end of the request, unless an end is noted, and
// returns the end noted.
func (rq *request) endRequest() time.Time {
	rq.mu.Lock()
	defer rq.mu.Unlock()

	if rq.requestEnd.IsZero() {
		rq.requestEnd = time.Now()
	}

	return rq.requestEnd
}

// stopTimers stops the request's time limits and keeps any from starting later:
// the app's answer is complete, the connection upgraded, or the request handled.
func (rq *request) stopTimers() {
	rq.mu.Lock()
	defer rq.mu.Unlock()

	rq.timersStopped = true
	stopTimer(rq.requestTimer)
	stopTimer(rq.responseTimer)
}

// stopTimer stops t, which is nil when its time limit is off.
func stopTimer(t *time.Timer) {
	if t != nil {
		t.Stop()
	}
}