github.com/kaisenlinux/docker.io@v0.0.0-20230510090727-ea55db55fac7/libnetwork/cmd/proxy/tcp_proxy.go (about)

     1  package main
     2  
     3  import (
     4  	"io"
     5  	"log"
     6  	"net"
     7  	"sync"
     8  )
     9  
    10  // TCPProxy is a proxy for TCP connections. It implements the Proxy interface to
    11  // handle TCP traffic forwarding between the frontend and backend addresses.
    12  type TCPProxy struct {
    13  	listener     *net.TCPListener
    14  	frontendAddr *net.TCPAddr
    15  	backendAddr  *net.TCPAddr
    16  }
    17  
    18  // NewTCPProxy creates a new TCPProxy.
    19  func NewTCPProxy(frontendAddr, backendAddr *net.TCPAddr) (*TCPProxy, error) {
    20  	// detect version of hostIP to bind only to correct version
    21  	ipVersion := ipv4
    22  	if frontendAddr.IP.To4() == nil {
    23  		ipVersion = ipv6
    24  	}
    25  	listener, err := net.ListenTCP("tcp"+string(ipVersion), frontendAddr)
    26  	if err != nil {
    27  		return nil, err
    28  	}
    29  	// If the port in frontendAddr was 0 then ListenTCP will have a picked
    30  	// a port to listen on, hence the call to Addr to get that actual port:
    31  	return &TCPProxy{
    32  		listener:     listener,
    33  		frontendAddr: listener.Addr().(*net.TCPAddr),
    34  		backendAddr:  backendAddr,
    35  	}, nil
    36  }
    37  
    38  func (proxy *TCPProxy) clientLoop(client *net.TCPConn, quit chan bool) {
    39  	backend, err := net.DialTCP("tcp", nil, proxy.backendAddr)
    40  	if err != nil {
    41  		log.Printf("Can't forward traffic to backend tcp/%v: %s\n", proxy.backendAddr, err)
    42  		client.Close()
    43  		return
    44  	}
    45  
    46  	var wg sync.WaitGroup
    47  	var broker = func(to, from *net.TCPConn) {
    48  		io.Copy(to, from)
    49  		from.CloseRead()
    50  		to.CloseWrite()
    51  		wg.Done()
    52  	}
    53  
    54  	wg.Add(2)
    55  	go broker(client, backend)
    56  	go broker(backend, client)
    57  
    58  	finish := make(chan struct{})
    59  	go func() {
    60  		wg.Wait()
    61  		close(finish)
    62  	}()
    63  
    64  	select {
    65  	case <-quit:
    66  	case <-finish:
    67  	}
    68  	client.Close()
    69  	backend.Close()
    70  	<-finish
    71  }
    72  
    73  // Run starts forwarding the traffic using TCP.
    74  func (proxy *TCPProxy) Run() {
    75  	quit := make(chan bool)
    76  	defer close(quit)
    77  	for {
    78  		client, err := proxy.listener.Accept()
    79  		if err != nil {
    80  			log.Printf("Stopping proxy on tcp/%v for tcp/%v (%s)", proxy.frontendAddr, proxy.backendAddr, err)
    81  			return
    82  		}
    83  		go proxy.clientLoop(client.(*net.TCPConn), quit)
    84  	}
    85  }
    86  
    87  // Close stops forwarding the traffic.
    88  func (proxy *TCPProxy) Close() { proxy.listener.Close() }
    89  
    90  // FrontendAddr returns the TCP address on which the proxy is listening.
    91  func (proxy *TCPProxy) FrontendAddr() net.Addr { return proxy.frontendAddr }
    92  
    93  // BackendAddr returns the TCP proxied address.
    94  func (proxy *TCPProxy) BackendAddr() net.Addr { return proxy.backendAddr }