package rproxy

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

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

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

// 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, "" until servewaf has been
	// asked.
	disposition string
	// answered is true when servewaf gave disposition, which makes a report due,
	// and byDefault when SWRP_DEFAULT_DISPOSITION did.
	answered  bool
	byDefault bool
	// asn and country are servewaf'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 the times below, 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
	// answerStarted is when the first byte of the app's answer arrived.
	answerStarted 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.out.Header().Set("Content-Type", "text/plain; charset=utf-8")
	_, _ = io.WriteString(rq.out, "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, and one that went away is not
// answered. A body over SWRP_MAX_REQUEST_SIZE goes on with the part read, so
// that servewaf 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, proxy.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:
		return false
	}
}

// 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. Left out are a
// target that is no path, such as *, and a path starting with //, which Go would
// send in Opaque as a scheme and host.
func setPathAsSent(out *url.URL, in *http.Request) {
	path := rules.Path(in)
	if strings.HasPrefix(path, "/") && !strings.HasPrefix(path, "//") {
		out.Opaque = path
	}
}

// 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,
		GotFirstResponseByte: rq.gotFirstResponseByte,
	})

	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)
}

// 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
// servewaf'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
	}

	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
	}

	res.Body = &responseBody{body: limitBody(res.Body, maxBytes), rq: rq}
	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 {
			return // the visitor has gone, and there is no one to answer
		}

		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)
	http.Error(rq.out, http.StatusText(status), status)
}

// 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.
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(limit))
	}
}

// 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 servewaf
// 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.stopReadingBody(rq.requestDeadline())
	}

	end := time.Now()
	requestEnd := rq.endRequest()
	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,
		DurationRequest:    requestlog.Milliseconds(requestEnd.Sub(rq.start)),
		DurationResponse:   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 := &proxy.Report{
		RequestID:     rq.id,
		Status:        rq.out.status,
		RequestBytes:  requestBytes,
		ResponseBytes: responseBytes,
		DurationTotal: requestlog.Milliseconds(end.Sub(rq.start)),
	}
	if !rq.upstreamStart.IsZero() {
		report.DurationUpstreamTotal = 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)
}

// 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() {
		// Once the app has answered, the transport, ReverseProxy and Go's server
		// each close the visitor's body, and closing it reads what is left of a
		// body under 256 KiB without requestBody.Read, so waiting is false while
		// that read waits for the visitor. That read ends now, before refuse:
		// should rproxy still answer itself, answer then sets the deadline it
		// gives its own answers.
		if !rq.answerStarted.IsZero() {
			rq.stopReadingBody(time.Now())
		}

		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())
}

// gotFirstResponseByte is called once the first byte of the app's answer has
// arrived.
func (rq *request) gotFirstResponseByte() {
	rq.mu.Lock()
	defer rq.mu.Unlock()

	rq.answerStarted = 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. An app may
// answer before the visitor has sent its whole body, as with a 404 it sends
// without reading it. The time limits then run on: until the transport's read of
// the body returns, Go's server cannot send the answer, and
// SWRP_MAX_REQUEST_TIME_MS is what ends that read for a visitor that stops
// sending.
func (rq *request) responseReceived() {
	if rq.body != nil && !rq.body.received.Load() {
		return
	}

	rq.stopTimers()
}

// setResponseDeadline sets SWRP_MAX_RESPONSE_TIME_MS on the connection to the
// visitor: the app's answer must reach the visitor within it of the end of the
// request.
func (rq *request) setResponseDeadline() {
	requestEnd := rq.endRequest()

	limit := rq.h.config.MaxResponseTime
	if limit > 0 {
		_ = rq.rc.SetWriteDeadline(requestEnd.Add(limit))
	}
}

// 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()
	}
}