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

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

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