123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990 |
- package proxy
- import (
- "io"
- "net"
- "syscall"
- log "github.com/Sirupsen/logrus"
- )
- type TCPProxy struct {
- listener *net.TCPListener
- frontendAddr *net.TCPAddr
- backendAddr *net.TCPAddr
- }
- 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 {
- log.Printf("Can't forward traffic to backend tcp/%v: %s\n", proxy.backendAddr, err)
- client.Close()
- return
- }
- event := make(chan int64)
- var broker = func(to, from *net.TCPConn) {
- written, err := io.Copy(to, from)
- if 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()
- event <- written
- }
- go broker(client, backend)
- go broker(backend, client)
- var transferred int64 = 0
- for i := 0; i < 2; i++ {
- select {
- case written := <-event:
- transferred += written
- case <-quit:
- // Interrupt the two brokers and "join" them.
- client.Close()
- backend.Close()
- for ; i < 2; i++ {
- transferred += <-event
- }
- return
- }
- }
- client.Close()
- backend.Close()
- }
- func (proxy *TCPProxy) Run() {
- quit := make(chan bool)
- defer close(quit)
- for {
- client, err := proxy.listener.Accept()
- if err != nil {
- log.Printf("Stopping proxy on tcp/%v for tcp/%v (%s)", proxy.frontendAddr, proxy.backendAddr, err)
- return
- }
- go proxy.clientLoop(client.(*net.TCPConn), quit)
- }
- }
- func (proxy *TCPProxy) Close() { proxy.listener.Close() }
- func (proxy *TCPProxy) FrontendAddr() net.Addr { return proxy.frontendAddr }
- func (proxy *TCPProxy) BackendAddr() net.Addr { return proxy.backendAddr }
|