add height option to restore command

Signed-off-by: Jun Kimura <[email protected]>
datachainlab · Nov 23, 2023 · a84ba51 · a84ba51
1 parent cb165d7
commit a84ba51
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 8 deletions.
diff --git a/relay/cmd.go b/relay/cmd.go
@@ -10,7 +10,8 @@ import (
 )
 
 const (
-	flagSrc = "src"
+	flagSrc    = "src"
+	flagHeight = "height"
 )
 
 func LCPCmd(ctx *config.Context) *cobra.Command {
@@ -111,7 +112,7 @@ func restoreELCStateCmd(ctx *config.Context) *cobra.Command {
 				verifier = c[src]
 			}
 			prover := target.Prover.(*Prover)
-			if err := prover.restoreELCState(context.TODO(), verifier); err != nil {
+			if err := prover.restoreELCState(context.TODO(), verifier, viper.GetUint64(flagHeight)); err != nil {
 				return err
 			}
 			if err := prover.removeEnclaveKeyInfos(context.TODO()); err != nil {
@@ -120,7 +121,7 @@ func restoreELCStateCmd(ctx *config.Context) *cobra.Command {
 			return nil
 		},
 	}
-	return srcFlag(cmd)
+	return heightFlag(srcFlag(cmd))
 }
 
 func removeEnclaveKeyInfoCmd(ctx *config.Context) *cobra.Command {
@@ -153,3 +154,11 @@ func srcFlag(cmd *cobra.Command) *cobra.Command {
 	}
 	return cmd
 }
+
+func heightFlag(cmd *cobra.Command) *cobra.Command {
+	cmd.Flags().Uint64P(flagHeight, "", 0, "a height to restore")
+	if err := viper.BindPFlag(flagHeight, cmd.Flags().Lookup(flagHeight)); err != nil {
+		panic(err)
+	}
+	return cmd
+}
diff --git a/relay/restore.go b/relay/restore.go
@@ -14,7 +14,7 @@ import (
 	"github.com/hyperledger-labs/yui-relayer/core"
 )
 
-func (pr *Prover) restoreELCState(ctx context.Context, counterparty core.FinalityAwareChain) error {
+func (pr *Prover) restoreELCState(ctx context.Context, counterparty core.FinalityAwareChain, height uint64) error {
 	if err := pr.initServiceClient(); err != nil {
 		return err
 	}
@@ -31,7 +31,6 @@ func (pr *Prover) restoreELCState(ctx context.Context, counterparty core.Finalit
 	if err != nil {
 		return err
 	}
-
 	counterpartyClientRes, err := counterparty.QueryClientState(core.NewQueryContext(context.TODO(), cplatestHeight))
 	if err != nil {
 		return err
@@ -41,7 +40,16 @@ func (pr *Prover) restoreELCState(ctx context.Context, counterparty core.Finalit
 		return err
 	}
 
-	counterpartyConsRes, err := counterparty.QueryClientConsensusState(core.NewQueryContext(context.TODO(), cplatestHeight), cs.GetLatestHeight())
+	var restoreHeight ibcexported.Height
+	if height == 0 {
+		restoreHeight = cs.GetLatestHeight()
+	} else {
+		restoreHeight = clienttypes.NewHeight(cs.GetLatestHeight().GetRevisionNumber(), height)
+	}
+
+	log.Printf("try to restore ELC state: height=%v", restoreHeight)
+
+	counterpartyConsRes, err := counterparty.QueryClientConsensusState(core.NewQueryContext(context.TODO(), cplatestHeight), restoreHeight)
 	if err != nil {
 		return err
 	}
@@ -109,8 +117,8 @@ func (pr *Prover) restoreELCState(ctx context.Context, counterparty core.Finalit
 	if !ucc.NewStateID.EqualBytes(consensusState.StateId) {
 		return fmt.Errorf("unexpected state id: expected %v, but got %v", ucc.NewStateID, consensusState.StateId)
 	}
-	if !ucc.NewHeight.EQ(clientState.LatestHeight) {
-		return fmt.Errorf("unexpected height: expected %v, but got %v", ucc.NewHeight, clientState.LatestHeight)
+	if !ucc.NewHeight.EQ(restoreHeight) {
+		return fmt.Errorf("unexpected height: expected %v, but got %v", restoreHeight, ucc.NewHeight)
 	}
 
 	// TODO relayer should update res.ClientId in the config