forked from eycorsican/go-tun2socks
-
Notifications
You must be signed in to change notification settings - Fork 0
/
udp_conn.go
154 lines (138 loc) · 3.18 KB
/
udp_conn.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
151
152
153
154
package core
/*
#cgo CFLAGS: -I./c/include
#include "lwip/udp.h"
*/
import "C"
import (
"errors"
"fmt"
"net"
"sync"
"unsafe"
)
type udpConnState uint
const (
udpConnecting udpConnState = iota
udpConnected
udpClosed
)
type udpPacket struct {
data []byte
addr *net.UDPAddr
}
type udpConn struct {
sync.Mutex
pcb *C.struct_udp_pcb
handler UDPConnHandler
localAddr *net.UDPAddr
localIP C.ip_addr_t
localPort C.u16_t
state udpConnState
pending chan *udpPacket
}
func newUDPConn(pcb *C.struct_udp_pcb, handler UDPConnHandler, localIP C.ip_addr_t, localPort C.u16_t, localAddr, remoteAddr *net.UDPAddr) (UDPConn, error) {
conn := &udpConn{
handler: handler,
pcb: pcb,
localAddr: localAddr,
localIP: localIP,
localPort: localPort,
state: udpConnecting,
pending: make(chan *udpPacket, 1), // To hold the first packet on the connection
}
go func() {
err := handler.Connect(conn, remoteAddr)
if err != nil {
conn.Close()
} else {
conn.Lock()
conn.state = udpConnected
conn.Unlock()
// Once connected, send all pending data.
DrainPending:
for {
select {
case pkt := <-conn.pending:
err := conn.handler.ReceiveTo(conn, pkt.data, pkt.addr)
if err != nil {
break DrainPending
}
continue DrainPending
default:
break DrainPending
}
}
}
}()
return conn, nil
}
func (conn *udpConn) LocalAddr() *net.UDPAddr {
return conn.localAddr
}
func (conn *udpConn) checkState() error {
conn.Lock()
defer conn.Unlock()
switch conn.state {
case udpClosed:
return errors.New("connection closed")
case udpConnected:
return nil
case udpConnecting:
return errors.New("not connected")
}
return nil
}
// If the connection isn't ready yet, and there is room in the queue, make a copy
// and hold onto it until the connection is ready.
func (conn *udpConn) enqueueEarlyPacket(data []byte, addr *net.UDPAddr) bool {
conn.Lock()
defer conn.Unlock()
if conn.state == udpConnecting {
pkt := &udpPacket{data: append([]byte(nil), data...), addr: addr}
select {
// Data will be dropped if pending is full.
case conn.pending <- pkt:
return true
default:
}
}
return false
}
func (conn *udpConn) ReceiveTo(data []byte, addr *net.UDPAddr) error {
if conn.enqueueEarlyPacket(data, addr) {
return nil
}
if err := conn.checkState(); err != nil {
return err
}
err := conn.handler.ReceiveTo(conn, data, addr)
if err != nil {
return errors.New(fmt.Sprintf("write proxy failed: %v", err))
}
return nil
}
func (conn *udpConn) WriteFrom(data []byte, addr *net.UDPAddr) (int, error) {
if err := conn.checkState(); err != nil {
return 0, err
}
// FIXME any memory leaks?
cremoteIP := C.struct_ip_addr{}
if err := ipAddrATON(addr.IP.String(), &cremoteIP); err != nil {
return 0, err
}
buf := C.pbuf_alloc_reference(unsafe.Pointer(&data[0]), C.u16_t(len(data)), C.PBUF_ROM)
defer C.pbuf_free(buf)
C.udp_sendto(conn.pcb, buf, &conn.localIP, conn.localPort, &cremoteIP, C.u16_t(addr.Port))
return len(data), nil
}
func (conn *udpConn) Close() error {
connId := udpConnId{
src: conn.LocalAddr().String(),
}
conn.Lock()
conn.state = udpClosed
conn.Unlock()
udpConns.Delete(connId)
return nil
}