github.com/zhuohuang-hust/src-cbuild@v0.0.0-20230105071821-c7aab3e7c840/mergeCode/libnetwork/cmd/proxy/udp_proxy.go (about)

     1  package main
     2  
     3  import (
     4  	"encoding/binary"
     5  	"log"
     6  	"net"
     7  	"strings"
     8  	"sync"
     9  	"syscall"
    10  	"time"
    11  )
    12  
    13  const (
    14  	// UDPConnTrackTimeout is the timeout used for UDP connection tracking
    15  	UDPConnTrackTimeout = 90 * time.Second
    16  	// UDPBufSize is the buffer size for the UDP proxy
    17  	UDPBufSize = 65507
    18  )
    19  
    20  // A net.Addr where the IP is split into two fields so you can use it as a key
    21  // in a map:
    22  type connTrackKey struct {
    23  	IPHigh uint64
    24  	IPLow  uint64
    25  	Port   int
    26  }
    27  
    28  func newConnTrackKey(addr *net.UDPAddr) *connTrackKey {
    29  	if len(addr.IP) == net.IPv4len {
    30  		return &connTrackKey{
    31  			IPHigh: 0,
    32  			IPLow:  uint64(binary.BigEndian.Uint32(addr.IP)),
    33  			Port:   addr.Port,
    34  		}
    35  	}
    36  	return &connTrackKey{
    37  		IPHigh: binary.BigEndian.Uint64(addr.IP[:8]),
    38  		IPLow:  binary.BigEndian.Uint64(addr.IP[8:]),
    39  		Port:   addr.Port,
    40  	}
    41  }
    42  
    43  type connTrackMap map[connTrackKey]*net.UDPConn
    44  
    45  // UDPProxy is proxy for which handles UDP datagrams. It implements the Proxy
    46  // interface to handle UDP traffic forwarding between the frontend and backend
    47  // addresses.
    48  type UDPProxy struct {
    49  	listener       *net.UDPConn
    50  	frontendAddr   *net.UDPAddr
    51  	backendAddr    *net.UDPAddr
    52  	connTrackTable connTrackMap
    53  	connTrackLock  sync.Mutex
    54  }
    55  
    56  // NewUDPProxy creates a new UDPProxy.
    57  func NewUDPProxy(frontendAddr, backendAddr *net.UDPAddr) (*UDPProxy, error) {
    58  	listener, err := net.ListenUDP("udp", frontendAddr)
    59  	if err != nil {
    60  		return nil, err
    61  	}
    62  	return &UDPProxy{
    63  		listener:       listener,
    64  		frontendAddr:   listener.LocalAddr().(*net.UDPAddr),
    65  		backendAddr:    backendAddr,
    66  		connTrackTable: make(connTrackMap),
    67  	}, nil
    68  }
    69  
    70  func (proxy *UDPProxy) replyLoop(proxyConn *net.UDPConn, clientAddr *net.UDPAddr, clientKey *connTrackKey) {
    71  	defer func() {
    72  		proxy.connTrackLock.Lock()
    73  		delete(proxy.connTrackTable, *clientKey)
    74  		proxy.connTrackLock.Unlock()
    75  		proxyConn.Close()
    76  	}()
    77  
    78  	readBuf := make([]byte, UDPBufSize)
    79  	for {
    80  		proxyConn.SetReadDeadline(time.Now().Add(UDPConnTrackTimeout))
    81  	again:
    82  		read, err := proxyConn.Read(readBuf)
    83  		if err != nil {
    84  			if err, ok := err.(*net.OpError); ok && err.Err == syscall.ECONNREFUSED {
    85  				// This will happen if the last write failed
    86  				// (e.g: nothing is actually listening on the
    87  				// proxied port on the container), ignore it
    88  				// and continue until UDPConnTrackTimeout
    89  				// expires:
    90  				goto again
    91  			}
    92  			return
    93  		}
    94  		for i := 0; i != read; {
    95  			written, err := proxy.listener.WriteToUDP(readBuf[i:read], clientAddr)
    96  			if err != nil {
    97  				return
    98  			}
    99  			i += written
   100  		}
   101  	}
   102  }
   103  
   104  // Run starts forwarding the traffic using UDP.
   105  func (proxy *UDPProxy) Run() {
   106  	readBuf := make([]byte, UDPBufSize)
   107  	for {
   108  		read, from, err := proxy.listener.ReadFromUDP(readBuf)
   109  		if err != nil {
   110  			// NOTE: Apparently ReadFrom doesn't return
   111  			// ECONNREFUSED like Read do (see comment in
   112  			// UDPProxy.replyLoop)
   113  			if !isClosedError(err) {
   114  				log.Printf("Stopping proxy on udp/%v for udp/%v (%s)", proxy.frontendAddr, proxy.backendAddr, err)
   115  			}
   116  			break
   117  		}
   118  
   119  		fromKey := newConnTrackKey(from)
   120  		proxy.connTrackLock.Lock()
   121  		proxyConn, hit := proxy.connTrackTable[*fromKey]
   122  		if !hit {
   123  			proxyConn, err = net.DialUDP("udp", nil, proxy.backendAddr)
   124  			if err != nil {
   125  				log.Printf("Can't proxy a datagram to udp/%s: %s\n", proxy.backendAddr, err)
   126  				proxy.connTrackLock.Unlock()
   127  				continue
   128  			}
   129  			proxy.connTrackTable[*fromKey] = proxyConn
   130  			go proxy.replyLoop(proxyConn, from, fromKey)
   131  		}
   132  		proxy.connTrackLock.Unlock()
   133  		for i := 0; i != read; {
   134  			written, err := proxyConn.Write(readBuf[i:read])
   135  			if err != nil {
   136  				log.Printf("Can't proxy a datagram to udp/%s: %s\n", proxy.backendAddr, err)
   137  				break
   138  			}
   139  			i += written
   140  		}
   141  	}
   142  }
   143  
   144  // Close stops forwarding the traffic.
   145  func (proxy *UDPProxy) Close() {
   146  	proxy.listener.Close()
   147  	proxy.connTrackLock.Lock()
   148  	defer proxy.connTrackLock.Unlock()
   149  	for _, conn := range proxy.connTrackTable {
   150  		conn.Close()
   151  	}
   152  }
   153  
   154  // FrontendAddr returns the UDP address on which the proxy is listening.
   155  func (proxy *UDPProxy) FrontendAddr() net.Addr { return proxy.frontendAddr }
   156  
   157  // BackendAddr returns the proxied UDP address.
   158  func (proxy *UDPProxy) BackendAddr() net.Addr { return proxy.backendAddr }
   159  
   160  func isClosedError(err error) bool {
   161  	/* This comparison is ugly, but unfortunately, net.go doesn't export errClosing.
   162  	 * See:
   163  	 * http://golang.org/src/pkg/net/net.go
   164  	 * https://code.google.com/p/go/issues/detail?id=4337
   165  	 * https://groups.google.com/forum/#!msg/golang-nuts/0_aaCvBmOcM/SptmDyX1XJMJ
   166  	 */
   167  	return strings.HasSuffix(err.Error(), "use of closed network connection")
   168  }