From ec64c2e0e1b563da8731b3b770fe5f8de9a072ec Mon Sep 17 00:00:00 2001
From: Jon Michael Aanes <jonjmaa@gmail.com>
Date: Wed, 9 Apr 2025 22:09:15 +0200
Subject: [PATCH] Asset contract is now upgradable

---
 python/notamon_viewer/app.py               |  1 -
 rust/notamon-asset-contract/src/lib.rs     |  7 +++----
 rust/notamon-asset-contract/src/upgrade.rs | 21 +++++++++++++++++++++
 rust/notamon-common/src/permission.rs      | 14 +++++++++-----
 4 files changed, 33 insertions(+), 10 deletions(-)
 create mode 100644 rust/notamon-asset-contract/src/upgrade.rs

diff --git a/python/notamon_viewer/app.py b/python/notamon_viewer/app.py
index 71642a4..480c835 100644
--- a/python/notamon_viewer/app.py
+++ b/python/notamon_viewer/app.py
@@ -1,6 +1,5 @@
 import flask
 import pbc_client
-import pbcabi
 import base64
 import dataclasses
 import requests_cache
diff --git a/rust/notamon-asset-contract/src/lib.rs b/rust/notamon-asset-contract/src/lib.rs
index c9e3893..66affb1 100644
--- a/rust/notamon-asset-contract/src/lib.rs
+++ b/rust/notamon-asset-contract/src/lib.rs
@@ -4,21 +4,20 @@
 #[macro_use]
 extern crate pbc_contract_codegen;
 
-use pbc_contract_common::address::Address;
 use pbc_contract_common::context::ContractContext;
 use pbc_contract_common::events::EventGroup;
 use pbc_contract_common::zk::{SecretVarId, ZkInputDef, ZkState, ZkStateChange};
-use pbc_zk::Sbi8;
+use pbc_zk::Sbu8;
 use read_write_rpc_derive::ReadWriteRPC;
 use read_write_state_derive::ReadWriteState;
 use pbc_contract_common::avl_tree_map::AvlTreeMap;
-use pbc_traits::{ReadRPC, WriteRPC, ReadWriteState};
 use create_type_spec_derive::CreateTypeSpec;
 use std::fmt::Debug;
 
 use notamon_common::{AssetId, Permission, Permissions};
 
 mod zk_compute;
+mod upgrade;
 
 /// Metadata for secret-shared assets.
 #[derive(ReadWriteState, ReadWriteRPC, Debug)]
@@ -70,7 +69,7 @@ pub fn set_asset(
 ) -> (
     AssetContractState,
     Vec<EventGroup>,
-    ZkInputDef<AssetMetadata, Vec<Sbi8>>,
+    ZkInputDef<AssetMetadata, Vec<Sbu8>>,
 ) {
     state.permissions.assert_has_permission(&context.sender, Role::UPLOADER {});
     let input_def = ZkInputDef::with_metadata_and_size(Some(SHORTNAME_SET_ASSET_INPUTTED), AssetMetadata { asset_id }, asset_length * 8);
diff --git a/rust/notamon-asset-contract/src/upgrade.rs b/rust/notamon-asset-contract/src/upgrade.rs
new file mode 100644
index 0000000..64cb80b
--- /dev/null
+++ b/rust/notamon-asset-contract/src/upgrade.rs
@@ -0,0 +1,21 @@
+//! Upgrade logic for allowing upgrade.
+
+use crate::{AssetContractState, Role};
+use pbc_contract_codegen::upgrade_is_allowed;
+use pbc_contract_common::context::ContractContext;
+use pbc_contract_common::upgrade::ContractHashes;
+
+/// Checks whether the upgrade is allowed.
+///
+/// This contract allows the [`ContractState::upgrader`] to upgrade the contract at any time.
+#[upgrade_is_allowed]
+pub fn is_upgrade_allowed(
+    context: ContractContext,
+    state: AssetContractState,
+    _old_contract_hashes: ContractHashes,
+    _new_contract_hashes: ContractHashes,
+    _new_contract_rpc: Vec<u8>,
+) -> bool {
+    state.permissions.assert_has_permission(&context.sender, Role::ADMIN {});
+    true
+}
diff --git a/rust/notamon-common/src/permission.rs b/rust/notamon-common/src/permission.rs
index 0ecc752..c20df67 100644
--- a/rust/notamon-common/src/permission.rs
+++ b/rust/notamon-common/src/permission.rs
@@ -34,7 +34,7 @@ pub struct Permissions<KeyT: ReadWriteState> {
     permissions: AvlTreeMap<KeyT, Permission>,
 }
 
-impl <KeyT: ReadWriteState> Permissions<KeyT> {
+impl <KeyT: ReadWriteState + Copy> Permissions<KeyT> {
 
     pub fn new () -> Self {
         Self { permissions: AvlTreeMap::new() }
@@ -44,15 +44,19 @@ impl <KeyT: ReadWriteState> Permissions<KeyT> {
         self.permissions.insert(permission_key, permission);
     }
 
-    pub fn get_permission(&self, permission: &KeyT) -> Permission {
-        self.permissions.get(permission).unwrap_or(Permission::NONE)
+    pub fn get_permission(&self, permission: KeyT) -> Permission {
+        self.permissions.get(&permission).unwrap_or(Permission::NONE)
     }
 
 }
 
-impl <KeyT: ReadWriteState + Debug> Permissions<KeyT> {
+impl <KeyT: ReadWriteState + Debug + Copy> Permissions<KeyT> {
+    pub fn has_permission(&self, address: &Address, permission: KeyT) -> bool {
+        self.get_permission(permission).allows(address)
+    }
+
     pub fn assert_has_permission(&self, address: &Address, permission: KeyT) {
-        if !self.get_permission(&permission).allows(address) {
+        if !self.has_permission(address, permission) {
             panic!("User {address} does not have permission: {permission:?}");
         }
     }