socks5/udp_netip.go at master · wzshiming/socks5 · GitHub

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
package socks5

import (
	"bytes"
	"net"
	"net/netip"
)

// ReadFromUDPAddrPort implements the net.UDPConn ReadFromUDPAddrPort method.
func (c *UDPConn) ReadFromUDPAddrPort(b []byte) (n int, addr netip.AddrPort, err error) {
	udpConn, ok := c.PacketConn.(*net.UDPConn)
	if !ok {
		return 0, addr, errUnsupportedMethod
	}
	n, addr, err = udpConn.ReadFromUDPAddrPort(c.bufRead[:])
	if err != nil {
		return 0, addr, err
	}
	if n < len(c.prefix) || addr.String() != c.proxyAddress.String() {
		return 0, addr, errBadHeader
	}
	buf := bytes.NewBuffer(c.bufRead[len(c.prefix):n])
	a, err := readAddr(buf)
	if err != nil {
		return 0, addr, err
	}
	n = copy(b, buf.Bytes())
	netipaddr, err := netip.ParseAddr(a.IP.String())
	if err != nil {
		return 0, addr, err
	}
	return n, netip.AddrPortFrom(netipaddr, uint16(a.Port)), nil
}

// ReadMsgUDPAddrPort implements the net.UDPConn ReadMsgUDPAddrPort method.
func (c *UDPConn) ReadMsgUDPAddrPort(b, oob []byte) (n, oobn, flags int, addr netip.AddrPort, err error) {
	udpConn, ok := c.PacketConn.(*net.UDPConn)
	if !ok {
		return 0, 0, 0, addr, errUnsupportedMethod
	}
	n, oobn, flags, addr, err = udpConn.ReadMsgUDPAddrPort(c.bufRead[:], oob)
	if err != nil {
		return 0, 0, 0, addr, err
	}
	if n < len(c.prefix) || addr.String() != c.proxyAddress.String() {
		return 0, 0, 0, addr, errBadHeader
	}
	buf := bytes.NewBuffer(c.bufRead[len(c.prefix):n])
	a, err := readAddr(buf)
	if err != nil {
		return 0, 0, 0, addr, err
	}
	n = copy(b, buf.Bytes())
	netipaddr, err := netip.ParseAddr(a.IP.String())
	if err != nil {
		return 0, 0, 0, addr, err
	}
	return n, oobn, flags, netip.AddrPortFrom(netipaddr, uint16(a.Port)), nil
}

// WriteToUDPAddrPort implements the net.UDPConn WriteToUDPAddrPort method.
func (c *UDPConn) WriteToUDPAddrPort(b []byte, addr netip.AddrPort) (n int, err error) {
	udpConn, ok := c.PacketConn.(*net.UDPConn)
	if !ok {
		return 0, errUnsupportedMethod
	}
	buf := bytes.NewBuffer(c.bufWrite[:0])
	buf.Write(c.prefix)
	err = writeAddrWithStr(buf, addr.String())
	if err != nil {
		return 0, err
	}
	_, err = buf.Write(b)
	if err != nil {
		return 0, err
	}

	netipaddrport, err := netip.ParseAddrPort(c.proxyAddress.String())
	if err != nil {
		return 0, err
	}

	data := buf.Bytes()
	_, err = udpConn.WriteToUDPAddrPort(data, netipaddrport)
	if err != nil {
		return 0, err
	}
	return len(b), nil
}

// WriteMsgUDPAddrPort implements the net.UDPConn WriteMsgUDPAddrPort method.
func (c *UDPConn) WriteMsgUDPAddrPort(b, oob []byte, addr netip.AddrPort) (n, oobn int, err error) {
	udpConn, ok := c.PacketConn.(*net.UDPConn)
	if !ok {
		return 0, 0, errUnsupportedMethod
	}

	buf := bytes.NewBuffer(c.bufWrite[:0])
	buf.Write(c.prefix)
	err = writeAddrWithStr(buf, addr.String())
	if err != nil {
		return 0, 0, err
	}

	_, err = buf.Write(b)
	if err != nil {
		return 0, 0, err
	}

	netipaddrport, err := netip.ParseAddrPort(c.proxyAddress.String())
	if err != nil {
		return 0, 0, err
	}

	data := buf.Bytes()
	_, _, err = udpConn.WriteMsgUDPAddrPort(data, oob, netipaddrport)
	if err != nil {
		return 0, 0, err
	}
	return len(b), len(oob), nil
}