traefik/vendor/github.com/vulcand/oxy/utils/netutils.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)
	}
}