From 17387b762114314f200156e7347ee7223d5ab3bf Mon Sep 17 00:00:00 2001 From: Abhinav Chordia Date: Mon, 12 Oct 2026 05:45:37 +0000 Subject: [PATCH] sbx-mgr: inject the template-manager resolver Every template-manager GetTemplate sandbox-manager issues now goes through one TemplateResolver, carried in SandboxManagerServiceOptions as template_resolver. That covers resolve_template_paths and resolve_template_paths_for_start, which serve create, start, wake, pool adoption and donor cloning, plus restore staging's parent-template lookup and the template network-policy cache. Both resolver methods take a tonic::Request and return the tonic::Response, so request and response metadata cross the seam; callers call into_inner() as before. The default, HostTemplateManager, makes the same calls the callers made directly. get_template dials pplx_space_template_mgr_client::connect() for each call, and get_template_on_shared_channel uses the client crate's process-wide lazy() channel. Connect and RPC failures stay distinct (TemplateLookupError), so every caller keeps its error codes and messages. TemplateNetworkPolicyCache no longer has a Default constructor that bypassed the injected resolver. Tests substitute FakeTemplateResolver. It records each lookup's channel, request metadata and message, and returns a scripted answer: a response, a refusal, or a connect failure. - create_template_lookup_tests covers create's auto-lifecycle authority checks. Forwarded peers' Local authority, Unspecified and unknown authorities are refused before any lookup. Cluster authority from a forwarded peer, and Local from a local peer, send one dialed lookup for the requested template and surface its refusal as "GetTemplate failed:". An unreachable template-manager surfaces as Internal "cannot connect to template-manager:". No row is created, and node-manager, a refusing fake, receives no call. - The latest-alias start test and the orphaned-sandbox prepare test assert that no template lookup is sent. Previously they relied on an absent template-manager. - The network-policy cache tests build the cache with a refusing fake. The hint-admission and custom-override tests assert that no owner lookup is sent. - grpc_core's isolated_service, which lists every option, gets a FakeTemplateResolver that refuses every lookup, so those tests can no longer reach the fixed template-manager socket. --- crates/libs/sbx/mgr/server/src/lib.rs | 1 + .../sbx/mgr/server/src/service/auth_lookup.rs | 10 +- .../libs/sbx/mgr/server/src/service/create.rs | 2 +- .../mgr/server/src/service/create_support.rs | 60 +++--- .../src/service/create_support/tests.rs | 22 ++- .../sbx/mgr/server/src/service/grpc_core.rs | 12 +- .../grpc_core/create_template_lookup_tests.rs | 178 ++++++++++++++++++ .../mgr/server/src/service/grpc_core/tests.rs | 9 +- .../src/service/grpc_pool_wait_clone.rs | 4 +- crates/libs/sbx/mgr/server/src/service/mod.rs | 6 + .../mgr/server/src/service/network_policy.rs | 41 ++-- .../src/service/network_policy/tests.rs | 36 +++- .../server/src/service/recovery_cross_node.rs | 11 +- .../src/service/route_identity_tests.rs | 1 + .../libs/sbx/mgr/server/src/service/start.rs | 11 +- .../libs/sbx/mgr/server/src/service/tests.rs | 7 +- .../tests/node_manager_transport_tests.rs | 25 ++- .../sbx/mgr/server/src/snapshot_restore.rs | 53 +++--- .../sbx/mgr/server/src/template_resolver.rs | 67 +++++++ .../mgr/server/src/template_resolver/fake.rs | 108 +++++++++++ 20 files changed, 561 insertions(+), 103 deletions(-) create mode 100644 crates/libs/sbx/mgr/server/src/service/grpc_core/create_template_lookup_tests.rs create mode 100644 crates/libs/sbx/mgr/server/src/template_resolver.rs create mode 100644 crates/libs/sbx/mgr/server/src/template_resolver/fake.rs diff --git a/crates/libs/sbx/mgr/server/src/lib.rs b/crates/libs/sbx/mgr/server/src/lib.rs index 5b6e8b1f53..54498db7a1 100644 --- a/crates/libs/sbx/mgr/server/src/lib.rs +++ b/crates/libs/sbx/mgr/server/src/lib.rs @@ -35,6 +35,7 @@ pub mod snapshot_scheduler; pub mod snapshot_upload; pub mod snapshot_upload_jobs; pub mod state; +pub mod template_resolver; mod transfer_directory; pub mod watchdog; diff --git a/crates/libs/sbx/mgr/server/src/service/auth_lookup.rs b/crates/libs/sbx/mgr/server/src/service/auth_lookup.rs index 3ac46671d9..792ae03797 100644 --- a/crates/libs/sbx/mgr/server/src/service/auth_lookup.rs +++ b/crates/libs/sbx/mgr/server/src/service/auth_lookup.rs @@ -1005,13 +1005,15 @@ impl SandboxManagerService { } if inherit_template_default { - let resolved_template = - resolve_template_paths(pplx_space_template_mgr_client::proto::TemplateRef { + let resolved_template = resolve_template_paths( + self.template_resolver.as_ref(), + pplx_space_template_mgr_client::proto::TemplateRef { name: sandbox.template.clone(), content_hash: sandbox.template_hash.clone(), tag: None, - }) - .await?; + }, + ) + .await?; policy = resolved_template.default_auto_lifecycle; } diff --git a/crates/libs/sbx/mgr/server/src/service/create.rs b/crates/libs/sbx/mgr/server/src/service/create.rs index b960533438..e46a361d44 100644 --- a/crates/libs/sbx/mgr/server/src/service/create.rs +++ b/crates/libs/sbx/mgr/server/src/service/create.rs @@ -628,7 +628,7 @@ impl SandboxManagerService { "restore.template_prepare_source", START_TEMPLATE_RESOLUTION_TIMEOUT, id, - template_source_for_sandbox_prepare(sandbox), + template_source_for_sandbox_prepare(self.template_resolver.as_ref(), sandbox), ) .await?; let prepare_template = prepare_source.template; diff --git a/crates/libs/sbx/mgr/server/src/service/create_support.rs b/crates/libs/sbx/mgr/server/src/service/create_support.rs index d8365eaad1..1b00b5af80 100644 --- a/crates/libs/sbx/mgr/server/src/service/create_support.rs +++ b/crates/libs/sbx/mgr/server/src/service/create_support.rs @@ -564,21 +564,29 @@ pub(super) fn checkpoint_start_response_for_existing_started_sandbox( } pub(super) async fn resolve_template_paths( + templates: &dyn crate::template_resolver::TemplateResolver, template_ref: pplx_space_template_mgr_client::proto::TemplateRef, ) -> Result { - let mut tmpl_client = pplx_space_template_mgr_client::connect() + use crate::template_resolver::TemplateLookupError; + let inner = templates + .get_template(tonic::Request::new( + pplx_space_template_mgr_client::proto::GetTemplateRequest { + caller_org_id: None, + template: Some(template_ref), + vcpu: None, + memory_mib: None, + }, + )) .await - .map_err(|e| tonic::Status::internal(format!("cannot connect to template-manager: {e}")))?; - let resp = tmpl_client - .get_template(pplx_space_template_mgr_client::proto::GetTemplateRequest { - caller_org_id: None, - template: Some(template_ref), - vcpu: None, - memory_mib: None, - }) - .await - .map_err(|e| tonic::Status::new(e.code(), format!("GetTemplate failed: {e}")))?; - let inner = resp.into_inner(); + .map_err(|error| match error { + TemplateLookupError::Connect(e) => { + tonic::Status::internal(format!("cannot connect to template-manager: {e}")) + } + TemplateLookupError::Rpc(e) => { + tonic::Status::new(e.code(), format!("GetTemplate failed: {e}")) + } + })? + .into_inner(); let tmpl = inner.template.unwrap_or_default(); let arch = template_metadata_arch(&inner.template_dir).unwrap_or(inner.arch); let spaced_compatibility = template_metadata_spaced_compatibility(&inner.template_dir); @@ -670,6 +678,7 @@ pub(super) fn stage_base_snapshot_runtime_metadata( } pub(super) async fn resolve_template_paths_for_start( + templates: &dyn crate::template_resolver::TemplateResolver, template_ref: pplx_space_template_mgr_client::proto::TemplateRef, ) -> Result { // Start restores an existing sandbox's identity. A name-only lookup could @@ -683,7 +692,7 @@ pub(super) async fn resolve_template_paths_for_start( "cannot recover sandbox artifact without its saved template hash", )); } - resolve_template_paths(template_ref) + resolve_template_paths(templates, template_ref) .await .map_err(start_template_resolution_failed_status) } @@ -701,6 +710,7 @@ fn start_template_resolution_failed_status(error: Status) -> Status { } pub(super) async fn template_runtime_defaults_for_sandbox_start( + templates: &dyn crate::template_resolver::TemplateResolver, sandbox: &crate::db::SandboxRow, prepared_template_dir: Option<&str>, ) -> TemplateRuntimeDefaults { @@ -708,11 +718,14 @@ pub(super) async fn template_runtime_defaults_for_sandbox_start( return template_runtime_defaults(template_dir); } - match resolve_template_paths(pplx_space_template_mgr_client::proto::TemplateRef { - name: sandbox.template.clone(), - content_hash: sandbox.template_hash.clone(), - tag: None, - }) + match resolve_template_paths( + templates, + pplx_space_template_mgr_client::proto::TemplateRef { + name: sandbox.template.clone(), + content_hash: sandbox.template_hash.clone(), + tag: None, + }, + ) .await { Ok(resolved) => resolved.runtime_defaults, @@ -742,6 +755,7 @@ pub(super) fn sandbox_rootfs_exists_in(cache_dir: &std::path::Path, sandbox_id: } pub(super) async fn template_source_for_sandbox_prepare( + templates: &dyn crate::template_resolver::TemplateResolver, sandbox: &crate::db::SandboxRow, ) -> Result { if sandbox_rootfs_exists(&sandbox.id) { @@ -764,13 +778,15 @@ pub(super) async fn template_source_for_sandbox_prepare( ))); } - let resolved_template = - resolve_template_paths_for_start(pplx_space_template_mgr_client::proto::TemplateRef { + let resolved_template = resolve_template_paths_for_start( + templates, + pplx_space_template_mgr_client::proto::TemplateRef { name: sandbox.template.clone(), content_hash: sandbox.template_hash.clone(), tag: None, - }) - .await?; + }, + ) + .await?; reject_incompatible_template_arch( &sandbox.template, resolved_template diff --git a/crates/libs/sbx/mgr/server/src/service/create_support/tests.rs b/crates/libs/sbx/mgr/server/src/service/create_support/tests.rs index 2790329b4e..425243bf02 100644 --- a/crates/libs/sbx/mgr/server/src/service/create_support/tests.rs +++ b/crates/libs/sbx/mgr/server/src/service/create_support/tests.rs @@ -476,17 +476,23 @@ fn restore_provenance_classifies_legacy_remote_metadata() { #[tokio::test] async fn existing_sandbox_cannot_resolve_a_missing_artifact_by_latest_alias() { - // With no template-manager available, this must reject locally rather than - // asking it to pick the newest build after the sandbox lost its identity. - let error = - resolve_template_paths_for_start(pplx_space_template_mgr_client::proto::TemplateRef { + // This must reject locally rather than asking template-manager to pick + // the newest build after the sandbox lost its identity. + let templates = crate::template_resolver::fake::FakeTemplateResolver::refusing( + tonic::Status::not_found("a latest-alias lookup must not be sent"), + ); + let error = resolve_template_paths_for_start( + &templates, + pplx_space_template_mgr_client::proto::TemplateRef { name: "app-runtime".into(), content_hash: None, tag: None, - }) - .await - .err() - .expect("missing identity cannot select a replacement artifact"); + }, + ) + .await + .err() + .expect("missing identity cannot select a replacement artifact"); assert_eq!(error.code(), tonic::Code::FailedPrecondition); assert!(error.message().contains("saved template hash")); + assert!(templates.requests().is_empty()); } diff --git a/crates/libs/sbx/mgr/server/src/service/grpc_core.rs b/crates/libs/sbx/mgr/server/src/service/grpc_core.rs index 559bace91a..bb018178bf 100644 --- a/crates/libs/sbx/mgr/server/src/service/grpc_core.rs +++ b/crates/libs/sbx/mgr/server/src/service/grpc_core.rs @@ -197,13 +197,15 @@ impl SandboxManagerService { // Resolve template to content-addressed paths via template-manager. // If caller pinned a content_hash, use it; otherwise resolve newest. - let resolved_template = - resolve_template_paths(pplx_space_template_mgr_client::proto::TemplateRef { + let resolved_template = resolve_template_paths( + self.template_resolver.as_ref(), + pplx_space_template_mgr_client::proto::TemplateRef { name: template.to_string(), content_hash: req.content_hash.clone(), tag: None, - }) - .await?; + }, + ) + .await?; if let Some(policy) = template_network_policy_from_proto( req.template_network_policy.as_ref(), req.template_network_policy_generation, @@ -1176,5 +1178,7 @@ fn project_template_catalog( .collect() } +#[cfg(test)] +mod create_template_lookup_tests; #[cfg(test)] mod tests; diff --git a/crates/libs/sbx/mgr/server/src/service/grpc_core/create_template_lookup_tests.rs b/crates/libs/sbx/mgr/server/src/service/grpc_core/create_template_lookup_tests.rs new file mode 100644 index 0000000000..9e28442e0e --- /dev/null +++ b/crates/libs/sbx/mgr/server/src/service/grpc_core/create_template_lookup_tests.rs @@ -0,0 +1,178 @@ +//! Create's auto-lifecycle authority checks, which run just before template +//! lookup, against a scripted template resolver. Template-manager and +//! node-manager are fakes, so a request that gets past a check stops at the +//! resolver's scripted answer. + +use super::*; +use crate::node_manager_transport::fake::FakeNodeManager; +use crate::template_resolver::TemplateLookupError; +use crate::template_resolver::fake::{FakeTemplateResolver, LookupChannel}; +use pplx_space_common_core::host::linux::identities::users; +use pplx_space_common_svc::PeerIdentity; +use pplx_space_net_common::proto::PolicyAuthority; +use pplx_space_template_mgr_client::proto::{GetTemplateRequest, TemplateRef}; + +/// The resolver's refusal, distinct from every refusal a check produces before +/// lookup. +const LOOKUP_REFUSAL: &str = "scripted template lookup refusal"; + +fn refusing_lookup() -> FakeTemplateResolver { + FakeTemplateResolver::refusing(Status::not_found(LOOKUP_REFUSAL)) +} + +fn unreachable_template_manager() -> FakeTemplateResolver { + FakeTemplateResolver::new(|| { + let invalid = "http://[".parse::().unwrap_err(); + Err(TemplateLookupError::Connect( + pplx_space_template_mgr_client::ConnectError::InvalidUri(invalid), + )) + }) +} + +#[tokio::test] +async fn create_checks_auto_lifecycle_authority_before_template_lookup() { + if !crate::service::tests::enter_route_test_process(&[ + ("EPPO_SDK_KEY", None), + ("SPACE_FLAGD_ENABLED", None), + ]) { + return; + } + let forwarded = PeerIdentity::new( + users::SPACE_NODEAPI.uid, + users::SPACE_NODEAPI.primary_group.gid, + None, + ); + let local = PeerIdentity::new(0, 0, None); + struct Case { + name: &'static str, + peer: PeerIdentity, + authority: i32, + templates: fn() -> FakeTemplateResolver, + code: tonic::Code, + message_prefix: &'static str, + message_contains: &'static str, + looks_up_template: bool, + } + let cases = [ + Case { + name: "local authority from a forwarded peer", + peer: forwarded.clone(), + authority: PolicyAuthority::Local as i32, + templates: refusing_lookup, + code: tonic::Code::PermissionDenied, + message_prefix: "authority=local requires a local operator peer", + message_contains: "", + looks_up_template: false, + }, + Case { + name: "unspecified authority", + peer: forwarded.clone(), + authority: PolicyAuthority::Unspecified as i32, + templates: refusing_lookup, + code: tonic::Code::InvalidArgument, + message_prefix: "policy authority proto value is unspecified", + message_contains: "", + looks_up_template: false, + }, + Case { + name: "unknown authority", + peer: forwarded.clone(), + authority: 99, + templates: refusing_lookup, + code: tonic::Code::InvalidArgument, + message_prefix: "policy authority proto value 99 is unknown", + message_contains: "", + looks_up_template: false, + }, + Case { + name: "cluster authority from a forwarded peer", + peer: forwarded.clone(), + authority: PolicyAuthority::Cluster as i32, + templates: refusing_lookup, + code: tonic::Code::NotFound, + message_prefix: "GetTemplate failed: ", + message_contains: LOOKUP_REFUSAL, + looks_up_template: true, + }, + Case { + name: "local authority from a local peer", + peer: local, + authority: PolicyAuthority::Local as i32, + templates: refusing_lookup, + code: tonic::Code::NotFound, + message_prefix: "GetTemplate failed: ", + message_contains: LOOKUP_REFUSAL, + looks_up_template: true, + }, + Case { + name: "accepted authority with template-manager unreachable", + peer: forwarded, + authority: PolicyAuthority::Cluster as i32, + templates: unreachable_template_manager, + code: tonic::Code::Internal, + message_prefix: "cannot connect to template-manager: ", + message_contains: "invalid gRPC endpoint URI", + looks_up_template: true, + }, + ]; + for case in cases { + let templates = Arc::new((case.templates)()); + let node_manager = FakeNodeManager::default().serve(); + let sockets = tempfile::tempdir().unwrap(); + let service = crate::service::tests::node_manager_transport_tests::isolated_service( + crate::db::test_conn(), + node_manager.connector(), + Arc::clone(&templates) as Arc, + sockets.path(), + ); + let sandbox_id = SandboxId::new().to_string(); + let mut request = Request::new(CreateRequest { + template: "seed".into(), + owner: "root@local".into(), + runtime: pplx_space_common_core::types::RUNTIME_FIRECRACKER_ALIAS.into(), + boot_mode: BootMode::Cold as i32, + sandbox_id: Some(sandbox_id.clone()), + auto_lifecycle: Some(Default::default()), + auto_lifecycle_authority: case.authority, + ..Default::default() + }); + request.extensions_mut().insert(case.peer); + + let error = service.create(request).await.expect_err(case.name); + + assert_eq!(error.code(), case.code, "{}: {error}", case.name); + assert!( + error.message().starts_with(case.message_prefix) + && error.message().contains(case.message_contains), + "{}: {error}", + case.name + ); + let lookups = templates.lookups(); + let expected_lookups = if case.looks_up_template { + vec![GetTemplateRequest { + caller_org_id: None, + template: Some(TemplateRef { + name: "seed".into(), + content_hash: None, + tag: None, + }), + vcpu: None, + memory_mib: None, + }] + } else { + Vec::new() + }; + assert_eq!(templates.requests(), expected_lookups, "{}", case.name); + // Create dials template-manager for its lookup and sends no metadata. + for lookup in &lookups { + assert_eq!(lookup.channel, LookupChannel::Dialed, "{}", case.name); + assert!(lookup.metadata.is_empty(), "{}", case.name); + } + assert!(node_manager.calls().is_empty(), "{}", case.name); + assert!( + !crate::db::sandbox_exists(&service.db_conn().unwrap(), &sandbox_id).unwrap(), + "{}", + case.name + ); + } +} diff --git a/crates/libs/sbx/mgr/server/src/service/grpc_core/tests.rs b/crates/libs/sbx/mgr/server/src/service/grpc_core/tests.rs index d7808ccfa3..04c5a61c8d 100644 --- a/crates/libs/sbx/mgr/server/src/service/grpc_core/tests.rs +++ b/crates/libs/sbx/mgr/server/src/service/grpc_core/tests.rs @@ -33,8 +33,8 @@ fn enter_isolated_test_process() -> bool { /// A service built from explicit options. `Default` options read the node /// identity from a fixed host file and probe runtimes under the fixed run /// directory; here the node identity is fixed, runtime liveness reads -/// `scratch`, and the host-control, credential-manager and node-manager -/// sockets are absent paths under `scratch`. +/// `scratch`, the host-control, credential-manager and node-manager sockets +/// are absent paths under `scratch`, and every template lookup is refused. fn isolated_service( conn: rusqlite::Connection, scratch: &std::path::Path, @@ -60,6 +60,11 @@ fn isolated_service( node_manager: crate::node_manager_transport::fake::unbound( &scratch.join("absent-node-manager.sock"), ), + template_resolver: Arc::new( + crate::template_resolver::fake::FakeTemplateResolver::refusing( + Status::unavailable("template-manager is outside this test"), + ), + ), }, ) } diff --git a/crates/libs/sbx/mgr/server/src/service/grpc_pool_wait_clone.rs b/crates/libs/sbx/mgr/server/src/service/grpc_pool_wait_clone.rs index ccdc17b1ee..7f10447b82 100644 --- a/crates/libs/sbx/mgr/server/src/service/grpc_pool_wait_clone.rs +++ b/crates/libs/sbx/mgr/server/src/service/grpc_pool_wait_clone.rs @@ -258,6 +258,7 @@ impl SandboxManagerService { let resolved_template = if req.auto_lifecycle.is_none() { Some( resolve_template_paths( + self.template_resolver.as_ref(), pplx_space_template_mgr_client::proto::TemplateRef { name: sandbox.template.clone(), content_hash: sandbox.template_hash.clone(), @@ -693,7 +694,8 @@ impl SandboxManagerService { let uuid8 = SandboxRuntimeLabel::from_sandbox_id(new_sandbox_id); // Resolve template and compute hash - let resolved_template = resolve_template_paths(template_ref).await?; + let resolved_template = + resolve_template_paths(self.template_resolver.as_ref(), template_ref).await?; reject_incompatible_template_arch( template, resolved_template diff --git a/crates/libs/sbx/mgr/server/src/service/mod.rs b/crates/libs/sbx/mgr/server/src/service/mod.rs index ce3c26f0f5..f42685dc99 100644 --- a/crates/libs/sbx/mgr/server/src/service/mod.rs +++ b/crates/libs/sbx/mgr/server/src/service/mod.rs @@ -403,6 +403,7 @@ pub struct SandboxManagerService { Arc>, network_credentials_enabled: bool, node_manager: NodeManagerConnector, + template_resolver: Arc, gateway_caller_client: pplx_space_node_mgr_client::NodeManagerGrpcClient, gateway_authority: gateway_authority::GatewayAuthority, /// S3 upload config for suspend snapshots. `None` means uploads are disabled. @@ -974,6 +975,8 @@ pub struct SandboxManagerServiceOptions { >, /// Every node-manager client the service and its subsystems use. pub node_manager: NodeManagerConnector, + /// Every template-manager `GetTemplate` the service and its subsystems ask. + pub template_resolver: Arc, } impl Default for SandboxManagerServiceOptions { @@ -993,6 +996,7 @@ impl Default for SandboxManagerServiceOptions { node_managed_by: None, route_ed25519_verification_bundle: None, node_manager: NodeManagerConnector::host(), + template_resolver: Arc::new(crate::template_resolver::HostTemplateManager), } } } @@ -1678,6 +1682,7 @@ impl SandboxManagerService { template_network_policies: TemplateNetworkPolicyCache::new( options.shutdown_token.clone(), options.task_tracker.clone(), + Arc::clone(&options.template_resolver), ), network_policy_reconcile_queue, network_policy_reconciler_started: Arc::new(AtomicBool::new(false)), @@ -1743,6 +1748,7 @@ impl SandboxManagerService { crate::retained_source::NodeManagerSourceInspector(options.node_manager.clone()), ), node_manager: options.node_manager, + template_resolver: options.template_resolver, drain_state, feature_flags, route_ed25519_verification_bundle: Arc::new(ArcSwapOption::from( diff --git a/crates/libs/sbx/mgr/server/src/service/network_policy.rs b/crates/libs/sbx/mgr/server/src/service/network_policy.rs index 1c7b614cee..9111e83f4c 100644 --- a/crates/libs/sbx/mgr/server/src/service/network_policy.rs +++ b/crates/libs/sbx/mgr/server/src/service/network_policy.rs @@ -130,20 +130,22 @@ pub(crate) fn network_policy_override_from_set_request( } pub(crate) async fn fetch_template_default_network_policy( + templates: &dyn crate::template_resolver::TemplateResolver, template: &str, ) -> Result { - let mut client = pplx_space_template_mgr_client::lazy(); - let response = match client - .get_template(pplx_space_template_mgr_client::proto::GetTemplateRequest { - caller_org_id: None, - template: Some(pplx_space_template_mgr_client::proto::TemplateRef { - name: template.to_string(), - content_hash: None, - tag: None, - }), - vcpu: None, - memory_mib: None, - }) + let response = match templates + .get_template_on_shared_channel(tonic::Request::new( + pplx_space_template_mgr_client::proto::GetTemplateRequest { + caller_org_id: None, + template: Some(pplx_space_template_mgr_client::proto::TemplateRef { + name: template.to_string(), + content_hash: None, + tag: None, + }), + vcpu: None, + memory_mib: None, + }, + )) .await { Ok(response) => response.into_inner(), @@ -175,15 +177,7 @@ pub(crate) struct TemplateNetworkPolicyCache { fetch_slots: Arc, shutdown_token: tokio_util::sync::CancellationToken, task_tracker: tokio_util::task::TaskTracker, -} - -impl Default for TemplateNetworkPolicyCache { - fn default() -> Self { - Self::new( - tokio_util::sync::CancellationToken::new(), - tokio_util::task::TaskTracker::new(), - ) - } + templates: Arc, } #[derive(Default)] @@ -216,6 +210,7 @@ impl TemplateNetworkPolicyCache { pub(crate) fn new( shutdown_token: tokio_util::sync::CancellationToken, task_tracker: tokio_util::task::TaskTracker, + templates: Arc, ) -> Self { Self { inner: Arc::default(), @@ -224,6 +219,7 @@ impl TemplateNetworkPolicyCache { )), shutdown_token, task_tracker, + templates, } } @@ -233,8 +229,9 @@ impl TemplateNetworkPolicyCache { &self, template: &str, ) -> Result { + let templates = Arc::clone(&self.templates); self.get_or_fetch_with(template, |template| async move { - fetch_template_default_network_policy(&template).await + fetch_template_default_network_policy(templates.as_ref(), &template).await }) .await } diff --git a/crates/libs/sbx/mgr/server/src/service/network_policy/tests.rs b/crates/libs/sbx/mgr/server/src/service/network_policy/tests.rs index 9063e9f87f..4c10af9fd4 100644 --- a/crates/libs/sbx/mgr/server/src/service/network_policy/tests.rs +++ b/crates/libs/sbx/mgr/server/src/service/network_policy/tests.rs @@ -1,4 +1,19 @@ use super::*; +use crate::template_resolver::fake::FakeTemplateResolver; + +/// A cache whose owner lookups are refused, with the resolver that refuses +/// them. Tests that supply their own fetch never reach it. +fn refusing_cache() -> (TemplateNetworkPolicyCache, Arc) { + let templates = Arc::new(FakeTemplateResolver::refusing(Status::unavailable( + "template-manager is outside this test", + ))); + let cache = TemplateNetworkPolicyCache::new( + tokio_util::sync::CancellationToken::new(), + tokio_util::task::TaskTracker::new(), + Arc::clone(&templates) as Arc, + ); + (cache, templates) +} #[test] fn legacy_policy_hash_is_stable() { @@ -27,7 +42,7 @@ async fn template_policy_refresh_uses_current_owner_across_local_changes_and_rec (Local, 0, Local, 0), // Old template-manager local updates reuse zero. (Local, 9, Cluster, 1), // Delete/recreate legitimately changes authority. ] { - let cache = TemplateNetworkPolicyCache::default(); + let cache = refusing_cache().0; cache .get_or_fetch_with("seed", move |_| async move { Ok(EffectiveNetworkPolicy::new( @@ -111,7 +126,7 @@ async fn template_policy_hints_are_admission_requirements_not_cache_authority() Some(tonic::Code::FailedPrecondition), ), ] { - let cache = TemplateNetworkPolicyCache::default(); + let (cache, templates) = refusing_cache(); cache .get_or_fetch_with("seed", move |_| async move { Ok(EffectiveNetworkPolicy::new( @@ -134,6 +149,8 @@ async fn template_policy_hints_are_admission_requirements_not_cache_authority() true, ); let admission_error = cache.require_hint("seed", &hint, true).await.err(); + // The hint is checked against the cached owner answer, not a new lookup. + assert!(templates.lookups().is_empty()); assert_eq!(admission_error.as_ref().map(|error| error.code()), expected); if let Some(error) = &admission_error { assert!(pplx_space_sbx_mgr_client::is_template_policy_dependency( @@ -173,7 +190,7 @@ async fn template_policy_hints_are_admission_requirements_not_cache_authority() #[tokio::test(start_paused = true)] async fn custom_policy_override_does_not_require_template_owner_progress() { - let cache = TemplateNetworkPolicyCache::default(); + let (cache, templates) = refusing_cache(); let hint = EffectiveNetworkPolicy::new( NetworkPolicyV1::deny_all(), 50, @@ -190,11 +207,12 @@ async fn custom_policy_override_does_not_require_template_owner_progress() { inherits_template.then_some(tonic::Code::DeadlineExceeded) ); } + assert!(templates.lookups().is_empty()); } #[tokio::test(start_paused = true)] async fn template_policy_invalidation_fences_fetch_callers_and_successor_claims() { - let cache = TemplateNetworkPolicyCache::default(); + let cache = refusing_cache().0; let mut old = Box::pin(cache.get_or_fetch_with("seed", |_| std::future::pending())); assert!(futures_util::poll!(&mut old).is_pending()); let old_claim = Arc::clone(&cache.inner.lock().await.pending_fetches["seed"].claim); @@ -258,7 +276,7 @@ async fn template_policy_invalidation_fences_fetch_callers_and_successor_claims( #[tokio::test(start_paused = true)] async fn template_policy_wait_deadline_includes_cache_admission() { for release_midway in [false, true] { - let cache = TemplateNetworkPolicyCache::default(); + let cache = refusing_cache().0; let guard = cache.inner.lock().await; let mut waiting = Box::pin(cache.get_or_fetch_with("seed", |_| std::future::pending())); assert!(futures_util::poll!(&mut waiting).is_pending()); @@ -280,7 +298,7 @@ async fn template_policy_wait_deadline_includes_cache_admission() { #[tokio::test(start_paused = true)] async fn template_policy_fetch_progress_survives_foreground_timeouts_and_cancellation() { - let cache = TemplateNetworkPolicyCache::default(); + let cache = refusing_cache().0; let (alive, cancelled) = tokio::sync::oneshot::channel::<()>(); let (release, released) = tokio::sync::oneshot::channel::<()>(); let mut first = Box::pin(cache.get_or_fetch_with("seed", |_| async move { @@ -324,7 +342,7 @@ async fn template_policy_fetch_progress_survives_foreground_timeouts_and_cancell #[tokio::test(start_paused = true)] async fn template_policy_fetch_capacity_and_deadline_release_claims() { for shutdown in [false, true] { - let cache = TemplateNetworkPolicyCache::default(); + let cache = refusing_cache().0; let names: Vec<_> = (0..TEMPLATE_NETWORK_POLICY_FETCH_CAPACITY) .map(|i| format!("template-{i}")) .collect(); @@ -391,7 +409,7 @@ async fn template_policy_fetch_capacity_and_deadline_release_claims() { // A ready Shared can be completed by its caller before its tracked owner // is scheduled. Those queued owner tasks still consume the same hard cap. - let cache = TemplateNetworkPolicyCache::default(); + let cache = refusing_cache().0; for i in 0..TEMPLATE_NETWORK_POLICY_FETCH_CAPACITY { let name = format!("ready-{i}"); let mut caller = Box::pin(tokio::task::unconstrained(cache.get_or_fetch_with( @@ -426,7 +444,7 @@ async fn template_policy_fetch_capacity_and_deadline_release_claims() { #[tokio::test(start_paused = true)] async fn template_policy_cache_bounds_churn_without_retiring_pending_claims() { - let cache = TemplateNetworkPolicyCache::default(); + let cache = refusing_cache().0; let mut waiting = Box::pin(cache.get_or_fetch_with("pending", |_| std::future::pending())); assert!(futures_util::poll!(&mut waiting).is_pending()); let claim = Arc::clone(&cache.inner.lock().await.pending_fetches["pending"].claim); diff --git a/crates/libs/sbx/mgr/server/src/service/recovery_cross_node.rs b/crates/libs/sbx/mgr/server/src/service/recovery_cross_node.rs index 90c182efbc..371aa3173b 100644 --- a/crates/libs/sbx/mgr/server/src/service/recovery_cross_node.rs +++ b/crates/libs/sbx/mgr/server/src/service/recovery_cross_node.rs @@ -1765,9 +1765,14 @@ impl SandboxManagerService { .ok_or_else(|| { RestoreError::Other("restore execution identity is missing".into()) })?; - let phase2_result = restore_pull_rootfs(&staged, execution_fence, &self.node_manager) - .instrument(phase2_span) - .await; + let phase2_result = restore_pull_rootfs( + &staged, + execution_fence, + &self.node_manager, + self.template_resolver.as_ref(), + ) + .instrument(phase2_span) + .await; if let Err(ref e) = phase2_result { emit_sandbox_event( &canonical_id, diff --git a/crates/libs/sbx/mgr/server/src/service/route_identity_tests.rs b/crates/libs/sbx/mgr/server/src/service/route_identity_tests.rs index 53c72fffd1..f9c9078a4f 100644 --- a/crates/libs/sbx/mgr/server/src/service/route_identity_tests.rs +++ b/crates/libs/sbx/mgr/server/src/service/route_identity_tests.rs @@ -86,6 +86,7 @@ impl RouteTestService { node_managed_by: None, route_ed25519_verification_bundle: None, node_manager: crate::node_manager_transport::NodeManagerConnector::host(), + template_resolver: Arc::new(crate::template_resolver::HostTemplateManager), }, ); Self { diff --git a/crates/libs/sbx/mgr/server/src/service/start.rs b/crates/libs/sbx/mgr/server/src/service/start.rs index fd82475fdf..8dc53f2d9b 100644 --- a/crates/libs/sbx/mgr/server/src/service/start.rs +++ b/crates/libs/sbx/mgr/server/src/service/start.rs @@ -514,6 +514,7 @@ async fn prepare_suspend_restore_dir( storage_uri, transfer_scope_id, &service.node_manager, + service.template_resolver.as_ref(), ) .await } @@ -523,6 +524,7 @@ async fn prepare_suspend_restore_dir( execution_fence, transfer_scope_id, &service.node_manager, + service.template_resolver.as_ref(), ) .await } @@ -1363,6 +1365,7 @@ impl SandboxManagerService { if !vmstate_path.exists() { let resolved_template = resolve_template_paths_for_start( + self.template_resolver.as_ref(), pplx_space_template_mgr_client::proto::TemplateRef { name: sandbox.template.clone(), content_hash: sandbox.template_hash.clone(), @@ -1645,7 +1648,7 @@ impl SandboxManagerService { "fresh.template_prepare_source", START_TEMPLATE_RESOLUTION_TIMEOUT, &sandbox.id, - template_source_for_sandbox_prepare(&sandbox), + template_source_for_sandbox_prepare(self.template_resolver.as_ref(), &sandbox), ) .await?; let selection_mode = if prepare_source.existing_artifact { @@ -1661,7 +1664,11 @@ impl SandboxManagerService { START_TEMPLATE_RESOLUTION_TIMEOUT, &sandbox.id, TemplateRuntimeDefaults::default(), - template_runtime_defaults_for_sandbox_start(&sandbox, prepare_template_dir.as_deref()), + template_runtime_defaults_for_sandbox_start( + self.template_resolver.as_ref(), + &sandbox, + prepare_template_dir.as_deref(), + ), ) .await; diff --git a/crates/libs/sbx/mgr/server/src/service/tests.rs b/crates/libs/sbx/mgr/server/src/service/tests.rs index 9270f1c5ae..cd8073a246 100644 --- a/crates/libs/sbx/mgr/server/src/service/tests.rs +++ b/crates/libs/sbx/mgr/server/src/service/tests.rs @@ -5409,10 +5409,14 @@ pub(crate) mod service_tests { let sandbox_id = SandboxId::new().to_string(); insert_orphaned_test_sandbox(&conn, &sandbox_id, SandboxState::Failed); let row = crate::db::get_sandbox(&conn, &sandbox_id).unwrap().unwrap(); + let templates = crate::template_resolver::fake::FakeTemplateResolver::refusing( + Status::not_found("orphaned sandboxes must not look up templates"), + ); - let error = template_source_for_sandbox_prepare(&row) + let error = template_source_for_sandbox_prepare(&templates, &row) .await .expect_err("orphaned sandbox without rootfs is not startable"); + assert!(templates.requests().is_empty()); assert_eq!(error.code(), tonic::Code::FailedPrecondition); assert!(error.message().contains("cannot start orphaned sandbox")); @@ -12402,6 +12406,7 @@ pub(crate) mod service_tests { node_managed_by: None, route_ed25519_verification_bundle: None, node_manager: crate::node_manager_transport::NodeManagerConnector::host(), + template_resolver: Arc::new(crate::template_resolver::HostTemplateManager), } } diff --git a/crates/libs/sbx/mgr/server/src/service/tests/node_manager_transport_tests.rs b/crates/libs/sbx/mgr/server/src/service/tests/node_manager_transport_tests.rs index 46a0de93ad..879bbdfc5d 100644 --- a/crates/libs/sbx/mgr/server/src/service/tests/node_manager_transport_tests.rs +++ b/crates/libs/sbx/mgr/server/src/service/tests/node_manager_transport_tests.rs @@ -14,10 +14,12 @@ use pplx_space_node_mgr_client::proto::{ CleanupSandboxRequest, CleanupSandboxResponse, StopProcessRequest, StopProcessResponse, }; -/// A service whose node-manager clients all come from `node_manager`. -pub(crate) fn service_with_node_manager( +/// A service whose node-manager clients all come from `node_manager` and whose +/// template lookups all go to `template_resolver`. +pub(crate) fn isolated_service( conn: rusqlite::Connection, node_manager: NodeManagerConnector, + template_resolver: Arc, sockets: &Path, ) -> SandboxManagerService { SandboxManagerService::new_with_options( @@ -35,10 +37,29 @@ pub(crate) fn service_with_node_manager( node_managed_by: None, route_ed25519_verification_bundle: None, node_manager, + template_resolver, }, ) } +/// [`isolated_service`] whose template lookups are all refused. +pub(crate) fn service_with_node_manager( + conn: rusqlite::Connection, + node_manager: NodeManagerConnector, + sockets: &Path, +) -> SandboxManagerService { + isolated_service( + conn, + node_manager, + Arc::new( + crate::template_resolver::fake::FakeTemplateResolver::refusing(Status::unavailable( + "template-manager is outside this test", + )), + ), + sockets, + ) +} + #[tokio::test] async fn fatal_kill_signal_failure_still_schedules_admitted_cleanup() { let stops = Arc::new(Mutex::new(Vec::new())); diff --git a/crates/libs/sbx/mgr/server/src/snapshot_restore.rs b/crates/libs/sbx/mgr/server/src/snapshot_restore.rs index 422d0c7ecc..83252de5b6 100644 --- a/crates/libs/sbx/mgr/server/src/snapshot_restore.rs +++ b/crates/libs/sbx/mgr/server/src/snapshot_restore.rs @@ -1257,6 +1257,7 @@ pub async fn prepare_restore_dir( execution_fence: pplx_space_node_mgr_client::execution_fence::ExecutionFence, transfer_scope_id: pplx_space_common_core::types::TransferScopeId, node_manager: &crate::node_manager_transport::NodeManagerConnector, + templates: &dyn crate::template_resolver::TemplateResolver, ) -> Result { let latest = latest_ready_snapshot(sandbox_id).await?; let staged = restore_pull_metadata_from_source( @@ -1265,7 +1266,7 @@ pub async fn prepare_restore_dir( RestoreMetadataDestination::InPlace(transfer_scope_id), ) .await?; - restore_pull_rootfs(&staged, execution_fence, node_manager).await?; + restore_pull_rootfs(&staged, execution_fence, node_manager, templates).await?; Ok(RestorePrepared { snapshot_dir: staged.snapshot_dir, storage_uri: staged.storage_uri, @@ -1280,6 +1281,7 @@ pub async fn prepare_restore_dir_from_uri( storage_uri: &str, transfer_scope_id: pplx_space_common_core::types::TransferScopeId, node_manager: &crate::node_manager_transport::NodeManagerConnector, + templates: &dyn crate::template_resolver::TemplateResolver, ) -> Result { let config = SnapshotUploadConfig::from_uri(storage_uri) .map_err(|e| RestoreError::Other(format!("parse storage_uri: {e}")))?; @@ -1292,7 +1294,7 @@ pub async fn prepare_restore_dir_from_uri( RestoreMetadataDestination::InPlace(transfer_scope_id), ) .await?; - restore_pull_rootfs(&staged, execution_fence, node_manager).await?; + restore_pull_rootfs(&staged, execution_fence, node_manager, templates).await?; Ok(RestorePrepared { snapshot_dir: staged.snapshot_dir, storage_uri: staged.storage_uri, @@ -1852,6 +1854,7 @@ pub async fn restore_pull_rootfs( stage: &RestoreMetadataStaged, execution_fence: pplx_space_node_mgr_client::execution_fence::ExecutionFence, node_manager: &crate::node_manager_transport::NodeManagerConnector, + templates: &dyn crate::template_resolver::TemplateResolver, ) -> Result<(), RestoreError> { use pplx_space_node_mgr_client::execution_fence::InstanceTarget; let target = stage.filesystem_scope_id.map_or( @@ -1898,6 +1901,7 @@ pub async fn restore_pull_rootfs( stage.transfer_scope_id, target, node_manager, + templates, ) .await } @@ -2506,7 +2510,7 @@ async fn stage_encrypted_restore( /// responsible for cleanup of the partial staging directory. #[expect( clippy::too_many_arguments, - reason = "each argument is an independent restore input; node-manager is the injected transport" + reason = "each argument is an independent restore input; the node-manager and template-manager clients are injected" )] async fn stage_btrfs_send_restore( identity: RestoreTransferIdentity<'_>, @@ -2517,6 +2521,7 @@ async fn stage_btrfs_send_restore( transfer_scope_id: Option, target: pplx_space_node_mgr_client::execution_fence::InstanceTarget, node_manager: &crate::node_manager_transport::NodeManagerConnector, + templates: &dyn crate::template_resolver::TemplateResolver, ) -> Result<(), RestoreError> { let sandbox_id = identity.sandbox_id; let sidecar_bytes = std::fs::read(sidecar_path).map_err(|e| { @@ -2568,7 +2573,7 @@ async fn stage_btrfs_send_restore( } run_btrfs_staging( Box::pin(async { - ensure_local_template(sandbox_id, template_name, content_hash) + ensure_local_template(templates, sandbox_id, template_name, content_hash) .await .map(drop) }), @@ -2654,31 +2659,35 @@ async fn run_btrfs_staging<'a>( /// tmplmgr — same RPC the sandbox-create path uses. Returns the local /// directory the template lives at. async fn ensure_local_template( + templates: &dyn crate::template_resolver::TemplateResolver, sandbox_id: &str, template_name: &str, content_hash: &str, ) -> Result { + use crate::template_resolver::TemplateLookupError; async { - let mut client = pplx_space_template_mgr_client::connect() + let resp = templates + .get_template(tonic::Request::new( + pplx_space_template_mgr_client::proto::GetTemplateRequest { + caller_org_id: None, + template: Some(pplx_space_template_mgr_client::proto::TemplateRef { + name: template_name.to_string(), + content_hash: Some(content_hash.to_string()), + tag: None, + }), + vcpu: None, + memory_mib: None, + }, + )) .await - .map_err(|e| { - RestoreError::Other(format!( + .map_err(|error| match error { + TemplateLookupError::Connect(e) => RestoreError::Other(format!( "tmplmgr connect for GetTemplate({template_name}): {e}" - )) - })?; - let resp = client - .get_template(pplx_space_template_mgr_client::proto::GetTemplateRequest { - caller_org_id: None, - template: Some(pplx_space_template_mgr_client::proto::TemplateRef { - name: template_name.to_string(), - content_hash: Some(content_hash.to_string()), - tag: None, - }), - vcpu: None, - memory_mib: None, - }) - .await - .map_err(|error| template_pull_error(template_name, content_hash, error))? + )), + TemplateLookupError::Rpc(error) => { + template_pull_error(template_name, content_hash, error) + } + })? .into_inner(); if resp.template_dir.is_empty() { return Err(RestoreError::Other(format!( diff --git a/crates/libs/sbx/mgr/server/src/template_resolver.rs b/crates/libs/sbx/mgr/server/src/template_resolver.rs new file mode 100644 index 0000000000..36d1ed34c9 --- /dev/null +++ b/crates/libs/sbx/mgr/server/src/template_resolver.rs @@ -0,0 +1,67 @@ +//! How sandbox-manager looks up template builds. +//! +//! Create, start, wake, pool cloning, restore staging and the template +//! network-policy cache ask template-manager's `GetTemplate` through one +//! [`TemplateResolver`], carried in the service options. Production uses +//! [`HostTemplateManager`], which reaches template-manager's fixed Unix socket +//! through the `pplx-space-template-mgr-client` constructors exactly as each +//! caller did directly. Callers keep their own error mapping. Tests substitute +//! a resolver that records requests and returns scripted answers. + +use pplx_space_template_mgr_client::ConnectError; +use pplx_space_template_mgr_client::proto::{GetTemplateRequest, GetTemplateResponse}; + +/// Why a [`TemplateResolver::get_template`] call failed. The connect and RPC +/// failures stay distinct because callers report them differently. +#[derive(Debug, thiserror::Error)] +pub enum TemplateLookupError { + #[error(transparent)] + Connect(#[from] ConnectError), + #[error(transparent)] + Rpc(#[from] tonic::Status), +} + +/// Requests and responses cross the resolver whole, so callers can attach and +/// read gRPC metadata as they could on the client directly. +#[async_trait::async_trait] +pub trait TemplateResolver: Send + Sync { + /// `GetTemplate` on a client dialed for this call under the default + /// connect-retry policy. + async fn get_template( + &self, + request: tonic::Request, + ) -> Result, TemplateLookupError>; + + /// `GetTemplate` on the process-wide reconnecting channel, which dials on + /// first use. Connection failures surface as the RPC status. + async fn get_template_on_shared_channel( + &self, + request: tonic::Request, + ) -> Result, tonic::Status>; +} + +/// Template-manager at [`pplx_space_template_mgr_client::SOCKET_PATH`]. +pub struct HostTemplateManager; + +#[async_trait::async_trait] +impl TemplateResolver for HostTemplateManager { + async fn get_template( + &self, + request: tonic::Request, + ) -> Result, TemplateLookupError> { + let mut client = pplx_space_template_mgr_client::connect().await?; + Ok(client.get_template(request).await?) + } + + async fn get_template_on_shared_channel( + &self, + request: tonic::Request, + ) -> Result, tonic::Status> { + pplx_space_template_mgr_client::lazy() + .get_template(request) + .await + } +} + +#[cfg(test)] +pub(crate) mod fake; diff --git a/crates/libs/sbx/mgr/server/src/template_resolver/fake.rs b/crates/libs/sbx/mgr/server/src/template_resolver/fake.rs new file mode 100644 index 0000000000..1fed44511c --- /dev/null +++ b/crates/libs/sbx/mgr/server/src/template_resolver/fake.rs @@ -0,0 +1,108 @@ +//! A scripted template resolver for tests. + +use std::sync::{Mutex, PoisonError}; + +use pplx_space_template_mgr_client::proto::{GetTemplateRequest, GetTemplateResponse}; + +use super::{TemplateLookupError, TemplateResolver}; + +type Answer = + dyn Fn() -> Result, TemplateLookupError> + Send + Sync; + +/// Which [`TemplateResolver`] method a lookup used. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum LookupChannel { + /// [`TemplateResolver::get_template`]: a client dialed for the call. + Dialed, + /// [`TemplateResolver::get_template_on_shared_channel`]. + Shared, +} + +/// One lookup as the resolver received it. +#[derive(Clone, Debug)] +pub(crate) struct Lookup { + pub(crate) channel: LookupChannel, + pub(crate) metadata: tonic::metadata::MetadataMap, + pub(crate) request: GetTemplateRequest, +} + +/// Answers every lookup with a scripted result and records each lookup in +/// arrival order. A scripted connect failure reaches a shared-channel lookup +/// as `Unavailable`, as a lazy channel reports it. +pub(crate) struct FakeTemplateResolver { + answer: Box, + lookups: Mutex>, +} + +impl FakeTemplateResolver { + pub(crate) fn new( + answer: impl Fn() -> Result, TemplateLookupError> + + Send + + Sync + + 'static, + ) -> Self { + Self { + answer: Box::new(answer), + lookups: Mutex::new(Vec::new()), + } + } + + /// Refuse every lookup with `status`. + pub(crate) fn refusing(status: tonic::Status) -> Self { + Self::new(move || Err(TemplateLookupError::Rpc(status.clone()))) + } + + pub(crate) fn lookups(&self) -> Vec { + self.lookups + .lock() + .unwrap_or_else(PoisonError::into_inner) + .clone() + } + + pub(crate) fn requests(&self) -> Vec { + self.lookups() + .into_iter() + .map(|lookup| lookup.request) + .collect() + } + + fn answer( + &self, + channel: LookupChannel, + request: tonic::Request, + ) -> Result, TemplateLookupError> { + let (metadata, _, request) = request.into_parts(); + self.lookups + .lock() + .unwrap_or_else(PoisonError::into_inner) + .push(Lookup { + channel, + metadata, + request, + }); + (self.answer)() + } +} + +#[async_trait::async_trait] +impl TemplateResolver for FakeTemplateResolver { + async fn get_template( + &self, + request: tonic::Request, + ) -> Result, TemplateLookupError> { + self.answer(LookupChannel::Dialed, request) + } + + async fn get_template_on_shared_channel( + &self, + request: tonic::Request, + ) -> Result, tonic::Status> { + self.answer(LookupChannel::Shared, request) + .map_err(|error| match error { + TemplateLookupError::Rpc(status) => status, + TemplateLookupError::Connect(error) => { + tonic::Status::unavailable(error.to_string()) + } + }) + } +}