hijacked flag for ResponseWriter, skipping sending headers to hijacked connections

This commit is contained in:
qz 2015-03-12 11:08:33 +03:00
parent 2b85363447
commit b6a995d26f

View File

@ -26,13 +26,15 @@ type (
Status() int Status() int
Size() int Size() int
Written() bool Written() bool
Hijacked() bool
WriteHeaderNow() WriteHeaderNow()
} }
responseWriter struct { responseWriter struct {
http.ResponseWriter http.ResponseWriter
status int status int
size int size int
hijacked bool
} }
) )
@ -52,7 +54,7 @@ func (w *responseWriter) WriteHeader(code int) {
} }
func (w *responseWriter) WriteHeaderNow() { func (w *responseWriter) WriteHeaderNow() {
if !w.Written() { if !w.Written() && !w.Hijacked() {
w.size = 0 w.size = 0
w.ResponseWriter.WriteHeader(w.status) w.ResponseWriter.WriteHeader(w.status)
} }
@ -77,12 +79,17 @@ func (w *responseWriter) Written() bool {
return w.size != NoWritten return w.size != NoWritten
} }
func (w *responseWriter) Hijacked() bool {
return w.hijacked
}
// Implements the http.Hijacker interface // Implements the http.Hijacker interface
func (w *responseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) { func (w *responseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
hijacker, ok := w.ResponseWriter.(http.Hijacker) hijacker, ok := w.ResponseWriter.(http.Hijacker)
if !ok { if !ok {
return nil, nil, errors.New("the ResponseWriter doesn't support the Hijacker interface") return nil, nil, errors.New("the ResponseWriter doesn't support the Hijacker interface")
} }
w.hijacked = true
return hijacker.Hijack() return hijacker.Hijack()
} }