From d33099ba17c6493addc1b320890cf8a98ec0e21d Mon Sep 17 00:00:00 2001
From: Rutger Broekhoff
Date: Sat, 22 May 2021 11:20:41 +0200
Subject: Make Base32 encoder more like the standard library

---
 src/base32.zig | 40 ++++++++++++++++++++++++----------------
 src/nkeys.zig  | 18 +++++++++---------
 2 files changed, 33 insertions(+), 25 deletions(-)

diff --git a/src/base32.zig b/src/base32.zig
index 26622d8..149e5f0 100644
--- a/src/base32.zig
+++ b/src/base32.zig
@@ -7,8 +7,26 @@ pub const Encoder = struct {
     index: ?usize,
     bit_off: u3,
 
-    fn bitmask(self: *const Self) u8 {
-        return @as(u8, 0b11111000) >> self.bit_off;
+    pub fn init(buffer: []const u8) Encoder {
+        return .{
+            .buffer = buffer,
+            .index = 0,
+            .bit_off = 0,
+        };
+    }
+
+    pub fn calcSize(source_len: usize) usize {
+        const source_len_bits = source_len * 8;
+        return source_len_bits / 5 + (if (source_len_bits % 5 > 0) @as(usize, 1) else 0);
+    }
+
+    pub fn encode(dest: []u8, source: []const u8) []const u8 {
+        const out_len = calcSize(source.len);
+        std.debug.assert(dest.len >= out_len);
+
+        var e = init(source);
+        for (dest) |*b| b.* = e.next() orelse unreachable;
+        return dest[0..out_len];
     }
 
     fn n_front_bits(self: *const Self) u3 {
@@ -34,7 +52,8 @@ pub const Encoder = struct {
         // 5         0b00000111   2           0b11100
         // 6         0b00000011   3           0b11000
         // 7         0b00000001   4           0b10000
-        const bits = self.buffer[index] & self.bitmask();
+        const bitmask = @as(u8, 0b11111000) >> self.bit_off;
+        const bits = self.buffer[index] & bitmask;
         if (self.bit_off >= 4) return @truncate(u5, bits << (self.bit_off - 3));
         return @truncate(u5, bits >> (3 - self.bit_off));
     }
@@ -66,10 +85,12 @@ pub const Encoder = struct {
         return front_bits | back_bits;
     }
 
+    // Returns the corresponding ASCII character for 5 bits of the input.
     fn char(unencoded: u5) u8 {
         return unencoded + (if (unencoded < 26) @as(u8, 'A') else '2' - 26);
     }
 
+    // Returns the next byte of the encoded buffer.
     pub fn next(self: *Self) ?u8 {
         const unencoded = self.next_u5() orelse return null;
         return char(unencoded);
@@ -121,24 +142,11 @@ pub const Decoder = struct {
     }
 };
 
-pub fn encodedLen(src_len: usize) usize {
-    const src_len_bits = src_len * 8;
-    return src_len_bits / 5 + (if (src_len_bits % 5 > 0) @as(usize, 1) else 0);
-}
-
 pub fn decodedLen(enc_len: usize) usize {
     const enc_len_bits = enc_len * 5;
     return enc_len_bits / 8;
 }
 
-pub fn encode(bs: []const u8, out: []u8) usize {
-    var e = Encoder{ .buffer = bs, .index = 0, .bit_off = 0 };
-    for (out) |*b, i| {
-        b.* = e.next() orelse return i;
-    }
-    return out.len;
-}
-
 pub fn decode(ps: []const u8, out: []u8) DecodeError!usize {
     var d = Decoder{};
     var i: usize = 0;
diff --git a/src/nkeys.zig b/src/nkeys.zig
index b8693ae..818fd98 100644
--- a/src/nkeys.zig
+++ b/src/nkeys.zig
@@ -211,9 +211,9 @@ const binary_public_size = 1 + Ed25519.public_length + 2;
 // Two prefix bytes, two CRC bytes
 const binary_seed_size = 2 + Ed25519.seed_length + 2;
 
-pub const text_private_len = base32.encodedLen(binary_private_size);
-pub const text_public_len = base32.encodedLen(binary_public_size);
-pub const text_seed_len = base32.encodedLen(binary_seed_size);
+pub const text_private_len = base32.Encoder.calcSize(binary_private_size);
+pub const text_public_len = base32.Encoder.calcSize(binary_public_size);
+pub const text_seed_len = base32.Encoder.calcSize(binary_seed_size);
 
 pub const text_private = [text_private_len]u8;
 pub const text_public = [text_public_len]u8;
@@ -227,8 +227,8 @@ pub fn encodePrivate(key: *const [Ed25519.secret_length]u8) !text_private {
     return encode(1, key.len, &[_]u8{@enumToInt(KeyTypePrefixByte.private)}, key);
 }
 
-fn EncodedKey(comptime prefix_len: usize, comptime data_len: usize) type {
-    return [base32.encodedLen(prefix_len + data_len + 2)]u8;
+fn encoded_key(comptime prefix_len: usize, comptime data_len: usize) type {
+    return [base32.Encoder.calcSize(prefix_len + data_len + 2)]u8;
 }
 
 fn encode(
@@ -236,7 +236,7 @@ fn encode(
     comptime data_len: usize,
     prefix: *const [prefix_len]u8,
     data: *const [data_len]u8,
-) !EncodedKey(prefix_len, data_len) {
+) !encoded_key(prefix_len, data_len) {
     var buf: [prefix_len + data_len + 2]u8 = undefined;
     defer wipeBytes(&buf);
 
@@ -246,8 +246,8 @@ fn encode(
     var checksum = crc16.make(buf[0..off]);
     mem.writeIntLittle(u16, buf[buf.len - 2 .. buf.len], checksum);
 
-    var text: EncodedKey(prefix_len, data_len) = undefined;
-    std.debug.assert(base32.encode(&buf, &text) == text.len);
+    var text: encoded_key(prefix_len, data_len) = undefined;
+    _ = base32.Encoder.encode(&text, &buf);
 
     return text;
 }
@@ -285,7 +285,7 @@ fn DecodedNKey(comptime prefix_len: usize, comptime data_len: usize) type {
 fn decode(
     comptime prefix_len: usize,
     comptime data_len: usize,
-    text: *const [base32.encodedLen(prefix_len + data_len + 2)]u8,
+    text: *const [base32.Encoder.calcSize(prefix_len + data_len + 2)]u8,
 ) !DecodedNKey(prefix_len, data_len) {
     var raw: [prefix_len + data_len + 2]u8 = undefined;
     defer wipeBytes(&raw);
-- 
cgit v1.2.3