1	# Key derivation using pseudorandom functions (NIST SP 800-108r1-upd1,
2	# "Recommendation for Key Derivation Using Pseudorandom Functions", L.
3	# Chen, August 2022), over any MAC (mac.bend) as the PRF:
4	#   counter mode (section 4.1):  K(i) = PRF(K_IN, [i]_2 || fixed),
5	#   feedback mode (section 4.2): K(i) = PRF(K_IN, K(i-1) {|| [i]_2} || fixed),
6	#                                K(0) = IV,
7	# for i = 1 to n, K_OUT the leftmost L bits of K(1) || K(2) || ..., with
8	# the fixed input data Label || 0x00 || Context || [L]_2 (section 4). The
9	# counter [i]_2 is r bits; the encoding of L is a field of its own length.
10	# Both are template arguments in bytes (~rlen, ~llen: 4 and 4, 32 bits, the
11	# usual choice); in feedback mode ~rlen = 0 leaves the counter out.
12	#
13	# The output is generated a byte at a time, each K(i) computed when the
14	# bytes before it are used up, so a stream has the length asked for, and a
15	# shorter stream is a prefix of a longer one with the same fixed input
16	# (LAWS.bend: kbkdf_*). A derivation of L bits binds L into every PRF input,
17	# so the derivations of two lengths are not prefixes of each other: the
18	# KDF interface's prefix law holds of the streams, not of the derivations.
19	#
20	# Refused: more than 2^r - 1 blocks in counter mode, or 2^32 - 1 in feedback
21	# mode (CrBadOutputLength); an L that does not fit its field
22	# (CrBadOutputLength); and the MAC's own refusal of K_IN. As a KDF
23	# (kdf.bend) in counter mode, the secret is K_IN, the salt the Label and the
24	# info the Context: ~CrKbkdf.derive(~K, ~key, ~compute, ~hlen, ~rlen, ~llen).
25	import Base
26	import bend-osdev-bytes@0.1.0.0/bytes.bend as B
27	import ./crypto.bend as Cr
28	
29	def cr.kb.head(xs: List<&2, U32>) -> U32:
30	  match xs:
31	    case Con{h, t}:
32	      h
33	    case Nil{}:
34	      0
35	
36	def cr.kb.tail(xs: List<&2, U32>) -> List<&2, U32>:
37	  match xs:
38	    case Con{h, t}:
39	      t
40	    case Nil{}:
41	      Nil{}
42	
43	# [i]_2 as rlen bytes (the low 8 rlen bits of i, big-endian)
44	def CrKbkdf.counter(+rlen: Nat, i: U32) -> B.Octets:
45	  B.Octets.drop(Nat.sub(4n, rlen), B.Octets.be32(i))
46	
47	# n more bytes in counter mode: the rest of the current block cur, then
48	# PRF(K_IN, [i + 1]_2 || fixed), ...
49	def cr.kb.ctr(~K: Data, ~compute: K -> B.Octets -> B.Octets, n: Nat, cur: List<&2, U32>, +k: K, +rlen: Nat, +fixed: B.Octets, +i: U32) -> List<&2, U32>:
50	  match n:
51	    case 0n:
52	      Nil{}
53	    case 1n+p:
54	      match cur:
55	        case Con{h, t}:
56	          Con{h, cr.kb.ctr(~K, ~compute, p, t, k, rlen, fixed, i)}
57	        case Nil{}:
58	          +blk = B.Octets.to_list(compute(k, B.Octets.append(CrKbkdf.counter(rlen, U32.add(i, 1)), fixed)))
59	          Con{cr.kb.head(blk), cr.kb.ctr(~K, ~compute, p, cr.kb.tail(blk), k, rlen, fixed, U32.add(i, 1))}
60	
61	# The first n bytes of counter mode under the prepared key k with the fixed
62	# input data, unchecked
63	def CrKbkdf.counter_stream(~K: Data, ~compute: K -> B.Octets -> B.Octets, ~rlen: Nat, k: K, fixed: B.Octets, n: Nat) -> B.Octets:
64	  B.Octets.of_list(cr.kb.ctr(~K, ~compute, n, Nil{}, k, rlen, fixed, 0))
65	
66	# n more bytes in feedback mode: the rest of the current block cur, then
67	# PRF(K_IN, K(i) || [i + 1]_2 || fixed), with prev = K(i)
68	def cr.kb.fb(~K: Data, ~compute: K -> B.Octets -> B.Octets, n: Nat, cur: List<&2, U32>, +k: K, +rlen: Nat, +fixed: B.Octets, +prev: B.Octets, +i: U32) -> List<&2, U32>:
69	  match n:
70	    case 0n:
71	      Nil{}
72	    case 1n+p:
73	      match cur:
74	        case Con{h, t}:
75	          Con{h, cr.kb.fb(~K, ~compute, p, t, k, rlen, fixed, prev, i)}
76	        case Nil{}:
77	          +blk = compute(k, B.Octets.append(prev, B.Octets.append(CrKbkdf.counter(rlen, U32.add(i, 1)), fixed)))
78	          +bl = B.Octets.to_list(blk)
79	          Con{cr.kb.head(bl), cr.kb.fb(~K, ~compute, p, cr.kb.tail(bl), k, rlen, fixed, blk, U32.add(i, 1))}
80	
81	# The first n bytes of feedback mode under the prepared key k from the IV,
82	# with the fixed input data (rlen 0: no counter), unchecked
83	def CrKbkdf.feedback_stream(~K: Data, ~compute: K -> B.Octets -> B.Octets, ~rlen: Nat, k: K, iv: B.Octets, fixed: B.Octets, n: Nat) -> B.Octets:
84	  B.Octets.of_list(cr.kb.fb(~K, ~compute, n, Nil{}, k, rlen, fixed, iv, 0))
85	
86	# [L]_2: L = 8 n bits as llen big-endian bytes
87	def CrKbkdf.length(+llen: Nat, +n: Nat) -> B.Octets:
88	  B.Octets.drop(Nat.sub(4n, llen), B.Octets.be32(U32.from_nat(Nat.mul(8n, n))))
89	
90	# the fixed input data Label || 0x00 || Context || [L]_2 for n bytes
91	def CrKbkdf.fixed(~llen: Nat, label: B.Octets, context: B.Octets, +n: Nat) -> B.Octets:
92	  B.Octets.concat([label, B.Octets.of_list([0]), context, CrKbkdf.length(llen, n)])
93	
94	# whether n bytes need at most max blocks of hlen bytes, and L = 8 n fits
95	# llen bytes (and 32 bits)
96	def cr.kb.fits(+hlen: Nat, +llen: Nat, +n: Nat, +max: Nat) -> Bool:
97	  Bool.and(Bool.and(Nat.is_lt(0n, hlen), Nat.is_le(Nat.div(Nat.add(n, Nat.sub(hlen, 1n)), hlen), max)), Nat.is_lt(Nat.mul(8n, n), Nat.pow(2n, Nat.mul(8n, Nat.min(llen, 4n)))))
98	
99	def cr.kb.key(~K: Data, ~compute: K -> B.Octets -> B.Octets, ~rlen: Nat, kr: Result<&2, &2, Cr.CrError, K>, fixed: B.Octets, n: Nat) -> Result<&2, &2, Cr.CrError, B.Octets>:
100	  match kr:
101	    case Fail{e}:
102	      Fail{e}
103	    case Done{k}:
104	      Done{CrKbkdf.counter_stream(~K, ~compute, ~rlen, k, fixed, n)}
105	
106	def cr.kb.checked(~K: Data, ~compute: K -> B.Octets -> B.Octets, ~rlen: Nat, fits: Bool, kr: Result<&2, &2, Cr.CrError, K>, fixed: B.Octets, +n: Nat) -> Result<&2, &2, Cr.CrError, B.Octets>:
107	  match fits:
108	    case False{}:
109	      Fail{Cr.CrBadOutputLength{n}}
110	    case True{}:
111	      cr.kb.key(~K, ~compute, ~rlen, kr, fixed, n)
112	
113	# K_OUT of n bytes in counter mode over the MAC (~K, ~key, ~compute) with
114	# hlen-byte tags, an r = 8 rlen bit counter and an llen-byte [L]_2
115	def CrKbkdf.counter_kdf(~K: Data, ~key: B.Octets -> Result<&2, &2, Cr.CrError, K>, ~compute: K -> B.Octets -> B.Octets, ~hlen: Nat, ~rlen: Nat, ~llen: Nat, ki: B.Octets, label: B.Octets, context: B.Octets, +n: Nat) -> Result<&2, &2, Cr.CrError, B.Octets>:
116	  +max = Nat.sub(Nat.pow(2n, Nat.mul(8n, rlen)), 1n)
117	  cr.kb.checked(~K, ~compute, ~rlen, cr.kb.fits(hlen, llen, n, max), key(ki), CrKbkdf.fixed(~llen, label, context, n), n)
118	
119	def cr.kb.fkey(~K: Data, ~compute: K -> B.Octets -> B.Octets, ~rlen: Nat, kr: Result<&2, &2, Cr.CrError, K>, iv: B.Octets, fixed: B.Octets, n: Nat) -> Result<&2, &2, Cr.CrError, B.Octets>:
120	  match kr:
121	    case Fail{e}:
122	      Fail{e}
123	    case Done{k}:
124	      Done{CrKbkdf.feedback_stream(~K, ~compute, ~rlen, k, iv, fixed, n)}
125	
126	def cr.kb.fchecked(~K: Data, ~compute: K -> B.Octets -> B.Octets, ~rlen: Nat, fits: Bool, kr: Result<&2, &2, Cr.CrError, K>, iv: B.Octets, fixed: B.Octets, +n: Nat) -> Result<&2, &2, Cr.CrError, B.Octets>:
127	  match fits:
128	    case False{}:
129	      Fail{Cr.CrBadOutputLength{n}}
130	    case True{}:
131	      cr.kb.fkey(~K, ~compute, ~rlen, kr, iv, fixed, n)
132	
133	# K_OUT of n bytes in feedback mode from the IV (rlen 0: no counter)
134	def CrKbkdf.feedback_kdf(~K: Data, ~key: B.Octets -> Result<&2, &2, Cr.CrError, K>, ~compute: K -> B.Octets -> B.Octets, ~hlen: Nat, ~rlen: Nat, ~llen: Nat, ki: B.Octets, iv: B.Octets, label: B.Octets, context: B.Octets, +n: Nat) -> Result<&2, &2, Cr.CrError, B.Octets>:
135	  cr.kb.fchecked(~K, ~compute, ~rlen, cr.kb.fits(hlen, llen, n, 4294967295n), key(ki), iv, CrKbkdf.fixed(~llen, label, context, n), n)
136	
137	# The KDF interface's ~derive in counter mode: secret K_IN, salt the Label,
138	# info the Context
139	def CrKbkdf.derive(~K: Data, ~key: B.Octets -> Result<&2, &2, Cr.CrError, K>, ~compute: K -> B.Octets -> B.Octets, ~hlen: Nat, ~rlen: Nat, ~llen: Nat, secret: B.Octets, salt: B.Octets, info: B.Octets, n: Nat) -> Result<&2, &2, Cr.CrError, B.Octets>:
140	  CrKbkdf.counter_kdf(~K, ~key, ~compute, ~hlen, ~rlen, ~llen, secret, salt, info, n)