github.com/wlan0/docker@v1.5.0/pkg/proxy/tcp_proxy.go (about)

     1  package proxy
     2  
     3  import (
     4  	"io"
     5  	"net"
     6  	"syscall"
     7  
     8  	log "github.com/Sirupsen/logrus"
     9  )
    10  
    11  type TCPProxy struct {
    12  	listener     *net.TCPListener
    13  	frontendAddr *net.TCPAddr
    14  	backendAddr  *net.TCPAddr
    15  }
    16  
    17  func NewTCPProxy(frontendAddr, backendAddr *net.TCPAddr) (*TCPProxy, error) {
    18  	listener, err := net.ListenTCP("tcp", frontendAddr)
    19  	if err != nil {
    20  		return nil, err
    21  	}
    22  	// If the port in frontendAddr was 0 then ListenTCP will have a picked
    23  	// a port to listen on, hence the call to Addr to get that actual port:
    24  	return &TCPProxy{
    25  		listener:     listener,
    26  		frontendAddr: listener.Addr().(*net.TCPAddr),
    27  		backendAddr:  backendAddr,
    28  	}, nil
    29  }
    30  
    31  func (proxy *TCPProxy) clientLoop(client *net.TCPConn, quit chan bool) {
    32  	backend, err := net.DialTCP("tcp", nil, proxy.backendAddr)
    33  	if err != nil {
    34  		log.Printf("Can't forward traffic to backend tcp/%v: %s\n", proxy.backendAddr, err)
    35  		client.Close()
    36  		return
    37  	}
    38  
    39  	event := make(chan int64)
    40  	var broker = func(to, from *net.TCPConn) {
    41  		written, err := io.Copy(to, from)
    42  		if err != nil {
    43  			// If the socket we are writing to is shutdown with
    44  			// SHUT_WR, forward it to the other end of the pipe:
    45  			if err, ok := err.(*net.OpError); ok && err.Err == syscall.EPIPE {
    46  				from.CloseWrite()
    47  			}
    48  		}
    49  		to.CloseRead()
    50  		event <- written
    51  	}
    52  
    53  	go broker(client, backend)
    54  	go broker(backend, client)
    55  
    56  	var transferred int64 = 0
    57  	for i := 0; i < 2; i++ {
    58  		select {
    59  		case written := <-event:
    60  			transferred += written
    61  		case <-quit:
    62  			// Interrupt the two brokers and "join" them.
    63  			client.Close()
    64  			backend.Close()
    65  			for ; i < 2; i++ {
    66  				transferred += <-event
    67  			}
    68  			return
    69  		}
    70  	}
    71  	client.Close()
    72  	backend.Close()
    73  }
    74  
    75  func (proxy *TCPProxy) Run() {
    76  	quit := make(chan bool)
    77  	defer close(quit)
    78  	for {
    79  		client, err := proxy.listener.Accept()
    80  		if err != nil {
    81  			log.Printf("Stopping proxy on tcp/%v for tcp/%v (%s)", proxy.frontendAddr, proxy.backendAddr, err)
    82  			return
    83  		}
    84  		go proxy.clientLoop(client.(*net.TCPConn), quit)
    85  	}
    86  }
    87  
    88  func (proxy *TCPProxy) Close()                 { proxy.listener.Close() }
    89  func (proxy *TCPProxy) FrontendAddr() net.Addr { return proxy.frontendAddr }
    90  func (proxy *TCPProxy) BackendAddr() net.Addr  { return proxy.backendAddr }