-
Notifications
You must be signed in to change notification settings - Fork 92
/
udp_proxy.go
150 lines (138 loc) · 3.85 KB
/
udp_proxy.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
// Package udpproxy is from https://raw.githubusercontent.com/docker/libnetwork/fec6476dfa21380bf8ee4d74048515d968c1ee63/cmd/proxy/udp_proxy.go
package udpproxy
import (
"encoding/binary"
"fmt"
"io"
"net"
"strings"
"sync"
"syscall"
"time"
)
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.
// From libnetwork udp_proxy.go .
type UDPProxy struct {
LogWriter io.Writer
Listener *net.UDPConn
BackendDial func() (*net.UDPConn, error)
connTrackTable connTrackMap
connTrackLock sync.Mutex
}
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() {
proxy.connTrackTable = make(connTrackMap)
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) {
fmt.Fprintf(proxy.LogWriter, "Stopping proxy on udp: %v\n", err)
}
break
}
fromKey := newConnTrackKey(from)
proxy.connTrackLock.Lock()
proxyConn, hit := proxy.connTrackTable[*fromKey]
if !hit {
proxyConn, err = proxy.BackendDial()
if err != nil {
fmt.Fprintf(proxy.LogWriter, "Can't proxy a datagram to udp: %v\n", 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 {
fmt.Fprintf(proxy.LogWriter, "Can't proxy a datagram to udp: %v\n", 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()
}
}
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")
}