1	# SCRAM-SHA-256 (RFC 5802, RFC 7677) client for PostgreSQL SASL authentication, without channel binding.
2	import Base
3	import bend-kit-bytes@0.3.1.0/bytes.bend as Bytes
4	import bend-kit-crypto@0.1.1.0/crypto.bend as Crypto
5	
6	# Messages and keys are byte strings, one Char per octet. The password is its octets
7	# (UTF-8), used as is: no SASLprep, which is the same as libpq for ASCII passwords.
8	#   (st, msg) : First & String <- first(user)          -> SASLInitialResponse
9	#   (fin, msg) : Final & String <- respond(st, sf, pw)  -> SASLResponse
10	#   verify(fin, server_final)                           -> Done{Unit{}} or Fail
11	# Errors: 22 for a malformed server message, 13 for a server error or a wrong server
12	# signature, else the crypto effect's errno.
13	
14	def mechanism() -> String:
15	  "SCRAM-SHA-256"
16	
17	# bare: client-first-message-bare. nonce: the client nonce.
18	type First is Data:
19	  First{bare: String, nonce: String}
20	
21	# sig: the ServerSignature the server-final-message must carry.
22	type Final is Data:
23	  Final{sig: String}
24	
25	# nonce: the combined nonce. salt: decoded octets. iters: at least 1.
26	type Parsed is Data:
27	  Parsed{nonce: String, salt: String, iters: U32}
28	
29	# saslname (RFC 5802 §5.1): '=' as "=3D", ',' as "=2C".
30	def saslname.of(+c: U32, +rest: String) -> String:
31	  Bool.pick(String, U32.is_eq(c, 61), "=3D" ++ rest, Bool.pick(String, U32.is_eq(c, 44), "=2C" ++ rest, SCon{Chr{c}, rest}))
32	
33	def saslname(s: String) -> String:
34	  match s:
35	    case SNil{}:
36	      SNil{}
37	    case SCon{Chr{+c}, t}:
38	      saslname.of(c, saslname(t))
39	
40	# The client-first-message for a nonce: gs2 header "n,," (no channel binding), then the bare part.
41	def first.with(+user: String, +nonce: String) -> First & String:
42	  +bare = "n=" ++ saslname(user) ++ ",r=" ++ nonce
43	  (First{bare, nonce}, "n,," ++ bare)
44	
45	def first.of(+user: String, r: Result<&1, &1, U32 & String, U32 & Array<U32>>) -> Result<&1, &1, U32 & String, First & String>:
46	  match r:
47	    case Fail{e}:
48	      Fail{e}
49	    case Done{(len, buf)}:
50	      Done{first.with(user, Bytes.to_base64(Bytes.Bytes{len, buf}))}
51	
52	# A fresh 18-octet nonce from the OS secure random source, base64 (24 printable chars, no ',').
53	def first(+user: String) -> IO(Result<&1, &1, U32 & String, First & String>):
54	  do IO<Result<&1, &1, U32 & String, First & String>>:
55	    r : Result<&1, &1, U32 & String, U32 & Array<U32>> <- Crypto.random.words(18)
56	    return first.of(user, r)
57	
58	# server-first-message (RFC 5802 §7): r=nonce,s=salt,i=iterations, then extensions, which are ignored.
59	
60	# The value of attribute name, when s is name=value.
61	def attr(+name: U32, s: String) -> Maybe<&2, String>:
62	  match s:
63	    case SCon{Chr{+c}, SCon{'=', v}}:
64	      Bool.pick(Maybe<&2, String>, U32.is_eq(c, name), Some{v}, None{})
65	    case _:
66	      None{}
67	
68	# printable (RFC 5802 §7): %x21-2B / %x2D-7E.
69	def printable(s: String) -> Bool:
70	  match s:
71	    case SNil{}:
72	      True{}
73	    case SCon{Chr{+c}, t}:
74	      Bool.and(Bool.and(U32.is_le(33, c), U32.is_le(c, 126)), Bool.and(U32.is_ne(c, 44), printable(t)))
75	
76	def digits(s: String) -> Bool:
77	  match s:
78	    case SNil{}:
79	      True{}
80	    case SCon{Chr{+c}, t}:
81	      Bool.and(Bool.and(U32.is_le(48, c), U32.is_le(c, 57)), digits(t))
82	
83	# The server nonce must be the client nonce plus at least one more printable char.
84	def nonce.ok(+client: String, +server: String) -> Bool:
85	  Bool.and(Bool.and(String.starts_with(server, client), Nat.is_lt(String.length(client), String.length(server))), printable(server))
86	
87	def parse.iters(ok: Bool, m: Maybe<&2, U32>, +nonce: String, +salt: String) -> Result<&1, &1, U32 & String, Parsed>:
88	  match ok m:
89	    case True{} Some{+i}:
90	      Bool.pick(Result<&1, &1, U32 & String, Parsed>, U32.is_eq(i, 0), Fail{(22, "SCRAM iteration count is 0")}, Done{Parsed{nonce, salt, i}})
91	    case _ _:
92	      Fail{(22, "SCRAM iteration count is not a decimal U32")}
93	
94	def parse.salt.of(empty: Bool, +nonce: String, +salt: String, +it: String) -> Result<&1, &1, U32 & String, Parsed>:
95	  match empty:
96	    case True{}:
97	      Fail{(22, "SCRAM salt is empty")}
98	    case False{}:
99	      parse.iters(Bool.and(Bool.not(String.is_empty(it)), digits(it)), U32.read(it), nonce, salt)
100	
101	def parse.salt(m: Maybe<&1, Bytes.Bytes>, +nonce: String, +it: String) -> Result<&1, &1, U32 & String, Parsed>:
102	  match m:
103	    case None{}:
104	      Fail{(22, "SCRAM salt is not base64")}
105	    case Some{b}:
106	      +salt = Bytes.to_string(b)
107	      parse.salt.of(String.is_empty(salt), nonce, salt, it)
108	
109	def parse.nonce(ok: Bool, +nonce: String, +salt: String, +it: String) -> Result<&1, &1, U32 & String, Parsed>:
110	  match ok:
111	    case False{}:
112	      Fail{(22, "SCRAM server nonce does not extend the client nonce")}
113	    case True{}:
114	      parse.salt(Bytes.from_base64(salt), nonce, it)
115	
116	def parse.fields(+client: String, r: Maybe<&2, String>, s: Maybe<&2, String>, i: Maybe<&2, String>) -> Result<&1, &1, U32 & String, Parsed>:
117	  match r s i:
118	    case Some{+nonce} Some{+salt} Some{+it}:
119	      parse.nonce(nonce.ok(client, nonce), nonce, salt, it)
120	    case _ _ _:
121	      Fail{(22, "server-first-message is not r=,s=,i=")}
122	
123	def parse.parts(+client: String, parts: List<&2, String>) -> Result<&1, &1, U32 & String, Parsed>:
124	  match parts:
125	    case Con{+r, Con{+s, Con{+i, rest}}}:
126	      parse.fields(client, attr(114, r), attr(115, s), attr(105, i))
127	    case _:
128	      Fail{(22, "server-first-message is not r=,s=,i=")}
129	
130	# The server-first-message checked against the client nonce.
131	def parse(+client: String, +msg: String) -> Result<&1, &1, U32 & String, Parsed>:
132	  parse.parts(client, String.split(msg, ','))
133	
134	# Crypto over byte strings.
135	
136	def octets(r: Result<&1, &1, U32 & String, U32 & Array<U32>>) -> Result<&1, &1, U32 & String, String>:
137	  match r:
138	    case Fail{e}:
139	      Fail{e}
140	    case Done{(len, buf)}:
141	      Done{Bytes.to_string(Bytes.Bytes{len, buf})}
142	
143	def hmac.of(k: Bytes.Bytes, d: Bytes.Bytes) -> IO(Result<&1, &1, U32 & String, U32 & Array<U32>>):
144	  Bytes.Bytes{kl, kw} = k
145	  Bytes.Bytes{dl, dw} = d
146	  Crypto.hmac.words("SHA256", kl, kw, dl, dw)
147	
148	def hmac(+key: String, +data: String) -> IO(Result<&1, &1, U32 & String, String>):
149	  do IO<Result<&1, &1, U32 & String, String>>:
150	    r : Result<&1, &1, U32 & String, U32 & Array<U32>> <- hmac.of(Bytes.from_string(key), Bytes.from_string(data))
151	    return octets(r)
152	
153	def sha.of(b: Bytes.Bytes) -> IO(Result<&1, &1, U32 & String, U32 & Array<U32>>):
154	  Bytes.Bytes{len, buf} = b
155	  Crypto.sha256.words(len, buf)
156	
157	def sha(+data: String) -> IO(Result<&1, &1, U32 & String, String>):
158	  do IO<Result<&1, &1, U32 & String, String>>:
159	    r : Result<&1, &1, U32 & String, U32 & Array<U32>> <- sha.of(Bytes.from_string(data))
160	    return octets(r)
161	
162	def pbkdf2.of(p: Bytes.Bytes, s: Bytes.Bytes, +iters: U32) -> IO(Result<&1, &1, U32 & String, U32 & Array<U32>>):
163	  Bytes.Bytes{pl, pw} = p
164	  Bytes.Bytes{sl, sw} = s
165	  Crypto.pbkdf2.words("SHA256", pl, pw, sl, sw, iters, 32)
166	
167	# SaltedPassword = Hi(pass, salt, iters) (RFC 5802 §2.2), which is PBKDF2 with HMAC-SHA-256.
168	def pbkdf2(+pass: String, +salt: String, +iters: U32) -> IO(Result<&1, &1, U32 & String, String>):
169	  do IO<Result<&1, &1, U32 & String, String>>:
170	    r : Result<&1, &1, U32 & String, U32 & Array<U32>> <- pbkdf2.of(Bytes.from_string(pass), Bytes.from_string(salt), iters)
171	    return octets(r)
172	
173	def eq.of(a: Bytes.Bytes, b: Bytes.Bytes) -> IO(Bool):
174	  Bytes.Bytes{al, aw} = a
175	  Bytes.Bytes{bl, bw} = b
176	  Crypto.eq.ct.words(al, aw, bl, bw)
177	
178	def xor(a: String, b: String) -> String:
179	  match a b:
180	    case SCon{Chr{+x}, at} SCon{Chr{+y}, bt}:
181	      SCon{Chr{(x .^. y : U32)}, xor(at, bt)}
182	    case _ _:
183	      SNil{}
184	
185	# client-final-message (RFC 5802 §3, §7).
186	
187	def step.of(r: Result<&1, &1, U32 & String, String>, k: String -> IO(Result<&1, &1, U32 & String, Final & String>)) -> IO(Result<&1, &1, U32 & String, Final & String>):
188	  match r:
189	    case Fail{e}:
190	      IO.pure(Result<&1, &1, U32 & String, Final & String>, Fail{e})
191	    case Done{s}:
192	      k(s)
193	
194	def step(m: IO(Result<&1, &1, U32 & String, String>), k: String -> IO(Result<&1, &1, U32 & String, Final & String>)) -> IO(Result<&1, &1, U32 & String, Final & String>):
195	  do IO<Result<&1, &1, U32 & String, Final & String>>:
196	    r : Result<&1, &1, U32 & String, String> <- m
197	    step.of(r, k)
198	
199	# c=biws is base64 of the gs2 header "n,,". AuthMessage joins the three bare messages.
200	def respond.of(r: Result<&1, &1, U32 & String, Parsed>, +bare: String, +sf: String, +pass: String) -> IO(Result<&1, &1, U32 & String, Final & String>):
201	  match r:
202	    case Fail{e}:
203	      IO.pure(Result<&1, &1, U32 & String, Final & String>, Fail{e})
204	    case Done{Parsed{+nonce, +salt, +iters}}:
205	      +fb = "c=biws,r=" ++ nonce
206	      +auth = bare ++ "," ++ sf ++ "," ++ fb
207	      step(pbkdf2(pass, salt, iters), +salted =>
208	        step(hmac(salted, "Client Key"), +ck =>
209	          step(sha(ck), +sk =>
210	            step(hmac(sk, auth), +cs =>
211	              step(hmac(salted, "Server Key"), +svk =>
212	                step(hmac(svk, auth), +ss =>
213	                  IO.pure(Result<&1, &1, U32 & String, Final & String>, Done{(Final{ss}, fb ++ ",p=" ++ Bytes.to_base64(Bytes.from_string(xor(ck, cs))))})))))))
214	
215	# The client-final-message with ClientProof for the server-first-message sf and the password.
216	def respond(st: First, +sf: String, +pass: String) -> IO(Result<&1, &1, U32 & String, Final & String>):
217	  First{+bare, +client} = st
218	  respond.of(parse(client, sf), bare, sf, pass)
219	
220	# server-final-message (RFC 5802 §7): v=signature or e=error, then extensions, which are ignored.
221	
222	def final.sig(m: Maybe<&1, Bytes.Bytes>) -> Result<&1, &1, U32 & String, String>:
223	  match m:
224	    case None{}:
225	      Fail{(22, "SCRAM server signature is not base64")}
226	    case Some{b}:
227	      Done{Bytes.to_string(b)}
228	
229	def final.attr(s: String) -> Result<&1, &1, U32 & String, String>:
230	  match s:
231	    case SCon{'v', SCon{'=', v}}:
232	      final.sig(Bytes.from_base64(v))
233	    case SCon{'e', SCon{'=', e}}:
234	      Fail{(13, "SCRAM server error: " ++ e)}
235	    case _:
236	      Fail{(22, "server-final-message is not v= or e=")}
237	
238	def final.parts(parts: List<&2, String>) -> Result<&1, &1, U32 & String, String>:
239	  match parts:
240	    case Con{f, rest}:
241	      final.attr(f)
242	    case Nil{}:
243	      Fail{(22, "server-final-message is not v= or e=")}
244	
245	def verify.ok(ok: Bool) -> Result<&1, &1, U32 & String, Unit>:
246	  match ok:
247	    case True{}:
248	      Done{Unit{}}
249	    case False{}:
250	      Fail{(13, "SCRAM server signature does not match")}
251	
252	def verify.of(+sig: String, r: Result<&1, &1, U32 & String, String>) -> IO(Result<&1, &1, U32 & String, Unit>):
253	  match r:
254	    case Fail{e}:
255	      IO.pure(Result<&1, &1, U32 & String, Unit>, Fail{e})
256	    case Done{+got}:
257	      do IO<Result<&1, &1, U32 & String, Unit>>:
258	        ok : Bool <- eq.of(Bytes.from_string(sig), Bytes.from_string(got))
259	        return verify.ok(ok)
260	
261	# Done when the server-final-message proves the server knows the password, compared in constant time.
262	def verify(st: Final, +msg: String) -> IO(Result<&1, &1, U32 & String, Unit>):
263	  Final{+sig} = st
264	  verify.of(sig, final.parts(String.split(msg, ',')))