45 lines
945 B
Go
45 lines
945 B
Go
package http
|
|
|
|
import (
|
|
"bufio"
|
|
"errors"
|
|
"net"
|
|
"net/http"
|
|
|
|
"github.com/felixge/httpsnoop"
|
|
)
|
|
|
|
// Custom RW implementation
|
|
type ResponseWriter struct {
|
|
http.ResponseWriter
|
|
code int
|
|
Hijacker http.Hijacker
|
|
}
|
|
|
|
// Write status code to header
|
|
func (w *ResponseWriter) WriteHeader(statusCode int) {
|
|
w.code = statusCode
|
|
w.ResponseWriter.WriteHeader(statusCode)
|
|
}
|
|
|
|
// GetStatusCode returns the status code that was written
|
|
func (w *ResponseWriter) GetStatusCode() int {
|
|
return w.code
|
|
}
|
|
|
|
// Get new RW
|
|
func NewResponseWriter(w http.ResponseWriter) *ResponseWriter {
|
|
hijacker, _ := w.(http.Hijacker)
|
|
return &ResponseWriter{
|
|
ResponseWriter: httpsnoop.Wrap(w, httpsnoop.Hooks{}),
|
|
Hijacker: hijacker,
|
|
}
|
|
}
|
|
|
|
func (w *ResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
|
|
if w.Hijacker == nil {
|
|
return nil, nil, errors.New("http.Hijacker not implemented by underlying http.ResponseWriter")
|
|
}
|
|
return w.Hijacker.Hijack()
|
|
}
|