Send derived secrets to funrun (#42319)

GitOrigin-RevId: 067e60624c6b0324df93c5fcb4c67345a9dc9703

Author: goffrie

crates/application/src/lib.rs

@@ -1018,7 +1018,12 @@ impl<RT: Runtime> Application<RT> {
         let url = self
             .file_storage
             .transactional_file_storage
-            .generate_upload_url(&mut tx, self.key_broker(), issued_ts, component)
+            .generate_upload_url(
+                &mut tx,
+                &self.key_broker().function_runner_keybroker(),
+                issued_ts,
+                component,
+            )
             .await?;
 
         Ok(url)

crates/application/src/test_helpers.rs

@@ -228,7 +228,7 @@ impl<RT: Runtime> ApplicationTestExt<RT> for Application<RT> {
         let fetch_client = Arc::new(StaticFetchClient::new());
         let function_runner = Arc::new(InProcessFunctionRunner::new(
             DEV_INSTANCE_NAME.into(),
-            DEV_SECRET.try_into()?,
+            KeyBroker::new(DEV_INSTANCE_NAME, DEV_SECRET.try_into()?)?.function_runner_keybroker(),
             convex_origin.clone(),
             rt.clone(),
             persistence.reader(),

crates/file_storage/src/core.rs

@@ -43,8 +43,8 @@ use headers::{
     ContentType,
 };
 use keybroker::{
+    FunctionRunnerKeyBroker,
     Identity,
-    KeyBroker,
 };
 use maplit::btreemap;
 use model::{
@@ -97,7 +97,7 @@ impl<RT: Runtime> TransactionalFileStorage<RT> {
     pub fn generate_upload_url_with_origin(
         &self,
         origin_override: Option<ConvexOrigin>,
-        key_broker: &KeyBroker,
+        key_broker: &FunctionRunnerKeyBroker,
         issued_ts: UnixTimestamp,
         component: ComponentId,
     ) -> anyhow::Result<String> {
@@ -110,7 +110,7 @@ impl<RT: Runtime> TransactionalFileStorage<RT> {
     pub async fn generate_upload_url(
         &self,
         tx: &mut Transaction<RT>,
-        key_broker: &KeyBroker,
+        key_broker: &FunctionRunnerKeyBroker,
         issued_ts: UnixTimestamp,
         component: ComponentId,
     ) -> anyhow::Result<String> {

crates/function_runner/src/in_process_function_runner.rs

@@ -49,8 +49,8 @@ use futures::{
 };
 use isolate::ActionCallbacks;
 use keybroker::{
+    FunctionRunnerKeyBroker,
     Identity,
-    InstanceSecret,
 };
 use model::{
     config::types::ModuleConfig,
@@ -100,7 +100,7 @@ pub struct InProcessFunctionRunner<RT: Runtime> {
 
     // Static information about the backend.
     instance_name: String,
-    instance_secret: InstanceSecret,
+    key_broker: FunctionRunnerKeyBroker,
     convex_origin: ConvexOrigin,
     database: Database<RT>,
     // Use Weak reference to avoid reference cycle between InProcessFunctionRunner
@@ -112,7 +112,7 @@ pub struct InProcessFunctionRunner<RT: Runtime> {
 impl<RT: Runtime> InProcessFunctionRunner<RT> {
     pub fn new(
         instance_name: String,
-        instance_secret: InstanceSecret,
+        keybroker: FunctionRunnerKeyBroker,
         convex_origin: ConvexOrigin,
         rt: RT,
         persistence_reader: Arc<dyn PersistenceReader>,
@@ -127,7 +127,7 @@ impl<RT: Runtime> InProcessFunctionRunner<RT> {
             server,
             persistence_reader,
             instance_name,
-            instance_secret,
+            key_broker: keybroker,
             convex_origin,
             database,
             action_callbacks: Arc::new(RwLock::new(None)),
@@ -239,7 +239,7 @@ impl<RT: Runtime> FunctionRunner<RT> for InProcessFunctionRunner<RT> {
 
         let request_metadata = RunRequestArgs {
             instance_name: self.instance_name.clone(),
-            instance_secret: self.instance_secret,
+            key_broker: self.key_broker.clone(),
             reader: self.persistence_reader.clone(),
             convex_origin: self.convex_origin.clone(),
             bootstrap_metadata: self.database.bootstrap_metadata.clone(),

crates/function_runner/src/server.rs

@@ -58,9 +58,8 @@ use isolate::{
     IsolateClient,
 };
 use keybroker::{
+    FunctionRunnerKeyBroker,
     Identity,
-    InstanceSecret,
-    KeyBroker,
 };
 use model::{
     config::types::ModuleConfig,
@@ -118,7 +117,7 @@ const MAX_ISOLATE_WORKERS: usize = 128;
 
 pub struct RunRequestArgs {
     pub instance_name: String,
-    pub instance_secret: InstanceSecret,
+    pub key_broker: FunctionRunnerKeyBroker,
     pub reader: Arc<dyn PersistenceReader>,
     pub convex_origin: ConvexOrigin,
     pub bootstrap_metadata: BootstrapMetadata,
@@ -289,7 +288,7 @@ impl<RT: Runtime, S: StorageForInstance<RT>> FunctionRunnerCore<RT, S> {
         &self,
         RunRequestArgs {
             instance_name,
-            instance_secret,
+            key_broker,
             reader,
             convex_origin,
             bootstrap_metadata,
@@ -355,7 +354,6 @@ impl<RT: Runtime, S: StorageForInstance<RT>> FunctionRunnerCore<RT, S> {
             .storage_for_instance(&mut transaction, StorageUseCase::Modules)
             .await?;
 
-        let key_broker = KeyBroker::new(&instance_name, instance_secret)?;
         let environment_data = EnvironmentData {
             key_broker,
             default_system_env_vars,

crates/isolate/src/client.rs

@@ -101,8 +101,8 @@ use futures::{
     },
 };
 use keybroker::{
+    FunctionRunnerKeyBroker,
     Identity,
-    KeyBroker,
 };
 use model::{
     config::types::ModuleConfig,
@@ -352,7 +352,7 @@ pub struct ActionRequestParams {
 
 #[derive(Clone)]
 pub struct EnvironmentData<RT: Runtime> {
-    pub key_broker: KeyBroker,
+    pub key_broker: FunctionRunnerKeyBroker,
     pub default_system_env_vars: BTreeMap<EnvVarName, EnvVarValue>,
     pub file_storage: TransactionalFileStorage<RT>,
     pub module_loader: Arc<dyn ModuleCache<RT>>,

crates/isolate/src/environment/action/task_executor.rs

@@ -32,8 +32,8 @@ use futures::{
     StreamExt,
 };
 use keybroker::{
+    FunctionRunnerKeyBroker,
     Identity,
-    KeyBroker,
 };
 use parking_lot::Mutex;
 use serde_json::Value as JsonValue;
@@ -72,7 +72,7 @@ pub struct TaskExecutor<RT: Runtime> {
     pub action_callbacks: Arc<dyn ActionCallbacks>,
     pub fetch_client: Arc<dyn FetchClient>,
     pub _module_loader: Arc<dyn ModuleCache<RT>>,
-    pub key_broker: KeyBroker,
+    pub key_broker: FunctionRunnerKeyBroker,
     pub task_order: TaskOrder,
     pub task_retval_sender: mpsc::UnboundedSender<TaskResponse>,
     pub usage_tracker: FunctionUsageTracker,

crates/isolate/src/environment/udf/async_syscall.rs

@@ -61,7 +61,7 @@ use errors::{
     ErrorMetadataAnyhowExt,
 };
 use itertools::Itertools;
-use keybroker::KeyBroker;
+use keybroker::FunctionRunnerKeyBroker;
 use model::{
     components::{
         handles::FunctionHandlesModel,
@@ -279,7 +279,7 @@ pub enum ManagedQuery<RT: Runtime> {
 pub trait AsyncSyscallProvider<RT: Runtime> {
     fn rt(&self) -> &RT;
     fn tx(&mut self) -> anyhow::Result<&mut Transaction<RT>>;
-    fn key_broker(&self) -> &KeyBroker;
+    fn key_broker(&self) -> &FunctionRunnerKeyBroker;
     fn context(&self) -> &ExecutionContext;
 
     fn observe_identity(&mut self) -> anyhow::Result<()>;
@@ -348,7 +348,7 @@ impl<RT: Runtime> AsyncSyscallProvider<RT> for DatabaseUdfEnvironment<RT> {
         self.phase.component()
     }
 
-    fn key_broker(&self) -> &KeyBroker {
+    fn key_broker(&self) -> &FunctionRunnerKeyBroker {
         &self.key_broker
     }
 

crates/isolate/src/environment/udf/mod.rs

@@ -85,7 +85,7 @@ use deno_core::{
 };
 use errors::ErrorMetadata;
 use file_storage::TransactionalFileStorage;
-use keybroker::KeyBroker;
+use keybroker::FunctionRunnerKeyBroker;
 use rand::Rng;
 use rand_chacha::ChaCha12Rng;
 use serde_json::Value as JsonValue;
@@ -177,7 +177,7 @@ pub struct DatabaseUdfEnvironment<RT: Runtime> {
     query_manager: QueryManager<RT>,
 
     persistence_version: PersistenceVersion,
-    key_broker: KeyBroker,
+    key_broker: FunctionRunnerKeyBroker,
     log_lines: LogLines,
 
     /// Journal from a previous computation of this UDF used as an input to this

crates/isolate/src/isolate2/runner.rs

@@ -43,7 +43,7 @@ use database::{
     Transaction,
 };
 use errors::ErrorMetadata;
-use keybroker::KeyBroker;
+use keybroker::FunctionRunnerKeyBroker;
 use model::{
     config::module_loader::ModuleLoader,
     environment_variables::{
@@ -483,7 +483,7 @@ async fn run_request<RT: Runtime>(
     path_and_args: ValidatedPathAndArgs,
     shared: UdfShared<RT>,
     mut log_line_receiver: spsc::Receiver<LogLine>,
-    key_broker: KeyBroker,
+    key_broker: FunctionRunnerKeyBroker,
     execution_context: ExecutionContext,
     query_journal: QueryJournal,
 ) -> anyhow::Result<UdfOutcome> {
@@ -808,7 +808,7 @@ struct Isolate2SyscallProvider<'a, RT: Runtime> {
 
     syscall_trace: SyscallTrace,
 
-    key_broker: KeyBroker,
+    key_broker: FunctionRunnerKeyBroker,
     context: ExecutionContext,
 }
 
@@ -820,7 +820,7 @@ impl<'a, RT: Runtime> Isolate2SyscallProvider<'a, RT> {
         prev_journal: QueryJournal,
         is_system: bool,
         shared: UdfShared<RT>,
-        key_broker: KeyBroker,
+        key_broker: FunctionRunnerKeyBroker,
         context: ExecutionContext,
     ) -> Self {
         Self {
@@ -853,7 +853,7 @@ impl<RT: Runtime> AsyncSyscallProvider<RT> for Isolate2SyscallProvider<'_, RT> {
         Ok(ComponentId::Root)
     }
 
-    fn key_broker(&self) -> &KeyBroker {
+    fn key_broker(&self) -> &FunctionRunnerKeyBroker {
         &self.key_broker
     }
 
@@ -989,7 +989,7 @@ async fn tokio_thread<RT: Runtime>(
     path_and_args: ValidatedPathAndArgs,
     shared: UdfShared<RT>,
     log_line_receiver: spsc::Receiver<LogLine>,
-    key_broker: KeyBroker,
+    key_broker: FunctionRunnerKeyBroker,
     execution_context: ExecutionContext,
     query_journal: QueryJournal,
 ) {
@@ -1026,7 +1026,7 @@ pub async fn run_isolate_v2_udf<RT: Runtime>(
     execution_time_seed: SeedData,
     udf_type: UdfType,
     path_and_args: ValidatedPathAndArgs,
-    key_broker: KeyBroker,
+    key_broker: FunctionRunnerKeyBroker,
     context: ExecutionContext,
     query_journal: QueryJournal,
 ) -> anyhow::Result<(Transaction<RT>, UdfOutcome)> {