From f7e9ec49be83e20a0dfb61d4dbf071d08b5ab25a Mon Sep 17 00:00:00 2001 From: Dmitry Ivanov Date: Fri, 19 May 2023 17:40:56 +0300 Subject: [PATCH] Add stub for auth info cache --- proxy/src/auth/backend/classic.rs | 9 +++++---- proxy/src/console.rs | 4 +++- proxy/src/console/provider.rs | 5 ++++- proxy/src/console/provider/mock.rs | 14 ++++++++------ proxy/src/console/provider/neon.rs | 10 +++++++--- 5 files changed, 27 insertions(+), 15 deletions(-) diff --git a/proxy/src/auth/backend/classic.rs b/proxy/src/auth/backend/classic.rs index f5ba54ae20..1893340b0b 100644 --- a/proxy/src/auth/backend/classic.rs +++ b/proxy/src/auth/backend/classic.rs @@ -2,7 +2,7 @@ use super::AuthSuccess; use crate::{ auth::{self, AuthFlow, ClientCredentials}, compute::{self, ComputeNode, Password}, - console::{self, AuthInfo, ConsoleReqExtra}, + console::{self, AuthInfo, CachedAuthInfo, ConsoleReqExtra}, sasl, scram, stream::PqStream, }; @@ -21,18 +21,19 @@ pub(super) async fn authenticate( // prevent malicious probing (possible due to missing protocol steps). // This mocked secret will never lead to successful authentication. info!("authentication info not found, mocking it"); - AuthInfo::Scram(scram::ServerSecret::mock(creds.user, rand::random())) + let info = scram::ServerSecret::mock(creds.user, rand::random()); + CachedAuthInfo::new_uncached(AuthInfo::Scram(info)) }); let flow = AuthFlow::new(client); - let keys = match info { + let keys = match &*info { AuthInfo::Md5(_) => { info!("auth endpoint chooses MD5"); return Err(auth::AuthError::bad_auth_method("MD5")); } AuthInfo::Scram(secret) => { info!("auth endpoint chooses SCRAM"); - let scram = auth::Scram(&secret); + let scram = auth::Scram(secret); let client_key = match flow.begin(scram).await?.authenticate().await? { sasl::Outcome::Success(key) => key, sasl::Outcome::Failure(reason) => { diff --git a/proxy/src/console.rs b/proxy/src/console.rs index 1f3ef99555..c67ff8a18c 100644 --- a/proxy/src/console.rs +++ b/proxy/src/console.rs @@ -6,7 +6,9 @@ pub mod messages; /// Wrappers for console APIs and their mocks. pub mod provider; -pub use provider::{errors, Api, AuthInfo, CachedNodeInfo, ConsoleReqExtra, NodeInfo}; +pub use provider::{errors, Api, ConsoleReqExtra}; +pub use provider::{AuthInfo, NodeInfo}; +pub use provider::{CachedAuthInfo, CachedNodeInfo}; /// Various cache-related types. pub mod caches { diff --git a/proxy/src/console/provider.rs b/proxy/src/console/provider.rs index ce84aa0b40..ebe8b2b3e2 100644 --- a/proxy/src/console/provider.rs +++ b/proxy/src/console/provider.rs @@ -166,6 +166,9 @@ pub struct NodeInfo { pub type NodeInfoCache = TimedLru, NodeInfo>; pub type CachedNodeInfo = Cached; +pub type AuthInfoCache = TimedLru, AuthInfo>; +pub type CachedAuthInfo = Cached; + /// This will allocate per each call, but the http requests alone /// already require a few allocations, so it should be fine. #[async_trait] @@ -175,7 +178,7 @@ pub trait Api { &self, extra: &ConsoleReqExtra<'_>, creds: &ClientCredentials<'_>, - ) -> Result, errors::GetAuthInfoError>; + ) -> Result, errors::GetAuthInfoError>; /// Wake up the compute node and return the corresponding connection info. async fn wake_compute( diff --git a/proxy/src/console/provider/mock.rs b/proxy/src/console/provider/mock.rs index 03eac37b26..d6fa216651 100644 --- a/proxy/src/console/provider/mock.rs +++ b/proxy/src/console/provider/mock.rs @@ -2,8 +2,10 @@ use super::{ errors::{ApiError, GetAuthInfoError, WakeComputeError}, - AuthInfo, CachedNodeInfo, ConsoleReqExtra, NodeInfo, PgEndpoint, + ConsoleReqExtra, PgEndpoint, }; +use super::{AuthInfo, NodeInfo}; +use super::{CachedAuthInfo, CachedNodeInfo}; use crate::{auth::ClientCredentials, error::io_error, scram, url::ApiUrl}; use async_trait::async_trait; use futures::TryFutureExt; @@ -102,8 +104,9 @@ impl super::Api for Api { &self, _extra: &ConsoleReqExtra<'_>, creds: &ClientCredentials<'_>, - ) -> Result, GetAuthInfoError> { - self.do_get_auth_info(creds).await + ) -> Result, GetAuthInfoError> { + let res = self.do_get_auth_info(creds).await?; + Ok(res.map(CachedAuthInfo::new_uncached)) } #[tracing::instrument(skip_all)] @@ -112,9 +115,8 @@ impl super::Api for Api { _extra: &ConsoleReqExtra<'_>, _creds: &ClientCredentials<'_>, ) -> Result { - self.do_wake_compute() - .map_ok(CachedNodeInfo::new_uncached) - .await + let res = self.do_wake_compute().await?; + Ok(CachedNodeInfo::new_uncached(res)) } } diff --git a/proxy/src/console/provider/neon.rs b/proxy/src/console/provider/neon.rs index 7cd2af9856..27631f546a 100644 --- a/proxy/src/console/provider/neon.rs +++ b/proxy/src/console/provider/neon.rs @@ -3,8 +3,10 @@ use super::{ super::messages::{ConsoleError, GetRoleSecret, WakeCompute}, errors::{ApiError, GetAuthInfoError, WakeComputeError}, - ApiCaches, AuthInfo, CachedNodeInfo, ConsoleReqExtra, NodeInfo, + ApiCaches, ConsoleReqExtra, }; +use super::{AuthInfo, NodeInfo}; +use super::{CachedAuthInfo, CachedNodeInfo}; use crate::{auth::ClientCredentials, http, scram}; use async_trait::async_trait; use futures::TryFutureExt; @@ -110,8 +112,10 @@ impl super::Api for Api { &self, extra: &ConsoleReqExtra<'_>, creds: &ClientCredentials<'_>, - ) -> Result, GetAuthInfoError> { - self.do_get_auth_info(extra, creds).await + ) -> Result, GetAuthInfoError> { + // FIXME: add cache! + let res = self.do_get_auth_info(extra, creds).await?; + Ok(res.map(CachedAuthInfo::new_uncached)) } #[tracing::instrument(skip_all)]