From ef1b1109e7b6dc24cf44a78c8835b0d6748b05a7 Mon Sep 17 00:00:00 2001 From: CooooolFrog Date: Thu, 13 Aug 2026 11:28:26 +0800 Subject: [PATCH] feat(auth): add JWT authentication for gRPC clients Signed-off-by: CooooolFrog --- Cargo.lock | 23 + Cargo.toml | 5 + ci/Dockerfile | 3 + ci/Dockerfile.debug | 3 + ci/Dockerfile.dfinit | 3 + ci/Dockerfile.profiling | 3 + dragonfly-client-auth/Cargo.toml | 26 + dragonfly-client-auth/src/lib.rs | 1073 ++++++++++++++++++ dragonfly-client-auth/testdata/interop.json | 11 + dragonfly-client-config/Cargo.toml | 1 + dragonfly-client-config/src/dfdaemon.rs | 5 + dragonfly-client-metric/Cargo.toml | 1 + dragonfly-client-metric/src/lib.rs | 11 + dragonfly-client-request/Cargo.toml | 1 + dragonfly-client-request/src/lib.rs | 36 +- dragonfly-client-request/src/selector.rs | 12 +- dragonfly-client/Cargo.toml | 1 + dragonfly-client/src/bin/dfctl/task.rs | 62 +- dragonfly-client/src/dynconfig/mod.rs | 14 +- dragonfly-client/src/grpc/dfdaemon_upload.rs | 26 +- dragonfly-client/src/grpc/interceptor.rs | 61 +- dragonfly-client/src/grpc/manager.rs | 21 +- dragonfly-client/src/grpc/scheduler.rs | 46 +- 23 files changed, 1401 insertions(+), 47 deletions(-) create mode 100644 dragonfly-client-auth/Cargo.toml create mode 100644 dragonfly-client-auth/src/lib.rs create mode 100644 dragonfly-client-auth/testdata/interop.json diff --git a/Cargo.lock b/Cargo.lock index 3ca8c0a6f..f0ca41bc4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1299,6 +1299,7 @@ dependencies = [ "console-subscriber", "dashmap", "dragonfly-api", + "dragonfly-client-auth", "dragonfly-client-backend", "dragonfly-client-config", "dragonfly-client-core", @@ -1363,6 +1364,23 @@ dependencies = [ "uuid", ] +[[package]] +name = "dragonfly-client-auth" +version = "1.4.9" +dependencies = [ + "base64 0.22.1", + "humantime-serde", + "jsonwebtoken 10.3.0", + "prometheus", + "serde", + "serde_json", + "serde_yaml", + "tempfile", + "thiserror 2.0.18", + "tonic", + "validator", +] + [[package]] name = "dragonfly-client-backend" version = "1.4.9" @@ -1408,6 +1426,7 @@ dependencies = [ "bytesize", "bytesize-serde", "clap", + "dragonfly-client-auth", "dragonfly-client-core", "dragonfly-client-util", "home", @@ -1473,6 +1492,7 @@ version = "1.4.9" dependencies = [ "bytes", "dragonfly-api", + "dragonfly-client-auth", "dragonfly-client-config", "dragonfly-client-util", "fs2", @@ -1491,6 +1511,7 @@ dependencies = [ "async-trait", "bytes", "dragonfly-api", + "dragonfly-client-auth", "dragonfly-client-core", "dragonfly-client-util", "futures", @@ -2813,9 +2834,11 @@ dependencies = [ "base64 0.22.1", "getrandom 0.2.12", "js-sys", + "pem", "serde", "serde_json", "signature", + "simple_asn1", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 6769c260c..5c2c63d5d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -2,6 +2,7 @@ resolver = "2" members = [ "dragonfly-client", + "dragonfly-client-auth", "dragonfly-client-backend", "dragonfly-client-backend/examples/plugin", "dragonfly-client-config", @@ -27,6 +28,7 @@ keywords = ["container", "docker-image", "dragonfly", "dragonfly-client", "p2p"] [workspace.dependencies] anyhow = "1.0.102" async-trait = "0.1" +base64 = "0.22.1" bytes = "1.11" bytesize = { version = "1.3.3", features = ["serde"] } bytesize-serde = "0.2.1" @@ -37,6 +39,7 @@ crc32fast = "1.5.0" dashmap = "6.1.0" dragonfly-api = "=2.2.32" dragonfly-client = { path = "dragonfly-client", version = "1.4.9" } +dragonfly-client-auth = { path = "dragonfly-client-auth", version = "1.4.9" } dragonfly-client-backend = { path = "dragonfly-client-backend", version = "1.4.9" } dragonfly-client-config = { path = "dragonfly-client-config", version = "1.4.9" } dragonfly-client-core = { path = "dragonfly-client-core", version = "1.4.9" } @@ -71,6 +74,7 @@ hyper-util = { version = "0.1.20", features = [ "server-auto", "tokio", ] } +jsonwebtoken = { version = "10.3.0", features = ["aws_lc_rs"] } leaky-bucket = "1.1.2" local-ip-address = "0.6.5" lru = "0.18.1" @@ -91,6 +95,7 @@ opendal = { version = "0.55.0", default-features = false, features = [ "services-webhdfs", ] } percent-encoding = "2.3.2" +prometheus = { version = "0.13.4", features = ["process"] } prost-wkt-types = "0.7" rcgen = { version = "0.12.1", features = ["x509-parser"] } regex = "1.12.3" diff --git a/ci/Dockerfile b/ci/Dockerfile index b331cc1f1..16283ccff 100644 --- a/ci/Dockerfile +++ b/ci/Dockerfile @@ -16,6 +16,9 @@ COPY dragonfly-client/src ./dragonfly-client/src COPY dragonfly-client-core/Cargo.toml ./dragonfly-client-core/Cargo.toml COPY dragonfly-client-core/src ./dragonfly-client-core/src +COPY dragonfly-client-auth/Cargo.toml ./dragonfly-client-auth/Cargo.toml +COPY dragonfly-client-auth/src ./dragonfly-client-auth/src + COPY dragonfly-client-config/Cargo.toml ./dragonfly-client-config/Cargo.toml COPY dragonfly-client-config/src ./dragonfly-client-config/src COPY dragonfly-client-config/build.rs ./dragonfly-client-config/build.rs diff --git a/ci/Dockerfile.debug b/ci/Dockerfile.debug index e175303f6..46be0bdbd 100644 --- a/ci/Dockerfile.debug +++ b/ci/Dockerfile.debug @@ -16,6 +16,9 @@ COPY dragonfly-client/src ./dragonfly-client/src COPY dragonfly-client-core/Cargo.toml ./dragonfly-client-core/Cargo.toml COPY dragonfly-client-core/src ./dragonfly-client-core/src +COPY dragonfly-client-auth/Cargo.toml ./dragonfly-client-auth/Cargo.toml +COPY dragonfly-client-auth/src ./dragonfly-client-auth/src + COPY dragonfly-client-config/Cargo.toml ./dragonfly-client-config/Cargo.toml COPY dragonfly-client-config/src ./dragonfly-client-config/src COPY dragonfly-client-config/build.rs ./dragonfly-client-config/build.rs diff --git a/ci/Dockerfile.dfinit b/ci/Dockerfile.dfinit index 30509eb0b..700d2f336 100644 --- a/ci/Dockerfile.dfinit +++ b/ci/Dockerfile.dfinit @@ -16,6 +16,9 @@ COPY dragonfly-client/src ./dragonfly-client/src COPY dragonfly-client-core/Cargo.toml ./dragonfly-client-core/Cargo.toml COPY dragonfly-client-core/src ./dragonfly-client-core/src +COPY dragonfly-client-auth/Cargo.toml ./dragonfly-client-auth/Cargo.toml +COPY dragonfly-client-auth/src ./dragonfly-client-auth/src + COPY dragonfly-client-config/Cargo.toml ./dragonfly-client-config/Cargo.toml COPY dragonfly-client-config/src ./dragonfly-client-config/src COPY dragonfly-client-config/build.rs ./dragonfly-client-config/build.rs diff --git a/ci/Dockerfile.profiling b/ci/Dockerfile.profiling index 99951c975..5dd63aad9 100644 --- a/ci/Dockerfile.profiling +++ b/ci/Dockerfile.profiling @@ -16,6 +16,9 @@ COPY dragonfly-client/src ./dragonfly-client/src COPY dragonfly-client-core/Cargo.toml ./dragonfly-client-core/Cargo.toml COPY dragonfly-client-core/src ./dragonfly-client-core/src +COPY dragonfly-client-auth/Cargo.toml ./dragonfly-client-auth/Cargo.toml +COPY dragonfly-client-auth/src ./dragonfly-client-auth/src + COPY dragonfly-client-config/Cargo.toml ./dragonfly-client-config/Cargo.toml COPY dragonfly-client-config/src ./dragonfly-client-config/src COPY dragonfly-client-config/build.rs ./dragonfly-client-config/build.rs diff --git a/dragonfly-client-auth/Cargo.toml b/dragonfly-client-auth/Cargo.toml new file mode 100644 index 000000000..c84321c94 --- /dev/null +++ b/dragonfly-client-auth/Cargo.toml @@ -0,0 +1,26 @@ +[package] +name = "dragonfly-client-auth" +description = "Inter-component gRPC authentication for the Dragonfly client" +version.workspace = true +authors.workspace = true +homepage.workspace = true +repository.workspace = true +keywords.workspace = true +license.workspace = true +readme.workspace = true +edition.workspace = true + +[dependencies] +base64.workspace = true +humantime-serde.workspace = true +jsonwebtoken.workspace = true +prometheus.workspace = true +serde.workspace = true +thiserror.workspace = true +tonic.workspace = true +validator.workspace = true + +[dev-dependencies] +serde_json.workspace = true +serde_yaml.workspace = true +tempfile.workspace = true diff --git a/dragonfly-client-auth/src/lib.rs b/dragonfly-client-auth/src/lib.rs new file mode 100644 index 000000000..2b385d7bb --- /dev/null +++ b/dragonfly-client-auth/src/lib.rs @@ -0,0 +1,1073 @@ +/* + * Copyright 2026 The Dragonfly Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +use base64::{engine::general_purpose::STANDARD, Engine as _}; +use jsonwebtoken::{ + decode, decode_header, encode, Algorithm, DecodingKey, EncodingKey, Header, Validation, +}; +use prometheus::{IntCounterVec, Opts}; +use serde::{Deserialize, Serialize}; +use std::borrow::Cow; +use std::collections::HashMap; +use std::fmt; +use std::fs; +use std::path::PathBuf; +use std::sync::{Arc, LazyLock, Mutex, OnceLock}; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; +use thiserror::Error; +use tonic::metadata::{Ascii, MetadataMap, MetadataValue}; +use tonic::service::Interceptor; +use tonic::{Request, Status}; +use validator::{Validate, ValidationError, ValidationErrors}; + +/// JWT type used for Dragonfly inter-component gRPC authentication. +pub const TOKEN_TYPE: &str = "dragonfly-grpc+jwt"; + +/// Audience expected by Manager gRPC servers. +pub const AUDIENCE_MANAGER: &str = "urn:dragonfly:grpc:manager"; + +/// Audience expected by Scheduler gRPC servers. +pub const AUDIENCE_SCHEDULER: &str = "urn:dragonfly:grpc:scheduler"; + +/// Audience expected by dfdaemon gRPC servers. +pub const AUDIENCE_DFDAEMON: &str = "urn:dragonfly:grpc:dfdaemon"; + +const DEFAULT_ISSUER: &str = "dragonfly-internal"; +const AUTHORIZATION_METADATA_KEY: &str = "authorization"; +const BEARER_SCHEME: &str = "Bearer"; +const MAX_CREDENTIAL_LENGTH: usize = 4 * 1024; +const MINIMUM_KEY_SIZE: usize = 32; +const AUTHENTICATION_ERROR_MESSAGE: &str = "invalid authentication credentials"; + +const REASON_NONE: &str = "none"; +const REASON_MISSING: &str = "missing"; +const REASON_MALFORMED: &str = "malformed"; +const REASON_UNSUPPORTED_ALG: &str = "unsupported_alg"; +const REASON_INVALID_TYPE: &str = "invalid_type"; +const REASON_UNKNOWN_KEY_ID: &str = "unknown_kid"; +const REASON_INVALID_SIGNATURE: &str = "invalid_signature"; +const REASON_INVALID_ISSUER: &str = "invalid_issuer"; +const REASON_INVALID_AUDIENCE: &str = "invalid_audience"; +const REASON_EXPIRED: &str = "expired"; +const REASON_INVALID_ISSUED_AT: &str = "invalid_iat"; +const REASON_TTL_EXCEEDED: &str = "ttl_exceeded"; + +/// Server authentication metrics shared with dragonfly-client-metric. +pub static GRPC_AUTH_REQUESTS: LazyLock = LazyLock::new(|| { + IntCounterVec::new( + Opts::new( + "requests_total", + "Total number of inter-component gRPC authentication attempts.", + ) + .namespace("dragonfly") + .subsystem("grpc_auth"), + &["audience", "mode", "result", "reason"], + ) + .expect("metric can be created") +}); + +/// Client token generation and cache metrics shared with dragonfly-client-metric. +pub static GRPC_AUTH_CLIENT_TOKEN_EVENTS: LazyLock = LazyLock::new(|| { + IntCounterVec::new( + Opts::new( + "client_token_events_total", + "Total number of inter-component gRPC client JWT cache and generation events.", + ) + .namespace("dragonfly") + .subsystem("grpc_auth"), + &["audience", "event"], + ) + .expect("metric can be created") +}); + +fn default_token_ttl() -> Duration { + Duration::from_secs(10 * 60) +} + +fn default_max_token_ttl() -> Duration { + Duration::from_secs(15 * 60) +} + +fn default_clock_skew() -> Duration { + Duration::from_secs(30) +} + +fn default_refresh_before() -> Duration { + Duration::from_secs(60) +} + +fn default_issuer() -> String { + DEFAULT_ISSUER.to_string() +} + +/// Inter-component gRPC authentication mode. +#[derive(Debug, Clone, Copy, Default, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "lowercase")] +pub enum Mode { + /// Do not send or verify JWTs. + #[default] + Disabled, + + /// Send JWTs and allow missing credentials during a rolling upgrade. + Permissive, + + /// Require a valid JWT on every protected RPC. + Required, +} + +impl Mode { + /// Returns the configuration representation of the mode. + pub fn as_str(self) -> &'static str { + match self { + Self::Disabled => "disabled", + Self::Permissive => "permissive", + Self::Required => "required", + } + } +} + +/// A shared HMAC key file. +#[derive(Debug, Clone, Default, Deserialize)] +#[serde(default, rename_all = "camelCase")] +pub struct KeyConfig { + /// Key identifier serialized in the JWT `kid` header. + pub id: String, + + /// File containing an RFC 4648 standard Base64-encoded key. + pub secret_file: PathBuf, +} + +/// JWT signing and validation configuration. +#[derive(Debug, Clone, Deserialize)] +#[serde(default, rename_all = "camelCase")] +pub struct JwtConfig { + /// Exact issuer accepted by the verifier. + #[serde(default = "default_issuer")] + pub issuer: String, + + /// Lifetime of generated tokens. + #[serde(default = "default_token_ttl", with = "humantime_serde")] + pub token_ttl: Duration, + + /// Maximum lifetime accepted by the verifier. + #[serde(default = "default_max_token_ttl", with = "humantime_serde")] + pub max_token_ttl: Duration, + + /// Clock skew accepted by the verifier. + #[serde(default = "default_clock_skew", with = "humantime_serde")] + pub clock_skew: Duration, + + /// Duration before expiration at which a cached token is refreshed. + #[serde(default = "default_refresh_before", with = "humantime_serde")] + pub refresh_before: Duration, + + /// Key used to sign new tokens. + #[serde(rename = "activeKeyID")] + pub active_key_id: String, + + /// Keyring trusted by the verifier. + pub keys: Vec, +} + +impl Default for JwtConfig { + fn default() -> Self { + Self { + issuer: default_issuer(), + token_ttl: default_token_ttl(), + max_token_ttl: default_max_token_ttl(), + clock_skew: default_clock_skew(), + refresh_before: default_refresh_before(), + active_key_id: String::new(), + keys: Vec::new(), + } + } +} + +#[derive(Clone)] +struct CachedToken { + value: String, + expires_at: u64, +} + +#[derive(Default)] +struct Runtime { + keys: OnceLock>>>, + tokens: Mutex>, +} + +/// Inter-component gRPC authentication configuration and shared runtime cache. +#[derive(Clone, Deserialize)] +#[serde(default, rename_all = "camelCase")] +pub struct GrpcAuth { + /// Client and server authentication mode. + pub mode: Mode, + + /// Refuse to attach bearer credentials to plaintext transports. + pub require_transport_security: bool, + + /// JWT configuration. + pub jwt: JwtConfig, + + #[serde(skip)] + runtime: Arc, +} + +impl Default for GrpcAuth { + fn default() -> Self { + Self { + mode: Mode::Disabled, + require_transport_security: false, + jwt: JwtConfig::default(), + runtime: Arc::new(Runtime::default()), + } + } +} + +impl fmt::Debug for GrpcAuth { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("GrpcAuth") + .field("mode", &self.mode) + .field( + "require_transport_security", + &self.require_transport_security, + ) + .field("jwt", &self.jwt) + .finish_non_exhaustive() + } +} + +impl Validate for GrpcAuth { + fn validate(&self) -> Result<(), ValidationErrors> { + self.ensure_keyring().map(|_| ()).map_err(|error| { + let mut validation_error = ValidationError::new("grpc_auth"); + validation_error.message = Some(Cow::Owned(error.to_string())); + let mut validation_errors = ValidationErrors::new(); + validation_errors.add("grpcAuth", validation_error); + validation_errors + }) + } +} + +/// Authentication decision for an incoming RPC. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum AuthenticationOutcome { + /// Authentication is disabled. + Disabled, + + /// The request carried a valid JWT. + Authenticated, + + /// Missing credentials were allowed in permissive mode. + PermissiveMissing, +} + +/// A bounded authentication failure. +#[derive(Debug, Error)] +#[error("{detail}")] +pub struct AuthError { + reason: &'static str, + detail: String, +} + +impl AuthError { + fn new(reason: &'static str, detail: impl Into) -> Self { + Self { + reason, + detail: detail.into(), + } + } + + /// Returns the bounded metric reason. + pub fn reason(&self) -> &'static str { + self.reason + } +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +struct Claims { + iss: String, + aud: String, + iat: i64, + exp: i64, +} + +impl GrpcAuth { + /// Returns whether authentication is enabled. + pub fn enabled(&self) -> bool { + self.mode != Mode::Disabled + } + + /// Validates that credentials may be sent on the selected transport. + pub fn ensure_transport_security(&self, secure: bool) -> Result<(), AuthError> { + if self.enabled() && self.require_transport_security && !secure { + return Err(AuthError::new( + REASON_MALFORMED, + "gRPC authentication requires transport security", + )); + } + + Ok(()) + } + + /// Generates or reuses a JWT for the target audience. + pub fn token(&self, audience: &str) -> Result, AuthError> { + self.token_at(audience, unix_time()?) + } + + fn token_at(&self, audience: &str, now: u64) -> Result, AuthError> { + if !self.enabled() { + return Ok(None); + } + + validate_audience(audience)?; + let keys = self.ensure_keyring()?; + let cache_key = format!("{}\0{}", audience, self.jwt.active_key_id); + let mut tokens = + self.runtime.tokens.lock().map_err(|_| { + AuthError::new(REASON_MALFORMED, "JWT token cache lock is poisoned") + })?; + + if let Some(cached) = tokens.get(&cache_key) { + if cached.expires_at.saturating_sub(now) > self.jwt.refresh_before.as_secs() { + GRPC_AUTH_CLIENT_TOKEN_EVENTS + .with_label_values(&[audience, "cache_hit"]) + .inc(); + return Ok(Some(cached.value.clone())); + } + } + + GRPC_AUTH_CLIENT_TOKEN_EVENTS + .with_label_values(&[audience, "cache_miss"]) + .inc(); + let key = keys.get(&self.jwt.active_key_id).ok_or_else(|| { + AuthError::new(REASON_UNKNOWN_KEY_ID, "active JWT key is unavailable") + })?; + let expires_at = now + .checked_add(self.jwt.token_ttl.as_secs()) + .ok_or_else(|| AuthError::new(REASON_TTL_EXCEEDED, "JWT expiration time overflows"))?; + let issued_at = i64::try_from(now).map_err(|_| { + AuthError::new(REASON_INVALID_ISSUED_AT, "JWT issued-at time overflows") + })?; + let expiration = i64::try_from(expires_at) + .map_err(|_| AuthError::new(REASON_TTL_EXCEEDED, "JWT expiration time overflows"))?; + let claims = Claims { + iss: self.jwt.issuer.clone(), + aud: audience.to_string(), + iat: issued_at, + exp: expiration, + }; + let mut header = Header::new(Algorithm::HS256); + header.typ = Some(TOKEN_TYPE.to_string()); + header.kid = Some(self.jwt.active_key_id.clone()); + let token = encode(&header, &claims, &EncodingKey::from_secret(key)).map_err(|error| { + AuthError::new(REASON_MALFORMED, format!("failed to sign JWT: {error}")) + })?; + + tokens.insert( + cache_key, + CachedToken { + value: token.clone(), + expires_at, + }, + ); + GRPC_AUTH_CLIENT_TOKEN_EVENTS + .with_label_values(&[audience, "generated"]) + .inc(); + Ok(Some(token)) + } + + /// Authenticates incoming gRPC metadata for the target audience. + pub fn authenticate( + &self, + metadata: &MetadataMap, + audience: &str, + ) -> Result { + self.authenticate_at(metadata, audience, unix_time()?) + } + + fn authenticate_at( + &self, + metadata: &MetadataMap, + audience: &str, + now: u64, + ) -> Result { + if !self.enabled() { + return Ok(AuthenticationOutcome::Disabled); + } + + validate_audience(audience)?; + let values: Vec<_> = metadata + .get_all(AUTHORIZATION_METADATA_KEY) + .iter() + .collect(); + if values.is_empty() { + if self.mode == Mode::Permissive { + return Ok(AuthenticationOutcome::PermissiveMissing); + } + + return Err(AuthError::new( + REASON_MISSING, + "authorization metadata is missing", + )); + } + + if values.len() != 1 { + return Err(AuthError::new( + REASON_MALFORMED, + "multiple authorization metadata values", + )); + } + + let credential = values[0] + .to_str() + .map_err(|_| AuthError::new(REASON_MALFORMED, "authorization metadata is not ASCII"))?; + if credential.len() > MAX_CREDENTIAL_LENGTH { + return Err(AuthError::new( + REASON_MALFORMED, + "authorization metadata is too large", + )); + } + + let parts: Vec<_> = credential.split_whitespace().collect(); + if parts.len() != 2 || !parts[0].eq_ignore_ascii_case(BEARER_SCHEME) || parts[1].is_empty() + { + return Err(AuthError::new( + REASON_MALFORMED, + "malformed bearer credential", + )); + } + + self.verify_at(parts[1], audience, now)?; + Ok(AuthenticationOutcome::Authenticated) + } + + fn verify_at(&self, raw_token: &str, audience: &str, now: u64) -> Result<(), AuthError> { + let header = decode_header(raw_token).map_err(|error| { + AuthError::new( + REASON_MALFORMED, + format!("failed to parse JWT header: {error}"), + ) + })?; + if header.alg != Algorithm::HS256 { + return Err(AuthError::new( + REASON_UNSUPPORTED_ALG, + "JWT algorithm is not HS256", + )); + } + + if header.typ.as_deref() != Some(TOKEN_TYPE) { + return Err(AuthError::new(REASON_INVALID_TYPE, "JWT type is invalid")); + } + + let key_id = header + .kid + .as_deref() + .filter(|key_id| !key_id.is_empty()) + .ok_or_else(|| AuthError::new(REASON_UNKNOWN_KEY_ID, "JWT key id is missing"))?; + let keys = self.ensure_keyring()?; + let key = keys + .get(key_id) + .ok_or_else(|| AuthError::new(REASON_UNKNOWN_KEY_ID, "JWT key id is not trusted"))?; + + let mut validation = Validation::new(Algorithm::HS256); + validation.validate_exp = false; + validation.validate_nbf = false; + validation.validate_aud = false; + validation.required_spec_claims.clear(); + let token = decode::(raw_token, &DecodingKey::from_secret(key), &validation) + .map_err(|_| AuthError::new(REASON_INVALID_SIGNATURE, "JWT signature is invalid"))?; + let claims = token.claims; + + if claims.iss != self.jwt.issuer { + return Err(AuthError::new( + REASON_INVALID_ISSUER, + "JWT issuer is invalid", + )); + } + + if claims.aud != audience { + return Err(AuthError::new( + REASON_INVALID_AUDIENCE, + "JWT audience is invalid", + )); + } + + if claims.iat <= 0 { + return Err(AuthError::new( + REASON_INVALID_ISSUED_AT, + "JWT issued-at time is missing", + )); + } + + if claims.exp <= 0 { + return Err(AuthError::new( + REASON_EXPIRED, + "JWT expiration time is missing", + )); + } + + let now = i64::try_from(now) + .map_err(|_| AuthError::new(REASON_INVALID_ISSUED_AT, "server time overflows"))?; + let skew = i64::try_from(self.jwt.clock_skew.as_secs()) + .map_err(|_| AuthError::new(REASON_TTL_EXCEEDED, "JWT clock skew overflows"))?; + if claims.iat > now.saturating_add(skew) { + return Err(AuthError::new( + REASON_INVALID_ISSUED_AT, + "JWT issued-at time is in the future", + )); + } + + if claims.exp <= now.saturating_sub(skew) { + return Err(AuthError::new(REASON_EXPIRED, "JWT is expired")); + } + + if claims.exp <= claims.iat { + return Err(AuthError::new( + REASON_TTL_EXCEEDED, + "JWT lifetime is not positive", + )); + } + + let max_token_ttl = i64::try_from(self.jwt.max_token_ttl.as_secs()) + .map_err(|_| AuthError::new(REASON_TTL_EXCEEDED, "JWT maximum lifetime overflows"))?; + if claims.exp - claims.iat > max_token_ttl { + return Err(AuthError::new( + REASON_TTL_EXCEEDED, + "JWT lifetime exceeds the configured maximum", + )); + } + + Ok(()) + } + + fn ensure_keyring(&self) -> Result>>, AuthError> { + if let Some(keys) = self.runtime.keys.get() { + return Ok(keys.clone()); + } + + let keys = Arc::new(self.load_keyring()?); + let _ = self.runtime.keys.set(keys.clone()); + Ok(self.runtime.keys.get().cloned().unwrap_or(keys)) + } + + fn load_keyring(&self) -> Result>, AuthError> { + if !self.enabled() { + return Ok(HashMap::new()); + } + + if self.jwt.issuer.is_empty() { + return Err(AuthError::new(REASON_MALFORMED, "JWT issuer is required")); + } + + validate_duration(self.jwt.token_ttl, "tokenTTL", true)?; + validate_duration(self.jwt.max_token_ttl, "maxTokenTTL", true)?; + validate_duration(self.jwt.clock_skew, "clockSkew", false)?; + validate_duration(self.jwt.refresh_before, "refreshBefore", true)?; + if self.jwt.token_ttl > self.jwt.max_token_ttl { + return Err(AuthError::new( + REASON_MALFORMED, + "JWT tokenTTL must not exceed maxTokenTTL", + )); + } + + if self.jwt.refresh_before >= self.jwt.token_ttl { + return Err(AuthError::new( + REASON_MALFORMED, + "JWT refreshBefore must be less than tokenTTL", + )); + } + + if self.jwt.active_key_id.is_empty() { + return Err(AuthError::new( + REASON_UNKNOWN_KEY_ID, + "JWT activeKeyID is required", + )); + } + + if self.jwt.keys.is_empty() { + return Err(AuthError::new( + REASON_UNKNOWN_KEY_ID, + "JWT keyring must not be empty", + )); + } + + let mut keys = HashMap::with_capacity(self.jwt.keys.len()); + for key_config in &self.jwt.keys { + if key_config.id.is_empty() { + return Err(AuthError::new( + REASON_UNKNOWN_KEY_ID, + "JWT key id is required", + )); + } + + if keys.contains_key(&key_config.id) { + return Err(AuthError::new( + REASON_UNKNOWN_KEY_ID, + format!("JWT key id {:?} is duplicated", key_config.id), + )); + } + + let metadata = fs::metadata(&key_config.secret_file).map_err(|error| { + AuthError::new( + REASON_UNKNOWN_KEY_ID, + format!("failed to stat JWT secret file: {error}"), + ) + })?; + if !metadata.is_file() { + return Err(AuthError::new( + REASON_UNKNOWN_KEY_ID, + "JWT secret file is not a regular file", + )); + } + + let encoded = fs::read_to_string(&key_config.secret_file).map_err(|error| { + AuthError::new( + REASON_UNKNOWN_KEY_ID, + format!("failed to read JWT secret file: {error}"), + ) + })?; + let secret = STANDARD.decode(encoded.trim()).map_err(|_| { + AuthError::new( + REASON_UNKNOWN_KEY_ID, + "JWT secret file is not valid standard Base64", + ) + })?; + if secret.len() < MINIMUM_KEY_SIZE { + return Err(AuthError::new( + REASON_UNKNOWN_KEY_ID, + format!("decoded JWT secret must contain at least {MINIMUM_KEY_SIZE} bytes"), + )); + } + + keys.insert(key_config.id.clone(), secret); + } + + if !keys.contains_key(&self.jwt.active_key_id) { + return Err(AuthError::new( + REASON_UNKNOWN_KEY_ID, + "JWT active key is not trusted", + )); + } + + Ok(keys) + } + + fn record_server_result(&self, audience: &str, result: &str, reason: &str) { + GRPC_AUTH_REQUESTS + .with_label_values(&[audience, self.mode.as_str(), result, reason]) + .inc(); + } +} + +/// Tonic client interceptor that attaches a JWT to an RPC. +#[derive(Clone)] +pub struct ClientInterceptor { + auth: GrpcAuth, + audience: &'static str, +} + +impl ClientInterceptor { + /// Creates a client interceptor for a fixed audience and transport. + pub fn new(auth: GrpcAuth, audience: &'static str, secure: bool) -> Result { + validate_audience(audience)?; + auth.ensure_transport_security(secure)?; + auth.ensure_keyring()?; + Ok(Self { auth, audience }) + } +} + +impl Interceptor for ClientInterceptor { + fn call(&mut self, mut request: Request<()>) -> Result, Status> { + match self.auth.token(self.audience) { + Ok(Some(token)) => { + let value: MetadataValue = format!("{BEARER_SCHEME} {token}") + .parse() + .map_err(|_| Status::unauthenticated(AUTHENTICATION_ERROR_MESSAGE))?; + request + .metadata_mut() + .insert(AUTHORIZATION_METADATA_KEY, value); + } + Ok(None) => {} + Err(_) => { + GRPC_AUTH_CLIENT_TOKEN_EVENTS + .with_label_values(&[self.audience, "generation_failed"]) + .inc(); + return Err(Status::unauthenticated(AUTHENTICATION_ERROR_MESSAGE)); + } + } + + Ok(request) + } +} + +/// Tonic server interceptor that validates a JWT on an RPC. +#[derive(Clone)] +pub struct ServerInterceptor { + auth: GrpcAuth, + audience: &'static str, +} + +impl ServerInterceptor { + /// Creates a server interceptor for a fixed audience. + pub fn new(auth: GrpcAuth, audience: &'static str) -> Self { + Self { auth, audience } + } +} + +impl Interceptor for ServerInterceptor { + fn call(&mut self, request: Request<()>) -> Result, Status> { + match self.auth.authenticate(request.metadata(), self.audience) { + Ok(AuthenticationOutcome::Disabled) => Ok(request), + Ok(AuthenticationOutcome::Authenticated) => { + self.auth + .record_server_result(self.audience, "success", REASON_NONE); + Ok(request) + } + Ok(AuthenticationOutcome::PermissiveMissing) => { + self.auth + .record_server_result(self.audience, "allowed", REASON_MISSING); + Ok(request) + } + Err(error) => { + self.auth + .record_server_result(self.audience, "failure", error.reason()); + Err(Status::unauthenticated(AUTHENTICATION_ERROR_MESSAGE)) + } + } + } +} + +fn validate_duration(duration: Duration, name: &str, positive: bool) -> Result<(), AuthError> { + if duration.subsec_nanos() != 0 { + return Err(AuthError::new( + REASON_MALFORMED, + format!("JWT {name} must use whole seconds"), + )); + } + + if positive && duration < Duration::from_secs(1) { + return Err(AuthError::new( + REASON_MALFORMED, + format!("JWT {name} must be at least one second"), + )); + } + + if duration.as_secs() > i64::MAX as u64 { + return Err(AuthError::new( + REASON_MALFORMED, + format!("JWT {name} is too large"), + )); + } + + Ok(()) +} + +fn validate_audience(audience: &str) -> Result<(), AuthError> { + if matches!( + audience, + AUDIENCE_MANAGER | AUDIENCE_SCHEDULER | AUDIENCE_DFDAEMON + ) { + return Ok(()); + } + + Err(AuthError::new( + REASON_INVALID_AUDIENCE, + "unsupported JWT audience", + )) +} + +fn unix_time() -> Result { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|duration| duration.as_secs()) + .map_err(|_| AuthError::new(REASON_INVALID_ISSUED_AT, "system time precedes Unix epoch")) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::io::Write; + use std::ops::{Deref, DerefMut}; + use tempfile::NamedTempFile; + + const NOW: u64 = 1_786_435_200; + + #[derive(Deserialize)] + struct InteroperabilityFixture { + #[serde(rename = "secretBase64")] + secret_base64: String, + #[serde(rename = "keyID")] + key_id: String, + issuer: String, + audience: String, + #[serde(rename = "issuedAt")] + issued_at: u64, + #[serde(rename = "expiresAt")] + expires_at: u64, + #[serde(rename = "goToken")] + go_token: String, + #[serde(rename = "rustToken")] + rust_token: String, + } + + struct TestAuth { + auth: GrpcAuth, + _secret_file: NamedTempFile, + } + + impl Deref for TestAuth { + type Target = GrpcAuth; + + fn deref(&self) -> &Self::Target { + &self.auth + } + } + + impl DerefMut for TestAuth { + fn deref_mut(&mut self) -> &mut Self::Target { + &mut self.auth + } + } + + fn enabled_auth(mode: Mode) -> TestAuth { + let mut secret_file = NamedTempFile::new().unwrap(); + writeln!( + secret_file, + "{}", + STANDARD.encode(b"0123456789abcdef0123456789abcdef") + ) + .unwrap(); + TestAuth { + auth: GrpcAuth { + mode, + jwt: JwtConfig { + active_key_id: "test-key".to_string(), + keys: vec![KeyConfig { + id: "test-key".to_string(), + secret_file: secret_file.path().to_path_buf(), + }], + ..Default::default() + }, + ..Default::default() + }, + _secret_file: secret_file, + } + } + + fn interoperability_fixture() -> InteroperabilityFixture { + serde_json::from_str(include_str!("../testdata/interop.json")).unwrap() + } + + #[test] + fn defaults_are_disabled() { + let auth = GrpcAuth::default(); + assert_eq!(auth.mode, Mode::Disabled); + assert_eq!(auth.jwt.issuer, DEFAULT_ISSUER); + assert_eq!(auth.jwt.token_ttl, default_token_ttl()); + assert!(auth.validate().is_ok()); + } + + #[test] + fn deserializes_shared_component_configuration() { + let mut secret_file = NamedTempFile::new().unwrap(); + writeln!( + secret_file, + "{}", + STANDARD.encode(b"0123456789abcdef0123456789abcdef") + ) + .unwrap(); + let yaml = format!( + r#" +mode: required +requireTransportSecurity: true +jwt: + issuer: dragonfly-internal + activeKeyID: test-key + tokenTTL: 10m + maxTokenTTL: 15m + clockSkew: 30s + refreshBefore: 1m + keys: + - id: test-key + secretFile: {} +"#, + secret_file.path().display() + ); + + let auth: GrpcAuth = serde_yaml::from_str(&yaml).unwrap(); + assert_eq!(auth.mode, Mode::Required); + assert!(auth.require_transport_security); + assert_eq!(auth.jwt.active_key_id, "test-key"); + assert!(auth.validate().is_ok()); + } + + #[test] + fn validates_go_interoperability_fixture() { + let fixture = interoperability_fixture(); + let auth = enabled_auth(Mode::Required); + assert_eq!(auth.jwt.active_key_id, fixture.key_id); + assert_eq!(auth.jwt.issuer, fixture.issuer); + assert_eq!(fixture.audience, AUDIENCE_SCHEDULER); + assert_eq!( + fixture.secret_base64, + STANDARD.encode(b"0123456789abcdef0123456789abcdef") + ); + assert_eq!( + fixture.expires_at - fixture.issued_at, + auth.jwt.token_ttl.as_secs() + ); + assert!(auth + .verify_at(&fixture.go_token, AUDIENCE_SCHEDULER, fixture.issued_at) + .is_ok()); + assert!(auth + .verify_at(&fixture.rust_token, AUDIENCE_SCHEDULER, fixture.issued_at) + .is_ok()); + assert_eq!( + auth.token_at(AUDIENCE_SCHEDULER, fixture.issued_at) + .unwrap() + .unwrap(), + fixture.rust_token + ); + } + + #[test] + fn generates_and_caches_token() { + let auth = enabled_auth(Mode::Required); + let first = auth.token_at(AUDIENCE_MANAGER, NOW).unwrap().unwrap(); + let second = auth.token_at(AUDIENCE_MANAGER, NOW).unwrap().unwrap(); + assert_eq!(first, second); + assert!(auth.verify_at(&first, AUDIENCE_MANAGER, NOW).is_ok()); + } + + #[test] + fn overlapping_keys_support_rotation() { + let mut auth = enabled_auth(Mode::Required); + let mut old_secret_file = NamedTempFile::new().unwrap(); + let old_secret = b"old-key-material-old-key-material-"; + writeln!(old_secret_file, "{}", STANDARD.encode(old_secret)).unwrap(); + auth.jwt.keys.push(KeyConfig { + id: "old-key".to_string(), + secret_file: old_secret_file.path().to_path_buf(), + }); + + let claims = Claims { + iss: DEFAULT_ISSUER.to_string(), + aud: AUDIENCE_SCHEDULER.to_string(), + iat: NOW as i64, + exp: (NOW + default_token_ttl().as_secs()) as i64, + }; + let mut old_header = Header::new(Algorithm::HS256); + old_header.typ = Some(TOKEN_TYPE.to_string()); + old_header.kid = Some("old-key".to_string()); + let old_token = + encode(&old_header, &claims, &EncodingKey::from_secret(old_secret)).unwrap(); + assert!(auth.verify_at(&old_token, AUDIENCE_SCHEDULER, NOW).is_ok()); + + let new_token = auth.token_at(AUDIENCE_SCHEDULER, NOW).unwrap().unwrap(); + assert_eq!( + decode_header(&new_token).unwrap().kid.as_deref(), + Some("test-key") + ); + } + + #[test] + fn caches_token_across_concurrent_clones() { + let auth = enabled_auth(Mode::Required); + let barrier = Arc::new(std::sync::Barrier::new(16)); + let handles: Vec<_> = (0..16) + .map(|_| { + let auth = auth.auth.clone(); + let barrier = barrier.clone(); + std::thread::spawn(move || { + barrier.wait(); + auth.token_at(AUDIENCE_SCHEDULER, NOW).unwrap().unwrap() + }) + }) + .collect(); + let tokens: Vec<_> = handles + .into_iter() + .map(|handle| handle.join().unwrap()) + .collect(); + + assert!(tokens.windows(2).all(|pair| pair[0] == pair[1])); + } + + #[test] + fn enforces_server_modes() { + let disabled = GrpcAuth::default(); + assert_eq!( + disabled + .authenticate_at(&MetadataMap::new(), AUDIENCE_DFDAEMON, NOW) + .unwrap(), + AuthenticationOutcome::Disabled + ); + + let permissive = enabled_auth(Mode::Permissive); + assert_eq!( + permissive + .authenticate_at(&MetadataMap::new(), AUDIENCE_DFDAEMON, NOW) + .unwrap(), + AuthenticationOutcome::PermissiveMissing + ); + + let required = enabled_auth(Mode::Required); + assert_eq!( + required + .authenticate_at(&MetadataMap::new(), AUDIENCE_DFDAEMON, NOW) + .unwrap_err() + .reason(), + REASON_MISSING + ); + } + + #[test] + fn rejects_wrong_audience_and_expired_token() { + let fixture = interoperability_fixture(); + let auth = enabled_auth(Mode::Required); + assert_eq!( + auth.verify_at(&fixture.go_token, AUDIENCE_MANAGER, NOW) + .unwrap_err() + .reason(), + REASON_INVALID_AUDIENCE + ); + assert_eq!( + auth.verify_at( + &fixture.go_token, + AUDIENCE_SCHEDULER, + NOW + default_token_ttl().as_secs() + default_clock_skew().as_secs() + ) + .unwrap_err() + .reason(), + REASON_EXPIRED + ); + } + + #[test] + fn rejects_invalid_key_configuration() { + let mut auth = enabled_auth(Mode::Required); + auth.jwt.active_key_id = "missing".to_string(); + assert!(auth.validate().is_err()); + } + + #[test] + fn rejects_plaintext_when_required() { + let mut auth = enabled_auth(Mode::Required); + auth.require_transport_security = true; + assert!(auth.ensure_transport_security(false).is_err()); + assert!(auth.ensure_transport_security(true).is_ok()); + } +} diff --git a/dragonfly-client-auth/testdata/interop.json b/dragonfly-client-auth/testdata/interop.json new file mode 100644 index 000000000..ff36f31c2 --- /dev/null +++ b/dragonfly-client-auth/testdata/interop.json @@ -0,0 +1,11 @@ +{ + "_comment": "TEST ONLY. This deterministic key must never be used in a deployment.", + "secretBase64": "MDEyMzQ1Njc4OWFiY2RlZjAxMjM0NTY3ODlhYmNkZWY=", + "keyID": "test-key", + "issuer": "dragonfly-internal", + "audience": "urn:dragonfly:grpc:scheduler", + "issuedAt": 1786435200, + "expiresAt": 1786435800, + "goToken": "eyJhbGciOiJIUzI1NiIsImtpZCI6InRlc3Qta2V5IiwidHlwIjoiZHJhZ29uZmx5LWdycGMrand0In0.eyJpc3MiOiJkcmFnb25mbHktaW50ZXJuYWwiLCJhdWQiOiJ1cm46ZHJhZ29uZmx5OmdycGM6c2NoZWR1bGVyIiwiaWF0IjoxNzg2NDM1MjAwLCJleHAiOjE3ODY0MzU4MDB9.-BwX7Ur59WOu00J90_FUGr7wFQ7FqKmmYbd8OVNf8s4", + "rustToken": "eyJ0eXAiOiJkcmFnb25mbHktZ3JwYytqd3QiLCJhbGciOiJIUzI1NiIsImtpZCI6InRlc3Qta2V5In0.eyJpc3MiOiJkcmFnb25mbHktaW50ZXJuYWwiLCJhdWQiOiJ1cm46ZHJhZ29uZmx5OmdycGM6c2NoZWR1bGVyIiwiaWF0IjoxNzg2NDM1MjAwLCJleHAiOjE3ODY0MzU4MDB9.GPft73L1Y8aPMedYlf_drg_vHGzwirKIzfn3lxH8ZeY" +} diff --git a/dragonfly-client-config/Cargo.toml b/dragonfly-client-config/Cargo.toml index dd747e8cf..07a5e4d6d 100644 --- a/dragonfly-client-config/Cargo.toml +++ b/dragonfly-client-config/Cargo.toml @@ -11,6 +11,7 @@ edition.workspace = true build = "build.rs" [dependencies] +dragonfly-client-auth.workspace = true dragonfly-client-core.workspace = true dragonfly-client-util.workspace = true local-ip-address.workspace = true diff --git a/dragonfly-client-config/src/dfdaemon.rs b/dragonfly-client-config/src/dfdaemon.rs index 1b9da78a5..692db0c5b 100644 --- a/dragonfly-client-config/src/dfdaemon.rs +++ b/dragonfly-client-config/src/dfdaemon.rs @@ -15,6 +15,7 @@ */ use bytesize::ByteSize; +use dragonfly_client_auth::GrpcAuth; use dragonfly_client_core::{ error::{ErrorType, OrErr}, Result, @@ -1653,6 +1654,10 @@ pub struct Config { #[validate] pub upload: Upload, + /// Inter-component gRPC authentication configuration. + #[validate] + pub grpc_auth: GrpcAuth, + /// The manager configuration for dfdaemon. #[validate] pub manager: Manager, diff --git a/dragonfly-client-metric/Cargo.toml b/dragonfly-client-metric/Cargo.toml index e17dae88d..4a2386aa0 100644 --- a/dragonfly-client-metric/Cargo.toml +++ b/dragonfly-client-metric/Cargo.toml @@ -12,6 +12,7 @@ edition.workspace = true [dependencies] dragonfly-api.workspace = true +dragonfly-client-auth.workspace = true dragonfly-client-config.workspace = true dragonfly-client-util.workspace = true bytes.workspace = true diff --git a/dragonfly-client-metric/src/lib.rs b/dragonfly-client-metric/src/lib.rs index 27d0f2fbf..536f8fa57 100644 --- a/dragonfly-client-metric/src/lib.rs +++ b/dragonfly-client-metric/src/lib.rs @@ -16,6 +16,7 @@ use bytes::Bytes; use dragonfly_api::common::v2::{Range, TrafficType}; +use dragonfly_client_auth::{GRPC_AUTH_CLIENT_TOKEN_EVENTS, GRPC_AUTH_REQUESTS}; use dragonfly_client_config::{ dfdaemon::Config, BUILD_PLATFORM, CARGO_PKG_VERSION, GIT_COMMIT_DATE, GIT_COMMIT_SHORT_HASH, }; @@ -740,6 +741,14 @@ fn register_custom_metrics() { REGISTRY .register(Box::new(UPLOAD_TASK_BLOCKED_COUNT.clone())) .expect("metric can be registered"); + + REGISTRY + .register(Box::new(GRPC_AUTH_REQUESTS.clone())) + .expect("metric can be registered"); + + REGISTRY + .register(Box::new(GRPC_AUTH_CLIENT_TOKEN_EVENTS.clone())) + .expect("metric can be registered"); } /// Resets all custom metrics. @@ -776,6 +785,8 @@ fn reset_custom_metrics() { DELETE_HOST_FAILURE_COUNT.reset(); DISK_SPACE.reset(); DISK_USAGE_SPACE.reset(); + GRPC_AUTH_REQUESTS.reset(); + GRPC_AUTH_CLIENT_TOKEN_EVENTS.reset(); } /// Represents the size of the task. diff --git a/dragonfly-client-request/Cargo.toml b/dragonfly-client-request/Cargo.toml index 8b752f425..28b5e80a7 100644 --- a/dragonfly-client-request/Cargo.toml +++ b/dragonfly-client-request/Cargo.toml @@ -15,6 +15,7 @@ default = [] preheat = ["dep:oci-client", "dep:oci-spec"] [dependencies] +dragonfly-client-auth.workspace = true dragonfly-client-core.workspace = true dragonfly-client-util.workspace = true dragonfly-api.workspace = true diff --git a/dragonfly-client-request/src/lib.rs b/dragonfly-client-request/src/lib.rs index b686c0b76..a843b3249 100644 --- a/dragonfly-client-request/src/lib.rs +++ b/dragonfly-client-request/src/lib.rs @@ -21,6 +21,9 @@ use dragonfly_api::dfdaemon::v2::{ dfdaemon_upload_client::DfdaemonUploadClient as DfdaemonUploadGRPCClient, DownloadTaskRequest, }; use dragonfly_api::scheduler::v2::scheduler_client::SchedulerClient; +use dragonfly_client_auth::{ + ClientInterceptor as AuthClientInterceptor, GrpcAuth, AUDIENCE_DFDAEMON, AUDIENCE_SCHEDULER, +}; use dragonfly_client_util::digest::is_blob_url; use dragonfly_client_util::http::{ headermap_to_hashmap, query_params::default_proxy_rule_filtered_query_params, @@ -408,6 +411,9 @@ pub struct Builder { /// The number of times to retry a request. max_retries: u8, + + /// Inter-component gRPC JWT authentication configuration. + grpc_auth: GrpcAuth, } /// Implements Default trait. @@ -419,6 +425,7 @@ impl Default for Builder { scheduler_request_timeout: DEFAULT_SCHEDULER_REQUEST_TIMEOUT, health_check_interval: Duration::from_secs(60), max_retries: 1, + grpc_auth: GrpcAuth::default(), } } } @@ -449,6 +456,12 @@ impl Builder { self } + /// Sets inter-component gRPC JWT authentication. + pub fn grpc_auth(mut self, grpc_auth: GrpcAuth) -> Self { + self.grpc_auth = grpc_auth; + self + } + /// Builds and returns a Proxy instance. pub async fn build(self) -> Result { // Validate input parameters. @@ -468,8 +481,18 @@ impl Builder { )) })?; - // Create scheduler client. - let scheduler_client = SchedulerClient::new(scheduler_channel); + // Create scheduler client. Health checks remain unauthenticated. + let secure_transport = url::Url::parse(&self.scheduler_endpoint) + .map(|url| url.scheme() == "https") + .unwrap_or(false); + let scheduler_interceptor = AuthClientInterceptor::new( + self.grpc_auth.clone(), + AUDIENCE_SCHEDULER, + secure_transport, + ) + .map_err(|error| Error::InvalidArgument(error.to_string()))?; + let scheduler_client = + SchedulerClient::with_interceptor(scheduler_channel, scheduler_interceptor); // Create seed peer selector. let seed_peer_selector = Arc::new( @@ -506,6 +529,7 @@ impl Builder { .idle_timeout(DEFAULT_CLIENT_POOL_IDLE_TIMEOUT) .build(), id_generator: Arc::new(id_generator), + grpc_auth: self.grpc_auth, }; Ok(proxy) @@ -552,6 +576,9 @@ pub struct Proxy { /// The task id generator. id_generator: Arc, + + /// Inter-component gRPC JWT authentication configuration. + grpc_auth: GrpcAuth, } /// Implements the proxy client that sends requests via Dragonfly. @@ -795,7 +822,10 @@ impl Request for Proxy { Error::Internal(format!("failed to connect to seed peer {addr}: {err}")) })?; - let mut client = DfdaemonUploadGRPCClient::new(channel) + let interceptor = + AuthClientInterceptor::new(self.grpc_auth.clone(), AUDIENCE_DFDAEMON, false) + .map_err(|err| Error::InvalidArgument(err.to_string()))?; + let mut client = DfdaemonUploadGRPCClient::with_interceptor(channel, interceptor) .max_decoding_message_size(usize::MAX) .max_encoding_message_size(usize::MAX); diff --git a/dragonfly-client-request/src/selector.rs b/dragonfly-client-request/src/selector.rs index b22fa08a3..ceef48212 100644 --- a/dragonfly-client-request/src/selector.rs +++ b/dragonfly-client-request/src/selector.rs @@ -17,6 +17,7 @@ use async_trait::async_trait; use dragonfly_api::common::v2::Host; use dragonfly_api::scheduler::v2::{scheduler_client::SchedulerClient, ListHostsRequest}; +use dragonfly_client_auth::ClientInterceptor as AuthClientInterceptor; use dragonfly_client_core::{Error, Result}; use dragonfly_client_util::hashring::VNodeHashRing; use dragonfly_client_util::net::format_url; @@ -27,6 +28,7 @@ use std::str::FromStr; use std::time::Duration; use tokio::sync::{Mutex, RwLock}; use tokio::task::JoinSet; +use tonic::service::interceptor::InterceptedService; use tonic::transport::{Channel, Endpoint}; use tonic_health::pb::{ health_client::HealthClient as HealthGRPCClient, HealthCheckRequest, HealthCheckResponse, @@ -61,7 +63,7 @@ pub struct SeedPeerSelector { health_check_interval: Duration, /// The client to communicate with the scheduler service. - scheduler_client: SchedulerClient, + scheduler_client: SchedulerClient>, /// The data of the seed peer selector. seed_peers: RwLock, @@ -74,7 +76,7 @@ pub struct SeedPeerSelector { impl SeedPeerSelector { /// Creates a new seed peer selector. pub async fn new( - scheduler_client: SchedulerClient, + scheduler_client: SchedulerClient>, health_check_interval: Duration, ) -> Result { let seed_peer_selector = Self { @@ -229,6 +231,7 @@ impl Selector for SeedPeerSelector { mod tests { use super::*; use dragonfly_api::common::v2::Host; + use dragonfly_client_auth::{GrpcAuth, AUDIENCE_SCHEDULER}; use dragonfly_client_util::net::format_socket_addr; use std::collections::HashMap; use std::sync::Arc; @@ -262,7 +265,10 @@ mod tests { SeedPeerSelector { health_check_interval: Duration::from_secs(10), - scheduler_client: SchedulerClient::new(channel), + scheduler_client: SchedulerClient::with_interceptor( + channel, + AuthClientInterceptor::new(GrpcAuth::default(), AUDIENCE_SCHEDULER, false).unwrap(), + ), seed_peers: RwLock::new(SeedPeers { hosts: HashMap::new(), hashring: VNodeHashRing::new(1), diff --git a/dragonfly-client/Cargo.toml b/dragonfly-client/Cargo.toml index 45b6f92f0..6a6b3d5ea 100644 --- a/dragonfly-client/Cargo.toml +++ b/dragonfly-client/Cargo.toml @@ -31,6 +31,7 @@ name = "dfctl" path = "src/bin/dfctl/main.rs" [dependencies] +dragonfly-client-auth.workspace = true dragonfly-client-core.workspace = true dragonfly-client-config.workspace = true dragonfly-client-storage.workspace = true diff --git a/dragonfly-client/src/bin/dfctl/task.rs b/dragonfly-client/src/bin/dfctl/task.rs index 444795bf1..97d1f5a4b 100644 --- a/dragonfly-client/src/bin/dfctl/task.rs +++ b/dragonfly-client/src/bin/dfctl/task.rs @@ -24,6 +24,7 @@ use dragonfly_api::scheduler::v2::{ scheduler_client::SchedulerClient as SchedulerGRPCClient, PreheatFileRequest, PreheatImageRequest, }; +use dragonfly_client_auth::{ClientInterceptor, GrpcAuth, AUDIENCE_SCHEDULER}; use dragonfly_client_backend::{hdfs, object_storage, oci}; use dragonfly_client_core::{ error::{ErrorType, OrErr}, @@ -38,6 +39,7 @@ use dragonfly_client_util::{ net::preferred_local_ip, }; use oci_client::Reference; +use serde::Deserialize; use std::path::PathBuf; use std::time::Duration; use tabled::{ @@ -818,6 +820,13 @@ pub struct PreheatCommand { )] scheduler_endpoint: String, + #[arg( + long, + env = "DFCTL_TASK_PREHEAT_GRPC_AUTH_CONFIG", + help = "Path to a YAML file containing the Dragonfly grpcAuth configuration" + )] + grpc_auth_config: Option, + #[arg( long, default_value_t = false, @@ -1248,16 +1257,37 @@ impl PreheatCommand { /// Run the preheat logic based on the URL type (image or file) and mode (gRPC or request SDK). async fn run(&self) -> Result<()> { + let grpc_auth = self.load_grpc_auth().await?; match (self.url.starts_with(oci::SCHEME), self.request_sdk) { - (true, true) => self.preheat_image_by_request_sdk().await, - (true, false) => self.preheat_image().await, - (false, true) => self.preheat_file_by_request_sdk().await, - (false, false) => self.preheat_file().await, + (true, true) => self.preheat_image_by_request_sdk(grpc_auth).await, + (true, false) => self.preheat_image(grpc_auth).await, + (false, true) => self.preheat_file_by_request_sdk(grpc_auth).await, + (false, false) => self.preheat_file(grpc_auth).await, } } + /// Loads the shared grpcAuth section used by dfdaemon and dfctl. + async fn load_grpc_auth(&self) -> Result { + #[derive(Deserialize)] + struct GrpcAuthFile { + #[serde(rename = "grpcAuth")] + grpc_auth: GrpcAuth, + } + + let Some(path) = &self.grpc_auth_config else { + return Ok(GrpcAuth::default()); + }; + + let content = tokio::fs::read(path) + .await + .map_err(|error| Error::Unknown(format!("read gRPC auth config: {error}")))?; + serde_yaml::from_slice::(&content) + .map(|config| config.grpc_auth) + .map_err(|error| Error::Unknown(format!("parse gRPC auth config: {error}"))) + } + /// Preheats an OCI image via the scheduler's gRPC PreheatImage RPC. - async fn preheat_image(&self) -> Result<()> { + async fn preheat_image(&self, grpc_auth: GrpcAuth) -> Result<()> { let reference: Reference = self .url .strip_prefix(&format!("{}://", oci::SCHEME)) @@ -1276,7 +1306,12 @@ impl PreheatCommand { .or_err(ErrorType::ParseError)? .connect() .await?; - let mut client = SchedulerGRPCClient::new(channel); + let secure_transport = Url::parse(&self.scheduler_endpoint) + .map(|url| url.scheme() == "https") + .unwrap_or(false); + let interceptor = ClientInterceptor::new(grpc_auth, AUDIENCE_SCHEDULER, secure_transport) + .map_err(|error| Error::Unknown(error.to_string()))?; + let mut client = SchedulerGRPCClient::with_interceptor(channel, interceptor); let filtered_query_params = self .filtered_query_params @@ -1321,12 +1356,17 @@ impl PreheatCommand { } /// Preheats a file via the scheduler's gRPC PreheatFile RPC. - async fn preheat_file(&self) -> Result<()> { + async fn preheat_file(&self, grpc_auth: GrpcAuth) -> Result<()> { let channel = Channel::from_shared(self.scheduler_endpoint.clone()) .or_err(ErrorType::ParseError)? .connect() .await?; - let mut client = SchedulerGRPCClient::new(channel); + let secure_transport = Url::parse(&self.scheduler_endpoint) + .map(|url| url.scheme() == "https") + .unwrap_or(false); + let interceptor = ClientInterceptor::new(grpc_auth, AUDIENCE_SCHEDULER, secure_transport) + .map_err(|error| Error::Unknown(error.to_string()))?; + let mut client = SchedulerGRPCClient::with_interceptor(channel, interceptor); let filtered_query_params = self .filtered_query_params @@ -1396,9 +1436,10 @@ impl PreheatCommand { } /// Preheats an OCI image via the Dragonfly request SDK. - async fn preheat_image_by_request_sdk(&self) -> Result<()> { + async fn preheat_image_by_request_sdk(&self, grpc_auth: GrpcAuth) -> Result<()> { let proxy = Proxy::builder() .scheduler_endpoint(self.scheduler_endpoint.clone()) + .grpc_auth(grpc_auth) .build() .await .map_err(|err| Error::Unknown(format!("failed to build proxy: {err}")))?; @@ -1450,9 +1491,10 @@ impl PreheatCommand { } /// Preheats a file via the Dragonfly request SDK. - async fn preheat_file_by_request_sdk(&self) -> Result<()> { + async fn preheat_file_by_request_sdk(&self, grpc_auth: GrpcAuth) -> Result<()> { let proxy = Proxy::builder() .scheduler_endpoint(self.scheduler_endpoint.clone()) + .grpc_auth(grpc_auth) .build() .await .map_err(|err| Error::Unknown(format!("failed to build proxy: {err}")))?; diff --git a/dragonfly-client/src/dynconfig/mod.rs b/dragonfly-client/src/dynconfig/mod.rs index e1efb60b0..5a078d0b7 100644 --- a/dragonfly-client/src/dynconfig/mod.rs +++ b/dragonfly-client/src/dynconfig/mod.rs @@ -219,11 +219,15 @@ impl Dynconfig { async fn backend(config: Arc, dynconfig_path: PathBuf) -> Result { match config.manager.addr { Some(ref addr) => { - let manager_client = ManagerClient::new(&config.manager, addr.clone()) - .await - .inspect_err(|err| { - error!("initialize manager client failed: {}", err); - })?; + let manager_client = ManagerClient::new_with_auth( + &config.manager, + config.grpc_auth.clone(), + addr.clone(), + ) + .await + .inspect_err(|err| { + error!("initialize manager client failed: {}", err); + })?; info!("refresh dynamic configuration from manager {}", addr); Ok(Backend::Remote(Remote::new( diff --git a/dragonfly-client/src/grpc/dfdaemon_upload.rs b/dragonfly-client/src/grpc/dfdaemon_upload.rs index 55f99d5db..52718adb6 100644 --- a/dragonfly-client/src/grpc/dfdaemon_upload.rs +++ b/dragonfly-client/src/grpc/dfdaemon_upload.rs @@ -36,6 +36,7 @@ use dragonfly_api::dfdaemon::v2::{ UpdatePersistentTaskRequest, }; use dragonfly_api::errordetails::v2::Backend; +use dragonfly_client_auth::AUDIENCE_DFDAEMON; use dragonfly_client_backend::StatRequest; use dragonfly_client_config::dfdaemon::Config; use dragonfly_client_core::{ @@ -160,13 +161,20 @@ impl DfdaemonUploadServer { persistent_cache_task: self.persistent_cache_task.clone(), system_monitor: self.system_monitor.clone(), }, - ExtractTracingInterceptor, + ExtractTracingInterceptor::with_auth(self.config.grpc_auth.clone(), AUDIENCE_DFDAEMON), ); // Register the reflection service. let reflection = tonic_reflection::server::Builder::configure() .register_encoded_file_descriptor_set(dragonfly_api::FILE_DESCRIPTOR_SET) .build_v1()?; + let reflection = InterceptedService::new( + reflection, + dragonfly_client_auth::ServerInterceptor::new( + self.config.grpc_auth.clone(), + AUDIENCE_DFDAEMON, + ), + ); // Clone the shutdown channel. let mut shutdown = self.shutdown.clone(); @@ -2500,12 +2508,14 @@ impl DfdaemonUploadClient { super::REQUEST_TIMEOUT }; - let channel = match config + let client_tls_config = config .upload .client .load_client_tls_config(domain_name.as_str()) - .await? - { + .await?; + let secure_transport = + client_tls_config.is_some() || Url::parse(addr.as_str())?.scheme() == "https"; + let channel = match client_tls_config { Some(client_tls_config) => { Channel::from_static(Box::leak(addr.clone().into_boxed_str())) .tls_config(client_tls_config)? @@ -2537,7 +2547,13 @@ impl DfdaemonUploadClient { .or_err(ErrorType::ConnectError)?, }; - let client = DfdaemonUploadGRPCClient::with_interceptor(channel, InjectTracingInterceptor) + let interceptor = InjectTracingInterceptor::with_auth( + config.grpc_auth.clone(), + AUDIENCE_DFDAEMON, + secure_transport, + ) + .map_err(|error| ClientError::Unknown(error.to_string()))?; + let client = DfdaemonUploadGRPCClient::with_interceptor(channel, interceptor) .max_decoding_message_size(usize::MAX) .max_encoding_message_size(usize::MAX); Ok(Self { client }) diff --git a/dragonfly-client/src/grpc/interceptor.rs b/dragonfly-client/src/grpc/interceptor.rs index 47fb40c59..e05913e77 100644 --- a/dragonfly-client/src/grpc/interceptor.rs +++ b/dragonfly-client/src/grpc/interceptor.rs @@ -14,6 +14,10 @@ * limitations under the License. */ +use dragonfly_client_auth::{ + AuthError, ClientInterceptor as AuthClientInterceptor, GrpcAuth, + ServerInterceptor as AuthServerInterceptor, +}; use tonic::{metadata, service::Interceptor, Request, Status}; use tracing_opentelemetry::OpenTelemetrySpanExt; @@ -51,9 +55,33 @@ impl opentelemetry::propagation::Injector for MetadataMap<'_> { } } -/// Auto-inject tracing gRPC interceptor. +/// Auto-inject tracing and optional component authentication gRPC interceptor. #[derive(Clone)] -pub struct InjectTracingInterceptor; +pub struct InjectTracingInterceptor { + auth: Option, +} + +/// Backward-compatible tracing-only interceptor value. +#[allow(non_upper_case_globals)] +pub const InjectTracingInterceptor: InjectTracingInterceptor = + InjectTracingInterceptor { auth: None }; + +impl InjectTracingInterceptor { + /// Creates an interceptor for the target component audience. + pub fn with_auth( + auth: GrpcAuth, + audience: &'static str, + secure_transport: bool, + ) -> Result { + Ok(Self { + auth: Some(AuthClientInterceptor::new( + auth, + audience, + secure_transport, + )?), + }) + } +} /// Implements the tonic Interceptor interface. impl Interceptor for InjectTracingInterceptor { @@ -64,18 +92,41 @@ impl Interceptor for InjectTracingInterceptor { prop.inject_context(&context, &mut MetadataMap(request.metadata_mut())); }); - Ok(request) + match self.auth.as_mut() { + Some(auth) => auth.call(request), + None => Ok(request), + } } } -/// Auto-extract tracing gRPC interceptor. +/// Auto-extract tracing and optional component authentication gRPC interceptor. #[derive(Clone)] -pub struct ExtractTracingInterceptor; +pub struct ExtractTracingInterceptor { + auth: Option, +} + +/// Backward-compatible tracing-only interceptor value. +#[allow(non_upper_case_globals)] +pub const ExtractTracingInterceptor: ExtractTracingInterceptor = + ExtractTracingInterceptor { auth: None }; + +impl ExtractTracingInterceptor { + /// Creates an authenticated interceptor for the target component audience. + pub fn with_auth(auth: GrpcAuth, audience: &'static str) -> Self { + Self { + auth: Some(AuthServerInterceptor::new(auth, audience)), + } + } +} /// Implements the tonic Interceptor interface. impl Interceptor for ExtractTracingInterceptor { /// Calls and injects tracing context into global propagator. fn call(&mut self, mut request: Request<()>) -> std::result::Result, Status> { + if let Some(auth) = self.auth.as_mut() { + request = auth.call(request)?; + } + let parent_cx = opentelemetry::global::get_text_map_propagator(|prop| { prop.extract(&MetadataMap(request.metadata_mut())) }); diff --git a/dragonfly-client/src/grpc/manager.rs b/dragonfly-client/src/grpc/manager.rs index 2680c7c37..ec52d2937 100644 --- a/dragonfly-client/src/grpc/manager.rs +++ b/dragonfly-client/src/grpc/manager.rs @@ -19,6 +19,7 @@ use dragonfly_api::manager::v2::{ manager_client::ManagerClient as ManagerGRPCClient, DeleteSeedPeerRequest, ListSchedulersRequest, ListSchedulersResponse, SeedPeer, UpdateSeedPeerRequest, }; +use dragonfly_client_auth::{GrpcAuth, AUDIENCE_MANAGER}; use dragonfly_client_config::dfdaemon::Manager; use dragonfly_client_core::{ error::{ErrorType, OrErr}, @@ -41,9 +42,19 @@ pub struct ManagerClient { /// Implements the grpc client of the manager. impl ManagerClient { - /// Creates a new manager client. + /// Creates a new manager client without inter-component authentication. pub async fn new(config: &Manager, addr: String) -> Result { - let domain_name = Url::parse(addr.as_str())? + Self::new_with_auth(config, GrpcAuth::default(), addr).await + } + + /// Creates a new manager client with inter-component authentication. + pub async fn new_with_auth( + config: &Manager, + grpc_auth: GrpcAuth, + addr: String, + ) -> Result { + let url = Url::parse(addr.as_str())?; + let domain_name = url .host_str() .ok_or(Error::InvalidParameter) .inspect_err(|_err| { @@ -52,6 +63,7 @@ impl ManagerClient { .to_string(); let client_tls_config = config.load_client_tls_config(domain_name.as_str()).await?; + let secure_transport = client_tls_config.is_some() || url.scheme() == "https"; let health_client = HealthClient::new(addr.as_str(), client_tls_config.clone()).await?; match health_client.check().await { Ok(resp) => { @@ -94,7 +106,10 @@ impl ManagerClient { .or_err(ErrorType::ConnectError)?, }; - let client = ManagerGRPCClient::with_interceptor(channel, InjectTracingInterceptor) + let interceptor = + InjectTracingInterceptor::with_auth(grpc_auth, AUDIENCE_MANAGER, secure_transport) + .map_err(|error| Error::Unknown(error.to_string()))?; + let client = ManagerGRPCClient::with_interceptor(channel, interceptor) .max_decoding_message_size(usize::MAX) .max_encoding_message_size(usize::MAX); Ok(Self { client }) diff --git a/dragonfly-client/src/grpc/scheduler.rs b/dragonfly-client/src/grpc/scheduler.rs index 7a7c23f23..e680d2e35 100644 --- a/dragonfly-client/src/grpc/scheduler.rs +++ b/dragonfly-client/src/grpc/scheduler.rs @@ -32,6 +32,7 @@ use dragonfly_api::scheduler::v2::{ UploadPersistentTaskFailedRequest, UploadPersistentTaskFinishedRequest, UploadPersistentTaskStartedRequest, }; +use dragonfly_client_auth::AUDIENCE_SCHEDULER; use dragonfly_client_config::dfdaemon::Config; use dragonfly_client_core::error::{ErrorType, OrErr}; use dragonfly_client_core::{Error, Result}; @@ -198,10 +199,12 @@ impl SchedulerClient { // Reuse the cached channel of the scheduler. let channel = scheduler_client.channel(addr).await?; - let mut client = - SchedulerGRPCClient::with_interceptor(channel, InjectTracingInterceptor) - .max_decoding_message_size(usize::MAX) - .max_encoding_message_size(usize::MAX); + let mut client = SchedulerGRPCClient::with_interceptor( + channel, + scheduler_client.auth_interceptor()?, + ) + .max_decoding_message_size(usize::MAX) + .max_encoding_message_size(usize::MAX); client.announce_host(request).await?; Ok(()) } @@ -244,10 +247,12 @@ impl SchedulerClient { // Reuse the cached channel of the scheduler. let channel = scheduler_client.channel(addr).await?; - let mut client = - SchedulerGRPCClient::with_interceptor(channel, InjectTracingInterceptor) - .max_decoding_message_size(usize::MAX) - .max_encoding_message_size(usize::MAX); + let mut client = SchedulerGRPCClient::with_interceptor( + channel, + scheduler_client.auth_interceptor()?, + ) + .max_decoding_message_size(usize::MAX) + .max_encoding_message_size(usize::MAX); client.announce_host(request).await?; Ok(()) } @@ -295,10 +300,12 @@ impl SchedulerClient { // Reuse the cached channel of the scheduler. let channel = scheduler_client.channel(addr).await?; - let mut client = - SchedulerGRPCClient::with_interceptor(channel, InjectTracingInterceptor) - .max_decoding_message_size(usize::MAX) - .max_encoding_message_size(usize::MAX); + let mut client = SchedulerGRPCClient::with_interceptor( + channel, + scheduler_client.auth_interceptor()?, + ) + .max_decoding_message_size(usize::MAX) + .max_encoding_message_size(usize::MAX); client.delete_host(request).await?; Ok(()) } @@ -584,12 +591,25 @@ impl SchedulerClient { let channel = self.channel(addr.addr).await?; Ok( - SchedulerGRPCClient::with_interceptor(channel, InjectTracingInterceptor) + SchedulerGRPCClient::with_interceptor(channel, self.auth_interceptor()?) .max_decoding_message_size(usize::MAX) .max_encoding_message_size(usize::MAX), ) } + /// Creates a JWT and tracing interceptor for scheduler RPCs. + fn auth_interceptor(&self) -> Result { + let secure_transport = self.config.scheduler.ca_cert.is_some() + && self.config.scheduler.cert.is_some() + && self.config.scheduler.key.is_some(); + InjectTracingInterceptor::with_auth( + self.config.grpc_auth.clone(), + AUDIENCE_SCHEDULER, + secure_transport, + ) + .map_err(|error| Error::Unknown(error.to_string())) + } + /// Returns the cached grpc channel of the scheduler address, connecting and /// caching it if it is absent. The channel balances requests over a pool of /// connections and reconnects on demand, so reusing it avoids a connection