Matrix protocol in OCaml, Eio specialised
1(** Server-side key backup and recovery.
2
3 This module implements Matrix server-side backup of room keys using the
4 m.megolm_backup.v1.curve25519-aes-sha2 backup algorithm.
5
6 Key backup allows clients to store encrypted room keys on the server so
7 they can be recovered on new devices or after data loss.
8
9 Note: Due to various known flaws in this algorithm, it is provided mainly
10 for backwards compatibility with existing backups. *)
11
12open Mirage_crypto_ec
13
14(** {1 Backup Key Types} *)
15
16(** Signature verification state *)
17type signature_state =
18 | Missing (** No signature found *)
19 | Invalid (** Signature is invalid *)
20 | Valid_but_not_trusted (** Valid but signer not trusted *)
21 | Valid_and_trusted (** Valid and signer is trusted *)
22
23let signature_state_trusted = function
24 | Valid_and_trusted -> true
25 | _ -> false
26
27(** Signature verification result *)
28type signature_verification = {
29 device_signature : signature_state;
30 user_identity_signature : signature_state;
31 other_signatures : (string * signature_state) list; (** device_id -> state *)
32}
33
34let signature_verification_trusted v =
35 signature_state_trusted v.device_signature ||
36 signature_state_trusted v.user_identity_signature ||
37 List.exists (fun (_, s) -> signature_state_trusted s) v.other_signatures
38
39let empty_signature_verification = {
40 device_signature = Missing;
41 user_identity_signature = Missing;
42 other_signatures = [];
43}
44
45(** Auth data for m.megolm_backup.v1.curve25519-aes-sha2 *)
46type megolm_v1_auth_data = {
47 public_key : string; (** Base64 Curve25519 public key *)
48 signatures : (string * (string * string) list) list; (** user_id -> (key_id, sig) *)
49}
50
51(** Convert signatures to JSON *)
52let signatures_to_json signatures =
53 let user_sigs = List.map (fun (user_id, key_sigs) ->
54 let inner = List.map (fun (key_id, sig_) ->
55 Printf.sprintf {|"%s":"%s"|} key_id sig_
56 ) key_sigs |> String.concat "," in
57 Printf.sprintf {|"%s":{%s}|} user_id inner
58 ) signatures |> String.concat "," in
59 "{" ^ user_sigs ^ "}"
60
61(** Encode auth data to JSON string *)
62let megolm_v1_auth_data_to_json auth_data =
63 Printf.sprintf {|{"public_key":"%s","signatures":%s}|}
64 auth_data.public_key
65 (signatures_to_json auth_data.signatures)
66
67(** Room key backup info - describes the backup algorithm and parameters *)
68type backup_info =
69 | Megolm_v1_curve25519_aes_sha2 of megolm_v1_auth_data
70 | Other of { algorithm : string; auth_data : Jsont.json }
71
72(** Backup version info from server *)
73type backup_version_info = {
74 version : string;
75 algorithm : string;
76 auth_data : Jsont.json;
77 count : int;
78 etag : string;
79}
80
81(** {1 Backup Encryption Key} *)
82
83(** Private key for decrypting backed up room keys *)
84type backup_decryption_key = {
85 private_key : string; (** Base64-encoded X25519 private key *)
86 public_key : string; (** Base64-encoded X25519 public key *)
87}
88
89(** Public key for encrypting room keys for backup *)
90type backup_encryption_key = {
91 public_key : string; (** Base64-encoded X25519 public key *)
92 mutable backup_version : string option;
93 mutable signatures : (string * (string * string) list) list;
94}
95
96(** Generate a new backup key pair *)
97let generate_backup_key () =
98 let priv, pub = X25519.gen_key () in
99 (* X25519.secret is a string, and gen_key returns (secret, public) where public is a string *)
100 let priv_bytes = X25519.secret_to_octets priv in
101 {
102 private_key = Base64.encode_string priv_bytes;
103 public_key = Base64.encode_string pub;
104 }
105
106(** Create encryption key from decryption key *)
107let encryption_key_of_decryption_key (decryption_key : backup_decryption_key) = {
108 public_key = decryption_key.public_key;
109 backup_version = None;
110 signatures = [];
111}
112
113(** Create encryption key from base64 public key *)
114let encryption_key_of_base64 public_key =
115 match Base64.decode public_key with
116 | Error _ -> Error "Invalid base64 encoding"
117 | Ok bytes ->
118 if String.length bytes <> 32 then
119 Error "Invalid key length"
120 else
121 Ok {
122 public_key;
123 backup_version = None;
124 signatures = [];
125 }
126
127(** Get the backup algorithm name *)
128let backup_algorithm = "m.megolm_backup.v1.curve25519-aes-sha2"
129
130(** {1 Room Key Encryption for Backup} *)
131
132(** Encrypted session data (matches Matrix spec) *)
133type encrypted_session_data = {
134 ephemeral : string; (** Base64 ephemeral public key *)
135 ciphertext : string; (** Base64 ciphertext *)
136 mac : string; (** Base64 MAC *)
137}
138
139(** Key backup data for a single session *)
140type key_backup_data = {
141 first_message_index : int;
142 forwarded_count : int;
143 is_verified : bool;
144 session_data : encrypted_session_data;
145}
146
147(** Room key backup request (for upload) *)
148type keys_backup_request = {
149 rooms : (string * (string * key_backup_data) list) list; (** room_id -> session_id -> data *)
150}
151
152(** Encrypt a room key for backup using X25519/AES-256/HMAC-SHA256 *)
153let encrypt_room_key (encryption_key : backup_encryption_key) ~session_key ~session_id ~room_id ~sender_key =
154 match Base64.decode encryption_key.public_key with
155 | Error _ -> Error "Invalid encryption key"
156 | Ok recipient_pub_bytes ->
157 if String.length recipient_pub_bytes <> 32 then
158 Error "Invalid X25519 public key length"
159 else
160 (* Generate ephemeral key pair *)
161 let ephemeral_priv, ephemeral_pub = X25519.gen_key () in
162
163 (* Perform X25519 key exchange *)
164 (match X25519.key_exchange ephemeral_priv recipient_pub_bytes with
165 | Error _ -> Error "Key exchange failed"
166 | Ok shared_secret ->
167 (* Derive encryption and MAC keys using HKDF *)
168 let prk = Hkdf.extract ~hash:`SHA256 ~salt:"" shared_secret in
169 let okm = Hkdf.expand ~hash:`SHA256 ~prk ~info:"" 64 in
170 let aes_key = String.sub okm 0 32 in
171 let mac_key = String.sub okm 32 32 in
172
173 (* Build the payload to encrypt *)
174 let payload = Printf.sprintf
175 {|{"algorithm":"m.megolm.v1.aes-sha2","room_id":"%s","sender_key":"%s","session_id":"%s","session_key":"%s"}|}
176 room_id sender_key session_id session_key
177 in
178
179 (* Pad payload to 16-byte boundary (PKCS#7) *)
180 let pad_len = 16 - (String.length payload mod 16) in
181 let padded = payload ^ String.make pad_len (Char.chr pad_len) in
182
183 (* Generate random IV *)
184 let iv = Mirage_crypto_rng.generate 16 in
185
186 (* Encrypt with AES-256-CBC *)
187 let key = Mirage_crypto.AES.CBC.of_secret aes_key in
188 let ciphertext = Mirage_crypto.AES.CBC.encrypt ~key ~iv padded in
189
190 (* Prepend IV to ciphertext for MAC calculation *)
191 let mac_input = iv ^ ciphertext in
192
193 (* Calculate HMAC-SHA256 *)
194 let mac = Digestif.SHA256.hmac_string ~key:mac_key mac_input in
195 let mac_bytes = Digestif.SHA256.to_raw_string mac in
196 (* Truncate to first 8 bytes as per spec *)
197 let mac_truncated = String.sub mac_bytes 0 8 in
198
199 Ok {
200 ephemeral = Base64.encode_string ephemeral_pub;
201 ciphertext = Base64.encode_string (iv ^ ciphertext);
202 mac = Base64.encode_string mac_truncated;
203 })
204
205(** Decrypt a room key from backup *)
206let decrypt_room_key (decryption_key : backup_decryption_key) (session_data : encrypted_session_data) =
207 match Base64.decode decryption_key.private_key,
208 Base64.decode session_data.ephemeral,
209 Base64.decode session_data.ciphertext,
210 Base64.decode session_data.mac with
211 | Error _, _, _, _ -> Error "Invalid private key encoding"
212 | _, Error _, _, _ -> Error "Invalid ephemeral key encoding"
213 | _, _, Error _, _ -> Error "Invalid ciphertext encoding"
214 | _, _, _, Error _ -> Error "Invalid MAC encoding"
215 | Ok priv_bytes, Ok ephemeral_bytes, Ok ciphertext_with_iv, Ok mac_bytes ->
216 match X25519.secret_of_octets priv_bytes with
217 | Error _ -> Error "Invalid X25519 private key"
218 | Ok (priv, _) ->
219 (* Perform X25519 key exchange *)
220 (match X25519.key_exchange priv ephemeral_bytes with
221 | Error _ -> Error "Key exchange failed"
222 | Ok shared_secret ->
223 (* Derive encryption and MAC keys using HKDF *)
224 let prk = Hkdf.extract ~hash:`SHA256 ~salt:"" shared_secret in
225 let okm = Hkdf.expand ~hash:`SHA256 ~prk ~info:"" 64 in
226 let aes_key = String.sub okm 0 32 in
227 let mac_key = String.sub okm 32 32 in
228
229 (* Verify MAC *)
230 let expected_mac = Digestif.SHA256.hmac_string ~key:mac_key ciphertext_with_iv in
231 let expected_mac_bytes = Digestif.SHA256.to_raw_string expected_mac in
232 let expected_mac_truncated = String.sub expected_mac_bytes 0 8 in
233 if mac_bytes <> expected_mac_truncated then
234 Error "MAC verification failed"
235 else begin
236 (* Extract IV and ciphertext *)
237 if String.length ciphertext_with_iv < 16 then
238 Error "Ciphertext too short"
239 else begin
240 let iv = String.sub ciphertext_with_iv 0 16 in
241 let ciphertext = String.sub ciphertext_with_iv 16 (String.length ciphertext_with_iv - 16) in
242
243 (* Decrypt with AES-256-CBC *)
244 let key = Mirage_crypto.AES.CBC.of_secret aes_key in
245 let plaintext = Mirage_crypto.AES.CBC.decrypt ~key ~iv ciphertext in
246
247 (* Remove PKCS#7 padding *)
248 let pad_len = Char.code (String.get plaintext (String.length plaintext - 1)) in
249 if pad_len < 1 || pad_len > 16 then
250 Error "Invalid padding"
251 else
252 let unpadded = String.sub plaintext 0 (String.length plaintext - pad_len) in
253 Ok unpadded
254 end
255 end)
256
257(** {1 Backup Machine State} *)
258
259(** State of the backup machine *)
260type backup_state =
261 | Disabled (** No backup configured *)
262 | Creating (** Creating a new backup *)
263 | Enabling (** Enabling existing backup *)
264 | Resuming (** Resuming existing backup *)
265 | Enabled (** Backup is active *)
266 | Downloading (** Downloading keys from backup *)
267 | Disabling (** Disabling backup *)
268
269(** Backup machine for managing room key backups *)
270type t = {
271 user_id : Matrix_proto.Id.User_id.t;
272 device_id : Matrix_proto.Id.Device_id.t;
273 mutable state : backup_state;
274 mutable encryption_key : backup_encryption_key option;
275 mutable decryption_key : backup_decryption_key option;
276 mutable backup_version : string option;
277 (* Pending sessions to backup *)
278 mutable pending_sessions : (string * string * string) list; (** (room_id, session_id, sender_key) *)
279 (* Sessions that have been backed up *)
280 mutable backed_up_sessions : (string * string) list; (** (room_id, session_id) *)
281}
282
283(** Create a new backup machine *)
284let create ~user_id ~device_id = {
285 user_id;
286 device_id;
287 state = Disabled;
288 encryption_key = None;
289 decryption_key = None;
290 backup_version = None;
291 pending_sessions = [];
292 backed_up_sessions = [];
293}
294
295(** Check if backup is enabled *)
296let is_enabled t = t.state = Enabled
297
298(** Get the current backup version *)
299let backup_version t = t.backup_version
300
301(** {1 Backup Setup} *)
302
303(** Enable backup with a new key *)
304let enable_with_new_key t =
305 let key = generate_backup_key () in
306 t.decryption_key <- Some key;
307 t.encryption_key <- Some (encryption_key_of_decryption_key key);
308 t.state <- Creating;
309 key
310
311(** Enable backup with an existing decryption key *)
312let enable_with_key t decryption_key =
313 t.decryption_key <- Some decryption_key;
314 t.encryption_key <- Some (encryption_key_of_decryption_key decryption_key);
315 t.state <- Enabling
316
317(** Enable backup with only an encryption key (upload-only mode) *)
318let enable_upload_only t encryption_key version =
319 t.encryption_key <- Some encryption_key;
320 t.decryption_key <- None;
321 t.backup_version <- Some version;
322 encryption_key.backup_version <- Some version;
323 t.state <- Enabled
324
325(** Set the backup version after creating *)
326let set_backup_version t version =
327 t.backup_version <- Some version;
328 (match t.encryption_key with
329 | Some key -> key.backup_version <- Some version
330 | None -> ());
331 t.state <- Enabled
332
333(** Disable backup *)
334let disable t =
335 t.state <- Disabling;
336 t.encryption_key <- None;
337 t.decryption_key <- None;
338 t.backup_version <- None;
339 t.backed_up_sessions <- [];
340 t.state <- Disabled
341
342(** {1 Session Management} *)
343
344(** Mark a session as needing backup *)
345let mark_session_for_backup t ~room_id ~session_id ~sender_key =
346 if not (List.mem (room_id, session_id) t.backed_up_sessions) then
347 t.pending_sessions <- (room_id, session_id, sender_key) :: t.pending_sessions
348
349(** Get number of pending sessions *)
350let pending_count t = List.length t.pending_sessions
351
352(** Check if a session has been backed up *)
353let is_session_backed_up t ~room_id ~session_id =
354 List.mem (room_id, session_id) t.backed_up_sessions
355
356(** Mark a session as backed up *)
357let mark_session_backed_up t ~room_id ~session_id =
358 t.pending_sessions <- List.filter (fun (r, s, _) -> r <> room_id || s <> session_id) t.pending_sessions;
359 if not (List.mem (room_id, session_id) t.backed_up_sessions) then
360 t.backed_up_sessions <- (room_id, session_id) :: t.backed_up_sessions
361
362(** {1 Room Key Recovery} *)
363
364(** Recovered room key data *)
365type recovered_room_key = {
366 room_id : string;
367 session_id : string;
368 session_key : string;
369 sender_key : string;
370 algorithm : string;
371 forwarded : bool;
372}
373
374(** Parse a recovered room key from decrypted JSON *)
375let parse_recovered_key json_str =
376 (* Simple JSON parsing - in production would use proper parser *)
377 let _get_field name str =
378 let _pattern = Printf.sprintf {|"%s":"|} name in
379 match String.split_on_char '"' str with
380 | _ -> None
381 in
382 (* Simplified - would use Jsont in real implementation *)
383 match Jsont_bytesrw.decode_string Jsont.json json_str with
384 | Error _ -> None
385 | Ok _json -> None (* Would extract fields from JSON *)
386
387(** Result of importing keys from backup *)
388type import_result = {
389 imported_count : int;
390 total_count : int;
391 keys : recovered_room_key list;
392}
393
394(** {1 Backup Creation API Helpers} *)
395
396(** Build auth data for a new backup *)
397let build_auth_data encryption_key = {
398 public_key = encryption_key.public_key;
399 signatures = encryption_key.signatures;
400}
401
402(** Create backup version request body *)
403let create_version_request_body t =
404 match t.encryption_key with
405 | None -> Error "No encryption key configured"
406 | Some key ->
407 let auth_data = build_auth_data key in
408 let auth_json = megolm_v1_auth_data_to_json auth_data in
409 Ok (Printf.sprintf
410 {|{"algorithm":"%s","auth_data":%s}|}
411 backup_algorithm auth_json)
412
413(** {1 Recovery Key Format} *)
414
415(** Encode a backup decryption key as a recovery key (human-readable) *)
416let encode_recovery_key (key : backup_decryption_key) =
417 match Base64.decode key.private_key with
418 | Error _ -> Error "Invalid private key"
419 | Ok bytes ->
420 (* Add header byte 0x8B, then calculate parity byte *)
421 let with_header = "\x8B" ^ bytes in
422 let parity = String.fold_left (fun acc c -> acc lxor Char.code c) 0 with_header in
423 let full = with_header ^ String.make 1 (Char.chr parity) in
424 (* Encode as base58 with spaces every 4 chars for readability *)
425 let _base58_alphabet = "123456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz" in
426 (* Simplified - would use proper base58 encoding; using base64 for now *)
427 let encoded = Base64.encode_string ~pad:false full in
428 (* Add spaces for readability *)
429 let with_spaces = String.to_seq encoded
430 |> Seq.mapi (fun i c -> if i > 0 && i mod 4 = 0 then [' '; c] else [c])
431 |> Seq.flat_map List.to_seq
432 |> String.of_seq
433 in
434 Ok with_spaces
435
436(** Decode a recovery key to a backup decryption key *)
437let decode_recovery_key recovery_key =
438 (* Remove spaces and decode *)
439 let cleaned = String.split_on_char ' ' recovery_key |> String.concat "" in
440 match Base64.decode cleaned with
441 | Error _ -> Error "Invalid recovery key format"
442 | Ok bytes ->
443 if String.length bytes < 3 then
444 Error "Recovery key too short"
445 else if String.get bytes 0 <> '\x8B' then
446 Error "Invalid recovery key header"
447 else begin
448 (* Verify parity *)
449 let key_bytes = String.sub bytes 1 (String.length bytes - 2) in
450 let expected_parity = Char.code (String.get bytes (String.length bytes - 1)) in
451 let actual_parity = String.fold_left (fun acc c -> acc lxor Char.code c) 0
452 (String.sub bytes 0 (String.length bytes - 1)) in
453 if expected_parity <> actual_parity then
454 Error "Recovery key parity check failed"
455 else begin
456 (* Derive public key from private key *)
457 match X25519.secret_of_octets key_bytes with
458 | Error _ -> Error "Invalid private key in recovery key"
459 | Ok (priv, pub_bytes) ->
460 ignore priv;
461 Ok {
462 private_key = Base64.encode_string key_bytes;
463 public_key = Base64.encode_string pub_bytes;
464 }
465 end
466 end