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 }