From 6c1f8a8a379ae64ebbe70400eb1d389ec49a16e2 Mon Sep 17 00:00:00 2001
From: Denozordec <shatskiy.blog@gmail.com>
Date: Tue, 31 Mar 2026 09:57:15 +0700
Subject: [PATCH] Add WebSocket support methods to statusWriter in gateway.go

- Implemented Hijack method to preserve WebSocket upgrade support through middleware.
- Added ReadFrom method to maintain fast-path behavior for io.Copy when the underlying writer supports it.
- Introduced Unwrap method to allow access to the original ResponseWriter, enhancing flexibility in response handling.
---
 internal/server/gateway.go | 23 +++++++++++++++++++++++
 1 file changed, 23 insertions(+)

diff --git a/internal/server/gateway.go b/internal/server/gateway.go
index 280b10f..8a00c44 100644
--- a/internal/server/gateway.go
+++ b/internal/server/gateway.go
@@ -1,9 +1,11 @@
 package server
 
 import (
+	"bufio"
 	"context"
 	"encoding/json"
 	"fmt"
+	"io"
 	"log/slog"
 	"net"
 	"net/http"
@@ -403,6 +405,27 @@ func (s *statusWriter) Flush() {
 	}
 }
 
+// Hijack preserves WebSocket upgrade support through middleware wrappers.
+func (s *statusWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	if h, ok := s.ResponseWriter.(http.Hijacker); ok {
+		return h.Hijack()
+	}
+	return nil, nil, http.ErrNotSupported
+}
+
+// ReadFrom keeps io.Copy fast-path behavior when the underlying writer supports it.
+func (s *statusWriter) ReadFrom(r io.Reader) (int64, error) {
+	if rf, ok := s.ResponseWriter.(io.ReaderFrom); ok {
+		return rf.ReadFrom(r)
+	}
+	return io.Copy(s.ResponseWriter, r)
+}
+
+// Unwrap lets net/http.ResponseController reach the original writer.
+func (s *statusWriter) Unwrap() http.ResponseWriter {
+	return s.ResponseWriter
+}
+
 // Shutdown idle connections on the shared transport.
 func (g *Gateway) Shutdown(ctx context.Context) error {
 	g.transport.CloseIdleConnections()