feat(sui-genesis-builder): Add Basic Output object validation (#237)

* implement output validation against

* whoops

* Add aliases to map

* Add struct to define potentially created objects. Update basic output validation to account for generated coins.

* move created objects map to migration struct

* Refactor module structure.

* Move migration logic out of mod.rs

* Add native token coin validation

* Fix native token validation

* Update crates/sui-genesis-builder/src/stardust/migration/verification/mod.rs

Co-authored-by: /alex/ <[email protected]>

* re-add test methods

* refactor: Fix some clippy

---------

Co-authored-by: /alex/ <[email protected]>
Co-authored-by: samuel_rufi <[email protected]>

crates/sui-genesis-builder/src/stardust/migration/mod.rs

@@ -1,6 +1,10 @@
 // Copyright (c) 2024 IOTA Stiftung
 // SPDX-License-Identifier: Apache-2.0
 
-pub mod migration;
+mod executor;
+mod migration;
 #[cfg(test)]
 mod tests;
+pub mod verification;
+
+pub use migration::*;

crates/sui-genesis-builder/src/stardust/migration/tests/alias.rs

@@ -1,7 +1,7 @@
 // Copyright (c) 2024 IOTA Stiftung
 // SPDX-License-Identifier: Apache-2.0
 
-use crate::stardust::migration::migration::Executor;
+use crate::stardust::migration::executor::Executor;
 use crate::stardust::migration::migration::{MIGRATION_PROTOCOL_VERSION, PACKAGE_DEPS};
 use crate::stardust::migration::tests::random_output_header;
 use crate::stardust::migration::tests::run_migration;
@@ -40,11 +40,7 @@ fn migrate_alias(
     let mut snapshot_buffer = Vec::new();
     Migration::new()
         .unwrap()
-        .run(
-            [].into_iter(),
-            [(header, stardust_alias.into())].into_iter(),
-            &mut snapshot_buffer,
-        )
+        .run([(header, stardust_alias.into())], &mut snapshot_buffer)
         .unwrap();
 
     let migrated_objects: Vec<Object> = bcs::from_bytes(&snapshot_buffer).unwrap();

crates/sui-genesis-builder/src/stardust/migration/tests/executor.rs

@@ -1,3 +1,6 @@
+// Copyright (c) 2024 IOTA Stiftung
+// SPDX-License-Identifier: Apache-2.0
+
 use iota_sdk::types::block::output::NativeTokens;
 use iota_sdk::types::block::{
     address::AliasAddress,
@@ -14,7 +17,9 @@ use sui_types::{
     object::Owner,
 };
 
-use crate::stardust::migration::migration::{Executor, NATIVE_TOKEN_BAG_KEY_TYPE};
+use crate::stardust::migration::executor::Executor;
+use crate::stardust::migration::migration::NATIVE_TOKEN_BAG_KEY_TYPE;
+use crate::stardust::migration::tests::random_output_header;
 use crate::stardust::native_token::package_builder;
 use crate::stardust::native_token::package_data::{NativeTokenModuleData, NativeTokenPackageData};
 
@@ -24,6 +29,7 @@ fn create_bag_with_pt() {
     let owner = AliasAddress::new(AliasId::new([0; AliasId::LENGTH]));
     let supply = 1_000_000;
     let token_scheme = SimpleTokenScheme::new(supply, 0, supply).unwrap();
+    let header = random_output_header();
     let foundry = FoundryOutputBuilder::new_with_amount(1000, 1, token_scheme.into())
         .with_unlock_conditions([UnlockCondition::from(
             ImmutableAliasAddressUnlockCondition::new(owner),
@@ -43,7 +49,7 @@ fn create_bag_with_pt() {
     let mut executor = Executor::new(ProtocolVersion::MAX).unwrap();
     let object_count = executor.store().objects().len();
     executor
-        .create_foundries([(foundry, foundry_package)].into_iter())
+        .create_foundries([(&header, &foundry, foundry_package)])
         .unwrap();
     // Foundry package publication creates four objects
     //
@@ -67,7 +73,7 @@ fn create_bag_with_pt() {
     let native_token = NativeToken::new(foundry_id.into(), token_amount).unwrap();
 
     // Create the bag
-    let (bag, _) = executor
+    let (bag, _, _) = executor
         .create_bag_with_pt(&NativeTokens::from_vec(vec![native_token]).unwrap())
         .unwrap();
     assert!(executor.store().get_object(bag.id.object_id()).is_none());

crates/sui-genesis-builder/src/stardust/migration/tests/mod.rs

@@ -22,27 +22,10 @@ fn random_output_header() -> OutputHeader {
 
 fn run_migration(outputs: impl IntoIterator<Item = (OutputHeader, Output)>) -> Vec<Object> {
     let mut snapshot_buffer = Vec::new();
-    let mut foundries = Vec::new();
-    let mut outputs_without_foundries = Vec::new();
-
-    for (header, output) in outputs.into_iter() {
-        match output {
-            Output::Foundry(foundry) => {
-                foundries.push((header, foundry));
-            }
-            other => {
-                outputs_without_foundries.push((header, other));
-            }
-        }
-    }
 
     Migration::new()
         .unwrap()
-        .run(
-            foundries.into_iter(),
-            outputs_without_foundries.into_iter(),
-            &mut snapshot_buffer,
-        )
+        .run(outputs, &mut snapshot_buffer)
         .unwrap();
 
     bcs::from_bytes(&snapshot_buffer).unwrap()

crates/sui-genesis-builder/src/stardust/migration/verification/alias.rs

@@ -0,0 +1,16 @@
+// Copyright (c) 2024 IOTA Stiftung
+// SPDX-License-Identifier: Apache-2.0
+
+use iota_sdk::types::block::output::AliasOutput;
+use sui_types::in_memory_storage::InMemoryStorage;
+
+use super::created_objects::CreatedObjects;
+
+pub fn verify_alias_output(
+    output: &AliasOutput,
+    created_objects: &CreatedObjects,
+    storage: &InMemoryStorage,
+) -> anyhow::Result<()> {
+    // TODO: Implementation. Returns Ok for now so the migration can be tested.
+    Ok(())
+}

crates/sui-genesis-builder/src/stardust/migration/verification/basic.rs

@@ -0,0 +1,141 @@
+// Copyright (c) 2024 IOTA Stiftung
+// SPDX-License-Identifier: Apache-2.0
+
+use anyhow::{anyhow, ensure, Result};
+use iota_sdk::types::block::output::BasicOutput;
+use sui_types::{balance::Balance, dynamic_field::Field, in_memory_storage::InMemoryStorage};
+
+use crate::stardust::migration::verification::{
+    created_objects::CreatedObjects,
+    util::{
+        verify_expiration_unlock_condition, verify_metadata_feature, verify_native_tokens,
+        verify_sender_feature, verify_storage_deposit_unlock_condition, verify_tag_feature,
+        verify_timelock_unlock_condition,
+    },
+};
+
+pub fn verify_basic_output(
+    output: &BasicOutput,
+    created_objects: &CreatedObjects,
+    storage: &InMemoryStorage,
+) -> Result<()> {
+    // If the output has multiple unlock conditions, then a genesis object should have been created.
+    if output.unlock_conditions().len() > 1 {
+        let created_output = created_objects
+            .output()
+            .and_then(|id| {
+                storage
+                    .get_object(id)
+                    .ok_or_else(|| anyhow!("missing object"))
+            })?
+            .to_rust::<crate::stardust::types::output::BasicOutput>()
+            .ok_or_else(|| anyhow!("invalid basic output object"))?;
+
+        // Amount
+        ensure!(
+            created_output.iota.value() == output.amount(),
+            "amount mismatch: found {}, expected {}",
+            created_output.iota.value(),
+            output.amount()
+        );
+
+        // Native Tokens
+        ensure!(
+            created_output.native_tokens.size == output.native_tokens().len() as u64,
+            "native tokens bag length mismatch: found {}, expected {}",
+            created_output.native_tokens.size,
+            output.native_tokens().len()
+        );
+        let created_native_token_fields = created_objects.native_tokens().and_then(|ids| {
+            ids.iter()
+                .map(|id| {
+                    let obj = storage
+                        .get_object(id)
+                        .ok_or_else(|| anyhow!("missing native token field for {id}"))?;
+                    obj.to_rust::<Field<String, Balance>>().ok_or_else(|| {
+                        anyhow!("expected a native token field, found {:?}", obj.type_())
+                    })
+                })
+                .collect::<Result<Vec<_>, _>>()
+        })?;
+        verify_native_tokens(output.native_tokens(), created_native_token_fields)?;
+
+        // Storage Deposit Return Unlock Condition
+        verify_storage_deposit_unlock_condition(
+            output.unlock_conditions().storage_deposit_return(),
+            created_output.storage_deposit_return.as_ref(),
+        )?;
+
+        // Timelock Unlock Condition
+        verify_timelock_unlock_condition(
+            output.unlock_conditions().timelock(),
+            created_output.timelock.as_ref(),
+        )?;
+
+        // Expiration Unlock Condition
+        verify_expiration_unlock_condition(
+            output.unlock_conditions().expiration(),
+            created_output.expiration.as_ref(),
+            output.address(),
+        )?;
+
+        // Metadata Feature
+        verify_metadata_feature(
+            output.features().metadata(),
+            created_output.metadata.as_ref(),
+        )?;
+
+        // Tag Feature
+        verify_tag_feature(output.features().tag(), created_output.tag.as_ref())?;
+
+        // Sender Feature
+        verify_sender_feature(output.features().sender(), created_output.sender)?;
+
+    // Otherwise the output contains only an address unlock condition and only a coin
+    // and possibly native tokens should have been created.
+    } else {
+        ensure!(
+            created_objects.output().is_err(),
+            "unexpected output object created for simple deposit"
+        );
+
+        // Coin value.
+        let created_coin = created_objects
+            .coin()
+            .and_then(|id| {
+                storage
+                    .get_object(id)
+                    .ok_or_else(|| anyhow!("missing coin"))
+            })?
+            .as_coin_maybe()
+            .ok_or_else(|| anyhow!("expected a coin"))?;
+        ensure!(
+            created_coin.value() == output.amount(),
+            "coin amount mismatch: found {}, expected {}",
+            created_coin.value(),
+            output.amount()
+        );
+
+        // Native Tokens
+        let created_native_token_coins = created_objects.native_tokens().and_then(|ids| {
+            ids.iter()
+                .map(|id| {
+                    let obj = storage
+                        .get_object(id)
+                        .ok_or_else(|| anyhow!("missing native token coin for {id}"))?;
+                    obj.as_coin_maybe().ok_or_else(|| {
+                        anyhow!("expected a native token coin, found {:?}", obj.type_())
+                    })
+                })
+                .collect::<Result<Vec<_>, _>>()
+        })?;
+        verify_native_tokens(output.native_tokens(), created_native_token_coins)?;
+    }
+
+    ensure!(
+        created_objects.package().is_err(),
+        "unexpected package found"
+    );
+
+    Ok(())
+}

crates/sui-genesis-builder/src/stardust/migration/verification/created_objects.rs

@@ -0,0 +1,72 @@
+// Copyright (c) 2024 IOTA Stiftung
+// SPDX-License-Identifier: Apache-2.0
+
+use anyhow::{anyhow, bail, Result};
+use sui_types::base_types::ObjectID;
+
+/// Defines objects that may have been created by migrating an [`Output`].
+#[derive(Default)]
+pub struct CreatedObjects {
+    output: Option<ObjectID>,
+    coin: Option<ObjectID>,
+    package: Option<ObjectID>,
+    native_tokens: Option<Vec<ObjectID>>,
+}
+
+impl CreatedObjects {
+    pub fn output(&self) -> Result<&ObjectID> {
+        self.output
+            .as_ref()
+            .ok_or_else(|| anyhow!("no created output object"))
+    }
+
+    pub(crate) fn set_output(&mut self, id: ObjectID) -> Result<()> {
+        if let Some(id) = self.output {
+            bail!("output already set: {id}")
+        }
+        self.output.replace(id);
+        Ok(())
+    }
+
+    pub fn coin(&self) -> Result<&ObjectID> {
+        self.coin
+            .as_ref()
+            .ok_or_else(|| anyhow!("no created coin object"))
+    }
+
+    pub(crate) fn set_coin(&mut self, id: ObjectID) -> Result<()> {
+        if let Some(id) = self.coin {
+            bail!("coin already set: {id}")
+        }
+        self.coin.replace(id);
+        Ok(())
+    }
+
+    pub fn package(&self) -> Result<&ObjectID> {
+        self.package
+            .as_ref()
+            .ok_or_else(|| anyhow!("no created package object"))
+    }
+
+    pub(crate) fn set_package(&mut self, id: ObjectID) -> Result<()> {
+        if let Some(id) = self.package {
+            bail!("package already set: {id}")
+        }
+        self.package.replace(id);
+        Ok(())
+    }
+
+    pub fn native_tokens(&self) -> Result<&[ObjectID]> {
+        self.native_tokens
+            .as_deref()
+            .ok_or_else(|| anyhow!("no created native token objects"))
+    }
+
+    pub(crate) fn set_native_tokens(&mut self, ids: Vec<ObjectID>) -> Result<()> {
+        if let Some(id) = &self.native_tokens {
+            bail!("native tokens already set: {id:?}")
+        }
+        self.native_tokens.replace(ids);
+        Ok(())
+    }
+}

crates/sui-genesis-builder/src/stardust/migration/verification/foundry.rs

@@ -0,0 +1,16 @@
+// Copyright (c) 2024 IOTA Stiftung
+// SPDX-License-Identifier: Apache-2.0
+
+use iota_sdk::types::block::output::FoundryOutput;
+use sui_types::in_memory_storage::InMemoryStorage;
+
+use super::created_objects::CreatedObjects;
+
+pub fn verify_foundry_output(
+    output: &FoundryOutput,
+    created_objects: &CreatedObjects,
+    storage: &InMemoryStorage,
+) -> anyhow::Result<()> {
+    // TODO: Implementation. Returns Ok for now so the migration can be tested.
+    Ok(())
+}

crates/sui-genesis-builder/src/stardust/migration/verification/mod.rs

@@ -0,0 +1,35 @@
+// Copyright (c) 2024 IOTA Stiftung
+// SPDX-License-Identifier: Apache-2.0
+
+//! The [`verification`] module contains the validation logic to make sure that the stardust outputs are correctly converted to the move objects.
+
+use iota_sdk::types::block::output::Output;
+use sui_types::in_memory_storage::InMemoryStorage;
+
+use crate::stardust::types::snapshot::OutputHeader;
+
+use self::created_objects::CreatedObjects;
+
+pub mod alias;
+pub mod basic;
+pub mod created_objects;
+pub mod foundry;
+pub mod nft;
+mod util;
+
+pub fn verify_output(
+    header: &OutputHeader,
+    output: &Output,
+    created_objects: &CreatedObjects,
+    storage: &InMemoryStorage,
+) -> anyhow::Result<()> {
+    match output {
+        Output::Alias(output) => alias::verify_alias_output(output, created_objects, storage),
+        Output::Basic(output) => basic::verify_basic_output(output, created_objects, storage),
+        Output::Foundry(output) => foundry::verify_foundry_output(output, created_objects, storage),
+        Output::Nft(output) => nft::verify_nft_output(output, created_objects, storage),
+        // Treasury outputs aren't used since Stardust, so no need to verify anything here.
+        Output::Treasury(_) => return Ok(()),
+    }
+    .map_err(|e| anyhow::anyhow!("error verifying output {}: {}", header.output_id(), e))
+}