| 138 | } |
| 139 | |
| 140 | func TestResponseWriterHijackAfterWrite(t *testing.T) { |
| 141 | tests := []struct { |
| 142 | name string |
| 143 | action func(w ResponseWriter) error // Action to perform before hijacking |
| 144 | expectWrittenBeforeHijack bool |
| 145 | expectHijackSuccess bool |
| 146 | expectWrittenAfterHijack bool |
| 147 | expectError error |
| 148 | }{ |
| 149 | { |
| 150 | name: "hijack before write should succeed", |
| 151 | action: func(w ResponseWriter) error { return nil }, |
| 152 | expectWrittenBeforeHijack: false, |
| 153 | expectHijackSuccess: true, |
| 154 | expectWrittenAfterHijack: true, // Hijack itself marks the writer as written |
| 155 | expectError: nil, |
| 156 | }, |
| 157 | { |
| 158 | name: "hijack after write should fail", |
| 159 | action: func(w ResponseWriter) error { |
| 160 | _, err := w.Write([]byte("test")) |
| 161 | return err |
| 162 | }, |
| 163 | expectWrittenBeforeHijack: true, |
| 164 | expectHijackSuccess: false, |
| 165 | expectWrittenAfterHijack: true, |
| 166 | expectError: errHijackAlreadyWritten, |
| 167 | }, |
| 168 | } |
| 169 | |
| 170 | for _, tc := range tests { |
| 171 | t.Run(tc.name, func(t *testing.T) { |
| 172 | hijacker := &mockHijacker{ResponseRecorder: httptest.NewRecorder()} |
| 173 | writer := &responseWriter{} |
| 174 | writer.reset(hijacker) |
| 175 | w := ResponseWriter(writer) |
| 176 | |
| 177 | // Check initial state |
| 178 | assert.False(t, w.Written(), "should not be written initially") |
| 179 | |
| 180 | // Perform pre-hijack action |
| 181 | require.NoError(t, tc.action(w), "unexpected error during pre-hijack action") |
| 182 | |
| 183 | // Check state before hijacking |
| 184 | assert.Equal(t, tc.expectWrittenBeforeHijack, w.Written(), "unexpected w.Written() state before hijack") |
| 185 | |
| 186 | // Attempt to hijack |
| 187 | _, _, hijackErr := w.Hijack() |
| 188 | |
| 189 | // Check results |
| 190 | require.ErrorIs(t, hijackErr, tc.expectError, "unexpected error from Hijack()") |
| 191 | assert.Equal(t, tc.expectHijackSuccess, hijacker.hijacked, "unexpected hijacker.hijacked state") |
| 192 | assert.Equal(t, tc.expectWrittenAfterHijack, w.Written(), "unexpected w.Written() state after hijack") |
| 193 | }) |
| 194 | } |
| 195 | } |
| 196 | |
| 197 | // Test: WebSocket compatibility - allow hijack after WriteHeaderNow(), but block after body data. |