Finish implementation of ChaCha20.

src/primitive/streamcipher/chacha20.zig

@@ -1,6 +1,10 @@
 const std = @import("std");
 const testing = std.testing;
 
+// ----------------------------------- ERROR DEFINITIONS -----------------------------------  //
+
+const KeyStreamDepleted = error.KeyStreamDepleted;
+
 // ----------------------------------- ChaCha20 CONSTANTS -----------------------------------  //
 
 const CHACHA20_BLOCK_SIZE = 512 / 8;
@@ -29,14 +33,16 @@ pub const ChaCha20_RFC7539_Parameters = struct {
     pub const COUNTER_SIZE = 32 / 8;
 };
 
-// ----------------------------------- ENCRYPTION/DECRYPTION -----------------------------------  //
+// ----------------------------------- CONTEXT MANAGEMENT -----------------------------------  //
 
 pub const ChaCha20Ctx = struct {
     key: [CHACHA20_KEY_WORDS]u32,
     nonce: [CHACHA20_NONCE_WORDS]u32,
     state: [CHACHA20_BLOCK_WORDS]u32,
     working_state: [CHACHA20_BLOCK_WORDS]u32,
-    keystream_idx: u6,
+    counter_idx_lsw: u8,
+    counter_idx_msw: u8,
+    keystream_idx: u8,
 };
 
 pub fn chacha20_new(
@@ -47,19 +53,21 @@ pub fn chacha20_new(
     nonce: *const [nonce_size]u8,
 ) ChaCha20Ctx {
     if (comptime counter_size + nonce_size != CHACHA20_NONCE_SIZE)
-        @panic("Invalid ChaCha initialization: The size of counter and nonce must add up to 16 bytes.");
+        @panic("Invalid ChaCha initialization: The lengths of the counter and nonce must add up to 16 bytes.");
+
+    const counter_words = comptime counter_size / 4;
+    const nonce_words = comptime nonce_size / 4;
 
     var ctx = ChaCha20Ctx{
         .key = undefined,
         .nonce = undefined,
         .state = undefined,
         .working_state = undefined,
+        .counter_idx_lsw = 0,
+        .counter_idx_msw = counter_words - 1,
         .keystream_idx = undefined,
     };
 
-    const counter_words = comptime counter_size / 4;
-    const nonce_words = comptime nonce_size / 4;
-
     chacha20_deserialize(CHACHA20_KEY_WORDS, key, &ctx.key);
     chacha20_deserialize(counter_words, counter, ctx.nonce[0..counter_words]);
     chacha20_deserialize(nonce_words, nonce, ctx.nonce[counter_words .. counter_words + nonce_words]);
@@ -71,9 +79,9 @@ pub fn chacha20_new(
 }
 
 pub fn chacha20_destroy(ctx: *ChaCha20Ctx) void {
-    @memset(ctx.state, 0);
+    @memset(&ctx.state, 0);
-    @memset(ctx.key, 0);
+    @memset(&ctx.key, 0);
-    @memset(ctx.nonce, 0);
+    @memset(&ctx.nonce, 0);
     ctx.keystream_idx = 0;
 }
 
@@ -105,13 +113,44 @@ pub fn chacha20_rfc7539_new(
     );
 }
 
-pub fn chacha20_quarter_round(
+// ----------------------------------- ENCRYPTION/DECRYPTION -----------------------------------  //
-    state: *[CHACHA20_BLOCK_WORDS]u32,
+
-    ia: u8,
+pub fn chacha20_encrypt_inplace(ctx: *ChaCha20Ctx, plaintext: []u8) !void {
-    ib: u8,
+    for (0..plaintext.len) |i| {
-    ic: u8,
+        try chacha20_ensure_keystream_expanded(ctx);
-    id: u8,
+        const keystream: [*]const u8 = @ptrCast(&ctx.state);
-) void {
+
+        plaintext[i] ^= keystream[ctx.keystream_idx];
+        ctx.keystream_idx += 1;
+    }
+}
+
+pub fn chacha20_decrypt_inplace(ctx: *ChaCha20Ctx, ciphertext: []u8) !void {
+    return chacha20_encrypt_inplace(ctx, ciphertext);
+}
+
+fn chacha20_ensure_keystream_expanded(ctx: *ChaCha20Ctx) !void {
+    if (ctx.keystream_idx == CHACHA20_BLOCK_SIZE) {
+        try chacha20_increment_counter(ctx);
+        chacha20_block_function(ctx);
+        ctx.keystream_idx = 0;
+    }
+}
+
+fn chacha20_increment_counter(ctx: *ChaCha20Ctx) !void {
+    for (ctx.counter_idx_lsw..ctx.counter_idx_msw + 1) |idx| {
+        const ov = @addWithOverflow(ctx.nonce[idx], 1);
+        ctx.nonce[idx] = ov[0];
+
+        if (ov[1] == 0)
+            return;
+    }
+    return KeyStreamDepleted;
+}
+
+// ----------------------------------- KEYSTREAM EXPANSION -----------------------------------  //
+
+pub fn chacha20_quarter_round(state: *[CHACHA20_BLOCK_WORDS]u32, ia: u8, ib: u8, ic: u8, id: u8) void {
     var a = state[ia];
     var b = state[ib];
     var c = state[ic];
@@ -235,21 +274,12 @@ test "ChaCha Quarter Round" {
         0x53372767, 0xb00a5631, 0x974c541a, 0x359e9963,
         0x5c971061, 0x3d631689, 0x2098d9d6, 0x91dbd320,
     };
-    // 0xe0, 0x31, 0x95, 0x87, 0x7d, 0xf3, 0xec, 0xc5, 0xb1, 0x61, 0x64, 0x51, 0x8a, 0x2f, 0xa6, 0xc9,
-    // 0xf3, 0x0e, 0xc2, 0x44, 0x7f, 0xaf, 0x90, 0x33, 0x0b, 0x69, 0xfc, 0xd9, 0x4c, 0x71, 0x5f, 0x2a,
-    // 0x67, 0x27, 0x37, 0x53, 0x31, 0x56, 0x0a, 0xb0, 0x1a, 0x54, 0x4c, 0x97, 0x63, 0x99, 0x9e, 0x35,
-    // 0x61, 0x10, 0x97, 0x5c, 0x89, 0x16, 0x63, 0x3d, 0xd6, 0xd9, 0x98, 0x20, 0x20, 0xd3, 0xdb, 0x91,
-
     const reference = [CHACHA20_BLOCK_WORDS]u32{
         0x879531e0, 0xc5ecf37d, 0xbdb886dc, 0xc9a62f8a,
         0x44c20ef3, 0x3390af7f, 0xd9fc690b, 0xcfacafd2,
         0xe46bea80, 0xb00a5631, 0x974c541a, 0x359e9963,
         0x5c971061, 0xccc07c79, 0x2098d9d6, 0x91dbd320,
     };
-    // 0xe0, 0x31, 0x95, 0x87, 0x7d, 0xf3, 0xec, 0xc5, 0xdc, 0x86, 0xb8, 0xbd, 0x8a, 0x2f, 0xa6, 0xc9,
-    // 0xf3, 0x0e, 0xc2, 0x44, 0x7f, 0xaf, 0x90, 0x33, 0x0b, 0x69, 0xfc, 0xd9, 0xd2, 0xaf, 0xac, 0xcf,
-    // 0x80, 0xea, 0x6b, 0xe4, 0x31, 0x56, 0x0a, 0xb0, 0x1a, 0x54, 0x4c, 0x97, 0x63, 0x99, 0x9e, 0x35,
-    // 0x61, 0x10, 0x97, 0x5c, 0x79, 0x7c, 0xc0, 0xcc, 0xd6, 0xd9, 0x98, 0x20, 0x20, 0xd3, 0xdb, 0x91,
 
     chacha20_quarter_round(&state, 2, 7, 8, 13);
     try testing.expectEqualSlices(u32, reference[0..], state[0..]);
@@ -257,14 +287,14 @@ test "ChaCha Quarter Round" {
 
 // https://www.rfc-editor.org/rfc/rfc7539#section-2.3.2
 test "ChaCha20 Block Function" {
-    var chacha = chacha20_rfc7539_new(
+    const key = [_]u8{
-        &.{
         0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f,
         0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, 0x18, 0x19, 0x1a, 0x1b, 0x1c, 0x1d, 0x1e, 0x1f,
-        },
+    };
-        &.{ 0x00, 0x00, 0x00, 0x09, 0x00, 0x00, 0x00, 0x4a, 0x00, 0x00, 0x00, 0x00 },
+    const nonce = [_]u8{
-        &.{ 0x01, 0x00, 0x00, 0x00 },
+        0x00, 0x00, 0x00, 0x09, 0x00, 0x00, 0x00, 0x4a, 0x00, 0x00, 0x00, 0x00,
-    );
+    };
+
     const reference = [CHACHA20_BLOCK_SIZE]u8{
         0x10, 0xf1, 0xe7, 0xe4, 0xd1, 0x3b, 0x59, 0x15, 0x50, 0x0f, 0xdd, 0x1f, 0xa3, 0x20, 0x71, 0xc4,
         0xc7, 0xd1, 0xf4, 0xc7, 0x33, 0xc0, 0x68, 0x03, 0x04, 0x22, 0xaa, 0x9a, 0xc3, 0xd4, 0x6c, 0x4e,
@@ -272,6 +302,9 @@ test "ChaCha20 Block Function" {
         0xb5, 0x12, 0x9c, 0xd1, 0xde, 0x16, 0x4e, 0xb9, 0xcb, 0xd0, 0x83, 0xe8, 0xa2, 0x50, 0x3c, 0x4e,
     };
 
+    var chacha = chacha20_rfc7539_new(&key, &nonce, &word_to_bytes_le(1));
+    defer chacha20_destroy(&chacha);
+
     chacha20_block_function(&chacha);
 
     var buffer: [CHACHA20_BLOCK_SIZE]u8 = undefined;
@@ -279,3 +312,35 @@ test "ChaCha20 Block Function" {
 
     try testing.expectEqualSlices(u8, reference[0..], buffer[0..]);
 }
+
+// https://www.rfc-editor.org/rfc/rfc7539#section-2.4.2
+test "ChaCha20 Cipher" {
+    const key = [CHACHA20_KEY_SIZE]u8{
+        0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f,
+        0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, 0x18, 0x19, 0x1a, 0x1b, 0x1c, 0x1d, 0x1e, 0x1f,
+    };
+    const nonce = [ChaCha20_RFC7539_Parameters.NONCE_SIZE]u8{
+        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x4a, 0x00, 0x00, 0x00, 0x00,
+    };
+    const counter = word_to_bytes_le(1);
+    const plaintext = "Ladies and Gentlemen of the class of '99: If I could offer you only one tip for the future, sunscreen would be it.";
+    const reference = [_]u8{
+        0x6e, 0x2e, 0x35, 0x9a, 0x25, 0x68, 0xf9, 0x80, 0x41, 0xba, 0x07, 0x28, 0xdd, 0x0d, 0x69, 0x81,
+        0xe9, 0x7e, 0x7a, 0xec, 0x1d, 0x43, 0x60, 0xc2, 0x0a, 0x27, 0xaf, 0xcc, 0xfd, 0x9f, 0xae, 0x0b,
+        0xf9, 0x1b, 0x65, 0xc5, 0x52, 0x47, 0x33, 0xab, 0x8f, 0x59, 0x3d, 0xab, 0xcd, 0x62, 0xb3, 0x57,
+        0x16, 0x39, 0xd6, 0x24, 0xe6, 0x51, 0x52, 0xab, 0x8f, 0x53, 0x0c, 0x35, 0x9f, 0x08, 0x61, 0xd8,
+        0x07, 0xca, 0x0d, 0xbf, 0x50, 0x0d, 0x6a, 0x61, 0x56, 0xa3, 0x8e, 0x08, 0x8a, 0x22, 0xb6, 0x5e,
+        0x52, 0xbc, 0x51, 0x4d, 0x16, 0xcc, 0xf8, 0x06, 0x81, 0x8c, 0xe9, 0x1a, 0xb7, 0x79, 0x37, 0x36,
+        0x5a, 0xf9, 0x0b, 0xbf, 0x74, 0xa3, 0x5b, 0xe6, 0xb4, 0x0b, 0x8e, 0xed, 0xf2, 0x78, 0x5e, 0x42,
+        0x87, 0x4d,
+    };
+
+    var chacha = chacha20_rfc7539_new(&key, &nonce, &counter);
+    defer chacha20_destroy(&chacha);
+
+    var buffer: [plaintext.len]u8 = undefined;
+    @memcpy(&buffer, plaintext);
+
+    try chacha20_encrypt_inplace(&chacha, &buffer);
+    try testing.expectEqualSlices(u8, reference[0..], buffer[0..]);
+}