Skip to content

Commit dbb8048

Browse files
committed
Add other parts of session init.
1 parent a9ea5d4 commit dbb8048

8 files changed

Lines changed: 134 additions & 51 deletions

File tree

internal/config/server.go

Lines changed: 54 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -18,21 +18,23 @@ import (
1818
)
1919

2020
type ServerConfig struct {
21-
ConfigDir string `toml:"-"`
22-
ConfigPath string `toml:"-"`
23-
UDPHost string `toml:"UDP_HOST"`
24-
UDPPort int `toml:"UDP_PORT"`
25-
UDPReaders int `toml:"UDP_READERS"`
26-
SocketBufferSize int `toml:"SOCKET_BUFFER_SIZE"`
27-
MaxConcurrentRequests int `toml:"MAX_CONCURRENT_REQUESTS"`
28-
DNSRequestWorkers int `toml:"DNS_REQUEST_WORKERS"`
29-
MaxPacketSize int `toml:"MAX_PACKET_SIZE"`
30-
DropLogIntervalSecs float64 `toml:"DROP_LOG_INTERVAL_SECONDS"`
31-
Domain []string `toml:"DOMAIN"`
32-
MinVPNLabelLength int `toml:"MIN_VPN_LABEL_LENGTH"`
33-
DataEncryptionMethod int `toml:"DATA_ENCRYPTION_METHOD"`
34-
EncryptionKeyFile string `toml:"ENCRYPTION_KEY_FILE"`
35-
LogLevel string `toml:"LOG_LEVEL"`
21+
ConfigDir string `toml:"-"`
22+
ConfigPath string `toml:"-"`
23+
UDPHost string `toml:"UDP_HOST"`
24+
UDPPort int `toml:"UDP_PORT"`
25+
UDPReaders int `toml:"UDP_READERS"`
26+
SocketBufferSize int `toml:"SOCKET_BUFFER_SIZE"`
27+
MaxConcurrentRequests int `toml:"MAX_CONCURRENT_REQUESTS"`
28+
DNSRequestWorkers int `toml:"DNS_REQUEST_WORKERS"`
29+
MaxPacketSize int `toml:"MAX_PACKET_SIZE"`
30+
DropLogIntervalSecs float64 `toml:"DROP_LOG_INTERVAL_SECONDS"`
31+
Domain []string `toml:"DOMAIN"`
32+
MinVPNLabelLength int `toml:"MIN_VPN_LABEL_LENGTH"`
33+
SupportedUploadCompressionTypes []int `toml:"SUPPORTED_UPLOAD_COMPRESSION_TYPES"`
34+
SupportedDownloadCompressionTypes []int `toml:"SUPPORTED_DOWNLOAD_COMPRESSION_TYPES"`
35+
DataEncryptionMethod int `toml:"DATA_ENCRYPTION_METHOD"`
36+
EncryptionKeyFile string `toml:"ENCRYPTION_KEY_FILE"`
37+
LogLevel string `toml:"LOG_LEVEL"`
3638
}
3739

3840
func defaultServerConfig() ServerConfig {
@@ -41,19 +43,21 @@ func defaultServerConfig() ServerConfig {
4143
readers := min(max(runtime.NumCPU()/2, 1), 4)
4244

4345
return ServerConfig{
44-
UDPHost: "0.0.0.0",
45-
UDPPort: 53,
46-
UDPReaders: readers,
47-
SocketBufferSize: 8 * 1024 * 1024,
48-
MaxConcurrentRequests: 4096,
49-
DNSRequestWorkers: workers,
50-
MaxPacketSize: 65535,
51-
DropLogIntervalSecs: 2.0,
52-
Domain: nil,
53-
MinVPNLabelLength: 3,
54-
DataEncryptionMethod: 1,
55-
EncryptionKeyFile: "encrypt_key.txt",
56-
LogLevel: "INFO",
46+
UDPHost: "0.0.0.0",
47+
UDPPort: 53,
48+
UDPReaders: readers,
49+
SocketBufferSize: 8 * 1024 * 1024,
50+
MaxConcurrentRequests: 4096,
51+
DNSRequestWorkers: workers,
52+
MaxPacketSize: 65535,
53+
DropLogIntervalSecs: 2.0,
54+
Domain: nil,
55+
MinVPNLabelLength: 3,
56+
SupportedUploadCompressionTypes: []int{0, 1, 2, 3},
57+
SupportedDownloadCompressionTypes: []int{0, 1, 2, 3},
58+
DataEncryptionMethod: 1,
59+
EncryptionKeyFile: "encrypt_key.txt",
60+
LogLevel: "INFO",
5761
}
5862
}
5963

@@ -110,6 +114,8 @@ func LoadServerConfig(filename string) (ServerConfig, error) {
110114
if cfg.MinVPNLabelLength <= 0 {
111115
cfg.MinVPNLabelLength = 3
112116
}
117+
cfg.SupportedUploadCompressionTypes = normalizeCompressionTypeList(cfg.SupportedUploadCompressionTypes)
118+
cfg.SupportedDownloadCompressionTypes = normalizeCompressionTypeList(cfg.SupportedDownloadCompressionTypes)
113119

114120
if cfg.DataEncryptionMethod < 0 || cfg.DataEncryptionMethod > 5 {
115121
cfg.DataEncryptionMethod = 1
@@ -143,3 +149,23 @@ func (c ServerConfig) EncryptionKeyPath() string {
143149
}
144150
return filepath.Join(c.ConfigDir, c.EncryptionKeyFile)
145151
}
152+
153+
func normalizeCompressionTypeList(values []int) []int {
154+
if len(values) == 0 {
155+
return []int{0}
156+
}
157+
158+
seen := [4]bool{}
159+
out := make([]int, 0, len(values))
160+
for _, value := range values {
161+
if value < 0 || value > 3 || seen[value] {
162+
continue
163+
}
164+
seen[value] = true
165+
out = append(out, value)
166+
}
167+
if len(out) == 0 {
168+
return []int{0}
169+
}
170+
return out
171+
}

internal/enums/dns.go

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,8 +14,6 @@ const (
1414
PacketMTUDownRes = 0x04
1515
PacketSessionInit = 0x05
1616
PacketSessionAccept = 0x06
17-
PacketSetMTUReq = 0x07
18-
PacketSetMTURes = 0x08
1917
PacketPing = 0x09
2018
PacketPong = 0x0A
2119
PacketStreamSyn = 0x0B

internal/enums/dns_test.go

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -32,8 +32,6 @@ func TestPacketEnumValuesAreUnique(t *testing.T) {
3232
PacketMTUDownRes,
3333
PacketSessionInit,
3434
PacketSessionAccept,
35-
PacketSetMTUReq,
36-
PacketSetMTURes,
3735
PacketPing,
3836
PacketPong,
3937
PacketStreamSyn,

internal/udpserver/server.go

Lines changed: 41 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -35,14 +35,16 @@ const (
3535
)
3636

3737
type Server struct {
38-
cfg config.ServerConfig
39-
log *logger.Logger
40-
codec *security.Codec
41-
domainMatcher *domainmatcher.Matcher
42-
sessions *sessionStore
43-
packetPool sync.Pool
44-
droppedPackets atomic.Uint64
45-
lastDropLogUnix atomic.Int64
38+
cfg config.ServerConfig
39+
log *logger.Logger
40+
codec *security.Codec
41+
domainMatcher *domainmatcher.Matcher
42+
sessions *sessionStore
43+
uploadCompressionMask uint8
44+
downloadCompressionMask uint8
45+
packetPool sync.Pool
46+
droppedPackets atomic.Uint64
47+
lastDropLogUnix atomic.Int64
4648
}
4749

4850
type request struct {
@@ -53,11 +55,13 @@ type request struct {
5355

5456
func New(cfg config.ServerConfig, log *logger.Logger, codec *security.Codec) *Server {
5557
return &Server{
56-
cfg: cfg,
57-
log: log,
58-
codec: codec,
59-
domainMatcher: domainmatcher.New(cfg.Domain, cfg.MinVPNLabelLength),
60-
sessions: newSessionStore(),
58+
cfg: cfg,
59+
log: log,
60+
codec: codec,
61+
domainMatcher: domainmatcher.New(cfg.Domain, cfg.MinVPNLabelLength),
62+
sessions: newSessionStore(),
63+
uploadCompressionMask: buildCompressionMask(cfg.SupportedUploadCompressionTypes),
64+
downloadCompressionMask: buildCompressionMask(cfg.SupportedDownloadCompressionTypes),
6165
packetPool: sync.Pool{
6266
New: func() any {
6367
return make([]byte, cfg.MaxPacketSize)
@@ -273,7 +277,11 @@ func (s *Server) handleSessionInitRequest(questionPacket []byte, decision domain
273277
if vpnPacket.SessionID != 0 {
274278
return nil
275279
}
276-
record, _, err := s.sessions.findOrCreate(vpnPacket.Payload)
280+
requestedUpload, requestedDownload := compression.SplitPair(vpnPacket.Payload[1])
281+
resolvedUpload := resolveCompressionType(requestedUpload, s.uploadCompressionMask)
282+
resolvedDownload := resolveCompressionType(requestedDownload, s.downloadCompressionMask)
283+
284+
record, _, err := s.sessions.findOrCreate(vpnPacket.Payload, resolvedUpload, resolvedDownload)
277285
if err != nil || record == nil {
278286
return nil
279287
}
@@ -295,6 +303,25 @@ func (s *Server) handleSessionInitRequest(questionPacket []byte, decision domain
295303
return response
296304
}
297305

306+
func buildCompressionMask(values []int) uint8 {
307+
var mask uint8 = 1 << compression.TypeOff
308+
for _, value := range values {
309+
if value < compression.TypeOff || value > compression.TypeZLIB {
310+
continue
311+
}
312+
mask |= 1 << uint8(value)
313+
}
314+
return mask
315+
}
316+
317+
func resolveCompressionType(requested uint8, allowedMask uint8) uint8 {
318+
requested = compression.NormalizeType(requested)
319+
if allowedMask&(1<<requested) != 0 {
320+
return requested
321+
}
322+
return compression.TypeOff
323+
}
324+
298325
func (s *Server) onDrop(addr *net.UDPAddr) {
299326
total := s.droppedPackets.Add(1)
300327

internal/udpserver/server_test.go

Lines changed: 31 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -249,11 +249,41 @@ func TestHandlePacketRejectsMalformedSessionInit(t *testing.T) {
249249
}
250250
}
251251

252+
func TestHandlePacketNegotiatesUnsupportedCompressionToOff(t *testing.T) {
253+
codec, err := security.NewCodec(0, "")
254+
if err != nil {
255+
t.Fatalf("NewCodec returned error: %v", err)
256+
}
257+
258+
srv := New(config.ServerConfig{
259+
MaxPacketSize: 65535,
260+
Domain: []string{"a.com"},
261+
MinVPNLabelLength: 3,
262+
SupportedUploadCompressionTypes: []int{0, 1},
263+
SupportedDownloadCompressionTypes: []int{0},
264+
}, nil, codec)
265+
266+
payload := []byte{1, 0x13, 0x00, 0x96, 0x00, 0xC8, 0x44, 0x33, 0x22, 0x11}
267+
query := buildTunnelQueryWithSessionID(t, codec, "a.com", 0, ENUMS.PacketSessionInit, payload)
268+
response := srv.handlePacket(query)
269+
if len(response) == 0 {
270+
t.Fatal("handlePacket should return a session accept response")
271+
}
272+
273+
packet, err := DnsParser.ExtractVPNResponse(response, true)
274+
if err != nil {
275+
t.Fatalf("ExtractVPNResponse returned error: %v", err)
276+
}
277+
if got := packet.Payload[2]; got != 0x10 {
278+
t.Fatalf("unexpected negotiated compression pair: got=%#x want=%#x", got, 0x10)
279+
}
280+
}
281+
252282
func TestSessionStoreExpiresReuseSignatureWithoutDroppingSession(t *testing.T) {
253283
store := newSessionStore()
254284
payload := []byte{1, 0x21, 0x00, 0x96, 0x00, 0xC8, 0x44, 0x33, 0x22, 0x11}
255285

256-
record, reused, err := store.findOrCreate(payload)
286+
record, reused, err := store.findOrCreate(payload, 2, 1)
257287
if err != nil {
258288
t.Fatalf("findOrCreate returned error: %v", err)
259289
}

internal/udpserver/session.go

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,7 @@ func newSessionStore() *sessionStore {
5454
}
5555
}
5656

57-
func (s *sessionStore) findOrCreate(payload []byte) (*sessionRecord, bool, error) {
57+
func (s *sessionStore) findOrCreate(payload []byte, uploadCompressionType uint8, downloadCompressionType uint8) (*sessionRecord, bool, error) {
5858
if len(payload) != sessionInitDataSize || !isValidSessionResponseMode(payload[0]) {
5959
return nil, false, nil
6060
}
@@ -89,7 +89,8 @@ func (s *sessionStore) findOrCreate(payload []byte) (*sessionRecord, bool, error
8989
ReuseUntil: now.Add(sessionInitTTL),
9090
Signature: signature,
9191
}
92-
record.UploadCompression, record.DownloadCompression = compression.SplitPair(payload[1])
92+
record.UploadCompression = compression.NormalizeType(uploadCompressionType)
93+
record.DownloadCompression = compression.NormalizeType(downloadCompressionType)
9394
record.UploadMTU = clampMTU(binary.BigEndian.Uint16(payload[2:4]))
9495
record.DownloadMTU = clampMTU(binary.BigEndian.Uint16(payload[4:6]))
9596
copy(record.VerifyCode[:], payload[6:10])

internal/vpnproto/parser.go

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -223,8 +223,6 @@ func buildPacketFlags() [256]uint8 {
223223
ENUMS.PacketMTUDownReq,
224224
ENUMS.PacketSessionInit,
225225
ENUMS.PacketSessionAccept,
226-
ENUMS.PacketSetMTUReq,
227-
ENUMS.PacketSetMTURes,
228226
ENUMS.PacketPing,
229227
ENUMS.PacketPong,
230228
ENUMS.PacketErrorDrop,

server_config.toml.simple

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,11 @@ DROP_LOG_INTERVAL_SECONDS = 2.0
3434
# Allowed tunnel domains.
3535
DOMAIN = ["v.domain.com"]
3636

37+
# Compression types allowed by server policy:
38+
# 0=OFF, 1=ZSTD, 2=LZ4, 3=ZLIB
39+
SUPPORTED_UPLOAD_COMPRESSION_TYPES = [0, 1, 2, 3]
40+
SUPPORTED_DOWNLOAD_COMPRESSION_TYPES = [0, 1, 2, 3]
41+
3742
# Encryption method:
3843
# 0=None, 1=XOR, 2=ChaCha20, 3=AES-128-GCM, 4=AES-192-GCM, 5=AES-256-GCM
3944
DATA_ENCRYPTION_METHOD = 1

0 commit comments

Comments
 (0)