Merge pull request #317 from Dstack-TEE/imp-sdk-tcb-info

IMP: Typing & schema for TcbInfo in SDK

Author: Leechael

sdk/js/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@phala/dstack-sdk",
-  "version": "0.5.4",
+  "version": "0.5.5",
   "description": "dstack SDK",
   "main": "dist/node/index.js",
   "types": "dist/node/index.d.ts",

sdk/js/src/index.ts

@@ -44,14 +44,26 @@ export interface TcbInfo {
   rtmr1: string
   rtmr2: string
   rtmr3: string
+  app_compose: string
   event_log: EventLog[]
 }
 
-export interface InfoResponse {
+export type TcbInfoV03x = TcbInfo & {
+  rootfs_hash: string
+}
+
+export type TcbInfoV05x = TcbInfo & {
+  mr_aggregated: string
+  os_image_hash: string
+  compose_hash: string
+  device_id: string
+}
+
+export interface InfoResponse<VersionTcbInfo extends TcbInfo> {
   app_id: string
   instance_id: string
   app_cert: string
-  tcb_info: TcbInfo
+  tcb_info: VersionTcbInfo
   app_name: string
   device_id: string
   os_image_hash?: string // Optional: empty if OS image is not measured by KMS
@@ -132,7 +144,7 @@ export interface TlsKeyOptions {
   usageClientAuth?: boolean;
 }
 
-export class DstackClient {
+export class DstackClient<T extends TcbInfo = TcbInfoV05x> {
   protected endpoint: string
 
   constructor(endpoint: string | undefined = undefined) {
@@ -210,11 +222,11 @@ export class DstackClient {
     return Object.freeze(result)
   }
 
-  async info(): Promise<InfoResponse> {
-    const result = await send_rpc_request<Omit<InfoResponse, 'tcb_info'> & { tcb_info: string }>(this.endpoint, '/Info', '{}')
+  async info(): Promise<InfoResponse<T>> {
+    const result = await send_rpc_request<Omit<InfoResponse<TcbInfo>, 'tcb_info'> & { tcb_info: string }>(this.endpoint, '/Info', '{}')
     return Object.freeze({
       ...result,
-      tcb_info: JSON.parse(result.tcb_info) as TcbInfo,
+      tcb_info: JSON.parse(result.tcb_info) as T,
     })
   }
 
@@ -283,7 +295,7 @@ export class DstackClient {
   }
 }
 
-export class TappdClient extends DstackClient {
+export class TappdClient extends DstackClient<TcbInfoV03x> {
   constructor(endpoint: string | undefined = undefined) {
     if (endpoint === undefined) {
       if (process.env.TAPPD_SIMULATOR_ENDPOINT) {

sdk/python/pyproject.toml

@@ -4,7 +4,7 @@
 
 [project]
 name = "dstack-sdk"
-version = "0.5.0"
+version = "0.5.1"
 description = "dstack SDK for Python"
 authors = [
     {name = "Leechael Yim", email = "[email protected]"},

sdk/python/src/dstack_sdk/dstack_client.py

@@ -11,8 +11,10 @@
 import os
 from typing import Any
 from typing import Dict
+from typing import Generic
 from typing import List
 from typing import Optional
+from typing import TypeVar
 from typing import cast
 import warnings
 
@@ -157,38 +159,63 @@ class EventLog(BaseModel):
 
 
 class TcbInfo(BaseModel):
+    """Base TCB (Trusted Computing Base) information structure."""
+
     mrtd: str
     rtmr0: str
     rtmr1: str
     rtmr2: str
     rtmr3: str
-    os_image_hash: str = ""
-    compose_hash: str
-    device_id: str
     app_compose: str
     event_log: List[EventLog]
 
 
-class InfoResponse(BaseModel):
+class TcbInfoV03x(TcbInfo):
+    """TCB information for dstack OS version 0.3.x."""
+
+    rootfs_hash: str
+
+
+class TcbInfoV05x(TcbInfo):
+    """TCB information for dstack OS version 0.5.x."""
+
+    mr_aggregated: str
+    os_image_hash: str
+    compose_hash: str
+    device_id: str
+
+
+# Type variable for TCB info versions
+T = TypeVar("T", bound=TcbInfo)
+
+
+class InfoResponse(BaseModel, Generic[T]):
     app_id: str
     instance_id: str
     app_cert: str
-    tcb_info: TcbInfo
+    tcb_info: T
     app_name: str
     device_id: str
     os_image_hash: str = ""
     key_provider_info: str
     compose_hash: str
 
     @classmethod
-    def parse_response(cls, obj: Any) -> "InfoResponse":
+    def parse_response(cls, obj: Any, tcb_info_type: type[T]) -> "InfoResponse[T]":
+        """Parse response from service, automatically deserializing tcb_info.
+
+        Args:
+            obj: Raw response object from service
+            tcb_info_type: The specific TcbInfo subclass to use for parsing
+
+        """
         if (
             isinstance(obj, dict)
             and "tcb_info" in obj
             and isinstance(obj["tcb_info"], str)
         ):
             obj = dict(obj)
-            obj["tcb_info"] = TcbInfo(**json.loads(obj["tcb_info"]))
+            obj["tcb_info"] = tcb_info_type(**json.loads(obj["tcb_info"]))
         return cls(**obj)
 
 
@@ -311,10 +338,10 @@ async def get_quote(
         result = await self._send_rpc_request("GetQuote", {"report_data": hex})
         return GetQuoteResponse(**result)
 
-    async def info(self) -> InfoResponse:
+    async def info(self) -> InfoResponse[TcbInfo]:
         """Fetch service information including parsed TCB info."""
         result = await self._send_rpc_request("Info", {})
-        return InfoResponse.parse_response(result)
+        return InfoResponse.parse_response(result, TcbInfoV05x)
 
     async def emit_event(
         self,
@@ -391,7 +418,7 @@ def get_quote(
         raise NotImplementedError
 
     @call_async
-    def info(self) -> InfoResponse:
+    def info(self) -> InfoResponse[TcbInfo]:
         """Fetch service information including parsed TCB info."""
         raise NotImplementedError
 
@@ -503,6 +530,11 @@ async def tdx_quote(
 
         return GetQuoteResponse(**result)
 
+    async def info(self) -> InfoResponse[TcbInfo]:
+        """Fetch service information including parsed TCB info."""
+        result = await self._send_rpc_request("Info", {})
+        return InfoResponse.parse_response(result, TcbInfoV03x)
+
 
 class TappdClient(DstackClient):
     """Deprecated client kept for backward compatibility.
@@ -537,6 +569,11 @@ def tdx_quote(
         """Use ``get_quote`` instead (deprecated)."""
         raise NotImplementedError
 
+    @call_async
+    def info(self) -> InfoResponse[TcbInfo]:
+        """Fetch service information including parsed TCB info."""
+        raise NotImplementedError
+
     @call_async
     def __enter__(self):
         raise NotImplementedError