X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/70111d82653ca633d2397c74eb994a8a9d718e80..6766d1b04bd238d05890f3ec221c65e84920dde6:/sdk/go/httpserver/responsewriter.go

diff --git a/sdk/go/httpserver/responsewriter.go b/sdk/go/httpserver/responsewriter.go
index 3941c13c22..049a3f1aae 100644
--- a/sdk/go/httpserver/responsewriter.go
+++ b/sdk/go/httpserver/responsewriter.go
@@ -8,10 +8,13 @@ import (
 	"net/http"
 )
 
+const sniffBytes = 1024
+
 type ResponseWriter interface {
 	http.ResponseWriter
 	WroteStatus() int
 	WroteBodyBytes() int
+	Sniffed() []byte
 }
 
 // responseWriter wraps http.ResponseWriter and exposes the status
@@ -19,35 +22,67 @@ type ResponseWriter interface {
 // error.
 type responseWriter struct {
 	http.ResponseWriter
-	wroteStatus    int   // Last status given to WriteHeader()
+	wroteStatus    int   // First status given to WriteHeader()
 	wroteBodyBytes int   // Bytes successfully written
 	err            error // Last error returned from Write()
+	sniffed        []byte
 }
 
 func WrapResponseWriter(orig http.ResponseWriter) ResponseWriter {
 	return &responseWriter{ResponseWriter: orig}
 }
 
-func (w responseWriter) WriteHeader(s int) {
-	w.wroteStatus = s
+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) {
+	if w.wroteStatus == 0 {
+		w.wroteStatus = s
+	}
+	// ...else it's too late to change the status seen by the
+	// client -- but we call the wrapped WriteHeader() anyway so
+	// it can log a warning.
 	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)
+	} else if w.wroteStatus >= 400 {
+		w.sniff(data)
+	}
 	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
 }
+
+func (w *responseWriter) sniff(data []byte) {
+	max := sniffBytes - len(w.sniffed)
+	if max <= 0 {
+		return
+	} else if max < len(data) {
+		data = data[:max]
+	}
+	w.sniffed = append(w.sniffed, data...)
+}
+
+func (w *responseWriter) Sniffed() []byte {
+	return w.sniffed
+}