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

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

internal/proxy/connect.go

@@ -65,8 +65,11 @@ func (s *Server) serveConnect(w http.ResponseWriter, r *http.Request) {
 		client = io.MultiReader(bytes.NewReader(bytes.Clone(head)), conn)
 	}
 
-	up := &stream{s: s, id: id, name: "c2g"}
+	// Frames go to the client whole, from the gateway or from the proxy answering a dropped message, so they
+	// never interleave.
+	toClient := &lockedWriter{w: conn}
 	down := &stream{s: s, id: id, name: "g2c"}
+	up := &stream{s: s, id: id, name: "c2g", filter: s.IMQFilter, reply: toClient}
 	// IMQ is a session: when either side closes, it's over, so both connections close.
 	var once sync.Once
 	closeBoth := func() {
@@ -85,21 +88,43 @@ func (s *Server) serveConnect(w http.ResponseWriter, r *http.Request) {
 	go func() {
 		defer wg.Done()
 		defer closeBoth()
-		down.copy(conn, upstream)
+		down.copy(toClient, upstream)
 	}()
 	wg.Wait()
 	s.logf("#%d CONNECT %s closed after %s: %s, %s", id, target, since(start), up.summary(), down.summary())
 }
 
-// stream is one direction of a CONNECT, counting bytes and IMQ frames as they pass.
+// IMQFilter sees the client's outgoing IMQ messages and can change or drop them.
+type IMQFilter interface {
+	// Outgoing returns the message to send instead of message, or drop to send nothing and tell the client it
+	// was sent.
+	Outgoing(queue, mount string, message []byte) (out []byte, drop bool)
+}
+
+// lockedWriter serializes writes, each a whole number of frames.
+type lockedWriter struct {
+	mu sync.Mutex
+	w  io.Writer
+}
+
+func (l *lockedWriter) Write(b []byte) (int, error) {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+	return l.w.Write(b)
+}
+
+// stream is one direction of a CONNECT. It passes IMQ frames on whole, counting them (and, outgoing, filtering
+// them), or bytes as they come once the stream turns out not to be IMQ.
 type stream struct {
-	s      *Server
-	id     int64
-	name   string
-	bytes  int64
-	frames int
-	other  bool // not IMQ framing
-	reader imq.Reader
+	s       *Server
+	id      int64
+	name    string
+	bytes   int64
+	frames  int
+	other   bool // not IMQ framing
+	pending []byte
+	filter  IMQFilter // outgoing only
+	reply   io.Writer // outgoing only: to the client
 }
 
 func (st *stream) copy(dst io.Writer, src io.Reader) {
@@ -107,8 +132,8 @@ func (st *stream) copy(dst io.Writer, src io.Reader) {
 	for {
 		n, err := src.Read(buf)
 		if n > 0 {
-			st.observe(buf[:n])
-			if _, werr := dst.Write(buf[:n]); werr != nil {
+			st.bytes += int64(n)
+			if werr := st.pass(dst, buf[:n]); werr != nil {
 				return
 			}
 		}
@@ -118,22 +143,65 @@ func (st *stream) copy(dst io.Writer, src io.Reader) {
 	}
 }
 
-func (st *stream) observe(data []byte) {
-	st.bytes += int64(len(data))
+// pass writes on the frames data completes.
+func (st *stream) pass(dst io.Writer, data []byte) error {
 	if st.other {
-		return
+		_, err := dst.Write(data)
+		return err
 	}
-	frames, err := st.reader.Write(data)
-	st.frames += len(frames)
-	if st.s.Trace {
-		for _, f := range frames {
+	st.pending = append(st.pending, data...)
+	var out []byte
+	for {
+		f, n, err := imq.Next(st.pending)
+		if err != nil {
+			st.other = true
+			st.s.logf("#%d %s isn't IMQ framing; passing it through", st.id, st.name)
+			out = append(out, st.pending...)
+			st.pending = nil
+			break
+		}
+		if n == 0 {
+			break
+		}
+		frame := st.pending[:n]
+		st.pending = st.pending[n:]
+		st.frames++
+		if st.s.Trace {
 			st.s.logf("#%d %s %s %dB", st.id, st.name, imq.TypeName(f.Type), f.Size)
 		}
+		out = append(out, st.filtered(f, frame)...)
 	}
+	if len(st.pending) == 0 {
+		st.pending = nil
+	}
+	if len(out) == 0 {
+		return nil
+	}
+	_, err := dst.Write(out)
+	return err
+}
+
+// filtered is the frame to send on: as is, changed by the filter, or none (answered here).
+func (st *stream) filtered(f imq.Frame, frame []byte) []byte {
+	if st.filter == nil || f.Type != imq.C2gSendMessage {
+		return frame
+	}
+	m, err := imq.ParseSendMessage(f.Data)
 	if err != nil {
-		st.other = true
-		st.s.logf("#%d %s isn't IMQ framing; passing it through", st.id, st.name)
+		return frame
 	}
+	out, drop := st.filter.Outgoing(m.Queue, m.Mount, m.Message)
+	if drop {
+		if st.reply != nil {
+			st.reply.Write(imq.EncodeFrame(imq.G2cResult, imq.Result(m.OpID, 0)))
+		}
+		return nil
+	}
+	if bytes.Equal(out, m.Message) {
+		return frame
+	}
+	m.Message = out
+	return imq.EncodeFrame(imq.C2gSendMessage, m.Marshal())
 }
 
 func (st *stream) summary() string {

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
 

internal/proxy/proxy_test.go

@@ -3,6 +3,7 @@ package proxy
 import (
 	"bufio"
 	"bytes"
+	"compress/gzip"
 	"context"
 	"crypto/tls"
 	"encoding/base64"
@@ -18,6 +19,8 @@ import (
 	"sync"
 	"testing"
 	"time"
+
+	"git.t5de.dev/t5de/t5de-proxy/internal/imq"
 )
 
 const token = "secret"
@@ -340,3 +343,134 @@ func TestFallsBackToHTTPWithTheBody(t *testing.T) {
 		t.Fatalf("%d %q", resp.StatusCode, got)
 	}
 }
+
+// suffixFilter appends "!" to messages on mount "messages" and drops those saying "drop".
+type suffixFilter struct{}
+
+func (suffixFilter) Outgoing(queue, mount string, message []byte) ([]byte, bool) {
+	if string(message) == "drop" {
+		return nil, true
+	}
+	if mount == "messages" {
+		return append(bytes.Clone(message), '!'), false
+	}
+	return message, false
+}
+
+func TestConnectFiltersOutgoingMessages(t *testing.T) {
+	gateway := httptest.NewUnstartedServer(nil)
+	gateway.StartTLS()
+	gateway.Close()
+	ln, err := tls.Listen("tcp", "127.0.0.1:0", &tls.Config{Certificates: gateway.TLS.Certificates})
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer ln.Close()
+	received := make(chan []byte, 1)
+	go func() {
+		c, err := ln.Accept()
+		if err != nil {
+			return
+		}
+		defer c.Close()
+		var got []byte
+		buf := make([]byte, 4096)
+		c.SetReadDeadline(time.Now().Add(2 * time.Second))
+		for {
+			n, err := c.Read(buf)
+			got = append(got, buf[:n]...)
+			if err != nil {
+				break
+			}
+		}
+		received <- got
+	}()
+	roots := gateway.Client().Transport.(*http.Transport).TLSClientConfig.RootCAs
+	p, logs := newProxy(t, &Server{TLSConfig: &tls.Config{RootCAs: roots}, IMQFilter: suffixFilter{},
+		Dial: func(ctx context.Context, network, addr string) (net.Conn, error) {
+			return net.Dial(network, ln.Addr().String())
+		}})
+	conn, err := net.Dial("tcp", p.Listener.Addr().String())
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer conn.Close()
+	send := func(op uint64, mount, msg string) []byte {
+		return imq.EncodeFrame(imq.C2gSendMessage, imq.SendMessage{OpID: op, Queue: "/chat/1", Mount: mount, Message: []byte(msg)}.Marshal())
+	}
+	// Sent in pieces, to check frames are put back together.
+	stream := append(append(send(1, "messages", "hi"), send(2, "messages", "drop")...), send(3, "control", "x")...)
+	fmt.Fprintf(conn, "CONNECT example.com:443 HTTP/1.1\r\nProxy-authorization: %s\r\n\r\n", auth)
+	r := bufio.NewReader(conn)
+	for line, _ := r.ReadString('\n'); line != "\r\n"; line, _ = r.ReadString('\n') {
+		if line == "" {
+			t.Fatalf("no response; log %s", logs)
+		}
+	}
+	for i := 0; i < len(stream); i += 5 {
+		conn.Write(stream[i:min(i+5, len(stream))])
+		time.Sleep(time.Millisecond)
+	}
+	// The dropped message is answered by the proxy.
+	conn.SetReadDeadline(time.Now().Add(2 * time.Second))
+	var fromProxy []byte
+	buf := make([]byte, 64)
+	for len(fromProxy) < len(imq.EncodeFrame(imq.G2cResult, imq.Result(2, 0))) {
+		n, err := r.Read(buf)
+		if err != nil {
+			t.Fatalf("%v; log %s", err, logs)
+		}
+		fromProxy = append(fromProxy, buf[:n]...)
+	}
+	if !bytes.Equal(fromProxy, imq.EncodeFrame(imq.G2cResult, imq.Result(2, 0))) {
+		t.Fatalf("client got %x", fromProxy)
+	}
+	conn.Close()
+	got := <-received
+	want := append(send(1, "messages", "hi!"), send(3, "control", "x")...)
+	if !bytes.Equal(got, want) {
+		t.Fatalf("gateway got %q, want %q", got, want)
+	}
+}
+
+// testInterceptor answers /local itself, uppercases request bodies to /change, and appends to their responses.
+type testInterceptor struct{}
+
+func (testInterceptor) Intercept(x *Exchange) error {
+	switch x.URL.Path {
+	case "/local":
+		x.Response = LocalResponse(200, "text/plain", []byte("from the proxy"))
+	case "/change":
+		x.Body = bytes.ToUpper(x.Body)
+		x.ModifyResponse = func(resp *http.Response, body []byte) ([]byte, error) {
+			return append(body, " (changed)"...), nil
+		}
+	}
+	return nil
+}
+
+func TestInterceptor(t *testing.T) {
+	var gotBody string
+	origin := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.URL.Path == "/local" {
+			t.Error("a request answered by the proxy reached the server")
+		}
+		b, _ := io.ReadAll(r.Body)
+		gotBody = string(b)
+		// Compressed, as IMVU's servers answer: the interceptor sees it decompressed.
+		w.Header().Set("Content-Encoding", "gzip")
+		zw := gzip.NewWriter(w)
+		io.WriteString(zw, "echo "+gotBody)
+		zw.Close()
+	}))
+	defer origin.Close()
+	p, _ := newProxy(t, &Server{Interceptor: testInterceptor{}})
+	resp := send(t, p, "GET "+origin.URL+"/local HTTP/1.1\r\nHost: x\r\nProxy-Authorization: "+auth+"\r\n\r\n")
+	if b := body(t, resp); b != "from the proxy" {
+		t.Fatalf("local: %q", b)
+	}
+	resp = send(t, p, "POST "+origin.URL+"/change HTTP/1.1\r\nHost: x\r\nProxy-Authorization: "+auth+"\r\nContent-Length: 5\r\n\r\nhello")
+	if b := body(t, resp); b != "echo HELLO (changed)" || gotBody != "HELLO" || resp.Header.Get("Content-Encoding") != "" {
+		t.Fatalf("changed: %q, server got %q, %v", b, gotBody, resp.Header)
+	}
+}