| package proxy |
| |
| import ( |
| "encoding/binary" |
| "net" |
| "strings" |
| "sync" |
| "syscall" |
| "time" |
| |
| "github.com/Sirupsen/logrus" |
| ) |
| |
| const ( |
| // UDPConnTrackTimeout is the timeout used for UDP connection tracking |
| UDPConnTrackTimeout = 90 * time.Second |
| // UDPBufSize is the buffer size for the UDP proxy |
| UDPBufSize = 65507 |
| ) |
| |
| // A net.Addr where the IP is split into two fields so you can use it as a key |
| // in a map: |
| type connTrackKey struct { |
| IPHigh uint64 |
| IPLow uint64 |
| Port int |
| } |
| |
| func newConnTrackKey(addr *net.UDPAddr) *connTrackKey { |
| if len(addr.IP) == net.IPv4len { |
| return &connTrackKey{ |
| IPHigh: 0, |
| IPLow: uint64(binary.BigEndian.Uint32(addr.IP)), |
| Port: addr.Port, |
| } |
| } |
| return &connTrackKey{ |
| IPHigh: binary.BigEndian.Uint64(addr.IP[:8]), |
| IPLow: binary.BigEndian.Uint64(addr.IP[8:]), |
| Port: addr.Port, |
| } |
| } |
| |
| type connTrackMap map[connTrackKey]*net.UDPConn |
| |
| // UDPProxy is proxy for which handles UDP datagrams. It implements the Proxy |
| // interface to handle UDP traffic forwarding between the frontend and backend |
| // addresses. |
| type UDPProxy struct { |
| listener *net.UDPConn |
| frontendAddr *net.UDPAddr |
| backendAddr *net.UDPAddr |
| connTrackTable connTrackMap |
| connTrackLock sync.Mutex |
| } |
| |
| // NewUDPProxy creates a new UDPProxy. |
| func NewUDPProxy(frontendAddr, backendAddr *net.UDPAddr) (*UDPProxy, error) { |
| listener, err := net.ListenUDP("udp", frontendAddr) |
| if err != nil { |
| return nil, err |
| } |
| return &UDPProxy{ |
| listener: listener, |
| frontendAddr: listener.LocalAddr().(*net.UDPAddr), |
| backendAddr: backendAddr, |
| connTrackTable: make(connTrackMap), |
| }, nil |
| } |
| |
| func (proxy *UDPProxy) replyLoop(proxyConn *net.UDPConn, clientAddr *net.UDPAddr, clientKey *connTrackKey) { |
| defer func() { |
| proxy.connTrackLock.Lock() |
| delete(proxy.connTrackTable, *clientKey) |
| proxy.connTrackLock.Unlock() |
| proxyConn.Close() |
| }() |
| |
| readBuf := make([]byte, UDPBufSize) |
| for { |
| proxyConn.SetReadDeadline(time.Now().Add(UDPConnTrackTimeout)) |
| again: |
| read, err := proxyConn.Read(readBuf) |
| if err != nil { |
| if err, ok := err.(*net.OpError); ok && err.Err == syscall.ECONNREFUSED { |
| // This will happen if the last write failed |
| // (e.g: nothing is actually listening on the |
| // proxied port on the container), ignore it |
| // and continue until UDPConnTrackTimeout |
| // expires: |
| goto again |
| } |
| return |
| } |
| for i := 0; i != read; { |
| written, err := proxy.listener.WriteToUDP(readBuf[i:read], clientAddr) |
| if err != nil { |
| return |
| } |
| i += written |
| } |
| } |
| } |
| |
| // Run starts forwarding the traffic using UDP. |
| func (proxy *UDPProxy) Run() { |
| readBuf := make([]byte, UDPBufSize) |
| for { |
| read, from, err := proxy.listener.ReadFromUDP(readBuf) |
| if err != nil { |
| // NOTE: Apparently ReadFrom doesn't return |
| // ECONNREFUSED like Read do (see comment in |
| // UDPProxy.replyLoop) |
| if !isClosedError(err) { |
| logrus.Printf("Stopping proxy on udp/%v for udp/%v (%s)", proxy.frontendAddr, proxy.backendAddr, err) |
| } |
| break |
| } |
| |
| fromKey := newConnTrackKey(from) |
| proxy.connTrackLock.Lock() |
| proxyConn, hit := proxy.connTrackTable[*fromKey] |
| if !hit { |
| proxyConn, err = net.DialUDP("udp", nil, proxy.backendAddr) |
| if err != nil { |
| logrus.Printf("Can't proxy a datagram to udp/%s: %s\n", proxy.backendAddr, err) |
| proxy.connTrackLock.Unlock() |
| continue |
| } |
| proxy.connTrackTable[*fromKey] = proxyConn |
| go proxy.replyLoop(proxyConn, from, fromKey) |
| } |
| proxy.connTrackLock.Unlock() |
| for i := 0; i != read; { |
| written, err := proxyConn.Write(readBuf[i:read]) |
| if err != nil { |
| logrus.Printf("Can't proxy a datagram to udp/%s: %s\n", proxy.backendAddr, err) |
| break |
| } |
| i += written |
| } |
| } |
| } |
| |
| // Close stops forwarding the traffic. |
| func (proxy *UDPProxy) Close() { |
| proxy.listener.Close() |
| proxy.connTrackLock.Lock() |
| defer proxy.connTrackLock.Unlock() |
| for _, conn := range proxy.connTrackTable { |
| conn.Close() |
| } |
| } |
| |
| // FrontendAddr returns the UDP address on which the proxy is listening. |
| func (proxy *UDPProxy) FrontendAddr() net.Addr { return proxy.frontendAddr } |
| |
| // BackendAddr returns the proxied UDP address. |
| func (proxy *UDPProxy) BackendAddr() net.Addr { return proxy.backendAddr } |
| |
| func isClosedError(err error) bool { |
| /* This comparison is ugly, but unfortunately, net.go doesn't export errClosing. |
| * See: |
| * http://golang.org/src/pkg/net/net.go |
| * https://code.google.com/p/go/issues/detail?id=4337 |
| * https://groups.google.com/forum/#!msg/golang-nuts/0_aaCvBmOcM/SptmDyX1XJMJ |
| */ |
| return strings.HasSuffix(err.Error(), "use of closed network connection") |
| } |