Line data    Source code

       1              : use std::{collections::HashMap, net::SocketAddr};
       2              : 
       3              : use anyhow::Context;
       4              : use arc_swap::ArcSwapOption;
       5              : 
       6              : use crate::{
       7              :     compute::ConnCfg,
       8              :     console::{
       9              :         messages::{ColdStartInfo, EndpointJwksResponse, MetricsAuxInfo},
      10              :         NodeInfo,
      11              :     },
      12              :     context::RequestMonitoring,
      13              :     intern::{BranchIdInt, BranchIdTag, EndpointIdTag, InternId, ProjectIdInt, ProjectIdTag},
      14              :     EndpointId, RoleName,
      15              : };
      16              : 
      17              : use super::jwt::{AuthRule, FetchAuthRules, JwkCache};
      18              : 
      19              : pub struct LocalBackend {
      20              :     pub(crate) jwks_cache: JwkCache,
      21              :     pub(crate) node_info: NodeInfo,
      22              : }
      23              : 
      24              : impl LocalBackend {
      25            0 :     pub fn new(postgres_addr: SocketAddr) -> Self {
      26            0 :         LocalBackend {
      27            0 :             jwks_cache: JwkCache::default(),
      28            0 :             node_info: NodeInfo {
      29            0 :                 config: {
      30            0 :                     let mut cfg = ConnCfg::new();
      31            0 :                     cfg.host(&postgres_addr.ip().to_string());
      32            0 :                     cfg.port(postgres_addr.port());
      33            0 :                     cfg
      34            0 :                 },
      35            0 :                 // TODO(conrad): make this better reflect compute info rather than endpoint info.
      36            0 :                 aux: MetricsAuxInfo {
      37            0 :                     endpoint_id: EndpointIdTag::get_interner().get_or_intern("local"),
      38            0 :                     project_id: ProjectIdTag::get_interner().get_or_intern("local"),
      39            0 :                     branch_id: BranchIdTag::get_interner().get_or_intern("local"),
      40            0 :                     cold_start_info: ColdStartInfo::WarmCached,
      41            0 :                 },
      42            0 :                 allow_self_signed_compute: false,
      43            0 :             },
      44            0 :         }
      45            0 :     }
      46              : }
      47              : 
      48              : #[derive(Clone, Copy)]
      49              : pub(crate) struct StaticAuthRules;
      50              : 
      51              : pub static JWKS_ROLE_MAP: ArcSwapOption<JwksRoleSettings> = ArcSwapOption::const_empty();
      52              : 
      53              : #[derive(Debug, Clone)]
      54              : pub struct JwksRoleSettings {
      55              :     pub roles: HashMap<RoleName, EndpointJwksResponse>,
      56              :     pub project_id: ProjectIdInt,
      57              :     pub branch_id: BranchIdInt,
      58              : }
      59              : 
      60              : impl FetchAuthRules for StaticAuthRules {
      61            0 :     async fn fetch_auth_rules(
      62            0 :         &self,
      63            0 :         _ctx: &RequestMonitoring,
      64            0 :         _endpoint: EndpointId,
      65            0 :         role_name: RoleName,
      66            0 :     ) -> anyhow::Result<Vec<AuthRule>> {
      67            0 :         let mappings = JWKS_ROLE_MAP.load();
      68            0 :         let role_mappings = mappings
      69            0 :             .as_deref()
      70            0 :             .and_then(|m| m.roles.get(&role_name))
      71            0 :             .context("JWKs settings for this role were not configured")?;
      72            0 :         let mut rules = vec![];
      73            0 :         for setting in &role_mappings.jwks {
      74            0 :             rules.push(AuthRule {
      75            0 :                 id: setting.id.clone(),
      76            0 :                 jwks_url: setting.jwks_url.clone(),
      77            0 :                 audience: setting.jwt_audience.clone(),
      78            0 :             });
      79            0 :         }
      80              : 
      81            0 :         Ok(rules)
      82            0 :     }
      83              : }