forked from Kuingsmile/clash-core
/
packet.go
79 lines (69 loc) · 1.94 KB
/
packet.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
package shadowstream
import (
"crypto/rand"
"errors"
"io"
"net"
"github.com/pfugwtg/clash-core/common/pool"
)
// ErrShortPacket means the packet is too short to be a valid encrypted packet.
var ErrShortPacket = errors.New("short packet")
// Pack encrypts plaintext using stream cipher s and a random IV.
// Returns a slice of dst containing random IV and ciphertext.
// Ensure len(dst) >= s.IVSize() + len(plaintext).
func Pack(dst, plaintext []byte, s Cipher) ([]byte, error) {
if len(dst) < s.IVSize()+len(plaintext) {
return nil, io.ErrShortBuffer
}
iv := dst[:s.IVSize()]
_, err := rand.Read(iv)
if err != nil {
return nil, err
}
s.Encrypter(iv).XORKeyStream(dst[len(iv):], plaintext)
return dst[:len(iv)+len(plaintext)], nil
}
// Unpack decrypts pkt using stream cipher s.
// Returns a slice of dst containing decrypted plaintext.
func Unpack(dst, pkt []byte, s Cipher) ([]byte, error) {
if len(pkt) < s.IVSize() {
return nil, ErrShortPacket
}
if len(dst) < len(pkt)-s.IVSize() {
return nil, io.ErrShortBuffer
}
iv := pkt[:s.IVSize()]
s.Decrypter(iv).XORKeyStream(dst, pkt[len(iv):])
return dst[:len(pkt)-len(iv)], nil
}
type PacketConn struct {
net.PacketConn
Cipher
}
// NewPacketConn wraps a net.PacketConn with stream cipher encryption/decryption.
func NewPacketConn(c net.PacketConn, ciph Cipher) *PacketConn {
return &PacketConn{PacketConn: c, Cipher: ciph}
}
const maxPacketSize = 64 * 1024
func (c *PacketConn) WriteTo(b []byte, addr net.Addr) (int, error) {
buf := pool.Get(maxPacketSize)
defer pool.Put(buf)
buf, err := Pack(buf, b, c.Cipher)
if err != nil {
return 0, err
}
_, err = c.PacketConn.WriteTo(buf, addr)
return len(b), err
}
func (c *PacketConn) ReadFrom(b []byte) (int, net.Addr, error) {
n, addr, err := c.PacketConn.ReadFrom(b)
if err != nil {
return n, addr, err
}
bb, err := Unpack(b[c.IVSize():], b[:n], c.Cipher)
if err != nil {
return n, addr, err
}
copy(b, bb)
return len(bb), addr, err
}