summaryrefslogtreecommitdiff
path: root/vendor/github.com/pion/transport/vnet/udpproxy.go
blob: c7749553925265dbc7d02222dd669aeb6d742696 (plain)
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
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
package vnet

import (
	"net"
	"sync"
	"time"
)

// UDPProxy is a proxy between real server(net.UDPConn) and vnet.UDPConn.
//
// High level design:
//                           ..............................................
//                           :         Virtual Network (vnet)             :
//                           :                                            :
//   +-------+ *         1 +----+         +--------+                      :
//   | :App  |------------>|:Net|--o<-----|:Router |          .............................
//   +-------+             +----+         |        |          :        UDPProxy           :
//                           :            |        |       +----+     +---------+     +---------+     +--------+
//                           :            |        |--->o--|:Net|-->o-| vnet.   |-->o-|  net.   |--->-| :Real  |
//                           :            |        |       +----+     | UDPConn |     | UDPConn |     | Server |
//                           :            |        |          :       +---------+     +---------+     +--------+
//                           :            |        |          ............................:
//                           :            +--------+                       :
//                           ...............................................
type UDPProxy struct {
	// The router bind to.
	router *Router

	// Each vnet source, bind to a real socket to server.
	// key is real server addr, which is net.Addr
	// value is *aUDPProxyWorker
	workers sync.Map

	// For each endpoint, we never know when to start and stop proxy,
	// so we stop the endpoint when timeout.
	timeout time.Duration

	// For utest, to mock the target real server.
	// Optional, use the address of received client packet.
	mockRealServerAddr *net.UDPAddr
}

// NewProxy create a proxy, the router for this proxy belongs/bind to. If need to proxy for
// please create a new proxy for each router. For all addresses we proxy, we will create a
// vnet.Net in this router and proxy all packets.
func NewProxy(router *Router) (*UDPProxy, error) {
	v := &UDPProxy{router: router, timeout: 2 * time.Minute}
	return v, nil
}

// Close the proxy, stop all workers.
func (v *UDPProxy) Close() error {
	// nolint:godox // TODO: FIXME: Do cleanup.
	return nil
}

// Proxy starts a worker for server, ignore if already started.
func (v *UDPProxy) Proxy(client *Net, server *net.UDPAddr) error {
	// Note that even if the worker exists, it's also ok to create a same worker,
	// because the router will use the last one, and the real server will see a address
	// change event after we switch to the next worker.
	if _, ok := v.workers.Load(server.String()); ok {
		// nolint:godox // TODO: Need to restart the stopped worker?
		return nil
	}

	// Not exists, create a new one.
	worker := &aUDPProxyWorker{
		router: v.router, mockRealServerAddr: v.mockRealServerAddr,
	}
	v.workers.Store(server.String(), worker)

	return worker.Proxy(client, server)
}

// A proxy worker for a specified proxy server.
type aUDPProxyWorker struct {
	router             *Router
	mockRealServerAddr *net.UDPAddr

	// Each vnet source, bind to a real socket to server.
	// key is vnet client addr, which is net.Addr
	// value is *net.UDPConn
	endpoints sync.Map
}

func (v *aUDPProxyWorker) Proxy(client *Net, serverAddr *net.UDPAddr) error { // nolint:gocognit
	// Create vnet for real server by serverAddr.
	nw := NewNet(&NetConfig{
		StaticIP: serverAddr.IP.String(),
	})
	if err := v.router.AddNet(nw); err != nil {
		return err
	}

	// We must create a "same" vnet.UDPConn as the net.UDPConn,
	// which has the same ip:port, to copy packets between them.
	vnetSocket, err := nw.ListenUDP("udp4", serverAddr)
	if err != nil {
		return err
	}

	// Got new vnet client, start a new endpoint.
	findEndpointBy := func(addr net.Addr) (*net.UDPConn, error) {
		// Exists binding.
		if value, ok := v.endpoints.Load(addr.String()); ok {
			// Exists endpoint, reuse it.
			return value.(*net.UDPConn), nil
		}

		// The real server we proxy to, for utest to mock it.
		realAddr := serverAddr
		if v.mockRealServerAddr != nil {
			realAddr = v.mockRealServerAddr
		}

		// Got new vnet client, create new endpoint.
		realSocket, err := net.DialUDP("udp4", nil, realAddr)
		if err != nil {
			return nil, err
		}

		// Bind address.
		v.endpoints.Store(addr.String(), realSocket)

		// Got packet from real serverAddr, we should proxy it to vnet.
		// nolint:godox // TODO: FIXME: Do cleanup.
		go func(vnetClientAddr net.Addr) {
			buf := make([]byte, 1500)
			for {
				n, _, err := realSocket.ReadFrom(buf)
				if err != nil {
					return
				}

				if n <= 0 {
					continue // Drop packet
				}

				if _, err := vnetSocket.WriteTo(buf[:n], vnetClientAddr); err != nil {
					return
				}
			}
		}(addr)

		return realSocket, nil
	}

	// Start a proxy goroutine.
	// nolint:godox // TODO: FIXME: Do cleanup.
	go func() {
		buf := make([]byte, 1500)

		for {
			n, addr, err := vnetSocket.ReadFrom(buf)
			if err != nil {
				return
			}

			if n <= 0 || addr == nil {
				continue // Drop packet
			}

			realSocket, err := findEndpointBy(addr)
			if err != nil {
				continue // Drop packet.
			}

			if _, err := realSocket.Write(buf[:n]); err != nil {
				return
			}
		}
	}()

	return nil
}