UDP support

src/shadowsocks.zig

@@ -1,6 +1,7 @@
 pub const crypto = @import("shadowsocks/crypto.zig");
 pub const headers = @import("shadowsocks/headers.zig");
 pub const server = @import("shadowsocks/server.zig");
+pub const udp_server = @import("shadowsocks/udp_server.zig");
 pub const client = @import("shadowsocks/client.zig");
 
 const std = @import("std");
@@ -9,6 +10,7 @@ const network = @import("network");
 test {
     _ = @import("shadowsocks/tests.zig");
     _ = @import("shadowsocks/salts.zig");
+    _ = @import("shadowsocks/udp_server.zig");
 }
 
 test "FixedLengthRequestHeader - derive, encode, encrypt, decrypt, decode" {

src/shadowsocks/crypto.zig

@@ -19,6 +19,8 @@ fn Crypto(comptime method_name: []const u8, comptime TAlg: type) type {
     std.debug.assert(TAlg.nonce_length == 12);
 
     return struct {
+        pub const algorithm = TAlg;
+
         pub const name = method_name;
         pub const tag_length = TAlg.tag_length;
         pub const key_length = TAlg.key_length;
@@ -36,6 +38,17 @@ fn Crypto(comptime method_name: []const u8, comptime TAlg: type) type {
             return session_subkey;
         }
 
+        pub fn deriveSessionSubkeyWithSaltUdp(key: [key_length]u8, salt: [8]u8) [key_length]u8 {
+            var key_and_salt: [key.len + salt.len]u8 = undefined;
+            std.mem.copy(u8, key_and_salt[0..key.len], &key);
+            std.mem.copy(u8, key_and_salt[key.len .. key.len + salt.len], &salt);
+
+            var session_subkey: [key_length]u8 = undefined;
+            deriveSessionSubkey(&key_and_salt, &session_subkey);
+
+            return session_subkey;
+        }
+
         pub fn generateRandomSalt() ![salt_length]u8 {
             var seed: [std.rand.DefaultCsprng.secret_seed_length]u8 = undefined;
             try std.os.getrandom(&seed);
@@ -69,6 +82,41 @@ fn Crypto(comptime method_name: []const u8, comptime TAlg: type) type {
                 self.nonce += 1;
             }
         };
+
+        const TBlockAlg = switch (key_length) {
+            16 => std.crypto.core.aes.Aes128,
+            32 => std.crypto.core.aes.Aes256,
+            else => unreachable,
+        };
+        const block_size: usize = 16;
+
+        pub const BlockEncryptor = struct {
+            block_alg: std.crypto.core.aes.AesEncryptCtx(TBlockAlg),
+
+            pub fn init(key: [key_length]u8) @This() {
+                return .{
+                    .block_alg = TBlockAlg.initEnc(key),
+                };
+            }
+
+            pub fn encrypt(self: @This(), message: *const [block_size]u8, encrypted: *[block_size]u8) void {
+                self.block_alg.encrypt(encrypted, message);
+            }
+        };
+
+        pub const BlockDecryptor = struct {
+            block_alg: std.crypto.core.aes.AesDecryptCtx(TBlockAlg),
+
+            pub fn init(key: [key_length]u8) @This() {
+                return .{
+                    .block_alg = TBlockAlg.initDec(key),
+                };
+            }
+
+            pub fn decrypt(self: @This(), message: *[block_size]u8, encrypted: *const [block_size]u8) void {
+                self.block_alg.decrypt(message, encrypted);
+            }
+        };
     };
 }
 
@@ -87,6 +135,8 @@ pub const Methods = [_]type{
 };
 
 pub const Encryptor = Blake3Aes256Gcm.Encryptor;
+pub const BlockEncryptor = Blake3Aes256Gcm.BlockEncryptor;
+pub const BlockDecryptor = Blake3Aes256Gcm.BlockDecryptor;
 pub const generateRandomSalt = Blake3Aes256Gcm.generateRandomSalt;
 pub const deriveSessionSubkeyWithSalt = Blake3Aes256Gcm.deriveSessionSubkeyWithSalt;
 
@@ -164,3 +214,18 @@ test "Test decrypt real data" {
     var message: [11]u8 = undefined;
     try Aes256Gcm.decrypt(&message, fixed, tag.*, "", nonce, session_subkey);
 }
+
+test "Block encrypt decrypt" {
+    const key = [_]u8{1} ** 32;
+    const message: []const u8 = "abcdefghijklmnop";
+
+    const encryptor = BlockEncryptor.init(key);
+    var encrypted_message: [16]u8 = undefined;
+    encryptor.encrypt(message[0..16], &encrypted_message);
+
+    const decryptor = BlockDecryptor.init(key);
+    var restored_message: [16]u8 = undefined;
+    decryptor.decrypt(&restored_message, &encrypted_message);
+
+    try std.testing.expectEqualStrings(message, &restored_message);
+}

src/shadowsocks/headers.zig

@@ -1,13 +1,38 @@
 const std = @import("std");
 
-fn DecodeResult(comptime T: type) type {
+pub fn readSocksAddress(address_type: u8, reader: anytype, allocator: std.mem.Allocator) ![]u8 {
+    return switch (address_type) {
+        1 => {
+            var addr: []u8 = try allocator.alloc(u8, 4);
+            errdefer allocator.free(addr);
+            try reader.readNoEof(addr);
+            return addr;
+        },
+        3 => {
+            const address_length = try reader.readIntBig(u8);
+            var addr: []u8 = try allocator.alloc(u8, address_length);
+            errdefer allocator.free(addr);
+            try reader.readNoEof(addr);
+            return addr;
+        },
+        4 => {
+            var addr: []u8 = try allocator.alloc(u8, 16);
+            errdefer allocator.free(addr);
+            try reader.readNoEof(addr);
+            return addr;
+        },
+        else => unreachable,
+    };
+}
+
+pub fn DecodeResult(comptime T: type) type {
     return struct {
         bytes_read: usize,
         result: T,
     };
 }
 
-fn DecodeResultWithDeinit(comptime T: type) type {
+pub fn DecodeResultWithDeinit(comptime T: type) type {
     return struct {
         bytes_read: usize,
         result: T,
@@ -73,30 +98,8 @@ pub const VariableLengthRequestHeader = struct {
         const start_pos = reader.context.pos;
 
         const address_type = try reader.readIntBig(u8);
-        var address: []u8 = add: {
-            switch (address_type) {
-                1 => {
-                    var addr: []u8 = try allocator.alloc(u8, 4);
-                    errdefer allocator.free(addr);
-                    try reader.readNoEof(addr);
-                    break :add addr;
-                },
-                3 => {
-                    const address_length = try reader.readIntBig(u8);
-                    var addr: []u8 = try allocator.alloc(u8, address_length);
-                    errdefer allocator.free(addr);
-                    try reader.readNoEof(addr);
-                    break :add addr;
-                },
-                4 => {
-                    var addr: []u8 = try allocator.alloc(u8, 16);
-                    errdefer allocator.free(addr);
-                    try reader.readNoEof(addr);
-                    break :add addr;
-                },
-                else => unreachable,
-            }
-        };
+        var address = try readSocksAddress(address_type, reader, allocator);
+        errdefer allocator.free(address);
 
         const port = try reader.readIntBig(u16);