Support websocket frames watcher in interceptor

Author: josexy

README.md

@@ -88,12 +88,20 @@ handler, err := mitmpgo.NewMitmProxyHandler(
 ### With WebSocket Interceptor
 
 ```go
-websocketInterceptor := func(ctx context.Context, dir metadata.WSDirection, msgType int, b *buf.Buffer, req *http.Request, invoker mitmpgo.WebsocketDelegatedInvoker) error {
+websocketInterceptor := func(ctx context.Context, req *http.Request, rsp *http.Response, fw mitmpgo.WebsocketFramesWatcher) {
     // Log WebSocket messages
-    fmt.Printf("WS [%s] %s: %d bytes\n", dir, req.URL, b.Len())
-
-    // Forward the message
-    return invoker.Invoke(msgType, b)
+    log.Printf("WS url: %s", req.URL.String())
+
+    for frame := range fw.GetFrame() {
+        dir := frame.Direction()
+        msgType := frame.MessageType()
+        dataBuf := frame.DataBuffer()
+        log.Printf("---> %s %d %s", dir, msgType, dataBuf.String())
+        if err := frame.Invoke(); err != nil {
+            log.Printf("frame invoke error: %v", err)
+        }
+        frame.Release()
+    }
 }
 
 handler, err := mitmpgo.NewMitmProxyHandler(
@@ -163,6 +171,9 @@ mitmpgo.WithCertCachePool(2048, 30, 15)
 mitmpgo.WithDialer(&net.Dialer{
     Timeout: 30 * time.Second,
 })
+
+// Maximum channel size of WebSocket frames
+mitmpgo.WithMaxWebsocketFramesPerForward(4096)
 ```
 
 ### Interceptor Options
@@ -177,9 +188,6 @@ mitmpgo.WithWebsocketInterceptor(websocketInterceptor)
 // Chain multiple HTTP interceptors (executed in order)
 mitmpgo.WithChainHTTPInterceptor(interceptor1, interceptor2, interceptor3)
 
-// Chain multiple WebSocket interceptors (executed in order)
-mitmpgo.WithChainWebsocketInterceptor(wsInterceptor1, wsInterceptor2)
-
 // Set error handler
 mitmpgo.WithErrorHandler(func(ec mitmpgo.ErrorContext) {
     log.Printf("Error: %v", ec.Error)

buf/buffer.go

@@ -339,6 +339,9 @@ func (b *Buffer) String() string {
 	if b == nil {
 		return "<nil>"
 	}
+	if b.Len() == 0 {
+		return ""
+	}
 	return unsafe.String(&b.Bytes()[0], b.Len())
 }
 

bufconn.go

@@ -5,8 +5,8 @@ import (
 	"io"
 	"net"
 
-	"github.com/gorilla/websocket"
 	"github.com/josexy/mitmpgo/buf"
+	"github.com/josexy/websocket"
 )
 
 type bufConn struct {

examples/chain-interceptors/main.go

@@ -7,10 +7,8 @@ import (
 	"log/slog"
 	"net/http"
 	"os"
-	"time"
 
 	"github.com/josexy/mitmpgo"
-	"github.com/josexy/mitmpgo/buf"
 )
 
 func main() {
@@ -29,7 +27,6 @@ func main() {
 		mitmpgo.WithCACertPath(caCertPath),
 		mitmpgo.WithCAKeyPath(caKeyPath),
 		mitmpgo.WithChainHTTPInterceptor(httpInterceptor1, httpInterceptor2, httpInterceptor3),
-		mitmpgo.WithChainWebsocketInterceptor(websocketInterceptor1, websocketInterceptor2),
 	)
 	if err != nil {
 		panic(err)
@@ -68,26 +65,3 @@ func httpInterceptor3(ctx context.Context, req *http.Request, invoker mitmpgo.HT
 	slog.Debug("httpInterceptor3 after", slog.String("status", rsp.Status), slog.String("protocol", rsp.Proto))
 	return rsp, err
 }
-
-func websocketInterceptor1(ctx context.Context, dir mitmpgo.WSDirection, msgType int, b *buf.Buffer, req *http.Request, invoker mitmpgo.WebsocketDelegatedInvoker) error {
-	slog.Debug("websocketInterceptor1 before", slog.String("dir", dir.String()), slog.Int("msgType", msgType), slog.Int("len", b.Len()))
-	if dir == mitmpgo.Send {
-		b.WriteString("->" + time.Now().Format(time.DateTime))
-	}
-	err := invoker.Invoke(msgType, b)
-	if err != nil {
-		return err
-	}
-	slog.Debug("websocketInterceptor1 after")
-	return nil
-}
-
-func websocketInterceptor2(ctx context.Context, dir mitmpgo.WSDirection, msgType int, b *buf.Buffer, req *http.Request, invoker mitmpgo.WebsocketDelegatedInvoker) error {
-	slog.Debug("websocketInterceptor2 before", slog.String("dir", dir.String()), slog.Int("msgType", msgType), slog.Int("len", b.Len()))
-	err := invoker.Invoke(msgType, b)
-	if err != nil {
-		return err
-	}
-	slog.Debug("websocketInterceptor2 after")
-	return err
-}

examples/dumper/main.go

@@ -13,14 +13,14 @@ import (
 	"log/slog"
 	"net"
 	"net/http"
+	"net/http/httputil"
 	"os"
 	"os/signal"
 	"strings"
 	"syscall"
 	"time"
 
 	"github.com/josexy/mitmpgo"
-	"github.com/josexy/mitmpgo/buf"
 	"github.com/josexy/mitmpgo/metadata"
 )
 
@@ -37,6 +37,9 @@ type bodyDecoder struct {
 }
 
 func newBodyDecoder(r io.ReadCloser, encoding string, chunkType int) (io.ReadCloser, error) {
+	if r == http.NoBody { // no body and no need to replace it
+		return r, nil
+	}
 	if encoding == "" {
 		return newChunkBodyReader(r, CHUNK_SIZE, chunkType), nil
 	}
@@ -102,7 +105,7 @@ func (r *chunkBodyReader) Read(p []byte) (n int, err error) {
 		// fmt.Printf("--> hex dump(chunk size/data size: %d/%d):\n%s\n", r.N, n, hex.Dump(p[:n]))
 	}
 	if err == io.EOF {
-		fmt.Printf("<<-- [%d]full data dump (%d bytes):\n%s\n", r.chunkType, r.buf.Len(), r.buf.Bytes())
+		fmt.Printf("<<-- [%d]full data dump (%d bytes):\n", r.chunkType, r.buf.Len())
 	}
 	return
 }
@@ -147,6 +150,7 @@ func main() {
 		// mitmpgo.WithDisableProxy(),
 		// mitmpgo.WithDisableHTTP2(),
 		// mitmpgo.WithSkipVerifySSLFromServer(),
+		// mitmpgo.WithMaxWebsocketFramesPerForward(4096),
 	)
 	if err != nil {
 		panic(err)
@@ -207,7 +211,6 @@ func httpInterceptor(ctx context.Context, req *http.Request, invoker mitmpgo.HTT
 		slog.String("host", req.Host),
 		slog.String("proto", req.Proto),
 		slog.String("method", req.Method),
-		slog.Bool("tls", req.TLS != nil),
 		slog.String("url", req.URL.String()),
 		slog.Any("headers", map[string][]string(req.Header)),
 	)
@@ -257,17 +260,28 @@ func httpInterceptor(ctx context.Context, req *http.Request, invoker mitmpgo.HTT
 	return rsp, err
 }
 
-func websocketInterceptor(ctx context.Context, dir mitmpgo.WSDirection, msgType int, b *buf.Buffer, req *http.Request, wdi mitmpgo.WebsocketDelegatedInvoker) error {
+func websocketInterceptor(ctx context.Context, req *http.Request, rsp *http.Response, fw mitmpgo.WebsocketFramesWatcher) {
 	_md, _ := metadata.FromContext(ctx)
 	md := _md.MD()
-	slog.Debug("websocket",
+	slog.Debug("request",
+		slog.Bool("stream_body", md.StreamBody),
 		slog.String("source", md.SourceAddr.String()),
 		slog.String("destination", md.DestinationAddr.String()),
 		slog.String("hostport", md.RequestHostport),
-		slog.String("uri", req.URL.String()),
-		slog.String("direction", dir.String()),
-		slog.Int("msg_type", msgType),
+		slog.String("host", req.Host),
+		slog.String("proto", req.Proto),
+		slog.String("method", req.Method),
+		slog.String("url", req.URL.String()),
+		slog.Int("status_code", rsp.StatusCode),
+		slog.Any("request_headers", map[string][]string(req.Header)),
+		slog.Any("response_headers", map[string][]string(rsp.Header)),
 	)
+
+	data, _ := httputil.DumpRequest(req, false)
+	fmt.Printf("%s\n", string(data))
+	data, _ = httputil.DumpResponse(rsp, false)
+	fmt.Printf("%s\n", string(data))
+
 	if md.TLSState != nil {
 		slog.Debug("tls state",
 			slog.String("server_name", md.TLSState.ServerName),
@@ -292,10 +306,25 @@ func websocketInterceptor(ctx context.Context, dir mitmpgo.WSDirection, msgType
 			slog.Any("ip", md.ServerCertificate.IPAddresses),
 		)
 	}
-	// if md.Direction == metadata.Receive {
-	// 	b.WriteString(time.Now().String())
-	// }
-	return wdi.Invoke(msgType, b)
+
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case frame, ok := <-fw.GetFrame():
+			if !ok {
+				return
+			}
+			dir := frame.Direction()
+			msgType := frame.MessageType()
+			dataBuf := frame.DataBuffer()
+			fmt.Printf("---> %s %d %s\n", dir, msgType, dataBuf.String())
+			if err := frame.Invoke(); err != nil {
+				slog.Error("frame invoke error", slog.String("error", err.Error()))
+			}
+			frame.Release()
+		}
+	}
 }
 
 func getDecodedReader(r io.Reader, encoding string) (io.ReadCloser, error) {

go.mod

@@ -3,8 +3,8 @@ module github.com/josexy/mitmpgo
 go 1.25.4
 
 require (
-	github.com/gorilla/websocket v1.5.3
-	golang.org/x/net v0.48.0
+	github.com/josexy/websocket v0.0.0-20260219083038-11b2ba10886b
+	golang.org/x/net v0.50.0
 )
 
-require golang.org/x/text v0.32.0 // indirect
+require golang.org/x/text v0.34.0 // indirect

go.sum

@@ -1,6 +1,6 @@
-github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg=
-github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
-golang.org/x/net v0.48.0 h1:zyQRTTrjc33Lhh0fBgT/H3oZq9WuvRR5gPC70xpDiQU=
-golang.org/x/net v0.48.0/go.mod h1:+ndRgGjkh8FGtu1w1FGbEC31if4VrNVMuKTgcAAnQRY=
-golang.org/x/text v0.32.0 h1:ZD01bjUt1FQ9WJ0ClOL5vxgxOI/sVCNgX1YtKwcY0mU=
-golang.org/x/text v0.32.0/go.mod h1:o/rUWzghvpD5TXrTIBuJU77MTaN0ljMWE47kxGJQ7jY=
+github.com/josexy/websocket v0.0.0-20260219083038-11b2ba10886b h1:ahmk86ulmaZRnXZFfLWPahG0FCcBw6Bco8NML+4UCNU=
+github.com/josexy/websocket v0.0.0-20260219083038-11b2ba10886b/go.mod h1:E1y5c7BPj8laYUGMExE3snEx4mLawJGeVztAKQQeNcg=
+golang.org/x/net v0.50.0 h1:ucWh9eiCGyDR3vtzso0WMQinm2Dnt8cFMuQa9K33J60=
+golang.org/x/net v0.50.0/go.mod h1:UgoSli3F/pBgdJBHCTc+tp3gmrU4XswgGRgtnwWTfyM=
+golang.org/x/text v0.34.0 h1:oL/Qq0Kdaqxa1KbNeMKwQq0reLCCaFtqu2eNuSeNHbk=
+golang.org/x/text v0.34.0/go.mod h1:homfLqTYRFyVYemLBFl5GgL/DWEiH5wcsQ5gSh1yziA=

interceptor.go

@@ -28,6 +28,17 @@ func (d WSDirection) String() string {
 	}
 }
 
+type WsFrame interface {
+	Direction() WSDirection
+	MessageType() int
+	DataBuffer() *buf.Buffer
+
+	// Forward the websocket message and release the data buffer
+	Invoke() error
+	// MUST be called to release the data buffer
+	Release()
+}
+
 type (
 	HTTPDelegatedInvoker interface {
 		Invoke(request *http.Request) (*http.Response, error)
@@ -36,14 +47,17 @@ type (
 	WebsocketDelegatedInvoker interface {
 		Invoke(msgType int, dataPtr *buf.Buffer) error
 	}
+	WebsocketFramesWatcher interface {
+		GetFrame() <-chan WsFrame
+	}
 )
 
 type (
 	HTTPDelegatedInvokerFunc      func(*http.Request) (*http.Response, error)
 	WebsocketDelegatedInvokerFunc func(int, *buf.Buffer) error
 
 	HTTPInterceptor      func(context.Context, *http.Request, HTTPDelegatedInvoker) (*http.Response, error)
-	WebsocketInterceptor func(context.Context, WSDirection, int, *buf.Buffer, *http.Request, WebsocketDelegatedInvoker) error
+	WebsocketInterceptor func(context.Context, *http.Request, *http.Response, WebsocketFramesWatcher)
 )
 
 func (f HTTPDelegatedInvokerFunc) Invoke(r *http.Request) (*http.Response, error) { return f(r) }
@@ -69,18 +83,3 @@ func getChainHTTPInterceptor(interceptors []HTTPInterceptor, curr int, ctx conte
 		return interceptors[curr+1](ctx, r, getChainHTTPInterceptor(interceptors, curr+1, ctx, finalInvoker))
 	})
 }
-
-func chainWebsocketInterceptors(interceptors []WebsocketInterceptor) WebsocketInterceptor {
-	return func(ctx context.Context, d WSDirection, i int, b *buf.Buffer, r *http.Request, wdi WebsocketDelegatedInvoker) error {
-		return interceptors[0](ctx, d, i, b, r, getChainWebsocketInterceptor(interceptors, 0, ctx, d, r, wdi))
-	}
-}
-
-func getChainWebsocketInterceptor(interceptors []WebsocketInterceptor, curr int, ctx context.Context, dir WSDirection, req *http.Request, finalInvoker WebsocketDelegatedInvoker) WebsocketDelegatedInvoker {
-	if curr == len(interceptors)-1 {
-		return finalInvoker
-	}
-	return WebsocketDelegatedInvokerFunc(func(i int, b *buf.Buffer) error {
-		return interceptors[curr+1](ctx, dir, i, b, req, getChainWebsocketInterceptor(interceptors, curr+1, ctx, dir, req, finalInvoker))
-	})
-}

option.go

@@ -34,6 +34,8 @@ type options struct {
 	rootCAs       []string    // Paths to additional root CA certificate files
 	dialer        *net.Dialer // Custom dialer for outbound connections
 
+	wsMaxFramesPerForward int // Max frames channel size per single websocket forward
+
 	clientCerts map[string]ClientCert // Client certificate configuration
 
 	// Certificate cache pool configuration
@@ -50,14 +52,14 @@ type options struct {
 	httpInt       HTTPInterceptor
 	wsInt         WebsocketInterceptor
 	chainHttpInts []HTTPInterceptor
-	chainWsInts   []WebsocketInterceptor
 }
 
 // newOptions creates a new options instance with default values.
 // Default dialer timeout is 15 seconds.
 func newOptions(opt ...Option) *options {
 	options := &options{
-		dialer: &net.Dialer{Timeout: 15 * time.Second},
+		dialer:                &net.Dialer{Timeout: 15 * time.Second},
+		wsMaxFramesPerForward: 2048,
 	}
 	for _, o := range opt {
 		o.apply(options)
@@ -258,6 +260,22 @@ func WithCertCachePool(capacity, intervalSecond, expireSecond int) Option {
 	})
 }
 
+// WithMaxWebsocketFramesPerForward specifies the maximum channel size of frames that can be buffered
+// per single websocket forward.
+//
+// If not specified, default value(2048) is used.
+//
+// Example:
+//
+//	handler, err := NewMitmProxyHandler(
+//	    WithMaxWebsocketFramesPerForward(2048),
+//	)
+func WithMaxWebsocketFramesPerForward(maxFrames int) Option {
+	return OptionFunc(func(o *options) {
+		o.wsMaxFramesPerForward = maxFrames
+	})
+}
+
 // WithIncludeHosts specifies a whitelist of hosts that should be intercepted.
 // Only traffic to these hosts will be intercepted; all other traffic will pass through
 // without interception (passthrough mode).
@@ -436,32 +454,3 @@ func WithChainHTTPInterceptor(interceptors ...HTTPInterceptor) Option {
 		o.chainHttpInts = append(o.chainHttpInts, interceptors...)
 	})
 }
-
-// WithChainWebsocketInterceptor chains multiple WebSocket interceptors together.
-// Interceptors are executed in the order they are provided, forming a middleware chain.
-// Each interceptor can inspect/modify the WebSocket message, call the next interceptor,
-// and handle the message forwarding. And The final interceptor will forwards the message.
-//
-// Example:
-//
-//	loggingInterceptor := func(ctx context.Context, dir metadata.WSDirection, msgType int, data *buf.Buffer, req *http.Request, invoker WebsocketDelegatedInvoker) error {
-//	    log.Printf("[%s] Message: type=%d, size=%d", dir, msgType, data.Len())
-//	    return invoker.Invoke(msgType, data)
-//	}
-//
-//	filterInterceptor := func(ctx context.Context, dir metadata.WSDirection, msgType int, data *buf.Buffer, req *http.Request, invoker WebsocketDelegatedInvoker) error {
-//	    // Drop ping messages
-//	    if msgType == websocket.PingMessage {
-//	        return nil
-//	    }
-//	    return invoker.Invoke(msgType, data)
-//	}
-//
-//	handler, err := NewMitmProxyHandler(
-//	    WithChainWebsocketInterceptor(loggingInterceptor, filterInterceptor),
-//	)
-func WithChainWebsocketInterceptor(interceptors ...WebsocketInterceptor) Option {
-	return OptionFunc(func(o *options) {
-		o.chainWsInts = append(o.chainWsInts, interceptors...)
-	})
-}