feat: refine hijack behavior for response lifecycle and add tests

- Clarify the error message for attempted hijack after response body data is written
- Modify hijack behavior: allow hijacking after headers are written (for better websocket compatibility), but block hijacking after any body data is sent
- Add comprehensive tests to validate allowed hijack after header write and disallowed hijack after body write

fix https://github.com/gin-gonic/gin/issues/4372

Signed-off-by: appleboy <appleboy.tw@gmail.com>
This commit is contained in:
appleboy 2025-09-24 23:19:08 +08:00
parent 048f6fb884
commit 92ce464992
No known key found for this signature in database
2 changed files with 62 additions and 2 deletions

View File

@ -17,7 +17,7 @@ const (
defaultStatus = http.StatusOK defaultStatus = http.StatusOK
) )
var errHijackAlreadyWritten = errors.New("gin: response already written") var errHijackAlreadyWritten = errors.New("gin: response body already written")
// ResponseWriter ... // ResponseWriter ...
type ResponseWriter interface { type ResponseWriter interface {
@ -109,7 +109,9 @@ func (w *responseWriter) Written() bool {
// Hijack implements the http.Hijacker interface. // Hijack implements the http.Hijacker interface.
func (w *responseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) { func (w *responseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
if w.Written() { // Allow hijacking after headers are written (size == 0), but not after body data (size > 0)
// For compatibility with websocket libraries (e.g., github.com/coder/websocket)
if w.size > 0 {
return nil, nil, errHijackAlreadyWritten return nil, nil, errHijackAlreadyWritten
} }
if w.size < 0 { if w.size < 0 {

View File

@ -194,6 +194,64 @@ func TestResponseWriterHijackAfterWrite(t *testing.T) {
} }
} }
// Test: WebSocket compatibility - allow hijack after WriteHeaderNow(), but block after body data.
func TestResponseWriterHijackAfterWriteHeaderNow(t *testing.T) {
tests := []struct {
name string
action func(w ResponseWriter) error
expectWrittenBeforeHijack bool
expectHijackSuccess bool
expectWrittenAfterHijack bool
expectError error
}{
{
name: "hijack after WriteHeaderNow only should succeed (websocket pattern)",
action: func(w ResponseWriter) error {
w.WriteHeaderNow() // Simulate websocket.Accept() behavior
return nil
},
expectWrittenBeforeHijack: true,
expectHijackSuccess: true, // NEW BEHAVIOR: allow hijack after just header write
expectWrittenAfterHijack: true,
expectError: nil,
},
{
name: "hijack after WriteHeaderNow + Write should fail",
action: func(w ResponseWriter) error {
w.WriteHeaderNow()
_, err := w.Write([]byte("test"))
return err
},
expectWrittenBeforeHijack: true,
expectHijackSuccess: false,
expectWrittenAfterHijack: true,
expectError: errHijackAlreadyWritten,
},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
hijacker := &mockHijacker{ResponseRecorder: httptest.NewRecorder()}
writer := &responseWriter{}
writer.reset(hijacker)
w := ResponseWriter(writer)
require.NoError(t, tc.action(w), "unexpected error during pre-hijack action")
assert.Equal(t, tc.expectWrittenBeforeHijack, w.Written(), "unexpected w.Written() state before hijack")
_, _, hijackErr := w.Hijack()
if tc.expectError == nil {
assert.NoError(t, hijackErr, "expected hijack to succeed")
} else {
assert.ErrorIs(t, hijackErr, tc.expectError, "unexpected error from Hijack()")
}
assert.Equal(t, tc.expectHijackSuccess, hijacker.hijacked, "unexpected hijacker.hijacked state")
assert.Equal(t, tc.expectWrittenAfterHijack, w.Written(), "unexpected w.Written() state after hijack")
})
}
}
func TestResponseWriterFlush(t *testing.T) { func TestResponseWriterFlush(t *testing.T) {
testServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { testServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
writer := &responseWriter{} writer := &responseWriter{}