Repository: aptos-labs/aptos-core
File: aptos-move/framework/src/natives/code.rs (Rust)

Code before:

const ARBITRARY_POLICY: u8 = 0;

/// The native code context.
#[derive(Tid)]
pub struct NativeCodeContext {
    /// If false, publish requests are ignored and any attempts to publish code result in runtime
    /// errors.
    enabled: bool,
    /// Possibly stores (if not [None]) the request to publish a module bundle. The request is made
    /// using the native code defined in this context. It is later extracted by the VM for further
    /// checks and processing the actual publish.
    requested_module_bundle: Option<PublishRequest>,
}

impl VersionControlledNativeExtension for NativeCodeContext {
    fn undo(&mut self) {
        // No-op: nothing to undo.
    }

    fn save(&mut self) {
        // No-op: nothing to save.
    }

    fn update(&mut self, _txn_hash: &[u8; 32], _script_hash: &[u8], _session_counter: u8) {
        self.enabled = true;
        self.requested_module_bundle = None;
    }
}

impl NativeCodeContext {
    #[allow(clippy::new_without_default)]
    pub fn new() -> Self {
        Self {
            enabled: true,
            requested_module_bundle: None,
        }
    }

    pub fn extract_publish_request(&mut self) -> Option<PublishRequest> {
        if !self.enabled {
            return None;
        }

        self.enabled = false;
        self.requested_module_bundle.take()
    }
}

/// Represents a request for code publishing made from a native call and to be processed
/// by the Aptos VM.

Review comment (suggestion):
Note: we probably want to disable module publish for AA: currently one can run `code::publish_txn` in prologue via AA, but AptosVM will ignore these.

Code after:
    pub fn new() -> Self {
        Self {
            enabled: true,
            requested_module_bundle: None,
        }
    }

    pub fn extract_publish_request(&mut self) -> Option<PublishRequest> {
        if !self.enabled {
            return None;
        }

        self.enabled = false;
        self.requested_module_bundle.take()
    }
}

/// Represents a request for code publishing made from a native call and to be processed
/// by the Aptos VM.
pub struct PublishRequest {
    pub destination: AccountAddress,
    pub bundle: ModuleBundle,
    pub expected_modules: BTreeSet<String>,
    /// Allowed module dependencies. Empty for no restrictions. An empty string in the set
    /// allows all modules from that address.
    pub allowed_deps: Option<BTreeMap<AccountAddress, BTreeSet<String>>>,
    pub check_compat: bool,
}

/// Gets the string value embedded in a Move `string::String` struct.
fn get_move_string(v: Value) -> PartialVMResult<String> {
    let bytes = v
        .value_as::<Struct>()?
        .unpack()?
        .next()
        .ok_or_else(|| PartialVMError::new(StatusCode::INTERNAL_TYPE_ERROR))?
        .value_as::<Vec<u8>>()?;
    String::from_utf8(bytes).map_err(|_| PartialVMError::new(StatusCode::INTERNAL_TYPE_ERROR))
}

/// Gets the fields of the `code::AllowedDep` helper structure.
fn unpack_allowed_dep(v: Value) -> PartialVMResult<(AccountAddress, String)> {
    let mut fields = v.value_as::<Struct>()?.unpack()?.collect::<Vec<_>>();
    if fields.len() != 2 {
        return Err(PartialVMError::new(StatusCode::INTERNAL_TYPE_ERROR));
    }
    let module_name = get_move_string(fields.pop().unwrap())?;
    let account = fields.pop().unwrap().value_as::<AccountAddress>()?;
    Ok((account, module_name))
}