Merge pull request #160 from Dstack-TEE/api-extend-rtmr3

Author: h4x3rotab

guest-agent/rpc/proto/agent_rpc.proto

@@ -39,6 +39,9 @@ service DstackGuest {
   // Generates a TDX quote with given report data.
   rpc GetQuote(RawQuoteArgs) returns (GetQuoteResponse) {}
 
+  // Emit an event. This extends the event to RTMR3 on TDX platform.
+  rpc EmitEvent(EmitEventArgs) returns (google.protobuf.Empty) {}
+
   // Get worker info
   rpc Info(google.protobuf.Empty) returns (WorkerInfo) {}
 }
@@ -161,6 +164,13 @@ message GetQuoteResponse {
   bytes report_data = 3;
 }
 
+message EmitEventArgs {
+  // The event name
+  string event = 1;
+  // The event data
+  bytes payload = 2;
+}
+
 // The request to derive a key
 message WorkerInfo {
   // App ID

guest-agent/src/rpc_service.rs

@@ -6,9 +6,9 @@ use dstack_guest_agent_rpc::{
     dstack_guest_server::{DstackGuestRpc, DstackGuestServer},
     tappd_server::{TappdRpc, TappdServer},
     worker_server::{WorkerRpc, WorkerServer},
-    DeriveK256KeyResponse, DeriveKeyArgs, GetKeyArgs, GetKeyResponse, GetQuoteResponse,
-    GetTlsKeyArgs, GetTlsKeyResponse, RawQuoteArgs, TdxQuoteArgs, TdxQuoteResponse, WorkerInfo,
-    WorkerVersion,
+    DeriveK256KeyResponse, DeriveKeyArgs, EmitEventArgs, GetKeyArgs, GetKeyResponse,
+    GetQuoteResponse, GetTlsKeyArgs, GetTlsKeyResponse, RawQuoteArgs, TdxQuoteArgs,
+    TdxQuoteResponse, WorkerInfo, WorkerVersion,
 };
 use dstack_types::AppKeys;
 use fs_err as fs;
@@ -161,6 +161,13 @@ impl DstackGuestRpc for InternalRpcHandler {
         })
     }
 
+    async fn emit_event(self, request: EmitEventArgs) -> Result<()> {
+        if self.state.config().simulator.enabled {
+            return Ok(());
+        }
+        tdx_attest::extend_rtmr3(&request.event, &request.payload)
+    }
+
     async fn info(self) -> Result<WorkerInfo> {
         ExternalRpcHandler { state: self.state }.info().await
     }

sdk/curl/api.md

@@ -165,6 +165,33 @@ curl --unix-socket /var/run/dstack.sock http://dstack/Info
 }
 ```
 
+### 5. Emit Event
+
+Emit an event to be extended to RTMR3 on TDX platform. This API requires Dstack OS 0.5.0 or later.
+
+**Endpoint:** `/EmitEvent`
+
+**Request Parameters:**
+
+| Field | Type | Description | Example |
+|-------|------|-------------|----------|
+| `event` | string | The event name | `"custom-event"` |
+| `payload` | string | Hex-encoded payload data | `"deadbeef"` |
+
+**Example:**
+```bash
+curl --unix-socket /var/run/dstack.sock -X POST \
+  http://dstack/EmitEvent \
+  -H 'Content-Type: application/json' \
+  -d '{
+    "event": "custom-event",
+    "payload": "deadbeef"
+  }'
+```
+
+**Response:**
+Empty response with HTTP 200 status code on success.
+
 ## Error Responses
 
 All endpoints may return the following HTTP status codes:

sdk/go/dstack/client.go

@@ -417,3 +417,15 @@ func (c *DstackClient) Info(ctx context.Context) (*InfoResponse, error) {
 
 	return &response, nil
 }
+
+// EmitEvent sends an event to be extended to RTMR3 on TDX platform.
+// The event will be extended to RTMR3 with the provided name and payload.
+//
+// Requires Dstack OS 0.5.0 or later.
+func (c *DstackClient) EmitEvent(ctx context.Context, event string, payload []byte) error {
+	_, err := c.sendRPCRequest(ctx, "/EmitEvent", map[string]interface{}{
+		"event":   event,
+		"payload": hex.EncodeToString(payload),
+	})
+	return err
+}

sdk/js/src/index.ts

@@ -314,4 +314,28 @@ export class DstackClient {
       tcb_info: JSON.parse(result.tcb_info) as TcbInfo,
     })
   }
+
+  /**
+   * Emit an event. This extends the event to RTMR3 on TDX platform.
+   *
+   * Requires Dstack OS 0.5.0 or later.
+   *
+   * @param event The event name
+   * @param payload The event data as string or Buffer or Uint8Array
+   */
+  async emitEvent(event: string, payload: string | Buffer | Uint8Array): Promise<void> {
+    if (!event) {
+      throw new Error('Event name cannot be empty')
+    }
+
+    const hexPayload = to_hex(payload)
+    await send_rpc_request(
+      this.endpoint,
+      '/EmitEvent',
+      JSON.stringify({
+        event: event,
+        payload: hexPayload
+      })
+    )
+  }
 }

sdk/python/src/dstack_sdk/dstack_client.py

@@ -203,6 +203,33 @@ def info(self) -> InfoResponse:
         result = self._send_rpc_request("/Info", {})
         return InfoResponse.model_validate(result)
 
+    def emit_event(
+        self,
+        event: str,
+        payload: str | bytes,
+    ) -> None:
+        """
+        Emit an event. This extends the event to RTMR3 on TDX platform.
+
+        Requires Dstack OS 0.5.0 or later.
+
+        Args:
+            event: The event name
+            payload: The event data as string or bytes
+
+        Returns:
+            None
+        """
+        if not event:
+            raise ValueError("event name cannot be empty")
+
+        if isinstance(payload, str):
+            payload = payload.encode()
+
+        hex_payload = binascii.hexlify(payload).decode()
+        self._send_rpc_request("/EmitEvent", {"event": event, "payload": hex_payload})
+        return None
+
     def get_tls_key(
         self,
         subject: str | None = None,
@@ -285,6 +312,33 @@ async def info(self) -> InfoResponse:
         result = await self._send_rpc_request("/Info", {})
         return InfoResponse.model_validate(result)
 
+    async def emit_event(
+        self,
+        event: str,
+        payload: str | bytes,
+    ) -> None:
+        """
+        Emit an event. This extends the event to RTMR3 on TDX platform.
+
+        Requires Dstack OS 0.5.0 or later.
+
+        Args:
+            event: The event name
+            payload: The event data as string or bytes
+
+        Returns:
+            None
+        """
+        if not event:
+            raise ValueError("event name cannot be empty")
+
+        if isinstance(payload, str):
+            payload = payload.encode()
+
+        hex_payload = binascii.hexlify(payload).decode()
+        await self._send_rpc_request("/EmitEvent", {"event": event, "payload": hex_payload})
+        return None
+
     async def get_tls_key(
         self,
         subject: str | None = None,

tdx-attest/src/lib.rs

@@ -20,3 +20,14 @@ pub type TdxReportData = [u8; 64];
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub struct TdxReport(pub [u8; 1024]);
+
+pub fn extend_rtmr3(event: &str, payload: &[u8]) -> anyhow::Result<()> {
+    use anyhow::Context;
+    // This code is not defined in the TCG specification.
+    // See https://trustedcomputinggroup.org/wp-content/uploads/PC-ClientSpecific_Platform_Profile_for_TPM_2p0_Systems_v51.pdf
+    let event_type = 0x08000001;
+    let index = 3;
+    let log = eventlog::TdxEventLog::new(index, event_type, event.to_string(), payload.to_vec());
+    extend_rtmr(index, event_type, log.digest).context("Failed to extend RTMR")?;
+    log_rtmr_event(&log).context("Failed to log RTMR event")
+}