From c44be73b4485214c478e426325a324909d91df1e Mon Sep 17 00:00:00 2001
From: "Bill.W" <0xbillw@gmail.com>
Date: Wed, 8 May 2024 12:04:58 +0800
Subject: [PATCH] feat: support `take_checkpoint` from `ceseal` after
 completing the first block synchronization in `cifrost` (#340)

---
 scripts/run/cifrost.sh                    | 5 +++++
 standalone/teeworker/cifrost/src/lib.rs   | 6 ++++++
 standalone/teeworker/cifrost/src/types.rs | 5 +++++
 3 files changed, 16 insertions(+)

diff --git a/scripts/run/cifrost.sh b/scripts/run/cifrost.sh
index c8869324..701a26d6 100755
--- a/scripts/run/cifrost.sh
+++ b/scripts/run/cifrost.sh
@@ -10,6 +10,10 @@ ceseal_port=$((${CESEAL_PORT:-8000} + $inst_seq))
 pub_port=$((${PUB_PORT:-19999} + $inst_seq))
 mnemonic=${MNEMONIC:-//Ferdie}
 inject_key=$(printf %064d $(($inst_seq + 1)))
+snapshot="--take-checkpoint"
+if [[ -z "${SNAPSHOT}" ]]; then
+    snapshot=
+fi
 
 bin="../cess/target/debug/cifrost"
 log_file="./target/cifrost-$inst_seq.log"
@@ -26,5 +30,6 @@ $bin \
     --mnemonic $mnemonic \
     --attestation-provider none \
     --longevity 16 \
+    $snapshot \
     --operator cXjHGCWMUM8gM9YFJUK2rqq2tiFWB4huBKWdQPkWdcXcZHhHA |&
     tee $log_file
diff --git a/standalone/teeworker/cifrost/src/lib.rs b/standalone/teeworker/cifrost/src/lib.rs
index c8739f44..18b20fad 100644
--- a/standalone/teeworker/cifrost/src/lib.rs
+++ b/standalone/teeworker/cifrost/src/lib.rs
@@ -784,6 +784,11 @@ async fn bridge(
                 handover_worker_key(&mut cc, &mut next_pr).await?;
             }
 
+            if !flags.checkpoint_taked && args.take_checkpoint {
+                cc.take_checkpoint(()).await?;
+                flags.checkpoint_taked = true;
+            }
+
             sleep(Duration::from_millis(args.dev_wait_block_ms)).await;
             continue;
         }
@@ -846,6 +851,7 @@ pub async fn run() {
         endpoint_registered: false,
         master_key_apply_sent: false,
         restart_failure_count: 0,
+        checkpoint_taked: false,
     };
 
     loop {
diff --git a/standalone/teeworker/cifrost/src/types.rs b/standalone/teeworker/cifrost/src/types.rs
index 8679666d..1d5110a7 100644
--- a/standalone/teeworker/cifrost/src/types.rs
+++ b/standalone/teeworker/cifrost/src/types.rs
@@ -190,6 +190,10 @@ pub struct Args {
     /// Load handover proof after blocks synced.
     #[arg(long)]
     pub load_handover_proof: bool,
+
+    /// Take checkpoint from Ceseal after completing the first block synchronization in Cifrost
+    #[arg(long)]
+    pub take_checkpoint: bool,
 }
 
 #[derive(clap::ValueEnum, Clone, Copy, Debug)]
@@ -212,6 +216,7 @@ pub struct RunningFlags {
     pub endpoint_registered: bool,
     pub master_key_apply_sent: bool,
     pub restart_failure_count: u32,
+    pub checkpoint_taked: bool,
 }
 
 pub struct BlockSyncState {