kms-contract: Add back registered apps

Author: kvinwang

docs/deployment.md

@@ -292,24 +292,26 @@ After the dstack-vmm is ready, you can deploy an app on it following the steps b
 
 ### 1. On-chain Registration
 
-In order to deploy an app on the dstack-vmm, you need to deploy an App's control contract AppAuth. Developers can develop their own or choose the reference contract from the Dstack repository. Custom contracts need to implement the IAppAuth interface.
+The on-chain registration process includes two steps:
 
-The Dstack repository provides scripts to complete the deployment:
+1. Deploy an App's control contract AppAuth. Developers can develop their own or choose the reference contract from the Dstack repository. Custom contracts need to implement the IAppAuth interface.
+2. Call KmsAuth.registerApp(appAuthAddress) to register the contract.
 
-**Option 1: Traditional deployment**
+The Dstack repository provides scripts to complete these two steps:
+
+**Option 1: Traditional deployment (2 transactions)**
 ```bash
 git clone https://github.com/Dstack-TEE/dstack
 cd dstack/kms/auth-eth
 npm install
 npx hardhat compile
 export PRIVATE_KEY=<your eth private key here>
+export KMS_CONTRACT_ADDRESS=0xFE6C45aE66344CAEF5E5D7e2cbD476286D651875
 npx hardhat app:deploy --allow-any-device --network phala
 ```
 
-**Option 2: Factory deployment (recommended)**
+**Option 2: Factory deployment (1 transaction, recommended)**
 ```bash
-export PRIVATE_KEY=<your eth private key here>
-export KMS_CONTRACT_ADDRESS=0xFE6C45aE66344CAEF5E5D7e2cbD476286D651875
 npx hardhat kms:create-app --allow-any-device --network phala
 ```
 

kms/auth-eth/contracts/KmsAuth.sol

@@ -30,6 +30,9 @@ contract KmsAuth is
     /// @custom:oz-renamed-from tproxyAppId
     string public gatewayAppId;
 
+    // Mapping of registered apps
+    mapping(address => bool) public registeredApps;
+
     // Mapping of allowed aggregated MR measurements for running KMS
     mapping(bytes32 => bool) public kmsAllowedAggregatedMrs;
 
@@ -43,6 +46,7 @@ contract KmsAuth is
     address public appAuthImplementation;
 
     // Events
+    event AppRegistered(address appId);
     event KmsInfoSet(bytes k256Pubkey);
     event KmsAggregatedMrAdded(bytes32 mrAggregated);
     event KmsAggregatedMrRemoved(bytes32 mrAggregated);
@@ -117,15 +121,22 @@ contract KmsAuth is
         emit GatewayAppIdSet(appId);
     }
 
+    // Function to register an app
+    function registerApp(address appId) public {
+        require(appId != address(0), "Invalid app ID");
+        registeredApps[appId] = true;
+        emit AppRegistered(appId);
+    }
+
     // Function to set AppAuth implementation contract address
     function setAppAuthImplementation(address _implementation) external onlyOwner {
         require(_implementation != address(0), "Invalid implementation address");
         appAuthImplementation = _implementation;
         emit AppAuthImplementationSet(_implementation);
     }
 
-    // Factory method: Deploy AppAuth in single transaction
-    function deployApp(
+    // Factory method: Deploy and register AppAuth in single transaction
+    function deployAndRegisterApp(
         address initialOwner,
         bool disableUpgrades,
         bool allowAnyDevice,
@@ -147,7 +158,8 @@ contract KmsAuth is
 
         // Deploy proxy contract
         appId = address(new ERC1967Proxy(appAuthImplementation, initData));
-
+        // Register to KMS
+        registerApp(appId);
         emit AppDeployedViaFactory(appId, msg.sender);
     }
 
@@ -221,6 +233,10 @@ contract KmsAuth is
     function isAppAllowed(
         AppBootInfo calldata bootInfo
     ) external view override returns (bool isAllowed, string memory reason) {
+        // Check if app is registered
+        if (!registeredApps[bootInfo.appId]) {
+            return (false, "App not registered");
+        }
         // Check if the OS image is allowed
         if (!allowedOsImages[bootInfo.osImageHash]) {
             return (false, "OS image is not allowed");

kms/auth-eth/foundry-cast-cheatsheet.md

@@ -181,17 +181,17 @@ cast send $KMS_CONTRACT_ADDRESS "removeKmsDevice(bytes32)" \
 
 ```bash
 # kms:create-app - Deploy and register AppAuth in single transaction
-cast send $KMS_CONTRACT_ADDRESS "deployApp(address,bool,bool,bytes32,bytes32)" \
+cast send $KMS_CONTRACT_ADDRESS "deployAndRegisterApp(address,bool,bool,bytes32,bytes32)" \
   "$DEPLOYER_ADDRESS" false true \
   "0x1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef" \
   "0xfedcba0987654321fedcba0987654321fedcba0987654321fedcba0987654321" \
   --private-key $PRIVATE_KEY --rpc-url $RPC_URL
 # Parameters: (owner, disableUpgrades, allowAnyDevice, initialDeviceId, initialComposeHash)
 # Use 0x0000...0000 for empty device/hash values
-# To decode return: cast abi-decode "deployApp(address,bool,bool,bytes32,bytes32)(address,address)" RETURN_DATA
+# To decode return: cast abi-decode "deployAndRegisterApp(address,bool,bool,bytes32,bytes32)(address,address)" RETURN_DATA
 
 # Example with no initial data:
-cast send $KMS_CONTRACT_ADDRESS "deployApp(address,bool,bool,bytes32,bytes32)" \
+cast send $KMS_CONTRACT_ADDRESS "deployAndRegisterApp(address,bool,bool,bytes32,bytes32)" \
   "$DEPLOYER_ADDRESS" false true \
   "0x0000000000000000000000000000000000000000000000000000000000000000" \
   "0x0000000000000000000000000000000000000000000000000000000000000000" \
@@ -379,7 +379,7 @@ cast abi-decode "kmsAllowedAggregatedMrs(bytes32)(bool)" RETURN_DATA
 cast abi-decode "isAppAllowed((address,bytes32,address,bytes32,bytes32,bytes32,bytes32,string,string[]))(bool,string)" RETURN_DATA
 
 # Decode factory deployment response
-cast abi-decode "deployApp(address,bool,bool,bytes32,bytes32)(address,address)" RETURN_DATA
+cast abi-decode "deployAndRegisterApp(address,bool,bool,bytes32,bytes32)(address,address)" RETURN_DATA
 ```
 
 ### Get Contract Information
@@ -459,7 +459,7 @@ cast send $KMS_CONTRACT_ADDRESS "addKmsAggregatedMr(bytes32)" "0x..." \
   --private-key $PRIVATE_KEY --rpc-url $RPC_URL
 
 # 3. Users can now deploy apps via factory immediately!
-cast send $KMS_CONTRACT_ADDRESS "deployApp(address,bool,bool,bytes32,bytes32)" \
+cast send $KMS_CONTRACT_ADDRESS "deployAndRegisterApp(address,bool,bool,bytes32,bytes32)" \
   "$USER_ADDRESS" false true "0x..." "0x..." \
   --private-key $USER_PRIVATE_KEY --rpc-url $RPC_URL
 ```
@@ -485,7 +485,7 @@ cast send $KMS_CONTRACT_ADDRESS "addKmsAggregatedMr(bytes32)" "0x..." \
   --private-key $PRIVATE_KEY --rpc-url $RPC_URL
 
 # 5. Users can now deploy apps via factory
-cast send $KMS_CONTRACT_ADDRESS "deployApp(address,bool,bool,bytes32,bytes32)" \
+cast send $KMS_CONTRACT_ADDRESS "deployAndRegisterApp(address,bool,bool,bytes32,bytes32)" \
   "$USER_ADDRESS" false true "0x..." "0x..." \
   --private-key $USER_PRIVATE_KEY --rpc-url $RPC_URL
 ```
@@ -521,7 +521,7 @@ cast call $KMS_CONTRACT_ADDRESS "owner()" --rpc-url $RPC_URL
 | `info:kms` | `cast call ... kmsInfo` | Returns struct |
 | `app:deploy` | Complex hardhat task | Multi-transaction deployment |
 | `app:deploy-with-data` | Complex hardhat task | Use initializeWithData |
-| `app:deploy-factory` | `cast send ... deployApp` | **Single transaction deployment** ⭐ |
+| `app:deploy-factory` | `cast send ... deployAndRegisterApp` | **Single transaction deployment** ⭐ |
 | `kms:set-app-implementation` | `cast send ... setAppAuthImplementation` | Manual setup (rarely needed now) |
 | `kms:get-app-implementation` | `cast call ... appAuthImplementation` | Query factory implementation |
 
@@ -547,7 +547,7 @@ mycast send $APP_AUTH_ADDRESS "addDevice(bytes32)" \
   "0x1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef"
 
 # Example: Factory deployment
-mycast send $KMS_CONTRACT_ADDRESS "deployApp(address,bool,bool,bytes32,bytes32)" \
+mycast send $KMS_CONTRACT_ADDRESS "deployAndRegisterApp(address,bool,bool,bytes32,bytes32)" \
   "$DEPLOYER_ADDRESS" false true \
   "0x1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef" \
   "0xfedcba0987654321fedcba0987654321fedcba0987654321fedcba0987654321"

kms/auth-eth/hardhat.config.ts

@@ -245,6 +245,8 @@ task("app:deploy", "Deploy AppAuth with a UUPS proxy")
     console.log("Deploying with account:", deployerAddress);
     console.log("Account balance:", await accountBalance(ethers, deployerAddress));
 
+    const kmsContract = await getKmsAuth(ethers);
+
     // Parse device and hash (convert to bytes32, use 0x0 if empty)
     const deviceId = taskArgs.device ? taskArgs.device.trim() : "0x0000000000000000000000000000000000000000000000000000000000000000";
     const composeHash = taskArgs.hash ? taskArgs.hash.trim() : "0x0000000000000000000000000000000000000000000000000000000000000000";
@@ -269,15 +271,54 @@ task("app:deploy", "Deploy AppAuth with a UUPS proxy")
     if (!appAuth) {
       return;
     }
+    
     await appAuth.waitForDeployment();
     const proxyAddress = await appAuth.getAddress();
-    console.log("✅ App deployed successfully!");
-    console.log("AppAuth deployed to (App Id):", proxyAddress);
+    console.log("AppAuth deployed to:", proxyAddress);
 
-    if (hasInitialData) {
-      const hasDevice = deviceId !== "0x0000000000000000000000000000000000000000000000000000000000000000";
-      const hasHash = composeHash !== "0x0000000000000000000000000000000000000000000000000000000000000000";
-      console.log(`Deployed with ${hasDevice ? "1" : "0"} initial device and ${hasHash ? "1" : "0"} initial compose hash`);
+    const tx = await kmsContract.registerApp(proxyAddress);
+    const receipt = await waitTx(tx);
+    
+    // Parse the AppRegistered event from the logs
+    let appRegisteredEvent = null;
+    for (const log of receipt.logs) {
+      try {
+        const parsedLog = kmsContract.interface.parseLog({
+          topics: log.topics,
+          data: log.data
+        });
+        
+        if (parsedLog?.name === 'AppRegistered') {
+          appRegisteredEvent = parsedLog.args;
+          break;
+        }
+      } catch (e) {
+        continue;
+      }
+    }
+
+    if (appRegisteredEvent) {
+      console.log("✅ App deployed and registered successfully!");
+      console.log("App ID:", appRegisteredEvent.appId);
+      console.log("Proxy Address:", proxyAddress);
+      console.log("Owner:", deployerAddress);
+      console.log("Transaction hash:", tx.hash);
+      
+      if (hasInitialData) {
+        const hasDevice = deviceId !== "0x0000000000000000000000000000000000000000000000000000000000000000";
+        const hasHash = composeHash !== "0x0000000000000000000000000000000000000000000000000000000000000000";
+        console.log(`Deployed with ${hasDevice ? "1" : "0"} initial device and ${hasHash ? "1" : "0"} initial compose hash`);
+      }
+    } else {
+      console.log("✅ App deployed and registered successfully!");
+      console.log("Proxy Address:", proxyAddress);
+      console.log("Transaction hash:", tx.hash);
+      
+      if (hasInitialData) {
+        const hasDevice = deviceId !== "0x0000000000000000000000000000000000000000000000000000000000000000";
+        const hasHash = composeHash !== "0x0000000000000000000000000000000000000000000000000000000000000000";
+        console.log(`Deployed with ${hasDevice ? "1" : "0"} initial device and ${hasHash ? "1" : "0"} initial compose hash`);
+      }
     }
   });
 
@@ -303,7 +344,7 @@ task("kms:create-app", "Create AppAuth via KMS factory method (single transactio
     console.log("Using factory method for single-transaction deployment...");
 
     // Single transaction deployment via factory
-    const tx = await kmsAuth.deployApp(
+    const tx = await kmsAuth.deployAndRegisterApp(
       deployerAddress,  // deployer owns the contract
       false,           // disableUpgrades
       taskArgs.allowAnyDevice,
@@ -315,6 +356,7 @@ task("kms:create-app", "Create AppAuth via KMS factory method (single transactio
 
     // Parse events using contract interface
     let factoryEvent = null;
+    let registeredEvent = null;
 
     for (const log of receipt.logs) {
       try {
@@ -325,14 +367,16 @@ task("kms:create-app", "Create AppAuth via KMS factory method (single transactio
 
         if (parsedLog?.name === 'AppDeployedViaFactory') {
           factoryEvent = parsedLog.args;
+        } else if (parsedLog?.name === 'AppRegistered') {
+          registeredEvent = parsedLog.args;
         }
       } catch (e) {
         // Skip logs that can't be parsed by this contract
         continue;
       }
     }
 
-    if (factoryEvent) {
+    if (factoryEvent && registeredEvent) {
       console.log("✅ App deployed and registered successfully!");
       console.log("Proxy Address (App Id):", factoryEvent.appId);
       console.log("Owner:", factoryEvent.deployer);

kms/auth-eth/test/AppAuth.test.ts

@@ -87,24 +87,6 @@ describe("AppAuth", function () {
       expect(isAllowed).to.be.true;
     });
 
-    it("Should reject invalid app ID", async function () {
-      const bootInfo = {
-        tcbStatus: "UpToDate",
-        advisoryIds: [],
-        appId: ethers.Wallet.createRandom().address,
-        composeHash,
-        instanceId,
-        deviceId,
-        mrAggregated,
-        osImageHash,
-        mrSystem
-      };
-
-      const [isAllowed, reason] = await appAuth.isAppAllowed(bootInfo);
-      expect(isAllowed).to.be.false;
-      expect(reason).to.equal("Wrong app controller");
-    });
-
     it("Should reject unallowed compose hash", async function () {
       const bootInfo = {
         tcbStatus: "UpToDate",

kms/auth-eth/test/ethereum.integration.test.ts

@@ -101,9 +101,8 @@ describe('Integration Tests', () => {
           ...mockBootInfo,
           appId: ethers.Wallet.createRandom().address
         };
-
         const result = await backend.checkBoot(badBootInfo, false);
-        expect(result.reason).to.equal('App not deployed or invalid address');
+        expect(result.reason).to.equal('App not registered');
         expect(result.isAllowed).to.equal(false);
       });
     });

kms/auth-eth/test/ethereum.test.ts

@@ -72,9 +72,8 @@ describe('EthereumBackend', () => {
         ...mockBootInfo,
         appId: ethers.Wallet.createRandom().address
       };
-      
       const result = await backend.checkBoot(badBootInfo, false);
-      expect(result.reason).toBe('App not deployed or invalid address');
+      expect(result.reason).toBe('App not registered');
       expect(result.isAllowed).toBe(false);
     });
   });

kms/auth-eth/test/setup.ts

@@ -35,6 +35,7 @@ beforeAll(async () => {
   ], true) as AppAuth;
 
   const appId = await appAuth.getAddress();
+  await kmsAuth.registerApp(appId);
 
   // Set up KMS info with the generated app ID
   await kmsAuth.setKmsInfo({