flush buffer on body copy

SenseUnit · May 22, 2020 · 53d049d · 53d049d
1 parent a837129
commit 53d049d
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 2 deletions.
diff --git a/handler.go b/handler.go
@@ -1,7 +1,6 @@
 package main
 
 import (
-    "io"
     "net"
     "fmt"
     "time"
@@ -68,7 +67,7 @@ func (s *ProxyHandler) HandleRequest(wr http.ResponseWriter, req *http.Request)
     copyHeader(wr.Header(), resp.Header)
     wr.WriteHeader(resp.StatusCode)
     flush(wr)
-    io.Copy(wr, resp.Body)
+    copyBody(wr, resp.Body)
 }
 
 func (s *ProxyHandler) ServeHTTP(wr http.ResponseWriter, req *http.Request) {

diff --git a/utils.go b/utils.go
@@ -11,6 +11,8 @@ import (
     "bufio"
 )
 
+const COPY_BUF = 128 * 1024
+
 func proxy(ctx context.Context, left, right net.Conn) {
     wg := sync.WaitGroup{}
     cpy := func (dst, src net.Conn) {
@@ -91,3 +93,18 @@ func flush(flusher interface{}) bool {
     f.Flush()
     return true
 }
+
+func copyBody(wr io.Writer, body io.Reader) {
+    for {
+        buf := make([]byte, COPY_BUF)
+        bread, read_err := body.Read(buf)
+        var write_err error
+        if bread > 0 {
+            _, write_err = wr.Write(buf[:bread])
+            flush(wr)
+        }
+        if read_err != nil || write_err != nil {
+            break
+        }
+    }
+}