1	import Base
2	import ../../src/crypto/keccak/types.bend as T
3	import ./keccak/permutation.bend as K
4	
5	# Executable specification of SHA3-256 (FIPS 202), byte oriented, in the
6	# shape of HACL*'s Spec.SHA3: the sponge construction (section 4) on
7	# KECCAK-p[1600, n_r] (section 3, spec/crypto/keccak/permutation.bend) with
8	# rate r = 1088 bits = 136 bytes, capacity 512 and a 256-bit output:
9	#
10	#   SHA3-256(M) = KECCAK[512](M || 01, 256)        (section 6.1)
11	#
12	# The permutation's round count is a parameter (24 for SHA3-256, section 3.4)
13	# so the proofs never unfold 24 rounds. Only the record types of
14	# src/crypto/keccak/types.bend are shared with the implementation: a lane is
15	# T.W{lo, hi}, the state is T.S{a0..a24} with a[x + 5y] = A[x, y].
16	# Bytes are U32 values, each < 256; bits are little-endian within a byte and
17	# bytes little-endian within a lane (section B.1).
18	
19	# ---------------------------------------------------------------- padding
20	
21	# M || 01 || pad10*1(r, |M| + 2) at the byte level, for q = 1..136 bytes to
22	# the end of the block: the suffix bits 0, 1 and the first pad bit 1 are the
23	# byte 0x06; the last pad bit is the top bit, 0x80, of the block's last byte;
24	# with one byte left both fall in it: 0x86.
25	def padding_bytes(q: Nat) -> List<&2, U32>:
26	  match q:
27	    case 0n:
28	      Nil{}
29	    case 1n:
30	      [U32.or(6, 128)]
31	    case 2n+k:
32	      6 <> List.append(&2, U32, List.replicate(U32, k, 0), [128])
33	
34	def padding(+n: Nat) -> List<&2, U32>:
35	  padding_bytes(Nat.sub(136n, Nat.mod(n, 136n)))
36	
37	def pad(+bytes: List<&2, U32>) -> List<&2, U32>:
38	  List.append(&2, U32, bytes, padding(List.length(&2, U32, bytes)))
39	
40	# ---------------------------------------------------------------- lanes
41	
42	def le32(a: U32, b: U32, c: U32, d: U32) -> U32:
43	  U32.or(U32.or(U32.or(U32.and(a, 255), U32.shln(U32.and(b, 255), 8n)),
44	    U32.shln(U32.and(c, 255), 16n)), U32.shln(U32.and(d, 255), 24n))
45	
46	# Little-endian 64-bit lanes, eight bytes each.
47	def lanes(bytes: List<&2, U32>) -> List<&2, T.Lane>:
48	  match bytes:
49	    case a <> b <> c <> d <> e <> f <> g <> h <> rest:
50	      T.W{le32(a, b, c, d), le32(e, f, g, h)} <> lanes(rest)
51	    case _:
52	      Nil{}
53	
54	# ---------------------------------------------------------------- the sponge
55	
56	# S = 0^1600.
57	def initial() -> T.State:
58	  T.S{T.W{0, 0}, T.W{0, 0}, T.W{0, 0}, T.W{0, 0}, T.W{0, 0},
59	    T.W{0, 0}, T.W{0, 0}, T.W{0, 0}, T.W{0, 0}, T.W{0, 0},
60	    T.W{0, 0}, T.W{0, 0}, T.W{0, 0}, T.W{0, 0}, T.W{0, 0},
61	    T.W{0, 0}, T.W{0, 0}, T.W{0, 0}, T.W{0, 0}, T.W{0, 0},
62	    T.W{0, 0}, T.W{0, 0}, T.W{0, 0}, T.W{0, 0}, T.W{0, 0}}
63	
64	def to_lanes(s: T.State) -> List<&2, T.Lane>:
65	  match s:
66	    case T.S{a0, a1, a2, a3, a4, a5, a6, a7, a8, a9, a10, a11, a12, a13, a14, a15, a16, a17, a18, a19, a20, a21, a22, a23, a24}:
67	      [a0, a1, a2, a3, a4, a5, a6, a7, a8, a9, a10, a11, a12, a13, a14, a15, a16, a17, a18, a19, a20, a21, a22, a23, a24]
68	
69	def of_lanes(ls: List<&2, T.Lane>) -> T.State:
70	  match ls:
71	    case a0 <> a1 <> a2 <> a3 <> a4 <> a5 <> a6 <> a7 <> a8 <> a9 <> a10 <> a11 <> a12 <> a13 <> a14 <> a15 <> a16 <> a17 <> a18 <> a19 <> a20 <> a21 <> a22 <> a23 <> a24 <> rest:
72	      T.S{a0, a1, a2, a3, a4, a5, a6, a7, a8, a9, a10, a11, a12, a13, a14, a15, a16, a17, a18, a19, a20, a21, a22, a23, a24}
73	    case _:
74	      initial()
75	
76	# Lane-wise XOR of a block into the leading lanes of the state.
77	def xor_lanes(xs: List<&2, T.Lane>, ys: List<&2, T.Lane>) -> List<&2, T.Lane>:
78	  match xs ys:
79	    case x <> xt y <> yt:
80	      K.xor(x, y) <> xor_lanes(xt, yt)
81	    case xs _:
82	      xs
83	
84	# Section 4, step 6: S = f(S XOR (P_i || 0^c)) for each r-bit block P_i.
85	def absorb(ls: List<&2, T.Lane>, +rounds: Nat, s: T.State) -> T.State:
86	  match ls:
87	    case l0 <> l1 <> l2 <> l3 <> l4 <> l5 <> l6 <> l7 <> l8 <> l9 <> l10 <> l11 <> l12 <> l13 <> l14 <> l15 <> l16 <> rest:
88	      absorb(rest, rounds, K.rounds(rounds, 0n, of_lanes(xor_lanes(to_lanes(s),
89	        [l0, l1, l2, l3, l4, l5, l6, l7, l8, l9, l10, l11, l12, l13, l14, l15, l16]))))
90	    case _:
91	      s
92	
93	# Section 4, steps 7-10: d = 256 <= r, so Z is the first 32 bytes of S.
94	def lane_bytes(l: T.Lane) -> List<&2, U32>:
95	  match l:
96	    case T.W{+lo, +hi}:
97	      [U32.and(lo, 255), U32.and(U32.shrn(lo, 8n), 255), U32.and(U32.shrn(lo, 16n), 255), U32.shrn(lo, 24n),
98	       U32.and(hi, 255), U32.and(U32.shrn(hi, 8n), 255), U32.and(U32.shrn(hi, 16n), 255), U32.shrn(hi, 24n)]
99	
100	def state_bytes(ls: List<&2, T.Lane>) -> List<&2, U32>:
101	  match ls:
102	    case Nil{}:
103	      Nil{}
104	    case l <> rest:
105	      List.append(&2, U32, lane_bytes(l), state_bytes(rest))
106	
107	def squeeze(s: T.State) -> List<&2, U32>:
108	  List.take(&2, U32, state_bytes(to_lanes(s)), 32n)
109	
110	def keccak(bytes: List<&2, U32>, +rounds: Nat) -> List<&2, U32>:
111	  squeeze(absorb(lanes(pad(bytes)), rounds, initial()))
112	
113	# SHA3-256(M): 32 bytes.
114	def sha3_256(bytes: List<&2, U32>) -> List<&2, U32>:
115	  keccak(bytes, 24n)