Skip to content

Commit

Permalink
migrate tee logs
Browse files Browse the repository at this point in the history
  • Loading branch information
lalalune committed Feb 11, 2025
1 parent 920f1a0 commit 0ad1264
Show file tree
Hide file tree
Showing 23 changed files with 2,120 additions and 0 deletions.
Binary file modified bun.lockb
Binary file not shown.
99 changes: 99 additions & 0 deletions packages/agent/src/api.ts
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,10 @@ import {
type Character,
getEnvVariable,
IAgentRuntime,
ITeeLogService,
logger,
ServiceType,
TeeLogQuery,
type UUID,
validateCharacterConfig,
validateUuid
Expand Down Expand Up @@ -352,5 +355,101 @@ export function createApiRouter(
}
});

router.get("/tee/agents", async (req, res) => {
try {
const allAgents = [];

for (const agentRuntime of agents.values()) {
const teeLogService = agentRuntime
.getService<ITeeLogService>(ServiceType.TEE)
.getInstance();

const agents = await teeLogService.getAllAgents();
allAgents.push(...agents);
}

const runtime: IAgentRuntime = agents.values().next().value;
const teeLogService = runtime
.getService<ITeeLogService>(ServiceType.TEE)
.getInstance() as ITeeLogService;
const attestation = await teeLogService.generateAttestation(
JSON.stringify(allAgents)
);
res.json({ agents: allAgents, attestation: attestation });
} catch (error) {
logger.error("Failed to get TEE agents:", error);
res.status(500).json({
error: "Failed to get TEE agents",
});
}
});

router.get("/tee/agents/:agentId", async (req, res) => {
try {
const agentId = req.params.agentId;
const agentRuntime = agents.get(agentId);
if (!agentRuntime) {
res.status(404).json({ error: "Agent not found" });
return;
}

const teeLogService = agentRuntime
.getService<ITeeLogService>(ServiceType.TEE)
.getInstance();

const teeAgent = await teeLogService.getAgent(agentId);
const attestation = await teeLogService.generateAttestation(
JSON.stringify(teeAgent)
);
res.json({ agent: teeAgent, attestation: attestation });
} catch (error) {
logger.error("Failed to get TEE agent:", error);
res.status(500).json({
error: "Failed to get TEE agent",
});
}
});

router.post(
"/tee/logs",
async (req: express.Request, res: express.Response) => {
try {
const query = req.body.query || {};
const page = Number.parseInt(req.body.page) || 1;
const pageSize = Number.parseInt(req.body.pageSize) || 10;

const teeLogQuery: TeeLogQuery = {
agentId: query.agentId || "",
roomId: query.roomId || "",
userId: query.userId || "",
type: query.type || "",
containsContent: query.containsContent || "",
startTimestamp: query.startTimestamp || undefined,
endTimestamp: query.endTimestamp || undefined,
};
const agentRuntime: IAgentRuntime = agents.values().next().value;
const teeLogService = agentRuntime
.getService<ITeeLogService>(ServiceType.TEE);
const pageQuery = await teeLogService.getLogs(
teeLogQuery,
page,
pageSize
);
const attestation = await teeLogService.generateAttestation(
JSON.stringify(pageQuery)
);
res.json({
logs: pageQuery,
attestation: attestation,
});
} catch (error) {
logger.error("Failed to get TEE logs:", error);
res.status(500).json({
error: "Failed to get TEE logs",
});
}
}
);

return router;
}
118 changes: 118 additions & 0 deletions packages/core/src/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -141,6 +141,7 @@ export enum ServiceType {
REMOTE_FILES = "aws_s3",
WEB_SEARCH = "web_search",
EMAIL = "email",
TEE = "tee",
}

/**
Expand Down Expand Up @@ -1105,6 +1106,23 @@ export interface IFileService extends Service {
generateSignedUrl(fileName: string, expiresIn: number): Promise<string>;
}

export interface ITeeLogService extends Service {
getInstance(): ITeeLogService;
log(
agentId: string,
roomId: string,
userId: string,
type: string,
content: string,
): Promise<boolean>;

generateAttestation<T>(reportData: string, hashAlgorithm?: T | any): Promise<string>;
getAllAgents(): Promise<TeeAgent[]>;
getAgent(agentId: string): Promise<TeeAgent | null>;
getLogs(query: TeeLogQuery, page: number, pageSize: number): Promise<TeePageQuery<TeeLog[]>>;

}

export interface TestCase {
name: string;
fn: (runtime: IAgentRuntime) => Promise<void> | void;
Expand All @@ -1113,4 +1131,104 @@ export interface TestCase {
export interface TestSuite {
name: string;
tests: TestCase[];
}

// Represents a log entry in the TeeLog table, containing details about agent activities.
export interface TeeLog {
id: string;
agentId: string;
roomId: string;
userId: string;
type: string;
content: string;
timestamp: number;
signature: string;
}

export interface TeeLogQuery {
agentId?: string;
roomId?: string;
userId?: string;
type?: string;
containsContent?: string;
startTimestamp?: number;
endTimestamp?: number;
}

// Represents an agent in the TeeAgent table, containing details about the agent.
export interface TeeAgent {
id: string; // Primary key
// Allow duplicate agentId.
// This is to support the case where the same agentId is registered multiple times.
// Each time the agent restarts, we will generate a new keypair and attestation.
agentId: string;
agentName: string;
createdAt: number;
publicKey: string;
attestation: string;
}

export interface TeePageQuery<Result = any> {
page: number;
pageSize: number;
total?: number;
data?: Result;
}

export abstract class TeeLogDAO<DB = any> {
db: DB;

abstract initialize(): Promise<void>;

abstract addLog(log: TeeLog): Promise<boolean>;

abstract getPagedLogs(
query: TeeLogQuery,
page: number,
pageSize: number
): Promise<TeePageQuery<TeeLog[]>>;

abstract addAgent(agent: TeeAgent): Promise<boolean>;

abstract getAgent(agentId: string): Promise<TeeAgent>;

abstract getAllAgents(): Promise<TeeAgent[]>;
}

export enum TEEMode {
OFF = "OFF",
LOCAL = "LOCAL", // For local development with simulator
DOCKER = "DOCKER", // For docker development with simulator
PRODUCTION = "PRODUCTION", // For production without simulator
}

export interface RemoteAttestationQuote {
quote: string;
timestamp: number;
}

export interface DeriveKeyAttestationData {
agentId: string;
publicKey: string;
subject?: string;
}

export interface RemoteAttestationMessage {
agentId: string;
timestamp: number;
message: {
userId: string;
roomId: string;
content: string;
}
}

export interface SgxAttestation {
quote: string;
timestamp: number;
}

export enum TeeType {
SGX_GRAMINE = "sgx_gramine",
TDX_DSTACK = "tdx_dstack",
}
6 changes: 6 additions & 0 deletions packages/plugin-tee/.npmignore
Original file line number Diff line number Diff line change
@@ -0,0 +1,6 @@
*

!dist/**
!package.json
!readme.md
!tsup.config.ts
136 changes: 136 additions & 0 deletions packages/plugin-tee/README.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,136 @@
# TEE Log Plugin for Eliza

The TEE Log Plugin for Eliza is designed to enhance the logging capabilities of the Eliza by providing a structured way to generate, store and verify TEE (Trusted Execution Environment) logs for agents. This plugin ensures that all sensitive interactions are securely logged, providing a transparent and tamper-resistant record of all sensitive activities.

## Background

As Eliza is a fully autonomous AI agent capable of running within a TEE, we need to demonstrate to the outside world that we are indeed operating within a TEE. This allows external parties to verify that our actions are protected by the TEE and that they are autonomously executed by Eliza, without any third-party interference. Therefore, it is necessary to leverage TEE's remote attestation and establish a TEE logging mechanism to prove that these operations are entirely and autonomously performed by Eliza within the TEE.

## Requirements

Since the TEE Logging is based on the TEE, it is necessary to have a TEE enabled environment. Currently, we support Intel SGX (Gramine) and Intel TDX (dstack).
- using Intel SGX (Gramine), you need to enable the plugin-sgx in the Eliza runtime, which is enabled in SGX env automatically.
- using Intel TDX (dstack), you need to enable the plugin-tee in the Eliza runtime.

## TEE Logging Mechanism

## TEE Logging Mechanism

1. **Key Pair Generation and Attestation**:
- During startup, each agent generates a key pair and creates a remote attestation for the public key. The private key is securely stored in the TEE's encrypted memory. The agent's relevant information, along with the public key and attestation, is recorded in a local database. A new key pair is generated each time the agent is updated or restarted to ensure key security.

2. **Log Recording**:
- For each log entry, basic information is recorded, including `agentId`, `roomId`, `userId`, `type`, `content`, and `timestamp`. This information is concatenated and signed using the agent's corresponding private key to ensure verifiability. The verification process follows this trust chain:
- Verify the attestation.
- Trust the public key contained in the attestation.
- Use the public key to verify the signature.
- Trust the complete log record.

3. **Data Storage**:
- All log data must be stored in the TEE's encrypted file system in production environments. Storing data in plaintext is prohibited to prevent tampering.

4. **Log Extraction for Verification**:
- Third parties can extract TEE logs for verification purposes. Two types of information can be extracted:
- **Agent Information**: This includes the agent's metadata, public key, and attestation, which can be used to verify the agent's public key.
- **Log Information**: Required logs can be extracted, with the agent's attestation and public key used to verify the signature, ensuring that each record remains untampered.

5. **Integrity Protection**:
- When users extract TEE logs via the REST API, the results are hashed, and an attestation is generated. After extraction, users can verify the attestation by comparing the hash value contained within it to the extracted results, thereby ensuring the integrity of the data.

## Services

- **[TeeLogService]**: This service is responsible for generating and storing TEE logs for agents.

### Class: TeeLogService

The `TeeLogService` class implements the `ITeeLogService` interface and extends the `Service` class. It manages the logging of sensitive interactions within a Trusted Execution Environment (TEE).

#### Methods

- **getInstance()**: `TeeLogService`
- Returns the singleton instance of the `TeeLogService`.

- **static get serviceType()**: `ServiceType`
- Returns the service type for TEE logging.

- **async initialize(runtime: IAgentRuntime): Promise<void>**
- Initializes the TEE log service. It checks the runtime settings to configure the TEE type and enables logging if configured.

- **async log(agentId: string, roomId: string, userId: string, type: string, content: string): Promise<boolean>**
- Logs an interaction with the specified parameters. Returns `false` if TEE logging is not enabled.

- **async getAllAgents(): Promise<TeeAgent[]>**
- Retrieves all agents that have been logged. Returns an empty array if TEE logging is not enabled.

- **async getAgent(agentId: string): Promise<TeeAgent | undefined>**
- Retrieves the details of a specific agent by their ID. Returns `undefined` if TEE logging is not enabled.

- **async getLogs(query: TeeLogQuery, page: number, pageSize: number): Promise<PageQuery<TeeLog[]>>**
- Retrieves logs based on the provided query parameters. Returns an empty result if TEE logging is not enabled.

- **async generateAttestation(userReport: string): Promise<string>**
- Generates an attestation based on the provided user report.

### Storage

The TEE logs are stored in a SQLite database, which is located at `./data/tee_log.sqlite`. The database is automatically created when the service is initialized.

Important: You need to use the encrypted file system to store the database file in production, otherwise the database will be compromised. Since TEE only protects memory-in-use, the disk is not protected by the TEE. However, Many TEE development tools support the encrypted file system, for example, you can refer to the [Gramine Encrypted files](https://gramine.readthedocs.io/en/latest/manifest-syntax.html#encrypted-files) documentation for more information.

### Usage

To use the `TeeLogService`, ensure that the TEE environment is properly configured and initialized.

Enable the TEE logging in the Eliza .env file:

```env
TEE_LOG_ENABLED=true
```

The logging isn't integrated for actions by default, you need to integrate the logging for the actions you want to log. For example, if you want to log the `Continue` action of plugin-bootstrap, you can do the following:

First, add plugin-tee-log to the dependencies of plugin-bootstrap:

```json
"@elizaos/plugin-tee-log": "workspace:*",
```

Then, add the following code to the `Continue` action:

```typescript
import {
ServiceType,
ITeeLogService,
} from "@elizaos/core";


// In the handler of the action
handler: async (
runtime: IAgentRuntime,
message: Memory,
state: State,
options: any,
callback: HandlerCallback
) => {
// Continue the action

// Log the action
const teeLogService = runtime
.getService<ITeeLogService>(ServiceType.TEE_LOG)
.getInstance();
if (teeLogService.log(
runtime.agentId,
message.roomId,
message.userId,
"The type of the log, for example, Action:CONTINUE",
"The content that you want to log"
)
) {
console.log("Logged TEE log successfully");
}

// Continue the action
}
```

After configuring the logging for the action, you can run the Eliza and see the logs through the client-direct REST API. See more details in the [Client-Direct REST API](../client-direct/src/README.md) documentation.
Loading

0 comments on commit 0ad1264

Please sign in to comment.