159 lines
3.9 KiB
Go
159 lines
3.9 KiB
Go
package utils
|
|
|
|
import (
|
|
"bufio"
|
|
"fmt"
|
|
"io"
|
|
"net"
|
|
"net/http"
|
|
"net/url"
|
|
"reflect"
|
|
|
|
log "github.com/Sirupsen/logrus"
|
|
)
|
|
|
|
// ProxyWriter helps to capture response headers and status code
|
|
// from the ServeHTTP. It can be safely passed to ServeHTTP handler,
|
|
// wrapping the real response writer.
|
|
type ProxyWriter struct {
|
|
W http.ResponseWriter
|
|
Code int
|
|
Length int64
|
|
}
|
|
|
|
func (p *ProxyWriter) StatusCode() int {
|
|
if p.Code == 0 {
|
|
// per contract standard lib will set this to http.StatusOK if not set
|
|
// by user, here we avoid the confusion by mirroring this logic
|
|
return http.StatusOK
|
|
}
|
|
return p.Code
|
|
}
|
|
|
|
func (p *ProxyWriter) Header() http.Header {
|
|
return p.W.Header()
|
|
}
|
|
|
|
func (p *ProxyWriter) Write(buf []byte) (int, error) {
|
|
p.Length = p.Length + int64(len(buf))
|
|
return p.W.Write(buf)
|
|
}
|
|
|
|
func (p *ProxyWriter) WriteHeader(code int) {
|
|
p.Code = code
|
|
p.W.WriteHeader(code)
|
|
}
|
|
|
|
func (p *ProxyWriter) Flush() {
|
|
if f, ok := p.W.(http.Flusher); ok {
|
|
f.Flush()
|
|
}
|
|
}
|
|
|
|
func (p *ProxyWriter) CloseNotify() <-chan bool {
|
|
if cn, ok := p.W.(http.CloseNotifier); ok {
|
|
return cn.CloseNotify()
|
|
}
|
|
log.Warningf("Upstream ResponseWriter of type %v does not implement http.CloseNotifier. Returning dummy channel.", reflect.TypeOf(p.W))
|
|
return make(<-chan bool)
|
|
}
|
|
|
|
func (p *ProxyWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
|
|
if hi, ok := p.W.(http.Hijacker); ok {
|
|
return hi.Hijack()
|
|
}
|
|
log.Warningf("Upstream ResponseWriter of type %v does not implement http.Hijacker. Returning dummy channel.", reflect.TypeOf(p.W))
|
|
return nil, nil, fmt.Errorf("The response writer that was wrapped in this proxy, does not implement http.Hijacker. It is of type: %v", reflect.TypeOf(p.W))
|
|
}
|
|
|
|
func NewBufferWriter(w io.WriteCloser) *BufferWriter {
|
|
return &BufferWriter{
|
|
W: w,
|
|
H: make(http.Header),
|
|
}
|
|
}
|
|
|
|
type BufferWriter struct {
|
|
H http.Header
|
|
Code int
|
|
W io.WriteCloser
|
|
}
|
|
|
|
func (b *BufferWriter) Close() error {
|
|
return b.W.Close()
|
|
}
|
|
|
|
func (b *BufferWriter) Header() http.Header {
|
|
return b.H
|
|
}
|
|
|
|
func (b *BufferWriter) Write(buf []byte) (int, error) {
|
|
return b.W.Write(buf)
|
|
}
|
|
|
|
// WriteHeader sets rw.Code.
|
|
func (b *BufferWriter) WriteHeader(code int) {
|
|
b.Code = code
|
|
}
|
|
|
|
func (b *BufferWriter) CloseNotify() <-chan bool {
|
|
if cn, ok := b.W.(http.CloseNotifier); ok {
|
|
return cn.CloseNotify()
|
|
}
|
|
log.Warningf("Upstream ResponseWriter of type %v does not implement http.CloseNotifier. Returning dummy channel.", reflect.TypeOf(b.W))
|
|
return make(<-chan bool)
|
|
}
|
|
|
|
func (b *BufferWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
|
|
if hi, ok := b.W.(http.Hijacker); ok {
|
|
return hi.Hijack()
|
|
}
|
|
log.Warningf("Upstream ResponseWriter of type %v does not implement http.Hijacker. Returning dummy channel.", reflect.TypeOf(b.W))
|
|
return nil, nil, fmt.Errorf("The response writer that was wrapped in this proxy, does not implement http.Hijacker. It is of type: %v", reflect.TypeOf(b.W))
|
|
}
|
|
|
|
type nopWriteCloser struct {
|
|
io.Writer
|
|
}
|
|
|
|
func (*nopWriteCloser) Close() error { return nil }
|
|
|
|
// NopCloser returns a WriteCloser with a no-op Close method wrapping
|
|
// the provided Writer w.
|
|
func NopWriteCloser(w io.Writer) io.WriteCloser {
|
|
return &nopWriteCloser{w}
|
|
}
|
|
|
|
// CopyURL provides update safe copy by avoiding shallow copying User field
|
|
func CopyURL(i *url.URL) *url.URL {
|
|
out := *i
|
|
if i.User != nil {
|
|
out.User = &(*i.User)
|
|
}
|
|
return &out
|
|
}
|
|
|
|
// CopyHeaders copies http headers from source to destination, it
|
|
// does not overide, but adds multiple headers
|
|
func CopyHeaders(dst http.Header, src http.Header) {
|
|
for k, vv := range src {
|
|
dst[k] = append(dst[k], vv...)
|
|
}
|
|
}
|
|
|
|
// HasHeaders determines whether any of the header names is present in the http headers
|
|
func HasHeaders(names []string, headers http.Header) bool {
|
|
for _, h := range names {
|
|
if headers.Get(h) != "" {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
// RemoveHeaders removes the header with the given names from the headers map
|
|
func RemoveHeaders(headers http.Header, names ...string) {
|
|
for _, h := range names {
|
|
headers.Del(h)
|
|
}
|
|
}
|