Add rebuild-mst subcommand

Author: futurGH

bin/main.ml

@@ -251,6 +251,19 @@ let migrate_blobs ?did () =
       print_endline "migrating all blobs to S3" ;
       S3.Blob_migration.migrate_all ()
 
+let rebuild_mst ~did () =
+  print_endline ("rebuilding MST for " ^ did) ;
+  let%lwt repo = Repository.load did in
+  match%lwt Repository.rebuild_mst repo with
+  | Ok (commit_cid, commit) ->
+      print_endline
+        (Printf.sprintf "MST rebuilt successfully, new commit: %s (rev: %s)"
+           (Cid.to_string commit_cid) commit.rev ) ;
+      Lwt.return_unit
+  | Error exn ->
+      print_endline ("error rebuilding MST: " ^ Printexc.to_string exn) ;
+      exit 1
+
 let print_usage () =
   print_endline
   @@ String.trim
@@ -262,6 +275,7 @@ commands:
   create-invite [uses]     create an invite code with an optional number of uses (default: 1)
   migrate-blobs            migrate all local blobs to S3
   migrate-blobs <did>      migrate blobs for a specific user to S3
+  rebuild-mst <did>        rebuild MST from records table (recovery tool)
 
 see also: gen-keys
 |}
@@ -280,6 +294,8 @@ let () =
       Lwt_main.run (migrate_blobs ())
   | ["migrate-blobs"; did] ->
       Lwt_main.run (migrate_blobs ~did ())
+  | ["rebuild-mst"; did] ->
+      Lwt_main.run (rebuild_mst ~did ())
   | ["help"] | ["--help"] | ["-h"] ->
       print_usage ()
   | cmd :: _ ->

pegasus/lib/repository.ml

@@ -589,3 +589,24 @@ let import_car t (stream : Car.stream) : (t, exn) Lwt_result.t =
     t.commit <- Some (root, commit) ;
     Lwt.return_ok t
   with exn -> Lwt.return_error exn
+
+let rebuild_mst t : (Cid.t * signed_commit, exn) Lwt_result.t =
+  try%lwt
+    let%lwt record_cids = User_store.get_all_record_cids t.db in
+    let record_count = List.length record_cids in
+    Logs.info (fun m -> m "rebuilding MST from %d records" record_count) ;
+    let%lwt () = User_store.clear_mst t.db in
+    Logs.info (fun m -> m "cleared existing MST blocks") ;
+    let%lwt mst_result = Mst.of_assoc t.db record_cids in
+    Logs.info (fun m ->
+        m "built new MST with root %s" (Cid.to_string mst_result.root) ) ;
+    let%lwt prev_commit = User_store.get_commit t.db in
+    let%lwt new_commit =
+      put_commit t mst_result.root ~previous:(Option.map snd prev_commit)
+    in
+    let commit_cid, commit = new_commit in
+    Logs.info (fun m ->
+        m "inserted new commit %s with rev %s" (Cid.to_string commit_cid)
+          commit.rev ) ;
+    Lwt.return_ok new_commit
+  with exn -> Lwt.return_error exn