| package main |
| |
| import ( |
| "io" |
| "net" |
| "sync" |
| "syscall" |
| |
| "github.com/Sirupsen/logrus" |
| ) |
| |
| // TCPProxy is a proxy for TCP connections. It implements the Proxy interface to |
| // handle TCP traffic forwarding between the frontend and backend addresses. |
| type TCPProxy struct { |
| listener *net.TCPListener |
| frontendAddr *net.TCPAddr |
| backendAddr *net.TCPAddr |
| } |
| |
| // NewTCPProxy creates a new TCPProxy. |
| func NewTCPProxy(frontendAddr, backendAddr *net.TCPAddr) (*TCPProxy, error) { |
| listener, err := net.ListenTCP("tcp", frontendAddr) |
| if err != nil { |
| return nil, err |
| } |
| // If the port in frontendAddr was 0 then ListenTCP will have a picked |
| // a port to listen on, hence the call to Addr to get that actual port: |
| return &TCPProxy{ |
| listener: listener, |
| frontendAddr: listener.Addr().(*net.TCPAddr), |
| backendAddr: backendAddr, |
| }, nil |
| } |
| |
| func (proxy *TCPProxy) clientLoop(client *net.TCPConn, quit chan bool) { |
| backend, err := net.DialTCP("tcp", nil, proxy.backendAddr) |
| if err != nil { |
| logrus.Printf("Can't forward traffic to backend tcp/%v: %s\n", proxy.backendAddr, err) |
| client.Close() |
| return |
| } |
| |
| var wg sync.WaitGroup |
| var broker = func(to, from *net.TCPConn) { |
| if _, err := io.Copy(to, from); err != nil { |
| // If the socket we are writing to is shutdown with |
| // SHUT_WR, forward it to the other end of the pipe: |
| if err, ok := err.(*net.OpError); ok && err.Err == syscall.EPIPE { |
| from.CloseWrite() |
| } |
| } |
| to.CloseRead() |
| wg.Done() |
| } |
| |
| wg.Add(2) |
| go broker(client, backend) |
| go broker(backend, client) |
| |
| finish := make(chan struct{}) |
| go func() { |
| wg.Wait() |
| close(finish) |
| }() |
| |
| select { |
| case <-quit: |
| case <-finish: |
| } |
| client.Close() |
| backend.Close() |
| <-finish |
| } |
| |
| // Run starts forwarding the traffic using TCP. |
| func (proxy *TCPProxy) Run() { |
| quit := make(chan bool) |
| defer close(quit) |
| for { |
| client, err := proxy.listener.Accept() |
| if err != nil { |
| logrus.Printf("Stopping proxy on tcp/%v for tcp/%v (%s)", proxy.frontendAddr, proxy.backendAddr, err) |
| return |
| } |
| go proxy.clientLoop(client.(*net.TCPConn), quit) |
| } |
| } |
| |
| // Close stops forwarding the traffic. |
| func (proxy *TCPProxy) Close() { proxy.listener.Close() } |
| |
| // FrontendAddr returns the TCP address on which the proxy is listening. |
| func (proxy *TCPProxy) FrontendAddr() net.Addr { return proxy.frontendAddr } |
| |
| // BackendAddr returns the TCP proxied address. |
| func (proxy *TCPProxy) BackendAddr() net.Addr { return proxy.backendAddr } |