mirror of
https://github.com/neondatabase/neon.git
synced 2026-05-19 06:00:38 +00:00
781bf4945d
A smaller version of #12066 that is somewhat easier to review. Now that I've been using https://crates.io/crates/top-type-sizes I've found a lot more of the low hanging fruit that can be tweaks to reduce the memory usage. Some context for the optimisations: Rust's stack allocation in futures is quite naive. Stack variables, even if moved, often still end up taking space in the future. Rearranging the order in which variables are defined, and properly scoping them can go a long way. `async fn` and `async move {}` have a consequence that they always duplicate the "upvars" (aka captures). All captures are permanently allocated in the future, even if moved. We can be mindful when writing futures to only capture as little as possible. TlsStream is massive. Needs boxing so it doesn't contribute to the above issue. ## Measurements from `top-type-sizes`: ### Before ``` 10328 {async block@proxy::proxy::task_main::{closure#0}::{closure#0}} align=8 6120 {async fn body of proxy::proxy::handle_client<proxy::protocol2::ChainRW<tokio::net::TcpStream>>()} align=8 ``` ### After ``` 4040 {async block@proxy::proxy::task_main::{closure#0}::{closure#0}} 4704 {async fn body of proxy::proxy::handle_client<proxy::protocol2::ChainRW<tokio::net::TcpStream>>()} align=8 ```
62 lines
2.3 KiB
Rust
62 lines
2.3 KiB
Rust
use tokio::io::{AsyncRead, AsyncWrite};
|
|
use tracing::{debug, info, warn};
|
|
|
|
use super::{ComputeCredentials, ComputeUserInfo};
|
|
use crate::auth::backend::ComputeCredentialKeys;
|
|
use crate::auth::{self, AuthFlow};
|
|
use crate::config::AuthenticationConfig;
|
|
use crate::context::RequestContext;
|
|
use crate::control_plane::AuthSecret;
|
|
use crate::stream::{PqStream, Stream};
|
|
use crate::{compute, sasl};
|
|
|
|
pub(super) async fn authenticate(
|
|
ctx: &RequestContext,
|
|
creds: ComputeUserInfo,
|
|
client: &mut PqStream<Stream<impl AsyncRead + AsyncWrite + Unpin>>,
|
|
config: &'static AuthenticationConfig,
|
|
secret: AuthSecret,
|
|
) -> auth::Result<ComputeCredentials> {
|
|
let scram_keys = match secret {
|
|
#[cfg(any(test, feature = "testing"))]
|
|
AuthSecret::Md5(_) => {
|
|
debug!("auth endpoint chooses MD5");
|
|
return Err(auth::AuthError::MalformedPassword("MD5 not supported"));
|
|
}
|
|
AuthSecret::Scram(secret) => {
|
|
debug!("auth endpoint chooses SCRAM");
|
|
|
|
let auth_outcome = tokio::time::timeout(
|
|
config.scram_protocol_timeout,
|
|
AuthFlow::new(client, auth::Scram(&secret, ctx)).authenticate(),
|
|
)
|
|
.await
|
|
.inspect_err(|_| warn!("error processing scram messages error = authentication timed out, execution time exceeded {} seconds", config.scram_protocol_timeout.as_secs()))
|
|
.map_err(auth::AuthError::user_timeout)?
|
|
.inspect_err(|error| warn!(?error, "error processing scram messages"))?;
|
|
|
|
let client_key = match auth_outcome {
|
|
sasl::Outcome::Success(key) => key,
|
|
sasl::Outcome::Failure(reason) => {
|
|
// TODO: warnings?
|
|
// TODO: should we get rid of this because double logging?
|
|
info!("auth backend failed with an error: {reason}");
|
|
return Err(auth::AuthError::password_failed(&*creds.user));
|
|
}
|
|
};
|
|
|
|
compute::ScramKeys {
|
|
client_key: client_key.as_bytes(),
|
|
server_key: secret.server_key.as_bytes(),
|
|
}
|
|
}
|
|
};
|
|
|
|
Ok(ComputeCredentials {
|
|
info: creds,
|
|
keys: ComputeCredentialKeys::AuthKeys(postgres_client::config::AuthKeys::ScramSha256(
|
|
scram_keys,
|
|
)),
|
|
})
|
|
}
|