13111: Avoid double WriteHeader after implicit 200 on first Write.

[arvados.git] / sdk / go / httpserver / responsewriter.go

diff --git a/sdk/go/httpserver/responsewriter.go b/sdk/go/httpserver/responsewriter.go
index 3941c13c222504710e3b6b32bd36ad747a60dc8a..8dea759ccb9b1772b816ad565a279975ab751c8a 100644 (file)
--- a/sdk/go/httpserver/responsewriter.go
+++ b/sdk/go/httpserver/responsewriter.go
@@ -28,26 +28,36 @@ func WrapResponseWriter(orig http.ResponseWriter) ResponseWriter {
        return &responseWriter{ResponseWriter: orig}
 }
 
-func (w responseWriter) WriteHeader(s int) {
+func (w *responseWriter) CloseNotify() <-chan bool {
+       if cn, ok := w.ResponseWriter.(http.CloseNotifier); ok {
+               return cn.CloseNotify()
+       }
+       return nil
+}
+
+func (w *responseWriter) WriteHeader(s int) {
        w.wroteStatus = s
        w.ResponseWriter.WriteHeader(s)
 }
 
-func (w responseWriter) Write(data []byte) (n int, err error) {
+func (w *responseWriter) Write(data []byte) (n int, err error) {
+       if w.wroteStatus == 0 {
+               w.WriteHeader(http.StatusOK)
+       }
        n, err = w.ResponseWriter.Write(data)
        w.wroteBodyBytes += n
        w.err = err
        return
 }
 
-func (w responseWriter) WroteStatus() int {
+func (w *responseWriter) WroteStatus() int {
        return w.wroteStatus
 }
 
-func (w responseWriter) WroteBodyBytes() int {
+func (w *responseWriter) WroteBodyBytes() int {
        return w.wroteBodyBytes
 }
 
-func (w responseWriter) Err() error {
+func (w *responseWriter) Err() error {
        return w.err
 }