github_mirrors/go-zero
1
0
Fork 0
mirror of https://github.com/zeromicro/go-zero.git synced 2025-02-02 16:28:39 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix: timeout handler not implementing http.Flusher (#3225)

Browse Source
This commit is contained in:
Kevin Wan 2023-05-08 18:07:02 +08:00 committed by GitHub
parent 8762a3b7ba
commit 9f42eda9ff
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 24 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
rest/handler/timeouthandler.go
View File

@ -127,6 +127,12 @@ type timeoutWriter struct {
var _ http.Pusher = (*timeoutWriter)(nil)
func (tw *timeoutWriter) Flush() {
if flusher, ok := tw.w.(http.Flusher); ok {
flusher.Flush()
}
}
func (tw *timeoutWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
if hijacked, ok := tw.w.(http.Hijacker); ok {
return hijacked.Hijack()

18
rest/handler/timeouthandler_test.go
View File

@ -164,6 +164,24 @@ func TestTimeoutHijack(t *testing.T) {
})
}
func TestTimeoutFlush(t *testing.T) {
timeoutHandler := TimeoutHandler(time.Minute)
handler := timeoutHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
flusher, ok := w.(http.Flusher)
if !ok {
http.Error(w, "Streaming unsupported!", http.StatusInternalServerError)
return
}
flusher.Flush()
}))
req := httptest.NewRequest(http.MethodGet, "http://localhost", http.NoBody)
resp := httptest.NewRecorder()
handler.ServeHTTP(resp, req)
assert.Equal(t, http.StatusOK, resp.Code)
}
func TestTimeoutPusher(t *testing.T) {
handler := &timeoutWriter{
w: mockedPusher{},