67 lines
1.3 KiB
Go
67 lines
1.3 KiB
Go
|
package hijack
|
||
|
|
|
||
|
|
import (
|
||
|
|
"bufio"
|
||
|
|
"net"
|
||
|
|
"net/http"
|
||
|
|
"sync"
|
||
|
|
)
|
||
|
|
|
||
|
|
func New(w http.ResponseWriter, commit func(http.ResponseWriter)) (http.ResponseWriter, func()) {
|
||
|
|
responseWriter := &responseWriter{
|
||
|
|
commit: commit,
|
||
|
|
w: w,
|
||
|
|
}
|
||
|
|
|
||
|
|
if _, ok := w.(http.Hijacker); ok {
|
||
|
|
return &responseWriterHijacker{
|
||
|
|
responseWriter: responseWriter,
|
||
|
|
}, responseWriter.ensureCommitted
|
||
|
|
} else {
|
||
|
|
return responseWriter, responseWriter.ensureCommitted
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
var _ http.ResponseWriter = new(responseWriter)
|
||
|
|
var _ http.ResponseWriter = new(responseWriterHijacker)
|
||
|
|
var _ http.Hijacker = new(responseWriterHijacker)
|
||
|
|
|
||
|
|
type responseWriter struct {
|
||
|
|
commit func(http.ResponseWriter)
|
||
|
|
w http.ResponseWriter
|
||
|
|
hijacked bool
|
||
|
|
writeOnce sync.Once
|
||
|
|
}
|
||
|
|
|
||
|
|
func (rw *responseWriter) ensureCommitted() {
|
||
|
|
rw.writeOnce.Do(func() {
|
||
|
|
if rw.hijacked {
|
||
|
|
return
|
||
|
|
}
|
||
|
|
rw.commit(rw.w)
|
||
|
|
})
|
||
|
|
}
|
||
|
|
|
||
|
|
func (rw *responseWriter) Header() http.Header {
|
||
|
|
return rw.w.Header()
|
||
|
|
}
|
||
|
|
|
||
|
|
func (rw *responseWriter) Write(data []byte) (int, error) {
|
||
|
|
rw.ensureCommitted()
|
||
|
|
return rw.w.Write(data)
|
||
|
|
}
|
||
|
|
|
||
|
|
func (rw *responseWriter) WriteHeader(status int) {
|
||
|
|
rw.ensureCommitted()
|
||
|
|
rw.w.WriteHeader(status)
|
||
|
|
}
|
||
|
|
|
||
|
|
type responseWriterHijacker struct {
|
||
|
|
*responseWriter
|
||
|
|
}
|
||
|
|
|
||
|
|
func (rw *responseWriterHijacker) Hijack() (net.Conn, *bufio.ReadWriter, error) {
|
||
|
|
rw.hijacked = true
|
||
|
|
return rw.w.(http.Hijacker).Hijack()
|
||
|
|
}
|