Frame-by-frame IMQ with an outgoing filter; HTTP interceptor hook

Co-Authored-By: Claude Opus 5.5 <[email protected]>

internal/proxy/proxy.go

@@ -14,6 +14,7 @@ package proxy
 
 import (
 	"bytes"
+	"compress/gzip"
 	"context"
 	"crypto/subtle"
 	"crypto/tls"
@@ -24,6 +25,7 @@ import (
 	"net"
 	"net/http"
 	"net/url"
+	"strconv"
 	"strings"
 	"sync/atomic"
 	"time"
@@ -45,6 +47,10 @@ type Server struct {
 	// TLSConfig is the base TLS configuration of CONNECT's upstream connections (tests set RootCAs); nil is the
 	// default, with the system's roots.
 	TLSConfig *tls.Config
+	// Interceptor, if set, sees every HTTP request (with a body the proxy could hold in memory).
+	Interceptor Interceptor
+	// IMQFilter, if set, sees the client's outgoing IMQ messages (on CONNECT streams) and can change or drop them.
+	IMQFilter IMQFilter
 	// Dial makes CONNECT's upstream connections (tests point it at local servers); nil is a net.Dialer.
 	Dial func(ctx context.Context, network, addr string) (net.Conn, error)
 
@@ -169,6 +175,29 @@ func (s *Server) forward(w http.ResponseWriter, r *http.Request) {
 		http.Error(w, "t5de-proxy: "+err.Error(), http.StatusBadRequest)
 		return
 	}
+	var modify func(*http.Response, []byte) ([]byte, error)
+	if s.Interceptor != nil && replayable {
+		x := &Exchange{URL: r.URL, Request: out}
+		if out.GetBody != nil {
+			body, _ := out.GetBody()
+			x.Body, _ = io.ReadAll(body)
+		}
+		original := x.Body
+		if err := s.Interceptor.Intercept(x); err != nil {
+			s.logf("#%d %s %s: %v", id, r.Method, s.logURL(r.URL), err)
+			w.Header().Set(ErrorHeader, err.Error())
+			http.Error(w, "t5de-proxy: "+err.Error(), http.StatusBadGateway)
+			return
+		}
+		if x.Response != nil {
+			s.writeLocal(w, id, r, x.Response, start)
+			return
+		}
+		if !bytes.Equal(x.Body, original) {
+			setBody(out, x.Body)
+		}
+		modify = x.ModifyResponse
+	}
 
 	resp, err := s.transport().RoundTrip(out)
 	if err != nil && upgraded && replayable && r.Context().Err() == nil {
@@ -191,6 +220,11 @@ func (s *Server) forward(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 	defer resp.Body.Close()
+	if modify != nil {
+		if err := modifyResponse(resp, modify); err != nil {
+			s.logf("#%d %s %s: changing the response: %v", id, r.Method, s.logURL(r.URL), err)
+		}
+	}
 
 	removeHopHeaders(resp.Header)
 	for name, values := range resp.Header {
@@ -208,6 +242,85 @@ func (s *Server) forward(w http.ResponseWriter, r *http.Request) {
 	s.logf("#%d %s %s %d %dB %s%s", id, r.Method, s.logURL(r.URL), resp.StatusCode, n, since(start), note)
 }
 
+// Interceptor can answer requests itself, change them before they're sent, and change their responses. The
+// t5de-proxy's community products are served this way.
+type Interceptor interface {
+	Intercept(x *Exchange) error
+}
+
+// Exchange is a request on its way, as an Interceptor sees it.
+type Exchange struct {
+	// URL is the URL the client asked for (before any upgrade to HTTPS).
+	URL *url.URL
+	// Request is the request about to be sent; its body is Body.
+	Request *http.Request
+	// Body is the request's body; an Interceptor may replace it.
+	Body []byte
+	// Response, if an Interceptor sets it, answers the request: nothing is sent.
+	Response *http.Response
+	// ModifyResponse, if set, changes the server's response body (decompressed) before the client gets it.
+	ModifyResponse func(resp *http.Response, body []byte) ([]byte, error)
+}
+
+// LocalResponse is a response made by the proxy.
+func LocalResponse(status int, contentType string, body []byte) *http.Response {
+	h := http.Header{}
+	h.Set("Content-Type", contentType)
+	h.Set("Content-Length", strconv.Itoa(len(body)))
+	return &http.Response{StatusCode: status, Header: h, ContentLength: int64(len(body)), Body: io.NopCloser(bytes.NewReader(body))}
+}
+
+func (s *Server) writeLocal(w http.ResponseWriter, id int64, r *http.Request, resp *http.Response, start time.Time) {
+	defer resp.Body.Close()
+	for name, values := range resp.Header {
+		w.Header()[name] = values
+	}
+	w.WriteHeader(resp.StatusCode)
+	n, _ := io.Copy(w, resp.Body)
+	s.logf("#%d %s %s %d %dB %s (answered by the proxy)", id, r.Method, s.logURL(r.URL), resp.StatusCode, n, since(start))
+}
+
+func setBody(out *http.Request, body []byte) {
+	out.ContentLength = int64(len(body))
+	out.Header.Set("Content-Length", strconv.Itoa(len(body)))
+	out.GetBody = func() (io.ReadCloser, error) { return io.NopCloser(bytes.NewReader(body)), nil }
+	out.Body, _ = out.GetBody()
+	if len(body) == 0 {
+		out.Body, out.GetBody = nil, nil
+	}
+}
+
+// modifyResponse replaces resp's body with modify's version of it, uncompressed.
+func modifyResponse(resp *http.Response, modify func(*http.Response, []byte) ([]byte, error)) error {
+	raw, err := io.ReadAll(io.LimitReader(resp.Body, maxBuffered))
+	resp.Body.Close()
+	if err != nil {
+		resp.Body = io.NopCloser(bytes.NewReader(raw))
+		return err
+	}
+	body := raw
+	if strings.EqualFold(resp.Header.Get("Content-Encoding"), "gzip") {
+		zr, err := gzip.NewReader(bytes.NewReader(raw))
+		if err == nil {
+			body, err = io.ReadAll(io.LimitReader(zr, maxBuffered))
+		}
+		if err != nil {
+			resp.Body = io.NopCloser(bytes.NewReader(raw))
+			return err
+		}
+	}
+	changed, err := modify(resp, body)
+	if err != nil {
+		changed = body
+	}
+	resp.Header.Del("Content-Encoding")
+	resp.Header.Del("Content-Md5")
+	resp.Header.Set("Content-Length", strconv.Itoa(len(changed)))
+	resp.ContentLength = int64(len(changed))
+	resp.Body = io.NopCloser(bytes.NewReader(changed))
+	return err
+}
+
 // maxBuffered bounds the request bodies the proxy holds in memory; larger ones (uploads) are streamed.
 const maxBuffered = 32 << 20