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:
1 parent
2551888ed6
commit
2190230f81
6 files changed
+485
-24
No files matched your search
@@ -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
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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
@@ -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 {
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user