State

pub fn State(comptime f: u11, comptime capacity: u11, comptime rounds: u5) type

File

Code

pub fn State(comptime f: u11, comptime capacity: u11, comptime rounds: u5) type {
    comptime assert(f >= 200 and f <= 1600 and f % 200 == 0); // invalid state size
    comptime assert(capacity < f and capacity % 8 == 0); // invalid capacity size

    // In debug mode, track transitions to prevent insecure ones.
    const Op = enum { uninitialized, initialized, updated, absorb, squeeze };
    const TransitionTracker = if (mode == .debug) struct {
        op: Op = .uninitialized,

        fn to(tracker: *@This(), next_op: Op) void {
            switch (next_op) {
                .updated => {
                    switch (tracker.op) {
                        .uninitialized => @panic("cannot permute before initializing"),
                        else => {},
                    }
                },
                .absorb => {
                    switch (tracker.op) {
                        .squeeze => @panic("cannot absorb right after squeezing"),
                        else => {},
                    }
                },
                .squeeze => {
                    switch (tracker.op) {
                        .uninitialized => @panic("cannot squeeze before initializing"),
                        .initialized => @panic("cannot squeeze right after initializing"),
                        .absorb => @panic("cannot squeeze right after absorbing"),
                        else => {},
                    }
                },
                .uninitialized => @panic("cannot transition to uninitialized"),
                .initialized => {},
            }
            tracker.op = next_op;
        }
    } else struct {
        // No-op in non-debug modes.
        inline fn to(tracker: *@This(), next_op: Op) void {
            _ = tracker; // no-op
            _ = next_op; // no-op
        }
    };

    return struct {
        const Self = @This();

        /// The block length, or rate, in bytes.
        pub const rate = KeccakF(f).block_bytes - capacity / 8;
        /// Keccak does not have any options.
        pub const Options = struct {};

        /// The input delimiter.
        delim: u8,

        offset: usize = 0,
        buf: [rate]u8 = undefined,

        st: KeccakF(f) = .{},

        transition: TransitionTracker = .{},

        /// Absorb a slice of bytes into the sponge.
        pub fn absorb(self: *Self, bytes: []const u8) void {
            self.transition.to(.absorb);
            var i: usize = 0;
            if (self.offset > 0) {
                const left = @min(rate - self.offset, bytes.len);
                @memcpy(self.buf[self.offset..][0..left], bytes[0..left]);
                self.offset += left;
                if (left == bytes.len) return;
                if (self.offset == rate) {
                    self.st.addBytes(self.buf[0..]);
                    self.st.permuteR(rounds);
                    self.offset = 0;
                }
                i = left;
            }
            while (i + rate < bytes.len) : (i += rate) {
                self.st.addBytes(bytes[i..][0..rate]);
                self.st.permuteR(rounds);
            }
            const left = bytes.len - i;
            if (left > 0) {
                @memcpy(self.buf[0..left], bytes[i..][0..left]);
            }
            self.offset = left;
        }

        /// Initialize the state from a slice of bytes.
        pub fn init(bytes: [f / 8]u8, delim: u8) Self {
            var st = Self{ .st = KeccakF(f).init(bytes), .delim = delim };
            st.transition.to(.initialized);
            return st;
        }

        /// Permute the state
        pub fn permute(self: *Self) void {
            if (mode == .debug) {
                if (self.transition.op == .absorb and self.offset > 0) {
                    @panic("cannot permute with pending input - call fillBlock() or pad() instead");
                }
            }
            self.transition.to(.updated);
            self.st.permuteR(rounds);
            self.offset = 0;
        }

        /// Align the input to the rate boundary and permute.
        pub fn fillBlock(self: *Self) void {
            self.transition.to(.absorb);
            self.st.addBytes(self.buf[0..self.offset]);
            self.st.permuteR(rounds);
            self.offset = 0;
            self.transition.to(.updated);
        }

        /// Mark the end of the input.
        pub fn pad(self: *Self) void {
            self.transition.to(.absorb);
            self.st.addBytes(self.buf[0..self.offset]);
            if (self.offset == rate) {
                self.st.permuteR(rounds);
                self.offset = 0;
            }
            self.st.addByte(self.delim, self.offset);
            self.st.addByte(0x80, rate - 1);
            self.st.permuteR(rounds);
            self.offset = 0;
            self.transition.to(.updated);
        }

        /// Squeeze a slice of bytes from the sponge.
        /// The function can be called multiple times.
        pub fn squeeze(self: *Self, out: []u8) void {
            self.transition.to(.squeeze);
            var i: usize = 0;
            if (self.offset == rate) {
                self.st.permuteR(rounds);
            } else if (self.offset > 0) {
                @branchHint(.unlikely);
                var buf: [rate]u8 = undefined;
                self.st.extractBytes(buf[0..]);
                const left = @min(rate - self.offset, out.len);
                @memcpy(out[0..left], buf[self.offset..][0..left]);
                self.offset += left;
                if (left == out.len) return;
                if (self.offset == rate) {
                    self.offset = 0;
                    self.st.permuteR(rounds);
                }
                i = left;
            }
            while (i + rate < out.len) : (i += rate) {
                self.st.extractBytes(out[i..][0..rate]);
                self.st.permuteR(rounds);
            }
            const left = out.len - i;
            if (left > 0) {
                self.st.extractBytes(out[i..][0..left]);
            }
            self.offset = left;
        }
    };
}