base64

pub const base64 = struct {
    /// The base64 variant to use.
    pub const Variant = packed struct {
        /// Use the URL-safe alphabet instead of the standard alphabet.
        urlsafe_alphabet: bool = false,
        /// Enable padding with '=' characters.
        padding: bool = true,

        /// The standard base64 variant.
        pub const standard: Variant = .{ .urlsafe_alphabet = false, .padding = true };
        /// The URL-safe base64 variant.
        pub const urlsafe: Variant = .{ .urlsafe_alphabet = true, .padding = true };
        /// The standard base64 variant without padding.
        pub const standard_nopad: Variant = .{ .urlsafe_alphabet = false, .padding = false };
        /// The URL-safe base64 variant without padding.
        pub const urlsafe_nopad: Variant = .{ .urlsafe_alphabet = true, .padding = false };
    };

    /// Returns the length of the encoded base64 string for a given length.
    pub fn encodedLen(bin_len: usize, variant: Variant) usize {
        if (variant.padding) {
            return (bin_len + 2) / 3 * 4;
        } else {
            const leftover = bin_len % 3;
            return bin_len / 3 * 4 + (leftover * 4 + 2) / 3;
        }
    }

    /// Returns the maximum possible decoded size for a given input length - The actual length may be less if the input includes padding.
    /// `InvalidPadding` is returned if the input length is not valid.
    pub fn decodedLen(b64_len: usize, variant: Variant) !usize {
        var result = b64_len / 4 * 3;
        const leftover = b64_len % 4;
        if (variant.padding) {
            if (leftover % 4 != 0) return error.InvalidPadding;
        } else {
            if (leftover % 4 == 1) return error.InvalidPadding;
            result += leftover * 3 / 4;
        }
        return result;
    }

    /// Encodes a binary buffer into a base64 string.
    /// The output buffer must be at least `encodedLen(bin.len)` bytes long.
    pub fn encode(encoded: []u8, bin: []const u8, comptime variant: Variant) error{NoSpaceLeft}![]const u8 {
        var acc_len: u4 = 0;
        var b64_pos: usize = 0;
        var acc: u16 = 0;
        const nibbles = bin.len / 3;
        const remainder = bin.len - 3 * nibbles;
        var b64_len = nibbles * 4;
        if (remainder != 0) {
            b64_len += if (variant.padding) 4 else 2 + (remainder >> 1);
        }
        if (encoded.len < b64_len) {
            return error.NoSpaceLeft;
        }
        const urlsafe = variant.urlsafe_alphabet;
        for (bin) |v| {
            acc = (acc << 8) + v;
            acc_len += 8;
            while (acc_len >= 6) {
                acc_len -= 6;
                encoded[b64_pos] = charFromByte(@as(u6, @truncate(acc >> acc_len)), urlsafe);
                b64_pos += 1;
            }
        }
        if (acc_len > 0) {
            encoded[b64_pos] = charFromByte(@as(u6, @truncate(acc << (6 - acc_len))), urlsafe);
            b64_pos += 1;
        }
        while (b64_pos < b64_len) {
            encoded[b64_pos] = '=';
            b64_pos += 1;
        }
        return encoded[0..b64_pos];
    }

    /// Decodes a base64 string into a binary buffer.
    /// The output buffer must be at least `decodedLenUpperBound(encoded.len)` bytes long.
    pub fn decode(bin: []u8, encoded: []const u8, comptime variant: Variant) error{ InvalidCharacter, InvalidPadding }![]const u8 {
        return decodeAny(bin, encoded, variant, null) catch |err| {
            switch (err) {
                error.InvalidCharacter, error.InvalidPadding => |e| return e,
                else => unreachable,
            }
        };
    }

    //// A decoder that ignores certain characters.
    pub const DecoderWithIgnore = struct {
        /// The characters to ignore.
        ignored_chars: StaticBitSet(256) = undefined,

        /// Decodes a base64 string into a binary buffer.
        /// The output buffer must be at least `decodedLenUpperBound(encoded.len)` bytes long.
        pub fn decode(
            self: DecoderWithIgnore,
            bin: []u8,
            encoded: []const u8,
            comptime variant: Variant,
        ) error{ NoSpaceLeft, InvalidCharacter, InvalidPadding }![]const u8 {
            return decodeAny(bin, encoded, variant, self.ignored_chars);
        }

        /// Returns the decoded length of a base64 string, ignoring any characters in the ignore list.
        /// This operation does not run in constant time, but it aims to avoid leaking information about the underlying base64 string.
        pub fn decodedLenForSlice(decoder: DecoderWithIgnore, encoded: []const u8, variant: Variant) !usize {
            var b64_len = encoded.len;
            for (encoded) |c| {
                if (decoder.ignored_chars.isSet(c)) b64_len -= 1;
            }
            return base64.decodedLen(b64_len, variant);
        }

        /// Returns the maximum possible decoded size for a given input length after skipping ignored characters.
        pub fn decodedLenUpperBound(b64_len: usize) usize {
            return b64_len / 3 * 4;
        }
    };

    /// Creates a new decoder that ignores certain characters.
    pub fn decoderWithIgnore(ignore_chars: []const u8) error{InvalidCharacter}!DecoderWithIgnore {
        var ignored_chars = StaticBitSet(256).empty;
        for (ignore_chars) |c| {
            switch (c) {
                'A'...'Z', 'a'...'z', '0'...'9' => return error.InvalidCharacter,
                else => if (ignored_chars.isSet(c)) return error.InvalidCharacter,
            }
            ignored_chars.set(c);
        }
        return DecoderWithIgnore{ .ignored_chars = ignored_chars };
    }

    fn eq(x: u8, y: u8) u8 {
        return ~@as(u8, @truncate((0 -% (@as(u16, x) ^ @as(u16, y))) >> 8));
    }

    fn gt(x: u8, y: u8) u8 {
        return @truncate((@as(u16, y) -% @as(u16, x)) >> 8);
    }

    fn ge(x: u8, y: u8) u8 {
        return ~gt(y, x);
    }

    fn lt(x: u8, y: u8) u8 {
        return gt(y, x);
    }

    fn le(x: u8, y: u8) u8 {
        return ge(y, x);
    }

    fn charFromByte(x: u8, comptime urlsafe: bool) u8 {
        return (lt(x, 26) & (x +% 'A')) |
            (ge(x, 26) & lt(x, 52) & (x +% 'a' -% 26)) |
            (ge(x, 52) & lt(x, 62) & (x +% '0' -% 52)) |
            (eq(x, 62) & if (urlsafe) '-' else '+') | (eq(x, 63) & if (urlsafe) '_' else '/');
    }

    fn byteFromChar(c: u8, comptime urlsafe: bool) u8 {
        const x =
            (ge(c, 'A') & le(c, 'Z') & (c -% 'A')) |
            (ge(c, 'a') & le(c, 'z') & (c -% 'a' +% 26)) |
            (ge(c, '0') & le(c, '9') & (c -% '0' +% 52)) |
            (eq(c, if (urlsafe) '-' else '+') & 62) | (eq(c, if (urlsafe) '_' else '/') & 63);
        return x | (eq(x, 0) & ~eq(c, 'A'));
    }

    fn skipPadding(
        encoded: []const u8,
        padding_len: usize,
        ignored_chars: ?StaticBitSet(256),
    ) error{InvalidPadding}![]const u8 {
        var b64_pos: usize = 0;
        var i = padding_len;
        while (i > 0) {
            if (b64_pos >= encoded.len) {
                return error.InvalidPadding;
            }
            const c = encoded[b64_pos];
            if (c == '=') {
                i -= 1;
            } else if (ignored_chars) |set| {
                if (!set.isSet(c)) {
                    return error.InvalidPadding;
                }
            }
            b64_pos += 1;
        }
        return encoded[b64_pos..];
    }

    fn decodeAny(
        bin: []u8,
        encoded: []const u8,
        comptime variant: Variant,
        ignored_chars: ?StaticBitSet(256),
    ) error{ NoSpaceLeft, InvalidCharacter, InvalidPadding }![]const u8 {
        var acc: u16 = 0;
        var acc_len: u4 = 0;
        var bin_pos: usize = 0;
        var premature_end: ?usize = null;
        const urlsafe = variant.urlsafe_alphabet;
        for (encoded, 0..) |c, b64_pos| {
            const d = byteFromChar(c, urlsafe);
            if (d == 0xff) {
                if (ignored_chars) |set| {
                    if (set.isSet(c)) continue;
                }
                premature_end = b64_pos;
                break;
            }
            acc = (acc << 6) + d;
            acc_len += 6;
            if (acc_len >= 8) {
                acc_len -= 8;
                if (bin_pos >= bin.len) {
                    return error.NoSpaceLeft;
                }
                bin[bin_pos] = @truncate(acc >> acc_len);
                bin_pos += 1;
            }
        }
        if (acc_len > 4 or (acc & ((@as(u16, 1) << acc_len) -% 1)) != 0) {
            return error.InvalidCharacter;
        }
        const padding_len = acc_len / 2;
        if (premature_end) |pos| {
            const remaining =
                if (variant.padding)
                    try skipPadding(encoded[pos..], padding_len, ignored_chars)
                else
                    encoded[pos..];
            if (ignored_chars) |set| {
                for (remaining) |c| {
                    if (!set.isSet(c)) {
                        return error.InvalidCharacter;
                    }
                }
            } else if (remaining.len != 0) {
                return error.InvalidCharacter;
            }
        } else if (variant.padding and padding_len != 0) {
            return error.InvalidPadding;
        }
        return bin[0..bin_pos];
    }
}