route/lib/tunnel/httpproxy.go

package tunnel

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"
	"sync"

	"git.xeserv.us/xena/route/lib/tunnel/proto"
	"github.com/koding/logging"
)

var (
	httpLog = logging.NewLogger("http")
)

// HTTPProxy forwards HTTP traffic.
//
// When tunnel server requests a connection it's proxied to 127.0.0.1:incomingPort
// where incomingPort is control message LocalPort.
// Usually this is tunnel server's public exposed Port.
// This behaviour can be changed by setting LocalAddr or FetchLocalAddr.
// FetchLocalAddr takes precedence over LocalAddr.
//
// When connection to local server cannot be established proxy responds with http error message.
type HTTPProxy struct {
	// LocalAddr defines the TCP address of the local server.
	// This is optional if you want to specify a single TCP address.
	LocalAddr string
	// ErrorResp is custom response send to tunnel server when client cannot
	// establish connection to local server. If not set a default "no local server"
	// response is sent.
	ErrorResp *http.Response
	// Log is a custom logger that can be used for the proxy.
	// If not set a "http" logger is used.
	Log logging.Logger

	hs *http.Server
	rp *httputil.ReverseProxy
}

// Proxy is a ProxyFunc.
func (p *HTTPProxy) Proxy(remote net.Conn, msg *proto.ControlMessage) {
	if msg.Protocol != proto.HTTP && msg.Protocol != proto.WS {
		panic("Proxy mismatch")
	}

	var log = p.log()

	var port = msg.LocalPort
	if port == 0 {
		port = 80
	}

	var localAddr = fmt.Sprintf("127.0.0.1:%d", port)
	if p.LocalAddr != "" {
		localAddr = p.LocalAddr
	}

	if p.hs == nil {
		su, _ := url.Parse(fmt.Sprintf("http://%s", p.LocalAddr))
		p.rp = httputil.NewSingleHostReverseProxy(su)
		p.hs = &http.Server{
			Handler: p.rp,
		}
	}

	log.Debug("Dialing local server %q", localAddr)

	sl := singleListener{
		conn: remote,
	}

	err := p.hs.Serve(sl)
	if err != nil {
		log.Error("Dialing local server %q failed: %s", localAddr, err)
		p.sendError(remote)
		return
	}
}

func (p *HTTPProxy) sendError(remote net.Conn) {
	var w = noLocalServer()
	if p.ErrorResp != nil {
		w = p.ErrorResp
	}

	buf := new(bytes.Buffer)
	w.Write(buf)
	if _, err := io.Copy(remote, buf); err != nil {
		var log = p.log()
		log.Debug("Copy in-mem response error: %s", err)
	}

	remote.Close()
}

func noLocalServer() *http.Response {
	body := bytes.NewBufferString("no local server")
	return &http.Response{
		Status:        http.StatusText(http.StatusServiceUnavailable),
		StatusCode:    http.StatusServiceUnavailable,
		Proto:         "HTTP/1.1",
		ProtoMajor:    1,
		ProtoMinor:    1,
		Body:          ioutil.NopCloser(body),
		ContentLength: int64(body.Len()),
	}
}

func (p *HTTPProxy) log() logging.Logger {
	if p.Log != nil {
		return p.Log
	}
	return httpLog
}

// A singleListener is a net.Listener that returns a single connection, then
// gives the error io.EOF.
type singleListener struct {
	conn net.Conn
	once sync.Once
}

func (s singleListener) Accept() (net.Conn, error) {
	var c net.Conn
	s.once.Do(func() {
		c = s.conn
	})
	if c != nil {
		return c, nil
	}
	return nil, io.EOF
}

func (s singleListener) Close() error {
	s.once.Do(func() {
		s.conn.Close()
	})
	return nil
}

func (s singleListener) Addr() net.Addr {
	return s.conn.LocalAddr()
}