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

Co-Authored-By: Claude Opus 5.5 <[email protected]>
This commit is contained in:
dhkatzandClaude Opus 5.5 committed 2026-10-29 02:47:48 -07:00
1 parent 2551888ed6
commit 2190230f81
6 files changed
+485 -24

No files matched your search

+9 -1
View File
@@ -40,10 +40,11 @@ func TypeName(t uint64) string {
return fmt.Sprintf("type %d", t)
}
// Frame is one message: its type and the size of its data.
// Frame is one message: its type and data (Size is len(Data)).
type Frame struct {
Type uint64
Size int
Data []byte
}
// ErrMalformed means the stream isn't IMQ framing; a Reader stops decoding after it.
@@ -76,6 +77,12 @@ func (r *Reader) Write(data []byte) ([]Frame, error) {
}
}
// Next reads the frame at the start of b: the frame and its length in b, a length of 0 if b holds only part of it,
// or ErrMalformed.
func Next(b []byte) (Frame, int, error) {
return parse(b)
}
// maxFrame bounds a frame's size, so a stream that isn't IMQ isn't buffered without end.
const maxFrame = 16 << 20
@@ -116,6 +123,7 @@ func parse(b []byte) (Frame, int, error) {
return f, 0, nil
}
f.Size = int(size)
f.Data = b[i : i+int(size) : i+int(size)]
return f, i + int(size), nil
default:
return f, 0, ErrMalformed
+20 -2
View File
@@ -36,12 +36,12 @@ func TestReaderSplitsFrames(t *testing.T) {
}
got = append(got, frames...)
}
want := []Frame{{1302, 5}, {3103, 300}, {1301, 0}}
want := []Frame{{1302, 5, []byte("hello")}, {3103, 300, bytes.Repeat([]byte{1}, 300)}, {1301, 0, nil}}
if len(got) != len(want) {
t.Fatalf("got %v, want %v", got, want)
}
for i := range want {
if got[i] != want[i] {
if got[i].Type != want[i].Type || got[i].Size != want[i].Size || !bytes.Equal(got[i].Data, want[i].Data) {
t.Fatalf("got %v, want %v", got, want)
}
}
@@ -63,3 +63,21 @@ func TestTypeName(t *testing.T) {
t.Fatal(TypeName(3104), TypeName(7))
}
}
func TestSendMessageRoundTrip(t *testing.T) {
m := SendMessage{OpID: 300, Queue: "/chat/1", Mount: "messages", Message: []byte(`{"message": "*use 80 1"}`)}
data := m.Marshal()
got, err := ParseSendMessage(data)
if err != nil || got.OpID != 300 || got.Queue != "/chat/1" || got.Mount != "messages" || string(got.Message) != string(m.Message) {
t.Fatalf("%+v %v", got, err)
}
// A frame of it reads back as one frame with the same data.
var r Reader
frames, err := r.Write(EncodeFrame(C2gSendMessage, data))
if err != nil || len(frames) != 1 || frames[0].Type != C2gSendMessage || !bytes.Equal(frames[0].Data, data) {
t.Fatalf("%v %v", frames, err)
}
if _, err := ParseSendMessage(Result(1, 0)); err == nil {
t.Fatal("a G2cResult parsed as a C2gSendMessage")
}
}
+120
View File
@@ -0,0 +1,120 @@
package imq
import "errors"
// Message types the proxy reads or writes.
const (
C2gSendMessage = 1304
G2cResult = 3102
)
// SendMessage is a C2gSendMessage: a message the client sends to a queue's mount.
type SendMessage struct {
OpID uint64
Queue string
Mount string
Message []byte
}
// field is one protobuf field: its number and either a varint or bytes.
type field struct {
num uint64
varint uint64
bytes []byte
isLen bool
}
func fields(b []byte) ([]field, error) {
var out []field
for i := 0; i < len(b); {
key, n := varint(b[i:])
if n <= 0 {
return nil, ErrMalformed
}
i += n
f := field{num: key >> 3}
switch key & 7 {
case 0:
v, n := varint(b[i:])
if n <= 0 {
return nil, ErrMalformed
}
f.varint, i = v, i+n
case 2:
l, n := varint(b[i:])
if n <= 0 || l > uint64(len(b)-i-n) {
return nil, ErrMalformed
}
i += n
f.bytes, f.isLen, i = b[i:i+int(l)], true, i+int(l)
default:
return nil, ErrMalformed
}
out = append(out, f)
}
return out, nil
}
// ParseSendMessage reads a C2gSendMessage's data.
func ParseSendMessage(data []byte) (SendMessage, error) {
fs, err := fields(data)
if err != nil {
return SendMessage{}, err
}
var m SendMessage
seen := 0
for _, f := range fs {
switch {
case f.num == 1 && !f.isLen:
m.OpID, seen = f.varint, seen|1
case f.num == 2 && f.isLen:
m.Queue, seen = string(f.bytes), seen|2
case f.num == 3 && f.isLen:
m.Mount, seen = string(f.bytes), seen|4
case f.num == 4 && f.isLen:
m.Message, seen = f.bytes, seen|8
default:
return SendMessage{}, errors.New("imq: unexpected field in C2gSendMessage")
}
}
if seen != 15 {
return SendMessage{}, errors.New("imq: C2gSendMessage is missing fields")
}
return m, nil
}
// Marshal encodes the message's data, fields in the client's order.
func (m SendMessage) Marshal() []byte {
var b []byte
b = appendVarintField(b, 1, m.OpID)
b = appendBytesField(b, 2, []byte(m.Queue))
b = appendBytesField(b, 3, []byte(m.Mount))
return appendBytesField(b, 4, m.Message)
}
// Result encodes a G2cResult's data: the answer to the operation opID (status 0 is success).
func Result(opID, status uint64) []byte {
return appendVarintField(appendVarintField(nil, 1, opID), 2, status)
}
// EncodeFrame wraps a message's data in a frame of the given type.
func EncodeFrame(t uint64, data []byte) []byte {
return appendBytesField(appendVarintField(nil, 1, t), 2, data)
}
func appendVarint(b []byte, v uint64) []byte {
for v >= 0x80 {
b = append(b, byte(v)|0x80)
v >>= 7
}
return append(b, byte(v))
}
func appendVarintField(b []byte, num, v uint64) []byte {
return appendVarint(appendVarint(b, num<<3), v)
}
func appendBytesField(b []byte, num uint64, v []byte) []byte {
b = appendVarint(appendVarint(b, num<<3|2), uint64(len(v)))
return append(b, v...)
}
+89 -21
View File
@@ -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 {
+113
View File
@@ -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
+134
View File
@@ -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)
}
}