AES CBC Mode

Implement AES CBC mode. Also fix key initialization, as it turns out,
Nettle AES keys cannot be initialized for encryption and decryption at
the same time. Hence, initialize the key when encryption/decryption is
invoked instead.

build.zig

@@ -62,9 +62,20 @@ pub fn build(b: *std.Build) void {
 
     const run_exe_unit_tests = b.addRunArtifact(exe_unit_tests);
 
+    const aes_unit_tests = b.addTest(.{
+        .root_source_file = b.path("src/aes.zig"),
+        .target = target,
+        .optimize = optimize,
+    });
+    aes_unit_tests.linkSystemLibrary2("nettle", .{ .needed = true });
+    aes_unit_tests.linkLibC();
+
+    const run_aes_unit_tests = b.addRunArtifact(aes_unit_tests);
+
     // Similar to creating the run step earlier, this exposes a `test` step to
     // the `zig build --help` menu, providing a way for the user to request
     // running the unit tests.
     const test_step = b.step("test", "Run unit tests");
     test_step.dependOn(&run_exe_unit_tests.step);
+    test_step.dependOn(&run_aes_unit_tests.step);
 }

src/aes.zig

@@ -2,50 +2,193 @@ const std = @import("std");
 const nettle = @cImport({
     @cInclude("nettle/aes.h");
 });
+const xor = @import("xor.zig");
+
+pub const Mode = enum { ECB, CBC };
 
 pub const AES = struct {
     allocator: std.mem.Allocator,
     key: []const u8,
+    iv: ?[]const u8 = null,
     nettle_ctx: nettle.struct_aes_ctx,
+    mode: Mode,
 
     pub fn init(allocator: std.mem.Allocator, key: []const u8) AES {
+        return init_ecb(allocator, key);
+    }
+
+    pub fn init_ecb(allocator: std.mem.Allocator, key: []const u8) AES {
         const owned_key = allocator.dupe(u8, key) catch {
             @panic("Could not allocate key memory");
         };
+        errdefer allocator.free(owned_key);
 
-        var aes = AES{
+        const aes = AES{
             .allocator = allocator,
             .key = owned_key,
             .nettle_ctx = nettle.struct_aes_ctx{},
+            .mode = Mode.ECB,
         };
 
-        // this works only in this way: first set encrypt then invert for enabling decrypt
+        return aes;
-        // this is not commutative
+    }
-        nettle.aes_set_encrypt_key(&aes.nettle_ctx, aes.key.len, aes.key.ptr);
+
-        nettle.aes_invert_key(&aes.nettle_ctx, &aes.nettle_ctx);
+    pub fn init_cbc(allocator: std.mem.Allocator, key: []const u8, iv: []const u8) AES {
+        if (key.len != iv.len) @panic("Key and IV length must match");
+
+        const owned_key = allocator.dupe(u8, key) catch {
+            @panic("Could not allocate key memory");
+        };
+        errdefer allocator.free(owned_key);
+
+        const owned_iv = allocator.dupe(u8, iv) catch {
+            @panic("Coult not allocate key memory");
+        };
+        errdefer allocator.free(owned_iv);
+
+        const aes = AES{
+            .allocator = allocator,
+            .key = owned_key,
+            .nettle_ctx = nettle.struct_aes_ctx{},
+            .mode = Mode.CBC,
+            .iv = owned_iv,
+        };
 
         return aes;
     }
 
     /// Encrypt buffer `buf`. Must be integer multiple of keysize.
     /// Caller must free result.
-    pub fn encrypt(self: @This(), allocator: std.mem.Allocator, buf: []const u8) ![]u8 {
+    pub fn encrypt(self: *@This(), allocator: std.mem.Allocator, buf: []const u8) ![]u8 {
+        nettle.aes_set_encrypt_key(&self.nettle_ctx, self.key.len, self.key.ptr);
+
+        switch (self.mode) {
+            .ECB => return self.encrypt_ecb(allocator, buf),
+            .CBC => return self.encrypt_cbc(allocator, buf),
+        }
+    }
+
+    fn encrypt_ecb(self: @This(), allocator: std.mem.Allocator, buf: []const u8) ![]u8 {
         const enc = try allocator.alloc(u8, buf.len);
+        errdefer allocator.free(enc);
 
         nettle.aes_encrypt(&self.nettle_ctx, buf.len, enc.ptr, buf.ptr);
         return enc;
     }
 
+    fn encrypt_cbc(self: @This(), allocator: std.mem.Allocator, buf: []const u8) ![]u8 {
+        const enc = try allocator.alloc(u8, buf.len);
+        errdefer allocator.free(enc);
+
+        var arena_alloc = std.heap.ArenaAllocator.init(allocator);
+        defer arena_alloc.deinit();
+        const arena = arena_alloc.allocator();
+
+        var cbc_key = self.iv.?;
+
+        const block_size = self.key.len;
+        var block_idx: usize = 0;
+
+        var it = std.mem.window(u8, buf, block_size, block_size);
+
+        while (it.next()) |block| {
+            const cbc_block = try xor.xor_buffers(arena, cbc_key, block);
+            nettle.aes_encrypt(&self.nettle_ctx, block_size, enc[block_idx..].ptr, cbc_block.ptr);
+
+            cbc_key = enc[block_idx..(block_idx + block_size)];
+            block_idx += block_size;
+        }
+
+        return enc;
+    }
+
     /// Decrypt buffer `buf`. Must be integer multiple of keysize.
     /// Caller must free result.
-    pub fn decrypt(self: @This(), allocator: std.mem.Allocator, buf: []const u8) ![]u8 {
+    pub fn decrypt(self: *@This(), allocator: std.mem.Allocator, buf: []const u8) ![]u8 {
+        nettle.aes_set_decrypt_key(&self.nettle_ctx, self.key.len, self.key.ptr);
+
+        switch (self.mode) {
+            .ECB => return self.decrypt_ecb(allocator, buf),
+            .CBC => return self.decrypt_cbc(allocator, buf),
+        }
+    }
+
+    fn decrypt_ecb(self: @This(), allocator: std.mem.Allocator, buf: []const u8) ![]u8 {
         const dec = try allocator.alloc(u8, buf.len);
+        errdefer allocator.free(dec);
 
         nettle.aes_decrypt(&self.nettle_ctx, buf.len, dec.ptr, buf.ptr);
         return dec;
     }
 
+    fn decrypt_cbc(self: @This(), allocator: std.mem.Allocator, buf: []const u8) ![]u8 {
+        const dec = try allocator.alloc(u8, buf.len);
+        errdefer allocator.free(dec);
+
+        var arena_alloc = std.heap.ArenaAllocator.init(allocator);
+        defer arena_alloc.deinit();
+        const arena = arena_alloc.allocator();
+
+        var cbc_key = self.iv.?;
+
+        const block_size = self.key.len;
+        var block_idx: usize = 0;
+
+        var it = std.mem.window(u8, buf, block_size, block_size);
+
+        while (it.next()) |block| {
+            nettle.aes_decrypt(&self.nettle_ctx, block_size, dec[block_idx..].ptr, block.ptr);
+
+            const cbc_block = try xor.xor_buffers(arena, cbc_key, dec[block_idx..(block_idx + block_size)]);
+            std.mem.copyForwards(u8, dec[block_idx..(block_idx + block_size)], cbc_block);
+
+            cbc_key = buf[block_idx..(block_idx + block_size)];
+            block_idx += block_size;
+        }
+
+        return dec;
+    }
+
     pub fn deinit(self: @This()) void {
         self.allocator.free(self.key);
+        if (self.iv) |iv| self.allocator.free(iv);
     }
 };
+
+test "ecb" {
+    const allocator = std.testing.allocator;
+
+    const key = [_]u8{ 0x00, 0x01, 0x02, 0x03, 0x05, 0x06, 0x07, 0x08, 0x0A, 0x0B, 0x0C, 0x0D, 0x0F, 0x10, 0x11, 0x12 };
+    var data = [_]u8{ 0x50, 0x68, 0x12, 0xA4, 0x5F, 0x08, 0xC8, 0x89, 0xB9, 0x7F, 0x59, 0x80, 0x03, 0x8B, 0x83, 0x59 };
+    const cipher = [_]u8{ 0xD8, 0xF5, 0x32, 0x53, 0x82, 0x89, 0xEF, 0x7D, 0x06, 0xB5, 0x06, 0xA4, 0xFD, 0x5B, 0xE9, 0xC9 };
+
+    var aes = AES.init(allocator, &key);
+    defer aes.deinit();
+
+    for (0..5) |_| {
+        const res = try aes.encrypt(allocator, &data);
+        defer allocator.free(res);
+        try std.testing.expectEqualSlices(u8, &cipher, res);
+
+        const res2 = try aes.decrypt(allocator, res);
+        defer allocator.free(res2);
+        try std.testing.expectEqualSlices(u8, &data, res2);
+    }
+}
+
+test "cbc" {
+    var allocator = std.testing.allocator;
+
+    var aes = AES.init_cbc(allocator, "YELLOW SUBMARINE", "YELLOW SUBMARINE");
+    defer aes.deinit();
+
+    for (0..5) |_| {
+        const enc = try aes.encrypt(allocator, "YELLOW SUBMARINEYELLOW SUBMARINEYELLOW SUBMARINE");
+        defer allocator.free(enc);
+
+        const dec = try aes.decrypt(allocator, enc);
+        defer allocator.free(dec);
+
+        try std.testing.expectEqualStrings("YELLOW SUBMARINEYELLOW SUBMARINEYELLOW SUBMARINE", dec);
+    }
+}

src/main.zig

@@ -48,8 +48,13 @@ pub fn main() !void {
         try s1c8(allocator, stdout);
     }
 
-    if (std.mem.eql(u8, args[1], "s1c9")) {
+    if (std.mem.eql(u8, args[1], "s2c09")) {
-        try s1c9(allocator, stdout);
+        try s2c09(allocator, stdout);
+    }
+
+    if (std.mem.eql(u8, args[1], "s2c10")) {
+        try s2c10(allocator, stdout);
+    }
     }
 }
 
@@ -217,7 +222,7 @@ fn s1c7(allocator: std.mem.Allocator, stdout: anytype) !void {
     // decrypt
     const key = "YELLOW SUBMARINE";
 
-    const cipher = aes.AES.init(allocator, key);
+    var cipher = aes.AES.init(allocator, key);
     defer cipher.deinit();
 
     const result = try cipher.decrypt(allocator, input_decoded);
@@ -258,7 +263,7 @@ fn s1c8(allocator: std.mem.Allocator, stdout: anytype) !void {
     }
 }
 
-fn s1c9(allocator: std.mem.Allocator, stdout: anytype) !void {
+fn s2c09(allocator: std.mem.Allocator, stdout: anytype) !void {
     const input = "YELLOW SUBMARINE";
 
     const result = try padding.pkcs7(allocator, input, 20);
@@ -266,3 +271,36 @@ fn s1c9(allocator: std.mem.Allocator, stdout: anytype) !void {
 
     try stdout.print("{s}", .{result});
 }
+
+fn s2c10(allocator: std.mem.Allocator, stdout: anytype) !void {
+    // prepare input
+    const f = @embedFile("res/10.txt");
+
+    var f_stream = std.io.fixedBufferStream(f);
+    const reader = f_stream.reader();
+
+    var joined_lines = std.ArrayList(u8).init(allocator);
+    defer joined_lines.deinit();
+    const joined_lines_writer = joined_lines.writer();
+
+    while (true) {
+        reader.streamUntilDelimiter(joined_lines_writer, '\n', null) catch |err| {
+            if (err == error.EndOfStream) break else return err;
+        };
+    }
+
+    const input_decoded = try b64.decode(allocator, joined_lines.items);
+    defer allocator.free(input_decoded);
+
+    // decrypt
+    const key = "YELLOW SUBMARINE";
+    const iv = [_]u8{0x00} ** 16;
+
+    var cipher = aes.AES.init_cbc(allocator, key, &iv);
+    defer cipher.deinit();
+
+    const result = try cipher.decrypt(allocator, input_decoded);
+    defer allocator.free(result);
+
+    try stdout.print("{s}", .{result});
+}

src/res/10.txt

@@ -0,0 +1,64 @@
+CRIwqt4+szDbqkNY+I0qbNXPg1XLaCM5etQ5Bt9DRFV/xIN2k8Go7jtArLIy
+P605b071DL8C+FPYSHOXPkMMMFPAKm+Nsu0nCBMQVt9mlluHbVE/yl6VaBCj
+NuOGvHZ9WYvt51uR/lklZZ0ObqD5UaC1rupZwCEK4pIWf6JQ4pTyPjyiPtKX
+g54FNQvbVIHeotUG2kHEvHGS/w2Tt4E42xEwVfi29J3yp0O/TcL7aoRZIcJj
+MV4qxY/uvZLGsjo1/IyhtQp3vY0nSzJjGgaLYXpvRn8TaAcEtH3cqZenBoox
+BH3MxNjD/TVf3NastEWGnqeGp+0D9bQx/3L0+xTf+k2VjBDrV9HPXNELRgPN
+0MlNo79p2gEwWjfTbx2KbF6htgsbGgCMZ6/iCshy3R8/abxkl8eK/VfCGfA6
+bQQkqs91bgsT0RgxXSWzjjvh4eXTSl8xYoMDCGa2opN/b6Q2MdfvW7rEvp5m
+wJOfQFDtkv4M5cFEO3sjmU9MReRnCpvalG3ark0XC589rm+42jC4/oFWUdwv
+kzGkSeoabAJdEJCifhvtGosYgvQDARUoNTQAO1+CbnwdKnA/WbQ59S9MU61Q
+KcYSuk+jK5nAMDot2dPmvxZIeqbB6ax1IH0cdVx7qB/Z2FlJ/U927xGmC/RU
+FwoXQDRqL05L22wEiF85HKx2XRVB0F7keglwX/kl4gga5rk3YrZ7VbInPpxU
+zgEaE4+BDoEqbv/rYMuaeOuBIkVchmzXwlpPORwbN0/RUL89xwOJKCQQZM8B
+1YsYOqeL3HGxKfpFo7kmArXSRKRHToXuBgDq07KS/jxaS1a1Paz/tvYHjLxw
+Y0Ot3kS+cnBeq/FGSNL/fFV3J2a8eVvydsKat3XZS3WKcNNjY2ZEY1rHgcGL
+5bhVHs67bxb/IGQleyY+EwLuv5eUwS3wljJkGcWeFhlqxNXQ6NDTzRNlBS0W
+4CkNiDBMegCcOlPKC2ZLGw2ejgr2utoNfmRtehr+3LAhLMVjLyPSRQ/zDhHj
+Xu+Kmt4elmTmqLgAUskiOiLYpr0zI7Pb4xsEkcxRFX9rKy5WV7NhJ1lR7BKy
+alO94jWIL4kJmh4GoUEhO+vDCNtW49PEgQkundV8vmzxKarUHZ0xr4feL1ZJ
+THinyUs/KUAJAZSAQ1Zx/S4dNj1HuchZzDDm/nE/Y3DeDhhNUwpggmesLDxF
+tqJJ/BRn8cgwM6/SMFDWUnhkX/t8qJrHphcxBjAmIdIWxDi2d78LA6xhEPUw
+NdPPhUrJcu5hvhDVXcceZLa+rJEmn4aftHm6/Q06WH7dq4RaaJePP6WHvQDp
+zZJOIMSEisApfh3QvHqdbiybZdyErz+yXjPXlKWG90kOz6fx+GbvGcHqibb/
+HUfcDosYA7lY4xY17llY5sibvWM91ohFN5jyDlHtngi7nWQgFcDNfSh77TDT
+zltUp9NnSJSgNOOwoSSNWadm6+AgbXfQNX6oJFaU4LQiAsRNa7vX/9jRfi65
+5uvujM4ob199CZVxEls10UI9pIemAQQ8z/3rgQ3eyL+fViyztUPg/2IvxOHv
+eexE4owH4Fo/bRlhZK0mYIamVxsRADBuBlGqx1b0OuF4AoZZgUM4d8v3iyUu
+feh0QQqOkvJK/svkYHn3mf4JlUb2MTgtRQNYdZKDRgF3Q0IJaZuMyPWFsSNT
+YauWjMVqnj0AEDHh6QUMF8bXLM0jGwANP+r4yPdKJNsoZMpuVoUBJYWnDTV+
+8Ive6ZgBi4EEbPbMLXuqDMpDi4XcLE0UUPJ8VnmO5fAHMQkA64esY2QqldZ+
+5gEhjigueZjEf0917/X53ZYWJIRiICnmYPoM0GSYJRE0k3ycdlzZzljIGk+P
+Q7WgeJhthisEBDbgTuppqKNXLbNZZG/VaTdbpW1ylBv0eqamFOmyrTyh1APS
+Gn37comTI3fmN6/wmVnmV4/FblvVwLuDvGgSCGPOF8i6FVfKvdESs+yr+1AE
+DJXfp6h0eNEUsM3gXaJCknGhnt3awtg1fSUiwpYfDKZxwpPOYUuer8Wi+VCD
+sWsUpkMxhhRqOBKaQaBDQG+kVJu6aPFlnSPQQTi1hxLwi0l0Rr38xkr+lHU7
+ix8LeJVgNsQdtxbovE3i7z3ZcTFY7uJkI9j9E0muDN9x8y/YN25rm6zULYaO
+jUoP/7FQZsSgxPIUvUiXkEq+FU2h0FqAC7H18cr3Za5x5dpw5nwawMArKoqG
+9qlhqc34lXV0ZYwULu58EImFIS8+kITFuu7jOeSXbBgbhx8zGPqavRXeiu0t
+bJd0gWs+YgMLzXtQIbQuVZENMxJSZB4aw5lPA4vr1fFBsiU4unjOEo/XAgwr
+Tc0w0UndJFPvXRr3Ir5rFoIEOdRo+6os5DSlk82SBnUjwbje7BWsxWMkVhYO
+6bOGUm4VxcKWXu2jU66TxQVIHy7WHktMjioVlWJdZC5Hq0g1LHg1nWSmjPY2
+c/odZqN+dBBC51dCt4oi5UKmKtU5gjZsRSTcTlfhGUd6DY4Tp3CZhHjQRH4l
+Zhg0bF/ooPTxIjLKK4r0+yR0lyRjqIYEY27HJMhZDXFDxBQQ1UkUIhAvXacD
+WB2pb3YyeSQjt8j/WSbQY6TzdLq8SreZiuMWcXmQk4EH3xu8bPsHlcvRI+B3
+gxKeLnwrVJqVLkf3m2cSGnWQhSLGbnAtgQPA6z7u3gGbBmRtP0KnAHWSK7q6
+onMoYTH+b5iFjCiVRqzUBVzRRKjAL4rcL2nYeV6Ec3PlnboRzJwZIjD6i7WC
+dcxERr4WVOjOBX4fhhKUiVvlmlcu8CkIiSnZENHZCpI41ypoVqVarHpqh2aP
+/PS624yfxx2N3C2ci7VIuH3DcSYcaTXEKhz/PRLJXkRgVlWxn7QuaJJzDvpB
+oFndoRu1+XCsup/AtkLidsSXMFTo/2Ka739+BgYDuRt1mE9EyuYyCMoxO/27
+sn1QWMMd1jtcv8Ze42MaM4y/PhAMp2RfCoVZALUS2K7XrOLl3s9LDFOdSrfD
+8GeMciBbfLGoXDvv5Oqq0S/OvjdID94UMcadpnSNsist/kcJJV0wtRGfALG2
++UKYzEj/2TOiN75UlRvA5XgwfqajOvmIIXybbdhxpjnSB04X3iY82TNSYTmL
+LAzZlX2vmV9IKRRimZ2SpzNpvLKeB8lDhIyGzGXdiynQjFMNcVjZlmWHsH7e
+ItAKWmCwNkeuAfFwir4TTGrgG1pMje7XA7kMT821cYbLSiPAwtlC0wm77F0T
+a7jdMrLjMO29+1958CEzWPdzdfqKzlfBzsba0+dS6mcW/YTHaB4bDyXechZB
+k/35fUg+4geMj6PBTqLNNWXBX93dFC7fNyda+Lt9cVJnlhIi/61fr0KzxOeX
+NKgePKOC3Rz+fWw7Bm58FlYTgRgN63yFWSKl4sMfzihaQq0R8NMQIOjzuMl3
+Ie5ozSa+y9g4z52RRc69l4n4qzf0aErV/BEe7FrzRyWh4PkDj5wy5ECaRbfO
+7rbs1EHlshFvXfGlLdEfP2kKpT9U32NKZ4h+Gr9ymqZ6isb1KfNov1rw0KSq
+YNP+EyWCyLRJ3EcOYdvVwVb+vIiyzxnRdugB3vNzaNljHG5ypEJQaTLphIQn
+lP02xcBpMNJN69bijVtnASN/TLV5ocYvtnWPTBKu3OyOkcflMaHCEUgHPW0f
+mGfld4i9Tu35zrKvTDzfxkJX7+KJ72d/V+ksNKWvwn/wvMOZsa2EEOfdCidm
+oql027IS5XvSHynQtvFmw0HTk9UXt8HdVNTqcdy/jUFmXpXNP2Wvn8PrU2Dh
+kkIzWhQ5Rxd/vnM2QQr9Cxa2J9GXEV3kGDiZV90+PCDSVGY4VgF8y7GedI1h