Matrix protocol in OCaml, Eio specialised
at main 466 lines 18 kB view raw
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