github.com/dlintw/docker@v1.5.0-rc4/pkg/proxy/udp_proxy.go (about)

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