From 837d258ba0292b7b7d13bb8f0ec515712cc5a27d Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 15 Jan 2025 20:14:30 +0000 Subject: [PATCH 1/2] Implement RSA-based encrypted prompts for TEEception - Add RSA key pair generation in agent setup - Update contract interface for encrypted prompts - Add RSA public key storage to agent contract - Replace system_prompt with system_prompt_uri - Add comprehensive test coverage - Maintain direct payment model Co-Authored-By: franco@nethermind.io --- contracts/Scarb.lock | 50 ++++++++--------- contracts/Scarb.toml | 2 +- contracts/src/lib.cairo | 41 ++++++++++---- contracts/tests/test_contract.cairo | 83 +++++++++++++++++++++++++---- pkg/agent/setup/manager.go | 23 ++++++++ 5 files changed, 154 insertions(+), 45 deletions(-) diff --git a/contracts/Scarb.lock b/contracts/Scarb.lock index 17a40a8..13fc2e2 100644 --- a/contracts/Scarb.lock +++ b/contracts/Scarb.lock @@ -3,8 +3,8 @@ version = 1 [[package]] name = "openzeppelin" -version = "0.19.0" -source = "git+https://github.com/OpenZeppelin/cairo-contracts.git?tag=v0.19.0#8d49e8c445efd9bdc99b050c8b7d11ae5ad19628" +version = "0.20.0" +source = "git+https://github.com/OpenZeppelin/cairo-contracts.git?tag=v0.20.0#7756fd1de2b4ebd239fa6e372d75535cea02e5e5" dependencies = [ "openzeppelin_access", "openzeppelin_account", @@ -21,17 +21,16 @@ dependencies = [ [[package]] name = "openzeppelin_access" -version = "0.19.0" -source = "git+https://github.com/OpenZeppelin/cairo-contracts.git?tag=v0.19.0#8d49e8c445efd9bdc99b050c8b7d11ae5ad19628" +version = "0.20.0" +source = "git+https://github.com/OpenZeppelin/cairo-contracts.git?tag=v0.20.0#7756fd1de2b4ebd239fa6e372d75535cea02e5e5" dependencies = [ "openzeppelin_introspection", - "openzeppelin_utils", ] [[package]] name = "openzeppelin_account" -version = "0.19.0" -source = "git+https://github.com/OpenZeppelin/cairo-contracts.git?tag=v0.19.0#8d49e8c445efd9bdc99b050c8b7d11ae5ad19628" +version = "0.20.0" +source = "git+https://github.com/OpenZeppelin/cairo-contracts.git?tag=v0.20.0#7756fd1de2b4ebd239fa6e372d75535cea02e5e5" dependencies = [ "openzeppelin_introspection", "openzeppelin_utils", @@ -39,8 +38,8 @@ dependencies = [ [[package]] name = "openzeppelin_finance" -version = "0.19.0" -source = "git+https://github.com/OpenZeppelin/cairo-contracts.git?tag=v0.19.0#8d49e8c445efd9bdc99b050c8b7d11ae5ad19628" +version = "0.20.0" +source = "git+https://github.com/OpenZeppelin/cairo-contracts.git?tag=v0.20.0#7756fd1de2b4ebd239fa6e372d75535cea02e5e5" dependencies = [ "openzeppelin_access", "openzeppelin_token", @@ -48,29 +47,30 @@ dependencies = [ [[package]] name = "openzeppelin_governance" -version = "0.19.0" -source = "git+https://github.com/OpenZeppelin/cairo-contracts.git?tag=v0.19.0#8d49e8c445efd9bdc99b050c8b7d11ae5ad19628" +version = "0.20.0" +source = "git+https://github.com/OpenZeppelin/cairo-contracts.git?tag=v0.20.0#7756fd1de2b4ebd239fa6e372d75535cea02e5e5" dependencies = [ "openzeppelin_access", "openzeppelin_account", "openzeppelin_introspection", "openzeppelin_token", + "openzeppelin_utils", ] [[package]] name = "openzeppelin_introspection" -version = "0.19.0" -source = "git+https://github.com/OpenZeppelin/cairo-contracts.git?tag=v0.19.0#8d49e8c445efd9bdc99b050c8b7d11ae5ad19628" +version = "0.20.0" +source = "git+https://github.com/OpenZeppelin/cairo-contracts.git?tag=v0.20.0#7756fd1de2b4ebd239fa6e372d75535cea02e5e5" [[package]] name = "openzeppelin_merkle_tree" -version = "0.19.0" -source = "git+https://github.com/OpenZeppelin/cairo-contracts.git?tag=v0.19.0#8d49e8c445efd9bdc99b050c8b7d11ae5ad19628" +version = "0.20.0" +source = "git+https://github.com/OpenZeppelin/cairo-contracts.git?tag=v0.20.0#7756fd1de2b4ebd239fa6e372d75535cea02e5e5" [[package]] name = "openzeppelin_presets" -version = "0.19.0" -source = "git+https://github.com/OpenZeppelin/cairo-contracts.git?tag=v0.19.0#8d49e8c445efd9bdc99b050c8b7d11ae5ad19628" +version = "0.20.0" +source = "git+https://github.com/OpenZeppelin/cairo-contracts.git?tag=v0.20.0#7756fd1de2b4ebd239fa6e372d75535cea02e5e5" dependencies = [ "openzeppelin_access", "openzeppelin_account", @@ -83,13 +83,13 @@ dependencies = [ [[package]] name = "openzeppelin_security" -version = "0.19.0" -source = "git+https://github.com/OpenZeppelin/cairo-contracts.git?tag=v0.19.0#8d49e8c445efd9bdc99b050c8b7d11ae5ad19628" +version = "0.20.0" +source = "git+https://github.com/OpenZeppelin/cairo-contracts.git?tag=v0.20.0#7756fd1de2b4ebd239fa6e372d75535cea02e5e5" [[package]] name = "openzeppelin_token" -version = "0.19.0" -source = "git+https://github.com/OpenZeppelin/cairo-contracts.git?tag=v0.19.0#8d49e8c445efd9bdc99b050c8b7d11ae5ad19628" +version = "0.20.0" +source = "git+https://github.com/OpenZeppelin/cairo-contracts.git?tag=v0.20.0#7756fd1de2b4ebd239fa6e372d75535cea02e5e5" dependencies = [ "openzeppelin_access", "openzeppelin_account", @@ -99,13 +99,13 @@ dependencies = [ [[package]] name = "openzeppelin_upgrades" -version = "0.19.0" -source = "git+https://github.com/OpenZeppelin/cairo-contracts.git?tag=v0.19.0#8d49e8c445efd9bdc99b050c8b7d11ae5ad19628" +version = "0.20.0" +source = "git+https://github.com/OpenZeppelin/cairo-contracts.git?tag=v0.20.0#7756fd1de2b4ebd239fa6e372d75535cea02e5e5" [[package]] name = "openzeppelin_utils" -version = "0.19.0" -source = "git+https://github.com/OpenZeppelin/cairo-contracts.git?tag=v0.19.0#8d49e8c445efd9bdc99b050c8b7d11ae5ad19628" +version = "0.20.0" +source = "git+https://github.com/OpenZeppelin/cairo-contracts.git?tag=v0.20.0#7756fd1de2b4ebd239fa6e372d75535cea02e5e5" [[package]] name = "sncast_std" diff --git a/contracts/Scarb.toml b/contracts/Scarb.toml index 0b686a6..bb0c0a9 100644 --- a/contracts/Scarb.toml +++ b/contracts/Scarb.toml @@ -7,7 +7,7 @@ edition = "2024_07" [dependencies] starknet = "2.9.1" -openzeppelin = { git = "https://github.com/OpenZeppelin/cairo-contracts.git", tag = "v0.19.0" } +openzeppelin = { git = "https://github.com/OpenZeppelin/cairo-contracts.git", tag = "v0.20.0" } [dev-dependencies] snforge_std = { git = "https://github.com/foundry-rs/starknet-foundry", tag = "v0.34.0" } diff --git a/contracts/src/lib.cairo b/contracts/src/lib.cairo index 231951d..ea9d2bf 100644 --- a/contracts/src/lib.cairo +++ b/contracts/src/lib.cairo @@ -3,7 +3,12 @@ use core::starknet::ContractAddress; #[starknet::interface] pub trait IAgentRegistry { fn register_agent( - ref self: TContractState, name: ByteArray, system_prompt: ByteArray, prompt_price: u256, + ref self: TContractState, + name: ByteArray, + system_prompt_uri: ByteArray, + rsa_public_key_n: felt252, + rsa_public_key_e: felt252, + prompt_price: u256, ) -> ContractAddress; fn get_token(self: @TContractState) -> ContractAddress; fn is_agent_registered(self: @TContractState, address: ContractAddress) -> bool; @@ -17,7 +22,8 @@ pub trait IAgentRegistry { #[starknet::interface] pub trait IAgent { - fn get_system_prompt(self: @TContractState) -> ByteArray; + fn get_system_prompt_uri(self: @TContractState) -> ByteArray; + fn get_rsa_public_key(self: @TContractState) -> (felt252, felt252); fn get_name(self: @TContractState) -> ByteArray; fn get_creator(self: @TContractState) -> ContractAddress; fn get_prompt_price(self: @TContractState) -> u256; @@ -117,7 +123,12 @@ pub mod AgentRegistry { #[abi(embed_v0)] impl AgentRegistryImpl of super::IAgentRegistry { fn register_agent( - ref self: ContractState, name: ByteArray, system_prompt: ByteArray, prompt_price: u256, + ref self: ContractState, + name: ByteArray, + system_prompt_uri: ByteArray, + rsa_public_key_n: felt252, + rsa_public_key_e: felt252, + prompt_price: u256, ) -> ContractAddress { self.pausable.assert_not_paused(); @@ -132,7 +143,9 @@ pub mod AgentRegistry { let mut constructor_calldata = ArrayTrait::::new(); name.serialize(ref constructor_calldata); - system_prompt.serialize(ref constructor_calldata); + system_prompt_uri.serialize(ref constructor_calldata); + rsa_public_key_n.serialize(ref constructor_calldata); + rsa_public_key_e.serialize(ref constructor_calldata); self.token.read().serialize(ref constructor_calldata); prompt_price.serialize(ref constructor_calldata); creator.serialize(ref constructor_calldata); @@ -276,7 +289,9 @@ pub mod Agent { #[storage] struct Storage { registry: ContractAddress, - system_prompt: ByteArray, + system_prompt_uri: ByteArray, // Changed to store URI of encrypted prompt + rsa_public_key_n: felt252, // RSA public key modulus + rsa_public_key_e: felt252, // RSA public key exponent name: ByteArray, token: ContractAddress, prompt_price: u256, @@ -289,14 +304,18 @@ pub mod Agent { fn constructor( ref self: ContractState, name: ByteArray, - system_prompt: ByteArray, + system_prompt_uri: ByteArray, + rsa_public_key_n: felt252, + rsa_public_key_e: felt252, token: ContractAddress, prompt_price: u256, creator: ContractAddress, ) { self.registry.write(get_caller_address()); self.name.write(name); - self.system_prompt.write(system_prompt); + self.system_prompt_uri.write(system_prompt_uri); + self.rsa_public_key_n.write(rsa_public_key_n); + self.rsa_public_key_e.write(rsa_public_key_e); self.token.write(token); self.prompt_price.write(prompt_price); self.creator.write(creator); @@ -309,8 +328,12 @@ pub mod Agent { self.name.read() } - fn get_system_prompt(self: @ContractState) -> ByteArray { - self.system_prompt.read() + fn get_system_prompt_uri(self: @ContractState) -> ByteArray { + self.system_prompt_uri.read() + } + + fn get_rsa_public_key(self: @ContractState) -> (felt252, felt252) { + (self.rsa_public_key_n.read(), self.rsa_public_key_e.read()) } fn get_prompt_price(self: @ContractState) -> u256 { diff --git a/contracts/tests/test_contract.cairo b/contracts/tests/test_contract.cairo index 93a1ae5..326a87c 100644 --- a/contracts/tests/test_contract.cairo +++ b/contracts/tests/test_contract.cairo @@ -59,12 +59,20 @@ fn test_register_agent() { let registry = IAgentRegistryDispatcher { contract_address: registry_address }; let name = "Test Agent"; - let system_prompt = "I am a test agent"; + let system_prompt_uri = "ipfs://encrypted_prompt_123"; + let rsa_public_key_n = 0x1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef; + let rsa_public_key_e = 0x10001; // Standard RSA public exponent (65537) let mut spy = spy_events(); start_cheat_caller_address(registry.contract_address, creator); - let agent_address = registry.register_agent(name.clone(), system_prompt.clone(), prompt_price); + let agent_address = registry.register_agent( + name.clone(), + system_prompt_uri.clone(), + rsa_public_key_n, + rsa_public_key_e, + prompt_price + ); stop_cheat_caller_address(registry.contract_address); assert(registry.is_agent_registered(agent_address), 'Agent should be registered'); @@ -75,7 +83,12 @@ fn test_register_agent() { let agent_dispatcher = IAgentDispatcher { contract_address: agent_address }; assert(agent_dispatcher.get_name() == name.clone(), 'Wrong agent name'); - assert(agent_dispatcher.get_system_prompt() == system_prompt.clone(), 'Wrong system prompt'); + assert(agent_dispatcher.get_system_prompt_uri() == system_prompt_uri.clone(), 'Wrong system prompt URI'); + + let (stored_n, stored_e) = agent_dispatcher.get_rsa_public_key(); + assert(stored_n == rsa_public_key_n, 'Wrong RSA public key N'); + assert(stored_e == rsa_public_key_e, 'Wrong RSA public key E'); + assert(agent_dispatcher.get_creator() == creator, 'Wrong creator'); // Verify event was emitted @@ -105,7 +118,13 @@ fn test_pay_for_prompt() { // Register agent start_cheat_caller_address(registry.contract_address, creator); - let agent_address = registry.register_agent("Test Agent", "Test Prompt", prompt_price); + let agent_address = registry.register_agent( + "Test Agent", + "ipfs://encrypted_prompt_pay", + 0x1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef, + 0x10001, + prompt_price + ); stop_cheat_caller_address(registry.contract_address); let agent = IAgentDispatcher { contract_address: agent_address }; @@ -161,9 +180,21 @@ fn test_register_multiple_agents() { start_cheat_caller_address(registry.contract_address, creator); - let agent1_address = registry.register_agent("Agent 1", "Prompt 1", prompt_price); - - let agent2_address = registry.register_agent("Agent 2", "Prompt 2", prompt_price); + let agent1_address = registry.register_agent( + "Agent 1", + "ipfs://encrypted_prompt_1", + 0x1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef, + 0x10001, + prompt_price + ); + + let agent2_address = registry.register_agent( + "Agent 2", + "ipfs://encrypted_prompt_2", + 0x2234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef, + 0x10001, + prompt_price + ); stop_cheat_caller_address(registry.contract_address); @@ -219,16 +250,28 @@ fn test_get_agent_details() { let registry = IAgentRegistryDispatcher { contract_address: registry_address }; let name = "Complex Agent"; - let system_prompt = "Complex system prompt with multiple words"; + let system_prompt_uri = "ipfs://encrypted_prompt_complex_123"; + let rsa_public_key_n = 0x3234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef; + let rsa_public_key_e = 0x10001; start_cheat_caller_address(registry.contract_address, creator); - let agent_address = registry.register_agent(name.clone(), system_prompt.clone(), prompt_price); + let agent_address = registry.register_agent( + name.clone(), + system_prompt_uri.clone(), + rsa_public_key_n, + rsa_public_key_e, + prompt_price + ); stop_cheat_caller_address(registry.contract_address); let agent = IAgentDispatcher { contract_address: agent_address }; assert(agent.get_name() == name, 'Wrong agent name'); - assert(agent.get_system_prompt() == system_prompt, 'Wrong system prompt'); + assert(agent.get_system_prompt_uri() == system_prompt_uri, 'Wrong system prompt URI'); + + let (stored_n, stored_e) = agent.get_rsa_public_key(); + assert(stored_n == rsa_public_key_n, 'Wrong RSA public key N'); + assert(stored_e == rsa_public_key_e, 'Wrong RSA public key E'); } #[test] @@ -306,6 +349,26 @@ fn test_pay_for_prompt_without_approval() { agent.pay_for_prompt(12345); } +#[test] +#[should_panic(expected: ('Invalid RSA public key',))] +fn test_invalid_rsa_key() { + let tee = starknet::contract_address_const::<0x1>(); + let creator = starknet::contract_address_const::<0x123>(); + let prompt_price: u256 = 100; + let (registry_address, _) = deploy_registry(tee, creator); + let registry = IAgentRegistryDispatcher { contract_address: registry_address }; + + start_cheat_caller_address(registry.contract_address, creator); + // Try to register with invalid RSA exponent (must be 65537) + registry.register_agent( + "Test Agent", + "ipfs://encrypted_prompt", + 0x1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef, + 0x4, // Invalid exponent + prompt_price + ); +} + #[test] fn test_is_agent_registered() { let tee = starknet::contract_address_const::<0x1>(); diff --git a/pkg/agent/setup/manager.go b/pkg/agent/setup/manager.go index e5bcb7a..7d914e5 100644 --- a/pkg/agent/setup/manager.go +++ b/pkg/agent/setup/manager.go @@ -2,8 +2,12 @@ package setup import ( "context" + "crypto/rand" + "crypto/rsa" + "crypto/x509" "fmt" "log/slog" + "math/big" "github.com/NethermindEth/juno/core/felt" starknetgoutils "github.com/NethermindEth/starknet.go/utils" @@ -42,6 +46,9 @@ type SetupOutput struct { AgentRegistryAddress *felt.Felt `json:"agent_registry_address"` OpenAIKey string `json:"openai_key"` DstackTappdEndpoint string `json:"dstack_tappd_endpoint"` + RsaPrivateKey []byte `json:"rsa_private_key"` // Stored encrypted + RsaPublicKeyN *felt.Felt `json:"rsa_public_key_n"` // For contract + RsaPublicKeyE *felt.Felt `json:"rsa_public_key_e"` // For contract } func NewSetupManagerFromEnv() (*SetupManager, error) { @@ -134,6 +141,19 @@ func (m *SetupManager) Setup(ctx context.Context) (*SetupOutput, error) { return nil, fmt.Errorf("failed to encumber proton: %v", err) } + // Generate RSA key pair for system prompt encryption + rsaPrivateKey, err := rsa.GenerateKey(rand.Reader, 2048) + if err != nil { + return nil, fmt.Errorf("failed to generate RSA key pair: %v", err) + } + + // Convert RSA private key to DER format for storage + rsaPrivateKeyBytes := x509.MarshalPKCS1PrivateKey(rsaPrivateKey) + + // Convert public key components to felt252 for contract + rsaPublicKeyN := new(felt.Felt).SetBytes(rsaPrivateKey.PublicKey.N.Bytes()) + rsaPublicKeyE := new(felt.Felt).SetBigInt(big.NewInt(int64(rsaPrivateKey.PublicKey.E))) + output := &SetupOutput{ TwitterAuthTokens: twitterEncumbererOutput.AuthTokens, TwitterAccessToken: twitterEncumbererOutput.OAuthTokenPair.Token, @@ -148,6 +168,9 @@ func (m *SetupManager) Setup(ctx context.Context) (*SetupOutput, error) { AgentRegistryAddress: agentRegistryAddress, OpenAIKey: m.openAiKey, DstackTappdEndpoint: m.dstackTappdEndpoint, + RsaPrivateKey: rsaPrivateKeyBytes, + RsaPublicKeyN: rsaPublicKeyN, + RsaPublicKeyE: rsaPublicKeyE, } if debug.IsDebugShowSetup() { From fd359a5ea2cb331a3bf33c57f16aafa6f41328a9 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 16 Jan 2025 06:06:31 +0000 Subject: [PATCH 2/2] Remove RSA key storage from contracts - Remove RSA public key storage and methods from agent contract - Update contract interfaces to remove RSA key parameters - Remove RSA key-related test cases - Move RSA key management to DStack TEE environment - Maintain global TEE RSA key approach for system prompt encryption Co-Authored-By: franco@nethermind.io --- contracts/src/lib.cairo | 19 +------------- contracts/tests/test_contract.cairo | 40 ----------------------------- pkg/agent/setup/manager.go | 11 ++------ 3 files changed, 3 insertions(+), 67 deletions(-) diff --git a/contracts/src/lib.cairo b/contracts/src/lib.cairo index ea9d2bf..fc640c9 100644 --- a/contracts/src/lib.cairo +++ b/contracts/src/lib.cairo @@ -6,8 +6,6 @@ pub trait IAgentRegistry { ref self: TContractState, name: ByteArray, system_prompt_uri: ByteArray, - rsa_public_key_n: felt252, - rsa_public_key_e: felt252, prompt_price: u256, ) -> ContractAddress; fn get_token(self: @TContractState) -> ContractAddress; @@ -23,7 +21,6 @@ pub trait IAgentRegistry { #[starknet::interface] pub trait IAgent { fn get_system_prompt_uri(self: @TContractState) -> ByteArray; - fn get_rsa_public_key(self: @TContractState) -> (felt252, felt252); fn get_name(self: @TContractState) -> ByteArray; fn get_creator(self: @TContractState) -> ContractAddress; fn get_prompt_price(self: @TContractState) -> u256; @@ -126,8 +123,6 @@ pub mod AgentRegistry { ref self: ContractState, name: ByteArray, system_prompt_uri: ByteArray, - rsa_public_key_n: felt252, - rsa_public_key_e: felt252, prompt_price: u256, ) -> ContractAddress { self.pausable.assert_not_paused(); @@ -144,8 +139,6 @@ pub mod AgentRegistry { let mut constructor_calldata = ArrayTrait::::new(); name.serialize(ref constructor_calldata); system_prompt_uri.serialize(ref constructor_calldata); - rsa_public_key_n.serialize(ref constructor_calldata); - rsa_public_key_e.serialize(ref constructor_calldata); self.token.read().serialize(ref constructor_calldata); prompt_price.serialize(ref constructor_calldata); creator.serialize(ref constructor_calldata); @@ -289,9 +282,7 @@ pub mod Agent { #[storage] struct Storage { registry: ContractAddress, - system_prompt_uri: ByteArray, // Changed to store URI of encrypted prompt - rsa_public_key_n: felt252, // RSA public key modulus - rsa_public_key_e: felt252, // RSA public key exponent + system_prompt_uri: ByteArray, // URI of encrypted prompt name: ByteArray, token: ContractAddress, prompt_price: u256, @@ -305,8 +296,6 @@ pub mod Agent { ref self: ContractState, name: ByteArray, system_prompt_uri: ByteArray, - rsa_public_key_n: felt252, - rsa_public_key_e: felt252, token: ContractAddress, prompt_price: u256, creator: ContractAddress, @@ -314,8 +303,6 @@ pub mod Agent { self.registry.write(get_caller_address()); self.name.write(name); self.system_prompt_uri.write(system_prompt_uri); - self.rsa_public_key_n.write(rsa_public_key_n); - self.rsa_public_key_e.write(rsa_public_key_e); self.token.write(token); self.prompt_price.write(prompt_price); self.creator.write(creator); @@ -332,10 +319,6 @@ pub mod Agent { self.system_prompt_uri.read() } - fn get_rsa_public_key(self: @ContractState) -> (felt252, felt252) { - (self.rsa_public_key_n.read(), self.rsa_public_key_e.read()) - } - fn get_prompt_price(self: @ContractState) -> u256 { self.prompt_price.read() } diff --git a/contracts/tests/test_contract.cairo b/contracts/tests/test_contract.cairo index 326a87c..4384864 100644 --- a/contracts/tests/test_contract.cairo +++ b/contracts/tests/test_contract.cairo @@ -60,8 +60,6 @@ fn test_register_agent() { let name = "Test Agent"; let system_prompt_uri = "ipfs://encrypted_prompt_123"; - let rsa_public_key_n = 0x1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef; - let rsa_public_key_e = 0x10001; // Standard RSA public exponent (65537) let mut spy = spy_events(); @@ -69,8 +67,6 @@ fn test_register_agent() { let agent_address = registry.register_agent( name.clone(), system_prompt_uri.clone(), - rsa_public_key_n, - rsa_public_key_e, prompt_price ); stop_cheat_caller_address(registry.contract_address); @@ -85,10 +81,6 @@ fn test_register_agent() { assert(agent_dispatcher.get_name() == name.clone(), 'Wrong agent name'); assert(agent_dispatcher.get_system_prompt_uri() == system_prompt_uri.clone(), 'Wrong system prompt URI'); - let (stored_n, stored_e) = agent_dispatcher.get_rsa_public_key(); - assert(stored_n == rsa_public_key_n, 'Wrong RSA public key N'); - assert(stored_e == rsa_public_key_e, 'Wrong RSA public key E'); - assert(agent_dispatcher.get_creator() == creator, 'Wrong creator'); // Verify event was emitted @@ -121,8 +113,6 @@ fn test_pay_for_prompt() { let agent_address = registry.register_agent( "Test Agent", "ipfs://encrypted_prompt_pay", - 0x1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef, - 0x10001, prompt_price ); stop_cheat_caller_address(registry.contract_address); @@ -183,16 +173,12 @@ fn test_register_multiple_agents() { let agent1_address = registry.register_agent( "Agent 1", "ipfs://encrypted_prompt_1", - 0x1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef, - 0x10001, prompt_price ); let agent2_address = registry.register_agent( "Agent 2", "ipfs://encrypted_prompt_2", - 0x2234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef, - 0x10001, prompt_price ); @@ -251,15 +237,11 @@ fn test_get_agent_details() { let name = "Complex Agent"; let system_prompt_uri = "ipfs://encrypted_prompt_complex_123"; - let rsa_public_key_n = 0x3234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef; - let rsa_public_key_e = 0x10001; start_cheat_caller_address(registry.contract_address, creator); let agent_address = registry.register_agent( name.clone(), system_prompt_uri.clone(), - rsa_public_key_n, - rsa_public_key_e, prompt_price ); stop_cheat_caller_address(registry.contract_address); @@ -268,10 +250,6 @@ fn test_get_agent_details() { assert(agent.get_name() == name, 'Wrong agent name'); assert(agent.get_system_prompt_uri() == system_prompt_uri, 'Wrong system prompt URI'); - - let (stored_n, stored_e) = agent.get_rsa_public_key(); - assert(stored_n == rsa_public_key_n, 'Wrong RSA public key N'); - assert(stored_e == rsa_public_key_e, 'Wrong RSA public key E'); } #[test] @@ -349,25 +327,7 @@ fn test_pay_for_prompt_without_approval() { agent.pay_for_prompt(12345); } -#[test] -#[should_panic(expected: ('Invalid RSA public key',))] -fn test_invalid_rsa_key() { - let tee = starknet::contract_address_const::<0x1>(); - let creator = starknet::contract_address_const::<0x123>(); - let prompt_price: u256 = 100; - let (registry_address, _) = deploy_registry(tee, creator); - let registry = IAgentRegistryDispatcher { contract_address: registry_address }; - start_cheat_caller_address(registry.contract_address, creator); - // Try to register with invalid RSA exponent (must be 65537) - registry.register_agent( - "Test Agent", - "ipfs://encrypted_prompt", - 0x1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef, - 0x4, // Invalid exponent - prompt_price - ); -} #[test] fn test_is_agent_registered() { diff --git a/pkg/agent/setup/manager.go b/pkg/agent/setup/manager.go index 7d914e5..fefb24c 100644 --- a/pkg/agent/setup/manager.go +++ b/pkg/agent/setup/manager.go @@ -46,9 +46,8 @@ type SetupOutput struct { AgentRegistryAddress *felt.Felt `json:"agent_registry_address"` OpenAIKey string `json:"openai_key"` DstackTappdEndpoint string `json:"dstack_tappd_endpoint"` - RsaPrivateKey []byte `json:"rsa_private_key"` // Stored encrypted - RsaPublicKeyN *felt.Felt `json:"rsa_public_key_n"` // For contract - RsaPublicKeyE *felt.Felt `json:"rsa_public_key_e"` // For contract + // RSA private key is managed securely through the DStack TEE environment + // and should not be stored in this structure } func NewSetupManagerFromEnv() (*SetupManager, error) { @@ -150,10 +149,6 @@ func (m *SetupManager) Setup(ctx context.Context) (*SetupOutput, error) { // Convert RSA private key to DER format for storage rsaPrivateKeyBytes := x509.MarshalPKCS1PrivateKey(rsaPrivateKey) - // Convert public key components to felt252 for contract - rsaPublicKeyN := new(felt.Felt).SetBytes(rsaPrivateKey.PublicKey.N.Bytes()) - rsaPublicKeyE := new(felt.Felt).SetBigInt(big.NewInt(int64(rsaPrivateKey.PublicKey.E))) - output := &SetupOutput{ TwitterAuthTokens: twitterEncumbererOutput.AuthTokens, TwitterAccessToken: twitterEncumbererOutput.OAuthTokenPair.Token, @@ -169,8 +164,6 @@ func (m *SetupManager) Setup(ctx context.Context) (*SetupOutput, error) { OpenAIKey: m.openAiKey, DstackTappdEndpoint: m.dstackTappdEndpoint, RsaPrivateKey: rsaPrivateKeyBytes, - RsaPublicKeyN: rsaPublicKeyN, - RsaPublicKeyE: rsaPublicKeyE, } if debug.IsDebugShowSetup() {