github.com/sijibomii/docker@v0.0.0-20231230191044-5cf6ca554647/pkg/proxy/tcp_proxy.go (about)

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