smallwebwaf/internal/rproxy/bodies.go

package rproxy

import (
	"bufio"
	"errors"
	"io"
	"net"
	"net/http"
	"sync"
	"sync/atomic"
)

// errResponseTooLarge ends an app's response body that is longer than
// SWRP_MAX_RESPONSE_SIZE.
var errResponseTooLarge = errors.New(
	"the response body is over SWRP_MAX_RESPONSE_SIZE")

// requestBody is the visitor's request body. rproxy first reads up to the first
// 64 KiB of it into held, and then, on its way to the app, the transport reads
// it on a goroutine of its own.
type requestBody struct {
	// body is the visitor's body, ending in an *http.MaxBytesError past
	// SWRP_MAX_REQUEST_SIZE.
	body io.ReadCloser
	rq   *request
	// held is the part read before the request goes to the app, which Read gives
	// first.
	held []byte
	// waiting is true while a Read waits for the visitor to send more.
	waiting atomic.Bool
	// received is true once the visitor has sent the whole body.
	received atomic.Bool
	// bytes is how much of the body has been read.
	bytes atomic.Int64

	// mu is held for the whole of each Read, and guards stopped.
	mu sync.Mutex
	// stopped is true once the body no longer goes to the app; Read then reads
	// no more.
	stopped bool
}

// Read reads from the visitor's body, after the part held, which has been
// counted already.
func (b *requestBody) Read(p []byte) (int, error) {
	b.mu.Lock()
	defer b.mu.Unlock()

	if b.stopped {
		return 0, http.ErrBodyReadAfterClose
	}

	if len(b.held) > 0 {
		n := copy(p, b.held)
		b.held = b.held[n:]

		return n, nil
	}

	b.waiting.Store(true)
	n, err := b.body.Read(p)
	b.waiting.Store(false)
	b.bytes.Add(int64(n))

	var tooLarge *http.MaxBytesError

	switch {
	case errors.Is(err, io.EOF):
		b.received.Store(true)
	case errors.As(err, &tooLarge):
		b.rq.refuse(http.StatusRequestEntityTooLarge)
	}

	return n, err
}

// Close does nothing: Go's server closes the visitor's body once the handler
// returns. Closing it earlier, as the transport and ReverseProxy do, would wait
// for a Read in progress and then read up to 256 KiB more of the body.
func (b *requestBody) Close() error {
	return nil
}

// responseBody is the app's response body on its way to the visitor.
type responseBody struct {
	// body is the app's body, ending in an *http.MaxBytesError past
	// SWRP_MAX_RESPONSE_SIZE.
	body io.ReadCloser
	rq   *request
}

// Read reads from the app's body.
func (b *responseBody) Read(p []byte) (int, error) {
	n, err := b.body.Read(p)
	if err == nil {
		return n, nil
	}

	var tooLarge *http.MaxBytesError

	switch {
	case errors.Is(err, io.EOF):
		b.rq.responseReceived()
	case errors.As(err, &tooLarge):
		b.rq.refuse(http.StatusBadGateway)

		return n, errResponseTooLarge
	case b.rq.in.Context().Err() == nil:
		// The answer broke off, not because the visitor went away. If a time
		// limit cut it, that refusal came first and is the one kept.
		b.rq.refuse(http.StatusBadGateway)
	}

	return n, err
}

// Close closes the app's body.
func (b *responseBody) Close() error {
	return b.body.Close()
}

// upgradedConn is the connection to the app once the app has switched
// protocols, as for a WebSocket. ReverseProxy writes to it what the visitor
// sends and reads from it what the app sends, on goroutines of its own, until
// the connection closes; it counts the bytes each way, for the report.
type upgradedConn struct {
	io.ReadWriteCloser

	// fromApp is how many bytes the app has sent, and toApp how many the
	// visitor has.
	fromApp atomic.Int64
	toApp   atomic.Int64
}

// Read reads what the app sends.
func (c *upgradedConn) Read(p []byte) (int, error) {
	n, err := c.ReadWriteCloser.Read(p)
	c.fromApp.Add(int64(n))

	return n, err
}

// Write sends the app what the visitor sent.
func (c *upgradedConn) Write(p []byte) (int, error) {
	n, err := c.ReadWriteCloser.Write(p)
	c.toApp.Add(int64(n))

	return n, err
}

// CloseWrite tells the app that the visitor sends no more, while what the app
// sends still passes. ReverseProxy calls it once the visitor has stopped
// sending, and closes the connection there if it is not supported.
func (c *upgradedConn) CloseWrite() error {
	conn, ok := c.ReadWriteCloser.(interface{ CloseWrite() error })
	if !ok {
		return http.ErrNotSupported
	}

	return conn.CloseWrite()
}

// limitBody returns body, cut off with an *http.MaxBytesError after maxBytes, or
// unchanged if maxBytes is zero, which is off.
func limitBody(body io.ReadCloser, maxBytes int64) io.ReadCloser {
	if maxBytes == 0 {
		return body
	}

	// Without a ResponseWriter, MaxBytesReader only counts and cuts off.
	return http.MaxBytesReader(nil, body, maxBytes)
}

// responseWriter is the response to the visitor. It notes the status and size
// for the log line and the report.
type responseWriter struct {
	http.ResponseWriter

	// status is the final status sent, or zero before one is.
	status int
	bytes  int64
	// afterStatus, if set, is called once the final status has been written.
	afterStatus func()
}

// WriteHeader sends the status and headers. An informational 1xx status is
// passed on and the final status still comes later.
func (w *responseWriter) WriteHeader(status int) {
	if status >= http.StatusOK && w.status == 0 {
		w.status = status
	}

	w.ResponseWriter.WriteHeader(status)

	if status >= http.StatusOK && w.afterStatus != nil {
		w.afterStatus()
	}
}

// Write sends part of the body.
func (w *responseWriter) Write(p []byte) (int, error) {
	if w.status == 0 {
		w.status = http.StatusOK
	}

	//nolint:gosec // the answer, as it was given
	n, err := w.ResponseWriter.Write(p)
	w.bytes += int64(n)

	return n, err
}

// FlushError sends what has been written so far. http.ResponseController calls
// it, as ReverseProxy does after each write.
func (w *responseWriter) FlushError() error {
	return http.NewResponseController(w.ResponseWriter).Flush()
}

// Hijack takes over the connection to the visitor, which ReverseProxy does only
// to pass on the app's 101 Switching Protocols: that is then the status sent.
func (w *responseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	conn, buffered, err := http.NewResponseController(w.ResponseWriter).
		Hijack()
	if err == nil {
		w.status = http.StatusSwitchingProtocols
	}

	return conn, buffered, err
}

// Unwrap lets http.ResponseController reach net/http's own ResponseWriter for
// what responseWriter does not do itself.
func (w *responseWriter) Unwrap() http.ResponseWriter {
	return w.ResponseWriter
}