traefik/vendor/github.com/vulcand/oxy/utils/netutils.go
2017-11-22 18:20:03 +01:00

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)
}
}