httplog/httplog.go

// httplog provides several utilities for logging http requests and responses.
//
// This package is needs testing, and is missing features. You should probably not be using it.

package httplog

import (
	"bytes"
	"cmp"
	"context"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"sync"
	"sync/atomic"
)

var (
	logMutex  sync.Mutex
	logOutput io.Writer = os.Stderr
)

func SetLogOutput(w io.Writer) {
	logMutex.Lock()
	defer logMutex.Unlock()
	logOutput = w
}

func emit(prefix string, text []byte) {
	if len(text) == 0 {
		return
	}

	logMutex.Lock()
	defer logMutex.Unlock()

	for i := 0; i < len(text); {
		var line []byte
		if j := bytes.IndexByte(text[i:], '\n'); j != -1 {
			line = text[i : i+j]
			if len(line) > 0 && line[len(line)-1] == '\r' {
				line = line[:len(line)-1]
			}
			i = i + j + 1
		} else {
			line = text[i:]
			i = len(text)
		}

		if len(line) > 0 {
			fmt.Fprintf(logOutput, "%s%s\n", prefix, line)
		}
	}

	if f, ok := logOutput.(interface{ Flush() error }); ok {
		f.Flush()
	}
}

type logWriter struct {
	prefix string
	buf    []byte
}

func (w *logWriter) Write(p []byte) (int, error) {
	i := bytes.LastIndexByte(p, '\n')

	if i == -1 {
		w.buf = append(w.buf, p...)
		return len(p), nil
	}

	w.buf = append(w.buf, p[:i+1]...)
	emit(w.prefix, w.buf)
	w.buf = append(w.buf[:0], p[i+1:]...)
	return len(p), nil
}

func (w *logWriter) Close() error {
	if len(w.buf) > 0 {
		emit(w.prefix, w.buf)
		w.buf = w.buf[:0]
	}

	return nil
}

type logRequestWriter struct {
	logWriter
	headersWritten, closed bool
	getStatus              func() string
	getHeader              func() http.Header
	getTrailer             func() http.Header
}

func (w *logRequestWriter) Write(p []byte) (int, error) {
	if w.headersWritten {
		return w.logWriter.Write(p)
	}

	w.headersWritten = true

	if w.getStatus != nil {
		if status := w.getStatus(); status != "" {
			fmt.Fprintf(&w.logWriter, "%s\n", status)
		}
	}

	if w.getHeader != nil {
		if h := w.getHeader(); len(h) > 0 {
			h.Write(&w.logWriter)
		}
	}

	return w.logWriter.Write(p)
}

func (w *logRequestWriter) Close() error {
	if w.closed {
		return nil
	}

	w.closed = true

	w.Write(nil) // force headers be written, if it hasn't already.

	if w.getTrailer != nil {
		if h := w.getTrailer(); len(h) > 0 {
			io.WriteString(&w.logWriter, "\n")
			h.Write(&w.logWriter)
		}
	}

	return nil
}

// this is more complicated than it arguably ought to be, because
// [http.Request] notes that closing can happen concurrently with reading,
// and more importantly, closing needs to unblock an in-progress read.
//
// the unblocking detail is the for the underlying reader to figure out,
// but we need to at least handle concurrently being closed safely.
type teeReadCloser struct {
	m      sync.Mutex
	closed bool
	err    error
	r      io.ReadCloser
	log    logRequestWriter
}

func (r *teeReadCloser) Read(p []byte) (n int, err error) {
	r.m.Lock()
	closed, err := r.closed, r.err
	r.m.Unlock()

	if err != nil {
		return 0, err
	}

	if closed {
		return 0, os.ErrClosed
	}

	n, err = r.r.Read(p)

	r.m.Lock()
	defer r.m.Unlock()

	if err != nil && err != io.EOF {
		r.err = err
	}

	if n > 0 && !r.log.closed {
		if _, err := r.log.Write(p[:n]); err != nil {
			r.log.Close()
			log.Printf("error teeing read to writer: %v", err)
		}
	}

	if err != nil && !r.log.closed {
		if err := r.log.Close(); err != nil {
			log.Printf("error closing tee writer: %v", err)
		}
	}

	return n, err
}

func (r *teeReadCloser) Close() error {
	r.m.Lock()
	closed, err := r.closed, r.err
	r.m.Unlock()

	if closed {
		return err
	}

	err = r.r.Close()

	r.m.Lock()
	defer r.m.Unlock()

	r.closed = true
	r.err = cmp.Or(err, os.ErrClosed)

	if !r.log.closed {
		if err := r.log.Close(); err != nil {
			log.Printf("error closing tee writer: %v", err)
		}
	}

	return err
}

type teeResponseWriter struct {
	http.ResponseWriter
	log           logRequestWriter
	statusCode    int
	headerWritten bool
}

func (w *teeResponseWriter) WriteHeader(code int) {
	if w.headerWritten {
		return
	}

	w.statusCode = code
	w.headerWritten = true
	w.log.Write(nil) // force headers to be written
	w.ResponseWriter.WriteHeader(code)
}

func (w *teeResponseWriter) Write(p []byte) (int, error) {
	if !w.headerWritten {
		w.WriteHeader(http.StatusOK)
	}

	n, err := w.ResponseWriter.Write(p)
	w.log.Write(p[:n])
	return n, err
}

func (w *teeResponseWriter) Flush() {
	if !w.headerWritten {
		w.WriteHeader(http.StatusOK)
	}

	if f, ok := w.ResponseWriter.(http.Flusher); ok {
		f.Flush()
	}
}

var requestID atomic.Uint32

func makeSendRecvPrefixPair() (string, string) {
	id := requestID.Add(1)
	return fmt.Sprintf("%d ➡ ⁞ ", id), fmt.Sprintf("⬅ %d ⁞ ", id)
}

func logResponseWriter(prefix string, w http.ResponseWriter, r *http.Request) *teeResponseWriter {
	trw := &teeResponseWriter{
		ResponseWriter: w,
		log: logRequestWriter{
			logWriter:  logWriter{prefix: prefix},
			getHeader:  func() http.Header { return w.Header() },
			getTrailer: func() http.Header { return w.Header() },
		},
	}

	trw.log.getStatus = func() string {
		return fmt.Sprintf("%s %d %s", r.Proto, trw.statusCode, http.StatusText(trw.statusCode))
	}

	return trw
}

// Error logs an error message.
//
// Currently, the only reason you'd use this would be if [http.Client] or [http.Transport] returned an error instead of a response.
func Error(prefix string, err error) {
	emit(prefix, fmt.Appendf(nil, "⚡ error: %v", err))
}

// Handler invokes the next http.Handler, logging both the request and the generated response.
func Handler(w http.ResponseWriter, req *http.Request, next http.Handler) {
	send, recv := makeSendRecvPrefixPair()

	body := &teeReadCloser{
		r: req.Body,
		log: logRequestWriter{
			logWriter: logWriter{
				prefix: recv,
			},
			getStatus: func() string {
				return fmt.Sprintf("%s %s\n", req.Method, req.RequestURI)
			},
			getHeader: func() http.Header {
				return req.Header
			},
			getTrailer: func() http.Header {
				return req.Trailer
			},
		},
	}

	body.log.Write(nil) // we can write status and headers immediately

	defer body.log.Close()
	req2 := req.Clone(req.Context())
	req2.Body = body

	lrw := logResponseWriter(send, w, req2)
	defer lrw.log.Close()
	next.ServeHTTP(lrw, req2)
}

// Middleware wraps the Handler function, turning it into http middleware.
func Middleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		Handler(w, r, next)
	})
}

// NewRequest is semantically equivalent to [http.NewRequest]. The request will be logged immediately, and if
// http.NewRequestWithContext returns an error, that will be logged as well.
//
// If this function does not return an error, it will also return a prefix that can be used with the [Response] or [Error] functions.
func NewRequest(ctx context.Context, method, url string, r io.Reader) (req *http.Request, recvPrefix string, err error) {
	send, recv := makeSendRecvPrefixPair()

	if r == nil {
		r = http.NoBody
	}

	rc, ok := r.(io.ReadCloser)
	if !ok {
		rc = io.NopCloser(r)
	}

	body := &teeReadCloser{
		r: rc,
		log: logRequestWriter{
			logWriter: logWriter{
				prefix: send,
			},
			getHeader: func() http.Header {
				return req.Header
			},
			getTrailer: func() http.Header {
				return req.Trailer
			},
		},
	}

	// We can print the status line now, but the headers won't be available until the request has been passed to an http client or transport.
	fmt.Fprintf(&body.log.logWriter, "%s %s\n", method, url)

	req, err = http.NewRequestWithContext(ctx, method, url, body)
	if err != nil {
		body.r.Close()
		Error(send, err)
		return nil, "", err
	}

	return req, recv, err
}

// Response immediately logs the status line and headers for the response,
// and then replaces its Body field with a proxy object that will log the
// body as it's read, and then print any trailers once EOF is reached.
func Response(prefix string, resp *http.Response) *http.Response {
	body := &teeReadCloser{
		r: resp.Body,
		log: logRequestWriter{
			logWriter: logWriter{prefix: prefix},
			getStatus: func() string {
				return fmt.Sprintf("%s %s", resp.Proto, resp.Status)
			},
			getHeader: func() http.Header {
				return resp.Header
			},
			getTrailer: func() http.Header {
				return resp.Trailer
			},
		},
	}

	body.log.Write(nil) // the headers and status line are can be printed immediately.

	resp.Body = body
	return resp
}