From a99a627daac71d7aff02d947c047d857434a6a32 Mon Sep 17 00:00:00 2001
From: Willem Wyndham <willem@ahalabs.dev>
Date: Wed, 19 Jul 2023 12:01:40 -0400
Subject: [PATCH] fix: check if change is contract code

---
 cmd/soroban-cli/src/commands/contract/bump.rs | 40 ++++++++++---------
 .../src/commands/contract/restore.rs          | 33 +++++++++------
 2 files changed, 43 insertions(+), 30 deletions(-)

diff --git a/cmd/soroban-cli/src/commands/contract/bump.rs b/cmd/soroban-cli/src/commands/contract/bump.rs
index e34712ce2..eb40946ae 100644
--- a/cmd/soroban-cli/src/commands/contract/bump.rs
+++ b/cmd/soroban-cli/src/commands/contract/bump.rs
@@ -6,12 +6,12 @@ use std::{
 
 use clap::{command, Parser};
 use soroban_env_host::xdr::{
-    BumpFootprintExpirationOp, ContractDataEntry, ContractEntryBodyType, Error as XdrError,
-    ExtensionPoint, Hash, LedgerEntry, LedgerEntryChange, LedgerEntryData, LedgerFootprint,
-    LedgerKey, LedgerKeyContractData, Memo, MuxedAccount, Operation, OperationBody, Preconditions,
-    ReadXdr, ScAddress, ScSpecTypeDef, ScVal, SequenceNumber, SorobanResources,
-    SorobanTransactionData, Transaction, TransactionExt, TransactionMeta, TransactionMetaV3,
-    Uint256,
+    BumpFootprintExpirationOp, ContractCodeEntry, ContractDataEntry, ContractEntryBodyType,
+    Error as XdrError, ExtensionPoint, Hash, LedgerEntry, LedgerEntryChange, LedgerEntryData,
+    LedgerFootprint, LedgerKey, LedgerKeyContractData, Memo, MuxedAccount, Operation,
+    OperationBody, Preconditions, ReadXdr, ScAddress, ScSpecTypeDef, ScVal, SequenceNumber,
+    SorobanResources, SorobanTransactionData, Transaction, TransactionExt, TransactionMeta,
+    TransactionMetaV3, Uint256,
 };
 use stellar_strkey::DecodeError;
 
@@ -182,20 +182,24 @@ impl Cmd {
             return Err(Error::LedgerEntryNotFound);
         }
 
-        let (
-            LedgerEntryChange::State(_state),
-            LedgerEntryChange::Updated(LedgerEntry{
-                data: LedgerEntryData::ContractData(ContractDataEntry{
-                    expiration_ledger_seq,
+        match (&operations[0].changes[0], &operations[0].changes[1]) {
+            (
+                LedgerEntryChange::State(_),
+                LedgerEntryChange::Updated(LedgerEntry {
+                    data:
+                        LedgerEntryData::ContractData(ContractDataEntry {
+                            expiration_ledger_seq,
+                            ..
+                        })
+                        | LedgerEntryData::ContractCode(ContractCodeEntry {
+                            expiration_ledger_seq,
+                            ..
+                        }),
                     ..
                 }),
-                ..
-            })
-        ) = (&operations[0].changes[0], &operations[0].changes[1]) else {
-            return Err(Error::LedgerEntryNotFound);
-        };
-
-        Ok(*expiration_ledger_seq)
+            ) => Ok(*expiration_ledger_seq),
+            _ => Err(Error::LedgerEntryNotFound),
+        }
     }
 
     fn run_in_sandbox(&self) -> Result<u32, Error> {
diff --git a/cmd/soroban-cli/src/commands/contract/restore.rs b/cmd/soroban-cli/src/commands/contract/restore.rs
index 8dc824152..f94e50a97 100644
--- a/cmd/soroban-cli/src/commands/contract/restore.rs
+++ b/cmd/soroban-cli/src/commands/contract/restore.rs
@@ -6,12 +6,12 @@ use std::{
 
 use clap::{command, Parser};
 use soroban_env_host::xdr::{
-    ContractDataDurability, ContractDataEntry, ContractEntryBodyType, Error as XdrError,
-    ExtensionPoint, Hash, LedgerEntry, LedgerEntryChange, LedgerEntryData, LedgerFootprint,
-    LedgerKey, LedgerKeyContractData, Memo, MuxedAccount, Operation, OperationBody, Preconditions,
-    ReadXdr, RestoreFootprintOp, ScAddress, ScSpecTypeDef, ScVal, SequenceNumber, SorobanResources,
-    SorobanTransactionData, Transaction, TransactionExt, TransactionMeta, TransactionMetaV3,
-    Uint256,
+    ContractCodeEntry, ContractDataDurability, ContractDataEntry, ContractEntryBodyType,
+    Error as XdrError, ExtensionPoint, Hash, LedgerEntry, LedgerEntryChange, LedgerEntryData,
+    LedgerFootprint, LedgerKey, LedgerKeyContractData, Memo, MuxedAccount, Operation,
+    OperationBody, Preconditions, ReadXdr, RestoreFootprintOp, ScAddress, ScSpecTypeDef, ScVal,
+    SequenceNumber, SorobanResources, SorobanTransactionData, Transaction, TransactionExt,
+    TransactionMeta, TransactionMetaV3, Uint256,
 };
 use stellar_strkey::DecodeError;
 
@@ -186,12 +186,21 @@ impl Cmd {
         if operations[0].changes.len() != 1 {
             return Err(Error::LedgerEntryNotFound);
         }
-
-        let LedgerEntryChange::Created(LedgerEntry{ data: LedgerEntryData::ContractData(ContractDataEntry{expiration_ledger_seq, ..}), ..}) = &operations[0].changes[0] else {
-            return Err(Error::LedgerEntryNotFound);
-        };
-
-        Ok(*expiration_ledger_seq)
+        match operations[0].changes[0] {
+            LedgerEntryChange::Updated(LedgerEntry {
+                data:
+                    LedgerEntryData::ContractData(ContractDataEntry {
+                        expiration_ledger_seq,
+                        ..
+                    })
+                    | LedgerEntryData::ContractCode(ContractCodeEntry {
+                        expiration_ledger_seq,
+                        ..
+                    }),
+                ..
+            }) => Ok(expiration_ledger_seq),
+            _ => Err(Error::LedgerEntryNotFound),
+        }
     }
 
     fn run_in_sandbox(&self) -> Result<u32, Error> {