소스 검색

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

Kevin Wan 2 년 전
부모
커밋
9f42eda9ff
2개의 변경된 파일24개의 추가작업 그리고 0개의 파일을 삭제
  1. 6 0
      rest/handler/timeouthandler.go
  2. 18 0
      rest/handler/timeouthandler_test.go

+ 6 - 0
rest/handler/timeouthandler.go

@@ -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 - 0
rest/handler/timeouthandler_test.go

@@ -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{},