1 files changed, 308 insertions, 0 deletions
diff --git a/tcpproxy.go b/tcpproxy.go
new file mode 100644
index 0000000..fde438f
--- /dev/null
+++ b/tcpproxy.go
@@ -0,0 +1,308 @@
+// Copyright 2017 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+// Package tcpproxy lets users build TCP, HTTP/1, and TLS+SNI proxies.
+//
+// Typical usage:
+//
+//     var p tcpproxy.Proxy
+//     p.AddHTTPHostRoute(":80", "foo.com", tcpproxy.To("10.0.0.1:8081"))
+//     p.AddHTTPHostRoute(":80", "bar.com", tcpproxy.To("10.0.0.2:8082"))
+//     p.AddRoute(":80", tcpproxy.To("10.0.0.1:8081")) // fallback
+//     p.AddSNIHostRoute(":443", "foo.com", tcpproxy.To("10.0.0.1:4431"))
+//     p.AddSNIHostRoute(":443", "bar.com", tcpproxy.To("10.0.0.2:4432"))
+//     p.AddRoute(":443", tcpproxy.To("10.0.0.1:4431")) // fallback
+//     log.Fatal(p.Run())
+package tcpproxy
+
+import (
+	"bufio"
+	"bytes"
+	"context"
+	"errors"
+	"io"
+	"log"
+	"net"
+	"time"
+)
+
+// Proxy is a proxy. Its zero value is a valid proxy that does
+// nothing. Call methods to add routes before calling Run.
+type Proxy struct {
+	routes map[string][]route // ip:port => route
+
+	lns   []net.Listener
+	donec chan struct{} // closed before err
+	err   error         // any error from listening
+
+	// ListenFunc optionally specifies an alternate listen
+	// function. If nil, net.Dial is used.
+	// The provided net is always "tcp".
+	ListenFunc func(net, laddr string) (net.Listener, error)
+}
+
+type route struct {
+	matcher matcher
+	target  Target
+}
+
+type matcher interface {
+	match(*bufio.Reader) bool
+}
+
+func (p *Proxy) netListen() func(net, laddr string) (net.Listener, error) {
+	if p.ListenFunc != nil {
+		return p.ListenFunc
+	}
+	return net.Listen
+}
+
+func (p *Proxy) addRoute(ipPort string, matcher matcher, dest Target) {
+	if p.routes == nil {
+		p.routes = make(map[string][]route)
+	}
+	p.routes[ipPort] = append(p.routes[ipPort], route{matcher, dest})
+}
+
+// AddRoute appends an always-matching route to the ipPort listener,
+// directing any connection to dest.
+//
+// This is generally used as either the only rule (for simple TCP
+// proxies), or as the final fallback rule for an ipPort.
+//
+// The ipPort is any valid net.Listen TCP address.
+func (p *Proxy) AddRoute(ipPort string, dest Target) {
+	p.addRoute(ipPort, alwaysMatch{}, dest)
+}
+
+type alwaysMatch struct{}
+
+func (alwaysMatch) match(*bufio.Reader) bool { return true }
+
+// Run is calls Start, and then Wait.
+//
+// It blocks until there's an error. The return value is always
+// non-nil.
+func (p *Proxy) Run() error {
+	if err := p.Start(); err != nil {
+		return err
+	}
+	return p.Wait()
+}
+
+// Wait waits for the Proxy to finish running. Currently this can only
+// happen if a Listener is closed, or Close is called on the proxy.
+//
+// It is only valid to call Wait after a successful call to Start.
+func (p *Proxy) Wait() error {
+	<-p.donec
+	return p.err
+}
+
+// Close closes all the proxy's self-opened listeners.
+func (p *Proxy) Close() error {
+	for _, c := range p.lns {
+		c.Close()
+	}
+	return nil
+}
+
+// Start creates a TCP listener for each unique ipPort from the
+// previously created routes and starts the proxy. It returns any
+// error from starting listeners.
+//
+// If it returns a non-nil error, any successfully opened listeners
+// are closed.
+func (p *Proxy) Start() error {
+	if p.donec != nil {
+		return errors.New("already started")
+	}
+	p.donec = make(chan struct{})
+	errc := make(chan error, len(p.routes))
+	p.lns = make([]net.Listener, 0, len(p.routes))
+	for ipPort, routes := range p.routes {
+		ln, err := p.netListen()("tcp", ipPort)
+		if err != nil {
+			p.Close()
+			return err
+		}
+		p.lns = append(p.lns, ln)
+		go p.serveListener(errc, ln, routes)
+	}
+	go p.awaitFirstError(errc)
+	return nil
+}
+
+func (p *Proxy) awaitFirstError(errc <-chan error) {
+	p.err = <-errc
+	close(p.donec)
+}
+
+func (p *Proxy) serveListener(ret chan<- error, ln net.Listener, routes []route) {
+	for {
+		c, err := ln.Accept()
+		if err != nil {
+			ret <- err
+			return
+		}
+		go p.serveConn(c, routes)
+	}
+}
+
+// serveConn runs in its own goroutine and matches c against routes.
+// It returns whether it matched purely for testing.
+func (p *Proxy) serveConn(c net.Conn, routes []route) bool {
+	br := bufio.NewReader(c)
+	for _, route := range routes {
+		if route.matcher.match(br) {
+			buffered, _ := br.Peek(br.Buffered())
+			route.target.HandleConn(changeReaderConn{
+				r:    io.MultiReader(bytes.NewReader(buffered), c),
+				Conn: c,
+			}, c)
+			return true
+		}
+	}
+	// TODO: hook for this?
+	log.Printf("tcpproxy: no routes matched conn %v/%v; closing", c.RemoteAddr().String(), c.LocalAddr().String())
+	c.Close()
+	return false
+}
+
+// changeReaderConn is a net.Conn wrapper with a separate reader function.
+type changeReaderConn struct {
+	r io.Reader
+	net.Conn
+}
+
+func (c changeReaderConn) Read(p []byte) (int, error) { return c.r.Read(p) }
+
+// Target is what an incoming matched connection is sent to.
+type Target interface {
+	// HandleConn is called when an incoming connection is
+	// matched. After the call to HandleConn, the tcpproxy
+	// package never touches the conn again. Implementations are
+	// responsible for closing the conn when needed.
+	//
+	// The c Conn acts like a new conn, without any bytes consumed,
+	// but it has an unexported concrete type and cannot be type
+	// asserted to *net.TCPConn, etc.
+	//
+	// The rawConn represents the underlying connections (with
+	// some bytes removed) and should only be used for type
+	// assertions and setting deadlines, not reading.
+	HandleConn(c net.Conn, rawConn net.Conn)
+}
+
+// To is shorthand way of writing &DialProxy{Addr: addr}.
+func To(addr string) *DialProxy {
+	return &DialProxy{Addr: addr}
+}
+
+// DialProxy implements Target by dialing a new connection to Addr
+// and then proxying data back and forth.
+//
+// The To func is a shorthand way of creating a DialProxy.
+type DialProxy struct {
+	// Addr is the TCP address to proxy to.
+	Addr string
+
+	// KeepAlivePeriod sets the period between TCP keep alives.
+	// If zero, a default is used. To disable, use a negative number.
+	// The keep-alive is used for both the client connection and
+	KeepAlivePeriod time.Duration
+
+	// DialTimeout optionally specifies a dial timeout.
+	// If zero, a default is used.
+	// If negative, the timeout is disabled.
+	DialTimeout time.Duration
+
+	// DialContext optionally specifies an alternate dial function
+	// for TCP targets. If nil, the standard
+	// net.Dialer.DialContext method is used.
+	DialContext func(ctx context.Context, network, address string) (net.Conn, error)
+
+	// OnDialError optionally specifies an alternate way to handle errors dialing Addr.
+	// If nil, the error is logged and src is closed.
+	// If non-nil, src is not closed automatically.
+	OnDialError func(src net.Conn, dstDialErr error)
+}
+
+func (dp *DialProxy) HandleConn(src net.Conn, rawSrc net.Conn) {
+	ctx := context.Background()
+	var cancel context.CancelFunc
+	if dp.DialTimeout >= 0 {
+		ctx, cancel = context.WithTimeout(ctx, dp.dialTimeout())
+	}
+	dst, err := dp.dialContext()(ctx, "tcp", dp.Addr)
+	if cancel != nil {
+		cancel()
+	}
+	if err != nil {
+		dp.onDialError()(src, err)
+		return
+	}
+	defer src.Close()
+	defer dst.Close()
+	if ka := dp.keepAlivePeriod(); ka > 0 {
+		if c, ok := rawSrc.(*net.TCPConn); ok {
+			c.SetKeepAlive(true)
+			c.SetKeepAlivePeriod(ka)
+		}
+		if c, ok := dst.(*net.TCPConn); ok {
+			c.SetKeepAlive(true)
+			c.SetKeepAlivePeriod(ka)
+		}
+	}
+	errc := make(chan error, 1)
+	go proxyCopy(errc, src, dst)
+	go proxyCopy(errc, dst, src)
+	<-errc
+}
+
+// proxyCopy is the function that copies bytes around.
+// It's a named function instead of a func literal so users get
+// named goroutines in debug goroutine stack dumps.
+func proxyCopy(errc chan<- error, dst io.Writer, src io.Reader) {
+	// TODO: make caller switch from src to rawSrc after N bytes (e.g. 4KB)
+	// if the io.Copy optimization to switch to Linux splice happens.
+	// TODO: if the runtime provides a way to wait for
+	// readability, use that to avoid stranding big blocks of
+	// memory blocked in idle reads.
+	_, err := io.Copy(dst, src)
+	errc <- err
+}
+
+func (dp *DialProxy) keepAlivePeriod() time.Duration {
+	if dp.KeepAlivePeriod != 0 {
+		return dp.KeepAlivePeriod
+	}
+	return time.Minute
+}
+
+func (dp *DialProxy) dialTimeout() time.Duration {
+	if dp.DialTimeout > 0 {
+		return dp.DialTimeout
+	}
+	return 10 * time.Second
+}
+
+var defaultDialer = new(net.Dialer)
+
+func (dp *DialProxy) dialContext() func(ctx context.Context, network, address string) (net.Conn, error) {
+	if dp.DialContext != nil {
+		return dp.DialContext
+	}
+	return defaultDialer.DialContext
+}
+
+func (dp *DialProxy) onDialError() func(src net.Conn, dstDialErr error) {
+	if dp.OnDialError != nil {
+		return dp.OnDialError
+	}
+	return func(src net.Conn, dstDialErr error) {
+		log.Printf("tcpproxy: for incoming conn %v, error dialing %q: %v", src.RemoteAddr().String(), dp.Addr, dstDialErr)
+		src.Close()
+	}
+}