Merge pull request #231 from Dstack-TEE/addr-as-appid

KmsAuth: Use the contract address as app-id

Author: kvinwang

docs/deployment.md

@@ -295,7 +295,7 @@ After the dstack-vmm is ready, you can deploy an app on it following the steps b
 The on-chain registration process includes two steps:
 
 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 and obtain the App Id
+2. Call KmsAuth.registerApp(appAuthAddress) to register the contract.
 
 The Dstack repository provides scripts to complete these two steps:
 
@@ -352,7 +352,6 @@ Call the hardhat command to add it to the whitelist (using AppAuth as an example
 
 ```bash
 export PRIVATE_KEY=<your eth private key here>
-export KMS_CONTRACT_ADDRESS=0xFE6C45aE66344CAEF5E5D7e2cbD476286D651875
 npx hardhat app:add-hash --network phala --app-id 0xA35b434eE853fdf9c2Bf48Fa1583Ac1332d50255 0x44d9cb98aaa6ab11f5729fc7d6fd58117585e0e3fbec621612dcee6b2dfbcde5
 ```
 

kms/auth-eth/contracts/AppAuth.sol

@@ -16,9 +16,6 @@ contract AppAuth is
     IAppAuth,
     IAppAuthBasicManagement
 {
-    // App ID this contract is managing
-    address public appId;
-
     // Mapping of allowed compose hashes for this app
     mapping(bytes32 => bool) public allowedComposeHashes;
 
@@ -43,31 +40,28 @@ contract AppAuth is
     // Initialize the contract
     function initialize(
         address initialOwner,
-        address _appId,
         bool _disableUpgrades,
         bool _allowAnyDevice,
         bytes32 initialDeviceId,
         bytes32 initialComposeHash
     ) public initializer {
         require(initialOwner != address(0), "invalid owner address");
-        require(_appId != address(0), "invalid app ID");
-        
-        appId = _appId;
+
         _upgradesDisabled = _disableUpgrades;
         allowAnyDevice = _allowAnyDevice;
-        
+
         // Add initial device if provided
         if (initialDeviceId != bytes32(0)) {
             allowedDeviceIds[initialDeviceId] = true;
             emit DeviceAdded(initialDeviceId);
         }
-        
+
         // Add initial compose hash if provided
         if (initialComposeHash != bytes32(0)) {
             allowedComposeHashes[initialComposeHash] = true;
             emit ComposeHashAdded(initialComposeHash);
         }
-        
+
         __Ownable_init(initialOwner);
         __UUPSUpgradeable_init();
         __ERC165_init();
@@ -131,11 +125,6 @@ contract AppAuth is
     function isAppAllowed(
         IAppAuth.AppBootInfo calldata bootInfo
     ) external view override returns (bool isAllowed, string memory reason) {
-        // Check if this controller is responsible for the app
-        if (bootInfo.appId != appId) {
-            return (false, "Wrong app controller");
-        }
-
         // Check if compose hash is allowed
         if (!allowedComposeHashes[bootInfo.composeHash]) {
             return (false, "Compose hash not allowed");

kms/auth-eth/contracts/KmsAuth.sol

@@ -30,14 +30,8 @@ contract KmsAuth is
     /// @custom:oz-renamed-from tproxyAppId
     string public gatewayAppId;
 
-    // Struct to store App configuration
-    struct AppConfig {
-        bool isRegistered;
-        address controller;
-    }
-
     // Mapping of registered apps
-    mapping(address => AppConfig) public apps;
+    mapping(address => bool) public registeredApps;
 
     // Mapping of allowed aggregated MR measurements for running KMS
     mapping(bytes32 => bool) public kmsAllowedAggregatedMrs;
@@ -48,9 +42,6 @@ contract KmsAuth is
     // Mapping of allowed image measurements
     mapping(bytes32 => bool) public allowedOsImages;
 
-    // Sequence number for app IDs - per user
-    mapping(address => uint256) public nextAppSequence;
-
     // AppAuth implementation contract address for factory deployment
     address public appAuthImplementation;
 
@@ -65,7 +56,7 @@ contract KmsAuth is
     event OsImageHashRemoved(bytes32 osImageHash);
     event GatewayAppIdSet(string gatewayAppId);
     event AppAuthImplementationSet(address implementation);
-    event AppDeployedViaFactory(address indexed appId, address indexed proxyAddress, address indexed deployer);
+    event AppDeployedViaFactory(address indexed appId, address indexed deployer);
 
     /// @custom:oz-upgrades-unsafe-allow constructor
     constructor() {
@@ -77,7 +68,7 @@ contract KmsAuth is
         __Ownable_init(initialOwner);
         __UUPSUpgradeable_init();
         __ERC165_init();
-        
+
         // Set AppAuth implementation if provided
         if (_appAuthImplementation != address(0)) {
             appAuthImplementation = _appAuthImplementation;
@@ -130,32 +121,11 @@ contract KmsAuth is
         emit GatewayAppIdSet(appId);
     }
 
-    // View next app id
-    function nextAppId() public view returns (address appId) {
-        bytes32 fullHash = keccak256(
-            abi.encodePacked(
-                address(this),
-                msg.sender,
-                nextAppSequence[msg.sender]
-            )
-        );
-        return address(uint160(uint256(fullHash)));
-    }
-
-    // Internal function to register an app with the given app ID and controller
-    function _registerAppInternal(address appId, address controller) private {
-        require(!apps[appId].isRegistered, "App already registered");
-        apps[appId].isRegistered = true;
-        apps[appId].controller = controller;
-        nextAppSequence[msg.sender]++;
-        emit AppRegistered(appId);
-    }
-
     // Function to register an app
-    function registerApp(address controller) external {
-        require(controller != address(0), "Invalid controller address");
-        address appId = nextAppId();
-        _registerAppInternal(appId, controller);
+    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
@@ -172,30 +142,25 @@ contract KmsAuth is
         bool allowAnyDevice,
         bytes32 initialDeviceId,
         bytes32 initialComposeHash
-    ) external returns (address appId, address proxyAddress) {
+    ) external returns (address appId) {
         require(appAuthImplementation != address(0), "AppAuth implementation not set");
         require(initialOwner != address(0), "Invalid owner address");
-        
-        // Calculate app ID
-        appId = nextAppId();
-        
+
         // Prepare initialization data
         bytes memory initData = abi.encodeWithSelector(
-            bytes4(keccak256("initialize(address,address,bool,bool,bytes32,bytes32)")),
+            bytes4(keccak256("initialize(address,bool,bool,bytes32,bytes32)")),
             initialOwner,
-            appId,
             disableUpgrades,
             allowAnyDevice,
             initialDeviceId,
             initialComposeHash
         );
-        
+
         // Deploy proxy contract
-        proxyAddress = address(new ERC1967Proxy(appAuthImplementation, initData));
-        
+        appId = address(new ERC1967Proxy(appAuthImplementation, initData));
         // Register to KMS
-        _registerAppInternal(appId, proxyAddress);
-        emit AppDeployedViaFactory(appId, proxyAddress, msg.sender);
+        registerApp(appId);
+        emit AppDeployedViaFactory(appId, msg.sender);
     }
 
     // Function to register an aggregated MR measurement
@@ -269,23 +234,31 @@ contract KmsAuth is
         AppBootInfo calldata bootInfo
     ) external view override returns (bool isAllowed, string memory reason) {
         // Check if app is registered
-        if (!apps[bootInfo.appId].isRegistered) {
+        if (!registeredApps[bootInfo.appId]) {
             return (false, "App not registered");
         }
-
-        // Check aggregated MR and image measurements
+        // Check if the OS image is allowed
         if (!allowedOsImages[bootInfo.osImageHash]) {
             return (false, "OS image is not allowed");
         }
 
-        // Ask the app controller if the app is allowed to boot
-        address controller = apps[bootInfo.appId].controller;
-        if (controller == address(0)) {
-            return (false, "App controller not set");
+        // Check if the contract exists at the appId address
+        if (!isContract(bootInfo.appId)) {
+            return (false, "App not deployed or invalid address");
         }
-        return IAppAuth(controller).isAppAllowed(bootInfo);
+
+        // Call the app's isAppAllowed function
+        return IAppAuth(bootInfo.appId).isAppAllowed(bootInfo);
     }
 
     // Add storage gap for upgradeable contracts
     uint256[50] private __gap;
 }
+
+function isContract(address addr) view returns (bool){
+    uint32 size;
+    assembly {
+        size := extcodesize(addr)
+    }
+    return (size > 0);
+}

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

@@ -259,10 +259,6 @@ cast call $KMS_CONTRACT_ADDRESS "isAppAllowed((address,bytes32,address,bytes32,b
 ### Query Operations
 
 ```bash
-# Get app ID
-cast call $APP_AUTH_ADDRESS "appId()" --rpc-url $RPC_URL
-# To decode: cast abi-decode "appId()(address)" RETURN_DATA
-
 # Get owner
 cast call $APP_AUTH_ADDRESS "owner()" --rpc-url $RPC_URL
 # To decode: cast abi-decode "owner()(address)" RETURN_DATA

kms/auth-eth/hardhat.config.ts

@@ -76,13 +76,7 @@ async function getKmsAuth(ethers: any) {
 }
 
 async function getAppAuth(ethers: any, appId: string) {
-  const kmsAuth = await getKmsAuth(ethers);
-  const controller = (await kmsAuth.apps(appId)).controller;
-  if (controller === ethers.ZeroAddress) {
-    throw new Error("AppAuth contract not found");
-  }
-  console.log("AppAuth address:", controller);
-  return await ethers.getContractAt("AppAuth", controller);
+  return await ethers.getContractAt("AppAuth", appId);
 }
 
 // KMS Contract Tasks
@@ -240,14 +234,6 @@ task("kms:get-app-implementation", "Get current AppAuth implementation address")
     console.log("AppAuth implementation:", impl);
   });
 
-task("app:show-controller", "Show the controller of an AppAuth contract")
-  .addPositionalParam("appId", "App ID")
-  .setAction(async ({ appId }, { ethers }) => {
-    const kmsAuth = await getKmsAuth(ethers);
-    const controller = await kmsAuth.apps(appId).then((app: any) => app.controller);
-    console.log("AppAuth controller:", controller);
-  });
-
 task("app:deploy", "Deploy AppAuth with a UUPS proxy")
   .addFlag("allowAnyDevice", "Allow any device to boot this app")
   .addOptionalParam("device", "Initial device ID", "", types.string)
@@ -260,8 +246,6 @@ task("app:deploy", "Deploy AppAuth with a UUPS proxy")
     console.log("Account balance:", await accountBalance(ethers, deployerAddress));
 
     const kmsContract = await getKmsAuth(ethers);
-    const appId = await kmsContract.nextAppId();
-    console.log("App ID:", appId);
 
     // Parse device and hash (convert to bytes32, use 0x0 if empty)
     const deviceId = taskArgs.device ? taskArgs.device.trim() : "0x0000000000000000000000000000000000000000000000000000000000000000";
@@ -278,7 +262,6 @@ task("app:deploy", "Deploy AppAuth with a UUPS proxy")
     // Use standard deployment - all cases use the same 6-parameter initializer
     const appAuth = await deployContract(hre, "AppAuth", [
       deployerAddress, 
-      appId, 
       false, 
       taskArgs.allowAnyDevice,
       deviceId,
@@ -395,8 +378,7 @@ task("kms:create-app", "Create AppAuth via KMS factory method (single transactio
 
     if (factoryEvent && registeredEvent) {
       console.log("✅ App deployed and registered successfully!");
-      console.log("App ID:", factoryEvent.appId);
-      console.log("Proxy Address:", factoryEvent.proxyAddress);
+      console.log("Proxy Address (App Id):", factoryEvent.appId);
       console.log("Owner:", factoryEvent.deployer);
       console.log("Transaction hash:", tx.hash);
 

kms/auth-eth/src/server.ts

@@ -99,7 +99,9 @@ export async function build(): Promise<FastifyInstance> {
     try {
       return await server.ethereum.checkBoot(request.body, true);
     } catch (error) {
-      console.error(error);
+      if (!(error instanceof Error && "Test backend error" == error.message)) {
+        console.error(error);
+      }
       reply.code(200).send({
         isAllowed: false,
         gatewayAppId: '',