move: automated address mgmt for pkg upgrades

crates/sui-move-build/Cargo.toml

@@ -18,6 +18,7 @@ sui-verifier = { path = "../../sui-execution/latest/sui-verifier", package = "su
 serde-reflection.workspace = true
 sui-types.workspace = true
 sui-protocol-config.workspace = true
+sui-package-management.workspace = true
 
 move-binary-format.workspace = true
 move-bytecode-utils.workspace = true

crates/sui-move-build/src/lib.rs

@@ -40,6 +40,7 @@ use move_package::{
 };
 use move_symbol_pool::Symbol;
 use serde_reflection::Registry;
+use sui_package_management::{resolve_published_id, PublishedAtError};
 use sui_protocol_config::{Chain, ProtocolConfig, ProtocolVersion};
 use sui_types::{
     base_types::ObjectID,
@@ -223,7 +224,7 @@ pub fn build_from_resolution_graph(
     run_bytecode_verifier: bool,
     print_diags_to_stderr: bool,
 ) -> SuiResult<CompiledPackage> {
-    let (published_at, dependency_ids) = gather_published_ids(&resolution_graph);
+    let (published_at, dependency_ids) = gather_published_ids(&resolution_graph, None);
 
     let result = if print_diags_to_stderr {
         BuildConfig::compile_package(resolution_graph, &mut std::io::stderr())
@@ -634,12 +635,8 @@ pub struct PackageDependencies {
     pub unpublished: BTreeSet<Symbol>,
     /// Set of dependencies with invalid `published-at` addresses.
     pub invalid: BTreeMap<Symbol, String>,
-}
-
-#[derive(Debug, Clone)]
-pub enum PublishedAtError {
-    Invalid(String),
-    NotPresent,
+    /// Set of dependencies that have conflicting `published-at` addresses in Move.lock and Move.toml.
+    pub conflicting: BTreeMap<Symbol, (String, String)>,
 }
 
 /// Partition packages in `resolution_graph` into one of four groups:
@@ -649,16 +646,18 @@ pub enum PublishedAtError {
 /// - The names of packages that have a `published-at` field that isn't filled with a valid address.
 pub fn gather_published_ids(
     resolution_graph: &ResolvedGraph,
+    chain_id: Option<String>,
 ) -> (Result<ObjectID, PublishedAtError>, PackageDependencies) {
     let root = resolution_graph.root_package();
 
     let mut published = BTreeMap::new();
     let mut unpublished = BTreeSet::new();
     let mut invalid = BTreeMap::new();
+    let mut conflicting = BTreeMap::new();
     let mut published_at = Err(PublishedAtError::NotPresent);
 
     for (name, package) in &resolution_graph.package_table {
-        let property = published_at_property(package);
+        let property = resolve_published_id(package, chain_id.clone());
         if name == &root {
             // Separate out the root package as a special case
             published_at = property;
@@ -675,6 +674,9 @@ pub fn gather_published_ids(
             Err(PublishedAtError::Invalid(value)) => {
                 invalid.insert(*name, value);
             }
+            Err(PublishedAtError::Conflict(id_lock, id_manifest)) => {
+                conflicting.insert(*name, (id_lock, id_manifest));
+            }
         };
     }
 
@@ -684,6 +686,7 @@ pub fn gather_published_ids(
             published,
             unpublished,
             invalid,
+            conflicting,
         },
     )
 }

crates/sui-package-management/Cargo.toml

@@ -17,3 +17,4 @@ sui-json-rpc-types.workspace = true
 sui-sdk.workspace = true
 
 move-package.workspace = true
+move-symbol-pool.workspace = true

crates/sui-package-management/src/lib.rs

@@ -3,16 +3,31 @@
 
 use anyhow::{bail, Context};
 use std::path::PathBuf;
+use std::str::FromStr;
 
-use move_package::lock_file::{self, LockFile};
+use move_package::{
+    lock_file::{self, schema::ManagedPackage, LockFile},
+    resolution::resolution_graph::Package,
+    source_package::layout::SourcePackageLayout,
+};
+use move_symbol_pool::Symbol;
 use sui_json_rpc_types::{get_new_package_obj_from_response, SuiTransactionBlockResponse};
-use sui_sdk::wallet_context::WalletContext;
+use sui_sdk::{types::base_types::ObjectID, wallet_context::WalletContext};
+
+const PUBLISHED_AT_MANIFEST_FIELD: &str = "published-at";
 
 pub enum LockCommand {
     Publish,
     Upgrade,
 }
 
+#[derive(Debug, Clone)]
+pub enum PublishedAtError {
+    Invalid(String),
+    Conflict(String, String),
+    NotPresent,
+}
+
 /// Update the `Move.lock` file with automated address management info.
 /// Expects a wallet context, the publish or upgrade command, its response.
 /// The `Move.lock` principally file records the published address (i.e., package ID) of
@@ -73,3 +88,69 @@ pub async fn update_lock_file(
     lock.commit(lock_file)?;
     Ok(())
 }
+
+/// Find the published on-chain ID in the `Move.lock` or
+/// `Move.toml` file for the chain ID of the current environment. The published ID is resolved from the
+/// `Move.toml` if the Move.lock does not exist. Else, we resolve from the
+/// `Move.lock`, where addresses are automatically managed. If conflicting IDs are found in the
+/// `Move.lock` vs. `Move.toml`, a "Conflict" error message returns.
+pub fn resolve_published_id(
+    package: &Package,
+    chain_id: Option<String>,
+) -> Result<ObjectID, PublishedAtError> {
+    // Look up a valid `published-at` in the `Move.toml` first, which we'll
+    // return if the Move.lock does not manage addresses.
+    let published_id_in_manifest = match published_at_property(package) {
+        Ok(v) => Some(v),
+        Err(PublishedAtError::NotPresent) => None,
+        Err(e) => return Err(e), // An existing but invalid `published-at` in `Move.toml` should fail early.
+    };
+
+    let lock = package.package_path.join(SourcePackageLayout::Lock.path());
+    let mut lock_file = match std::fs::File::open(lock.clone()) {
+        Ok(f) => f,
+        Err(_) => match published_id_in_manifest {
+            Some(v) => {
+                return ObjectID::from_str(v.as_str())
+                    .map_err(|_| PublishedAtError::Invalid(v.to_owned()))
+            }
+            None => return Err(PublishedAtError::NotPresent),
+        },
+    };
+    let managed_packages = ManagedPackage::read(&mut lock_file).ok();
+    // Find the environment and ManagedPackage data for this chain_id.
+    let env_for_chain_id = managed_packages
+        .and_then(|m| {
+            m.into_iter().find(|(_, v)| {
+                if let Some(chain_id) = &chain_id {
+                    v.chain_id == *chain_id
+                } else {
+                    false
+                }
+            })
+        })
+        .map(|(k, v)| (k, v.latest_published_id));
+
+    let package_id = match (env_for_chain_id, published_id_in_manifest) {
+        (Some((_env, id_lock)), Some(id_manifest)) if id_lock != id_manifest => {
+            return Err(PublishedAtError::Conflict(id_lock, id_manifest))
+        }
+        (Some((_, id_lock)), _) => id_lock,
+        (None, Some(id_manifest)) => id_manifest, /* No info in Move.lock: Fall back to manifest */
+        _ => return Err(PublishedAtError::NotPresent), /* Neither in Move.toml nor Move.lock */
+    };
+    ObjectID::from_str(package_id.as_str())
+        .map_err(|_| PublishedAtError::Invalid(package_id.to_owned()))
+}
+
+fn published_at_property(package: &Package) -> Result<String, PublishedAtError> {
+    let Some(value) = package
+        .source_package
+        .package
+        .custom_properties
+        .get(&Symbol::from(PUBLISHED_AT_MANIFEST_FIELD))
+    else {
+        return Err(PublishedAtError::NotPresent);
+    };
+    Ok(value.to_string())
+}

crates/sui/src/client_commands.rs

@@ -47,9 +47,9 @@ use sui_json_rpc_types::{
 use sui_keys::keystore::AccountKeystore;
 use sui_move_build::{
     build_from_resolution_graph, check_invalid_dependencies, check_unpublished_dependencies,
-    gather_published_ids, BuildConfig, CompiledPackage, PackageDependencies, PublishedAtError,
+    gather_published_ids, BuildConfig, CompiledPackage, PackageDependencies,
 };
-use sui_package_management::LockCommand;
+use sui_package_management::{LockCommand, PublishedAtError};
 use sui_replay::ReplayToolCommand;
 use sui_sdk::{
     apis::ReadApi,
@@ -864,6 +864,11 @@ impl SuiClientCommands {
                         .map_err(|e| SuiError::ModulePublishFailure {
                             error: format!("Failed to canonicalize package path: {}", e),
                         })?;
+                let env_alias = context
+                    .config
+                    .get_active_env()
+                    .map(|e| e.alias.clone())
+                    .ok();
                 let (package_id, compiled_modules, dependencies, package_digest, upgrade_policy) =
                     upgrade_package(
                         client.read_api(),
@@ -872,6 +877,7 @@ impl SuiClientCommands {
                         upgrade_capability,
                         with_unpublished_dependencies,
                         skip_dependency_verification,
+                        env_alias,
                     )
                     .await?;
                 let tx_kind = client
@@ -1622,6 +1628,7 @@ pub(crate) async fn upgrade_package(
     upgrade_capability: ObjectID,
     with_unpublished_dependencies: bool,
     skip_dependency_verification: bool,
+    env_alias: Option<String>,
 ) -> Result<(ObjectID, Vec<Vec<u8>>, PackageDependencies, [u8; 32], u8), anyhow::Error> {
     let (dependencies, compiled_modules, compiled_package, package_id) = compile_package(
         read_api,
@@ -1634,12 +1641,25 @@ pub(crate) async fn upgrade_package(
 
     let package_id = package_id.map_err(|e| match e {
         PublishedAtError::NotPresent => {
-            anyhow!("No 'published-at' field in manifest for package to be upgraded.")
+            anyhow!("No 'published-at' field in Move.toml or Move.lock for package to be upgraded.")
         }
         PublishedAtError::Invalid(v) => anyhow!(
-            "Invalid 'published-at' field in manifest of package to be upgraded. \
+            "Invalid 'published-at' field in Move.toml or Move.lock of package to be upgraded. \
                          Expected an on-chain address, but found: {v:?}"
         ),
+        PublishedAtError::Conflict(id_lock, id_manifest) => {
+            let env_alias = format!("(currently {})", env_alias.unwrap_or_default());
+            anyhow!(
+                "Conflicting published package address: `Move.toml` contains published-at address \
+                 {id_manifest} but `Move.lock` file contains published-at address {id_lock}. \
+                 You may want to:
+
+                 - delete the published-at address in the `Move.toml` if the `Move.lock` address is correct; OR
+                 - update the `Move.lock` address using the `sui manage-package` command to be the same as the `Move.toml`; OR
+                 - check that your `sui active-env` {env_alias} corresponds to the chain on which the package is published (i.e., devnet, testnet, mainnet); OR
+                 - contact the maintainer if this package is a dependency and request resolving the conflict."
+            )
+        }
     })?;
 
     let resp = read_api
@@ -1700,7 +1720,8 @@ pub(crate) async fn compile_package(
         print_diags_to_stderr,
     };
     let resolution_graph = config.resolution_graph(&package_path)?;
-    let (package_id, dependencies) = gather_published_ids(&resolution_graph);
+    let chain_id = read_api.get_chain_identifier().await.ok();
+    let (package_id, dependencies) = gather_published_ids(&resolution_graph, chain_id);
     check_invalid_dependencies(&dependencies.invalid)?;
     if !with_unpublished_dependencies {
         check_unpublished_dependencies(&dependencies.unpublished)?;

crates/sui/src/client_ptb/builder.rs

@@ -974,6 +974,7 @@ impl<'a> PTBBuilder<'a> {
                         ObjectID::from_address(upgrade_cap_id.into_inner()),
                         false, /* with_unpublished_dependencies */
                         false, /* skip_dependency_verification */
+                        None,
                     )
                     .await
                     .map_err(|e| err!(path_loc, "{e}"))?;