laskoviymishka commented on code in PR #3170:
URL: https://github.com/apache/iceberg-rust/pull/3170#discussion_r4232872491
##########
crates/catalog/rest/src/auth/oauth2.rs:
##########
@@ -134,15 +167,111 @@ impl AuthManager for OAuth2Manager {
client: &HttpClient,
props: &HashMap<String, String>,
) -> Result<Box<dyn AuthSession>> {
- Ok(Box::new(self.session_from(client, props).await?))
+ let (session, _) = self.build_session_from_properties(client,
props).await?;
+ Ok(Box::new(session))
}
async fn catalog_session(
&self,
client: &HttpClient,
props: &HashMap<String, String>,
) -> Result<Arc<dyn AuthSession>> {
- Ok(Arc::new(self.session_from(client, props).await?))
+ if self.contextual_state.get().is_some() {
+ return Err(Error::new(
+ ErrorKind::PreconditionFailed,
+ "OAuth2Manager catalog session already initialized",
+ ));
+ }
+
+ let (session, token_exchange) =
self.build_session_from_properties(client, props).await?;
+
+ self.contextual_state
+ .set(ContextualAuthState {
+ token_exchange,
+ session_cache: Cache::builder()
+ .time_to_idle(session_idle_timeout(props)?)
+ .build(),
+ })
+ .map_err(|_| {
+ Error::new(
+ ErrorKind::PreconditionFailed,
+ "OAuth2Manager catalog session already initialized
concurrently",
+ )
+ })?;
+
+ Ok(Arc::new(session))
+ }
+
+ async fn contextual_session(
+ &self,
+ context: &SessionContext,
+ catalog_session: Arc<dyn AuthSession>,
+ ) -> Result<Arc<dyn AuthSession>> {
+ let credential = if let Some(token) =
context.credentials().get("token") {
+ ContextualCredential::Token(token.clone())
+ } else if let Some(credential) =
context.credentials().get("credential") {
+
ContextualCredential::ClientCredentials(parse_contextual_credential(credential))
+ } else {
+ return Ok(catalog_session);
+ };
+
+ let contextual_state = self.contextual_state.get().ok_or_else(|| {
+ Error::new(
+ ErrorKind::PreconditionFailed,
+ "OAuth2 catalog session must be initialized before contextual
sessions",
+ )
+ })?;
+
+ if let Some(session) =
contextual_state.session_cache.get(context.session_id()) {
Review Comment:
A cache hit returns the cached session without ever comparing the incoming
token/credential to the one it was built from, so a `session_id` reused with a
different or rotated credential silently gets the old session. For a static
context token that means it's never refreshed — any session touched once per
idle window (1h default) keeps sending the stale token until it 401s — and if
two contexts ever collide on an id, one tenant gets the other's bearer.
`test_context_token_takes_precedence_and_is_cached_by_session_id` locks this in
as the intended behavior.
I'd fold a fingerprint of the credential (hash of kind + value) into the
key, or store the credential in the entry and rebuild on mismatch. Java gets
away with id-only keying because its contract is that `session_id` is unique
per principal — if we're relying on that too, let's state it explicitly in the
trait doc.
##########
crates/catalog/rest/src/auth/oauth2.rs:
##########
@@ -134,15 +167,111 @@ impl AuthManager for OAuth2Manager {
client: &HttpClient,
props: &HashMap<String, String>,
) -> Result<Box<dyn AuthSession>> {
- Ok(Box::new(self.session_from(client, props).await?))
+ let (session, _) = self.build_session_from_properties(client,
props).await?;
+ Ok(Box::new(session))
}
async fn catalog_session(
&self,
client: &HttpClient,
props: &HashMap<String, String>,
) -> Result<Arc<dyn AuthSession>> {
- Ok(Arc::new(self.session_from(client, props).await?))
+ if self.contextual_state.get().is_some() {
Review Comment:
The `get().is_some()` pre-check isn't atomic with the later `set()`, and
more importantly `build_session_from_properties` has already written
`self.token` (when `props` carries `token`) by the time `set()` loses the race
or `session_idle_timeout(props)?` errors — so a failed or duplicate call
mutates the shared token cell past the error. I'd pull
`session_idle_timeout(props)?` to the first line, build the state fully, then
commit with a single `set()` as the only gate.
Separately, this now fails permanently on a second call — worth confirming
no builder path (`with_auth_manager` followed by multiple `load()`s, or a
retried `RestClient::init`) can reach it twice, and documenting the one-shot
contract.
##########
crates/catalog/rest/src/auth/oauth2.rs:
##########
@@ -134,15 +167,111 @@ impl AuthManager for OAuth2Manager {
client: &HttpClient,
props: &HashMap<String, String>,
) -> Result<Box<dyn AuthSession>> {
- Ok(Box::new(self.session_from(client, props).await?))
+ let (session, _) = self.build_session_from_properties(client,
props).await?;
+ Ok(Box::new(session))
}
async fn catalog_session(
&self,
client: &HttpClient,
props: &HashMap<String, String>,
) -> Result<Arc<dyn AuthSession>> {
- Ok(Arc::new(self.session_from(client, props).await?))
+ if self.contextual_state.get().is_some() {
+ return Err(Error::new(
+ ErrorKind::PreconditionFailed,
+ "OAuth2Manager catalog session already initialized",
+ ));
+ }
+
+ let (session, token_exchange) =
self.build_session_from_properties(client, props).await?;
+
+ self.contextual_state
+ .set(ContextualAuthState {
+ token_exchange,
+ session_cache: Cache::builder()
+ .time_to_idle(session_idle_timeout(props)?)
+ .build(),
+ })
+ .map_err(|_| {
+ Error::new(
+ ErrorKind::PreconditionFailed,
+ "OAuth2Manager catalog session already initialized
concurrently",
+ )
+ })?;
+
+ Ok(Arc::new(session))
+ }
+
+ async fn contextual_session(
+ &self,
+ context: &SessionContext,
+ catalog_session: Arc<dyn AuthSession>,
+ ) -> Result<Arc<dyn AuthSession>> {
+ let credential = if let Some(token) =
context.credentials().get("token") {
+ ContextualCredential::Token(token.clone())
+ } else if let Some(credential) =
context.credentials().get("credential") {
+
ContextualCredential::ClientCredentials(parse_contextual_credential(credential))
+ } else {
+ return Ok(catalog_session);
+ };
+
+ let contextual_state = self.contextual_state.get().ok_or_else(|| {
+ Error::new(
+ ErrorKind::PreconditionFailed,
+ "OAuth2 catalog session must be initialized before contextual
sessions",
+ )
+ })?;
+
+ if let Some(session) =
contextual_state.session_cache.get(context.session_id()) {
+ return Ok(session.clone());
+ }
+
+ let token_exchange = contextual_state.token_exchange.clone();
+
+ let candidate_session = Arc::new(match credential {
+ ContextualCredential::Token(token) => OAuth2Session {
+ token:
Arc::new(Mutex::new(Some(CachedToken::static_token(token)))),
+ token_source: TokenSource::StaticToken,
+ parent: Some(catalog_session),
+ },
+ ContextualCredential::ClientCredentials(credential) =>
OAuth2Session {
+ token: Arc::new(Mutex::new(None)),
+ token_source:
TokenSource::ClientCredentials(Box::new(ClientCredentialsConfig {
+ token_exchange: TokenExchangeConfig {
+ client: token_exchange
+ .client
+ .with_auth_session(catalog_session.clone()),
+ ..token_exchange
+ },
+ credential,
+ })),
+ parent: Some(catalog_session),
+ },
+ });
+
+ let session = contextual_state
+ .session_cache
+ .entry(context.session_id().to_string())
+ .or_insert(candidate_session) // If a concurrent insert won, we'll
use it instead because it may already be in use.
Review Comment:
This comment makes a concurrency claim — winner is reused, single exchange —
that nothing tests. I'd add a `join_all` test on a multi-thread runtime with
`.expect(1)` on the token mock so the single-exchange guarantee is actually
pinned.
A few other gaps worth closing this pass: a second `catalog_session` call
hitting `PreconditionFailed`, two different credential contexts exchanging
separately (we cover token isolation but not credential isolation), and a
token-endpoint failure propagating and then retrying.
##########
crates/catalog/rest/src/auth/oauth2.rs:
##########
@@ -134,15 +167,111 @@ impl AuthManager for OAuth2Manager {
client: &HttpClient,
props: &HashMap<String, String>,
) -> Result<Box<dyn AuthSession>> {
- Ok(Box::new(self.session_from(client, props).await?))
+ let (session, _) = self.build_session_from_properties(client,
props).await?;
+ Ok(Box::new(session))
}
async fn catalog_session(
&self,
client: &HttpClient,
props: &HashMap<String, String>,
) -> Result<Arc<dyn AuthSession>> {
- Ok(Arc::new(self.session_from(client, props).await?))
+ if self.contextual_state.get().is_some() {
+ return Err(Error::new(
+ ErrorKind::PreconditionFailed,
+ "OAuth2Manager catalog session already initialized",
+ ));
+ }
+
+ let (session, token_exchange) =
self.build_session_from_properties(client, props).await?;
+
+ self.contextual_state
+ .set(ContextualAuthState {
+ token_exchange,
+ session_cache: Cache::builder()
+ .time_to_idle(session_idle_timeout(props)?)
+ .build(),
+ })
+ .map_err(|_| {
+ Error::new(
+ ErrorKind::PreconditionFailed,
+ "OAuth2Manager catalog session already initialized
concurrently",
+ )
+ })?;
+
+ Ok(Arc::new(session))
+ }
+
+ async fn contextual_session(
+ &self,
+ context: &SessionContext,
+ catalog_session: Arc<dyn AuthSession>,
+ ) -> Result<Arc<dyn AuthSession>> {
+ let credential = if let Some(token) =
context.credentials().get("token") {
+ ContextualCredential::Token(token.clone())
+ } else if let Some(credential) =
context.credentials().get("credential") {
+
ContextualCredential::ClientCredentials(parse_contextual_credential(credential))
+ } else {
+ return Ok(catalog_session);
Review Comment:
A context carrying only an `id_token` / `access_token` / `jwt` — the RFC
8693 token-type keys Java's `maybeCreateChildSession` handles — falls through
here and runs with the shared catalog identity, with no error or log. That's a
fail-open: a caller who expected tenant identity silently gets the catalog
credential. At minimum I'd document the supported keys; better would be to fail
closed (or warn) when a context has credentials but none are recognized.
##########
crates/catalog/rest/src/auth/oauth2.rs:
##########
@@ -46,15 +49,37 @@ struct OAuth2Params {
/// Iceberg REST catalogs.
///
/// A configured `token` is used directly; otherwise `credential` is exchanged
-/// for a token at the token endpoint and cached. The cached token is shared
-/// across sessions so it survives the config handshake.
+/// for a token at the token endpoint and cached. Client-credential tokens are
+/// refreshed on demand after the server-provided `expires_in` duration;
+/// responses without `expires_in` remain cached. The init and catalog sessions
+/// share that cached token so it survives the config handshake; contextual
+/// sessions keep their tokens isolated from it.
pub struct OAuth2Manager {
- token: Arc<Mutex<Option<SensitiveString>>>,
- init_params: OAuth2Params,
+ token: Arc<Mutex<Option<CachedToken>>>,
+ initial_config: OAuth2Config,
/// True when the token endpoint was derived from the catalog URI (not
/// explicitly configured): it is then recomputed from the merged URI in
/// [`Self::catalog_session`], since `/v1/config` may override the URI.
endpoint_is_default: bool,
+ /// Installed by `catalog_session`; replacing it atomically invalidates
Review Comment:
`contextual_state` is a `OnceLock`, so it can never be replaced — this
comment describes a design that was dropped. Reword to something like "set once
by `catalog_session`; contextual sessions require it."
##########
crates/catalog/rest/src/auth/oauth2.rs:
##########
@@ -178,33 +309,74 @@ impl OAuth2Manager {
_ if self.endpoint_is_default => props
.get(REST_CATALOG_PROP_URI)
.map(|uri| default_token_endpoint(uri))
- .unwrap_or_else(|| self.init_params.token_endpoint.clone()),
- _ => self.init_params.token_endpoint.clone(),
+ .unwrap_or_else(|| self.initial_config.token_endpoint.clone()),
+ _ => self.initial_config.token_endpoint.clone(),
};
let credential = credential_from_props(props)
.map(|(id, secret)| (id, secret.into()))
- .or_else(|| self.init_params.credential.clone());
+ .or_else(|| self.initial_config.credential.clone());
- Ok(OAuth2Session {
+ let token_exchange = TokenExchangeConfig {
+ client: client.clone(),
+ extra_headers,
+ token_endpoint,
+ extra_oauth_params,
+ };
+
+ let session = OAuth2Session {
token: self.token.clone(),
// A configured token takes precedence over the credential: the
// token cell is pre-seeded, and the credential only comes into
// play once that token is gone.
token_source: match credential {
Some(credential) => {
TokenSource::ClientCredentials(Box::new(ClientCredentialsConfig {
- client: client.clone(),
+ token_exchange: token_exchange.clone(),
credential,
- token_endpoint,
- extra_headers,
- extra_oauth_params,
}))
}
None => TokenSource::StaticToken,
},
- })
+ parent: None,
+ };
+ Ok((session, token_exchange))
+ }
+}
+
+const AUTH_SESSION_TIMEOUT_MS_PROP: &str = "auth.session-timeout-ms";
+const DEFAULT_CONTEXTUAL_SESSION_IDLE_TIMEOUT: Duration =
Duration::from_hours(1);
+// Moka's cache builder panics for expiration durations longer than 1,000
+// years, so we validate the property before constructing the cache.
+const MAX_CONTEXTUAL_SESSION_IDLE_TIMEOUT: Duration =
+ Duration::from_secs(1_000 * 365 * 24 * 60 * 60);
+
+fn session_idle_timeout(props: &HashMap<String, String>) -> Result<Duration> {
+ let Some(timeout_prop) = props.get(AUTH_SESSION_TIMEOUT_MS_PROP) else {
+ return Ok(DEFAULT_CONTEXTUAL_SESSION_IDLE_TIMEOUT);
+ };
+
+ let timeout_ms = timeout_prop.parse::<u64>().map_err(|e| {
+ Error::new(
+ ErrorKind::PreconditionFailed,
Review Comment:
A malformed `auth.session-timeout-ms` is a bad user-supplied property — the
rest of the crate (and `CachedToken::from_response` just below) uses
`DataInvalid` for that, not `PreconditionFailed`. I'd switch both arms in
`session_idle_timeout` to `DataInvalid` unless `PreconditionFailed` is
deliberate. While we're here, the second `format!` in this function already
inlines its arg, so this one can be `format!("Property
{AUTH_SESSION_TIMEOUT_MS_PROP} not an integer")` to satisfy
`uninlined_format_args`.
##########
Cargo.toml:
##########
@@ -118,6 +118,7 @@ mimalloc = "0.1.46"
minijinja = "2.12.0"
mockall = "0.13.1"
mockito = "1"
+moka = { version = "0.12.10" }
Review Comment:
This can just be `moka = "0.12.10"` — the braces are redundant.
`crates/iceberg` and `integrations/cache-moka` still pin `0.12.10` directly
with their own features; optionally move them onto `workspace = true` while
you're here so there's a single source.
##########
crates/catalog/rest/src/auth/oauth2.rs:
##########
@@ -318,18 +537,28 @@ impl AuthSession for OAuth2Session {
/// for one, caches it, then uses it. Without a credential and without a
/// token, no authentication is attached.
async fn authenticate(&self, req: &mut HttpRequest) -> Result<()> {
+ if let Some(parent) = &self.parent {
+ parent.authenticate(req).await?;
Review Comment:
Every contextual request runs `parent.authenticate(req)` first and then
overwrites `Authorization`. For a static-token context that's a wasted parent
call; if the catalog session uses client credentials it can trigger or wait on
a parent token exchange on each request, and a parent failure then breaks a
request that carried a perfectly valid contextual token. There's also a latent
leak — if the contextual token is ever absent, the `None => Ok(())` path leaves
the parent's tenant token on the request.
The credential exchange already authenticates through
`with_auth_session(catalog_session)`, so the session-level parent is only
needed to carry the parent's non-auth headers. I'd set `parent` only where
that's actually required, or strip `Authorization` before calling the parent.
##########
crates/catalog/rest/src/auth/oauth2.rs:
##########
@@ -134,15 +167,111 @@ impl AuthManager for OAuth2Manager {
client: &HttpClient,
props: &HashMap<String, String>,
) -> Result<Box<dyn AuthSession>> {
- Ok(Box::new(self.session_from(client, props).await?))
+ let (session, _) = self.build_session_from_properties(client,
props).await?;
+ Ok(Box::new(session))
}
async fn catalog_session(
&self,
client: &HttpClient,
props: &HashMap<String, String>,
) -> Result<Arc<dyn AuthSession>> {
- Ok(Arc::new(self.session_from(client, props).await?))
+ if self.contextual_state.get().is_some() {
+ return Err(Error::new(
+ ErrorKind::PreconditionFailed,
+ "OAuth2Manager catalog session already initialized",
+ ));
+ }
+
+ let (session, token_exchange) =
self.build_session_from_properties(client, props).await?;
+
+ self.contextual_state
+ .set(ContextualAuthState {
+ token_exchange,
+ session_cache: Cache::builder()
+ .time_to_idle(session_idle_timeout(props)?)
+ .build(),
+ })
+ .map_err(|_| {
+ Error::new(
+ ErrorKind::PreconditionFailed,
+ "OAuth2Manager catalog session already initialized
concurrently",
+ )
+ })?;
+
+ Ok(Arc::new(session))
+ }
+
+ async fn contextual_session(
+ &self,
+ context: &SessionContext,
+ catalog_session: Arc<dyn AuthSession>,
+ ) -> Result<Arc<dyn AuthSession>> {
+ let credential = if let Some(token) =
context.credentials().get("token") {
+ ContextualCredential::Token(token.clone())
+ } else if let Some(credential) =
context.credentials().get("credential") {
+
ContextualCredential::ClientCredentials(parse_contextual_credential(credential))
+ } else {
+ return Ok(catalog_session);
+ };
+
+ let contextual_state = self.contextual_state.get().ok_or_else(|| {
+ Error::new(
+ ErrorKind::PreconditionFailed,
+ "OAuth2 catalog session must be initialized before contextual
sessions",
+ )
+ })?;
+
+ if let Some(session) =
contextual_state.session_cache.get(context.session_id()) {
+ return Ok(session.clone());
+ }
+
+ let token_exchange = contextual_state.token_exchange.clone();
+
+ let candidate_session = Arc::new(match credential {
+ ContextualCredential::Token(token) => OAuth2Session {
+ token:
Arc::new(Mutex::new(Some(CachedToken::static_token(token)))),
+ token_source: TokenSource::StaticToken,
+ parent: Some(catalog_session),
+ },
+ ContextualCredential::ClientCredentials(credential) =>
OAuth2Session {
+ token: Arc::new(Mutex::new(None)),
+ token_source:
TokenSource::ClientCredentials(Box::new(ClientCredentialsConfig {
+ token_exchange: TokenExchangeConfig {
+ client: token_exchange
+ .client
+ .with_auth_session(catalog_session.clone()),
+ ..token_exchange
+ },
+ credential,
+ })),
+ parent: Some(catalog_session),
+ },
+ });
+
+ let session = contextual_state
+ .session_cache
+ .entry(context.session_id().to_string())
+ .or_insert(candidate_session) // If a concurrent insert won, we'll
use it instead because it may already be in use.
+ .value()
+ .clone();
+
+ Ok(session)
+ }
+}
+
+enum ContextualCredential {
+ Token(SensitiveString),
+ ClientCredentials((Option<String>, SensitiveString)),
+}
+
+fn parse_contextual_credential(credential: &SensitiveString) ->
(Option<String>, SensitiveString) {
+ match credential.expose().split_once(':') {
Review Comment:
`split_once(':')` on `":secret"` yields `Some("")` as the client_id, which
then goes out as `client_id=` in the form body (Java behaves the same, so it's
not a divergence). Worth rejecting an empty client_id — and either way the
colon-less / `:secret` / `id:` cases have no coverage, so a couple of tests
around this helper would be good.
##########
crates/catalog/rest/src/auth/oauth2.rs:
##########
@@ -134,15 +167,111 @@ impl AuthManager for OAuth2Manager {
client: &HttpClient,
props: &HashMap<String, String>,
) -> Result<Box<dyn AuthSession>> {
- Ok(Box::new(self.session_from(client, props).await?))
+ let (session, _) = self.build_session_from_properties(client,
props).await?;
+ Ok(Box::new(session))
}
async fn catalog_session(
&self,
client: &HttpClient,
props: &HashMap<String, String>,
) -> Result<Arc<dyn AuthSession>> {
- Ok(Arc::new(self.session_from(client, props).await?))
+ if self.contextual_state.get().is_some() {
+ return Err(Error::new(
+ ErrorKind::PreconditionFailed,
+ "OAuth2Manager catalog session already initialized",
+ ));
+ }
+
+ let (session, token_exchange) =
self.build_session_from_properties(client, props).await?;
+
+ self.contextual_state
+ .set(ContextualAuthState {
+ token_exchange,
+ session_cache: Cache::builder()
+ .time_to_idle(session_idle_timeout(props)?)
Review Comment:
The cache sets only `time_to_idle` — no `max_capacity` and no `time_to_live`
— and moka's sync cache only runs maintenance on access, so an idle manager
never frees anything. In a multi-tenant server with a per-user or per-request
`session_id`, memory grows with every distinct id seen inside the idle window
(1h default, and the property goes up to the 1000-year cap). I'd add a
`max_capacity` (configurable, something like 1k–10k) and ideally a
`time_to_live` to bound absolute age.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
