From ba9e6b06b6663207ad6b0841ccfa078d2569e526 Mon Sep 17 00:00:00 2001
From: Andreas Huber <andreas.huber@lucares.de>
Date: Sun, 8 Sep 2024 11:48:28 +0200
Subject: [PATCH] fix some memory leaks and add better error handling

---
 src/base32.zig |  3 ++-
 src/main.zig   | 52 +++++++++++++++++++++++++++++++++++++++++++++++---
 2 files changed, 51 insertions(+), 4 deletions(-)

diff --git a/src/base32.zig b/src/base32.zig
index d83a258..2a01af8 100644
--- a/src/base32.zig
+++ b/src/base32.zig
@@ -3,7 +3,7 @@ const expect = std.testing.expect;
 const expectEqualStrings = std.testing.expectEqualStrings;
 const Allocator = std.mem.Allocator;
 
-const Base32Error = error{ InvalidLength, InvalidCharacter, InvalidPadding };
+pub const Base32Error = error{ InvalidCharacter, InvalidPadding };
 
 pub const Base32 = struct {
     ///
@@ -11,6 +11,7 @@ pub const Base32 = struct {
         const alphabet = "ABCDEFGHIJKLMNOPQRSTUVWXYZ234567";
         const expectedLength = data.len / 8 * 5 + (data.len % 8) * 5 / 8;
         const result = try allocator.alloc(u8, expectedLength);
+        errdefer allocator.free(result);
 
         var dest_index: usize = 0;
         var bytes: u40 = 0;
diff --git a/src/main.zig b/src/main.zig
index e1b5cdf..8fd6690 100644
--- a/src/main.zig
+++ b/src/main.zig
@@ -8,11 +8,32 @@ const testing = std.testing;
 const expect = testing.expect;
 const ArrayList = std.ArrayList;
 const Base32 = @import("base32.zig").Base32;
+const Base32Error = @import("base32.zig").Base32Error;
 
 // otpauth://totp/AWS+Dev?secret=47STA47VFCMMLLWOLHWO3KY7MYNC36MLCDTHOLIYKJCTTSSAMKVM7YA3VWT2AJEP&digits=6&icon=Amazon
 // otpauth://totp/Gitea%20%28git.lucares.de%29:andi?algorithm=SHA1&digits=6&issuer=Gitea&period=30&secret=MSP53Q672UJMSCLQVRCJKMMZKK7MWSMYFL77OQ24JBM65RZWY7F2Y45FMTNLYM36
 
 pub fn main() !void {
+    mainInternal() catch |err| {
+        switch (err) {
+            Base32Error.InvalidCharacter => {
+                try std.io.getStdErr().writer().print("The secret is invalid.\n", .{});
+            },
+            Base32Error.InvalidPadding => {
+                try std.io.getStdErr().writer().print("The secret is invalid.\n", .{});
+            },
+            ArgumentError.UnknownParameter => {
+                try std.io.getStdErr().writer().print("Unknown argument\n", .{});
+            },
+            else => |leftover_err| {
+                try std.io.getStdErr().writer().print("{?}\n", .{leftover_err});
+            },
+        }
+        std.process.exit(1);
+    };
+}
+
+fn mainInternal() !void {
     var gpa = std.heap.GeneralPurposeAllocator(.{}){};
     const allocator = gpa.allocator();
 
@@ -28,12 +49,11 @@ pub fn main() !void {
         }
     } else if (arg.show != null) {
         const authenticator = try getAuthenticator(allocator, arg.show.?, arg.config_location);
-        try std.io.getStdOut().writer().print("{s}\n", .{try authenticator.code(allocator, std.time.timestamp)});
+        const code = try authenticator.code(allocator, std.time.timestamp);
+        try std.io.getStdOut().writer().print("{s}\n", .{code});
     } else {
         printHelp();
     }
-
-    //std.process.exit(returnValue);
 }
 
 fn printHelp() void {
@@ -383,3 +403,29 @@ fn _closure_return_2000000000() i64 {
 fn _closure_return_20000000000() i64 {
     return 20000000000;
 }
+
+test "secret is not Base32 - invalid padding" {
+    const authenticator = Authenticator{ .url = OtpAuthUrl{ .name = "", .secretEncoded = "H", .url = "", .period = 60, .digits = 6 } };
+
+    const actualError = authenticator.code(std.testing.allocator, _closure_return_1725695340);
+    try std.testing.expectError(Base32Error.InvalidPadding, actualError);
+}
+test "secret is not Base32 - invalid padding, input has too many bytes" {
+    const authenticator = Authenticator{ .url = OtpAuthUrl{ .name = "", .secretEncoded = "HAX", .url = "", .period = 60, .digits = 6 } };
+
+    const actualError = authenticator.code(std.testing.allocator, _closure_return_1725695340);
+    try std.testing.expectError(Base32Error.InvalidPadding, actualError);
+}
+test "secret is not Base32 - invalid padding, trailing bits are not 0" {
+    const authenticator = Authenticator{ .url = OtpAuthUrl{ .name = "", .secretEncoded = "HX", .url = "", .period = 60, .digits = 6 } };
+
+    const actualError = authenticator.code(std.testing.allocator, _closure_return_1725695340);
+    try std.testing.expectError(Base32Error.InvalidPadding, actualError);
+}
+
+test "secret is not Base32 - invalid character" {
+    const authenticator = Authenticator{ .url = OtpAuthUrl{ .name = "", .secretEncoded = "0", .url = "", .period = 60, .digits = 6 } };
+
+    const actualError = authenticator.code(std.testing.allocator, _closure_return_1725695340);
+    try std.testing.expectError(Base32Error.InvalidCharacter, actualError);
+}