From 9f42eda9ffefbc1e6b3183bc8a1efc94f7cc2b71 Mon Sep 17 00:00:00 2001
From: Kevin Wan <wanjunfeng@gmail.com>
Date: Mon, 8 May 2023 18:07:02 +0800
Subject: [PATCH] fix: timeout handler not implementing http.Flusher (#3225)

---
 rest/handler/timeouthandler.go      |  6 ++++++
 rest/handler/timeouthandler_test.go | 18 ++++++++++++++++++
 2 files changed, 24 insertions(+)

diff --git a/rest/handler/timeouthandler.go b/rest/handler/timeouthandler.go
index ec34180a..52d82a80 100644
--- a/rest/handler/timeouthandler.go
+++ b/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()
diff --git a/rest/handler/timeouthandler_test.go b/rest/handler/timeouthandler_test.go
index 3fd8bdd2..acbda030 100644
--- a/rest/handler/timeouthandler_test.go
+++ b/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{},