Skip to content

Commit

Permalink
proxy: format imports with nightly rustfmt (#9414)
Browse files Browse the repository at this point in the history
```shell
cargo +nightly fmt -p proxy -- -l --config imports_granularity=Module,group_imports=StdExternalCrate,reorder_imports=true
```

These rust-analyzer settings for VSCode should help retain this style:
```json
  "rust-analyzer.imports.group.enable": true,
  "rust-analyzer.imports.prefix": "crate",
  "rust-analyzer.imports.merge.glob": false,
  "rust-analyzer.imports.granularity.group": "module",
  "rust-analyzer.imports.granularity.enforce": true,
```
  • Loading branch information
cloneable authored Oct 16, 2024
1 parent 89a65a9 commit f14e45f
Show file tree
Hide file tree
Showing 73 changed files with 723 additions and 832 deletions.
19 changes: 9 additions & 10 deletions proxy/src/auth/backend/classic.rs
Original file line number Diff line number Diff line change
@@ -1,16 +1,15 @@
use super::{ComputeCredentials, ComputeUserInfo};
use crate::{
auth::{self, backend::ComputeCredentialKeys, AuthFlow},
compute,
config::AuthenticationConfig,
context::RequestMonitoring,
control_plane::AuthSecret,
sasl,
stream::{PqStream, Stream},
};
use tokio::io::{AsyncRead, AsyncWrite};
use tracing::{info, warn};

use super::{ComputeCredentials, ComputeUserInfo};
use crate::auth::backend::ComputeCredentialKeys;
use crate::auth::{self, AuthFlow};
use crate::config::AuthenticationConfig;
use crate::context::RequestMonitoring;
use crate::control_plane::AuthSecret;
use crate::stream::{PqStream, Stream};
use crate::{compute, sasl};

pub(super) async fn authenticate(
ctx: &RequestMonitoring,
creds: ComputeUserInfo,
Expand Down
21 changes: 9 additions & 12 deletions proxy/src/auth/backend/console_redirect.rs
Original file line number Diff line number Diff line change
@@ -1,15 +1,3 @@
use crate::{
auth,
cache::Cached,
compute,
config::AuthenticationConfig,
context::RequestMonitoring,
control_plane::{self, provider::NodeInfo, CachedNodeInfo},
error::{ReportableError, UserFacingError},
proxy::connect_compute::ComputeConnectBackend,
stream::PqStream,
waiters,
};
use async_trait::async_trait;
use pq_proto::BeMessage as Be;
use thiserror::Error;
Expand All @@ -18,6 +6,15 @@ use tokio_postgres::config::SslMode;
use tracing::{info, info_span};

use super::ComputeCredentialKeys;
use crate::cache::Cached;
use crate::config::AuthenticationConfig;
use crate::context::RequestMonitoring;
use crate::control_plane::provider::NodeInfo;
use crate::control_plane::{self, CachedNodeInfo};
use crate::error::{ReportableError, UserFacingError};
use crate::proxy::connect_compute::ComputeConnectBackend;
use crate::stream::PqStream;
use crate::{auth, compute, waiters};

#[derive(Debug, Error)]
pub(crate) enum WebAuthError {
Expand Down
19 changes: 9 additions & 10 deletions proxy/src/auth/backend/hacks.rs
Original file line number Diff line number Diff line change
@@ -1,16 +1,15 @@
use super::{ComputeCredentials, ComputeUserInfo, ComputeUserInfoNoEndpoint};
use crate::{
auth::{self, AuthFlow},
config::AuthenticationConfig,
context::RequestMonitoring,
control_plane::AuthSecret,
intern::EndpointIdInt,
sasl,
stream::{self, Stream},
};
use tokio::io::{AsyncRead, AsyncWrite};
use tracing::{info, warn};

use super::{ComputeCredentials, ComputeUserInfo, ComputeUserInfoNoEndpoint};
use crate::auth::{self, AuthFlow};
use crate::config::AuthenticationConfig;
use crate::context::RequestMonitoring;
use crate::control_plane::AuthSecret;
use crate::intern::EndpointIdInt;
use crate::sasl;
use crate::stream::{self, Stream};

/// Compared to [SCRAM](crate::scram), cleartext password auth saves
/// one round trip and *expensive* computations (>= 4096 HMAC iterations).
/// These properties are benefical for serverless JS workers, so we
Expand Down
39 changes: 19 additions & 20 deletions proxy/src/auth/backend/jwt.rs
Original file line number Diff line number Diff line change
@@ -1,22 +1,22 @@
use std::{
future::Future,
sync::Arc,
time::{Duration, SystemTime},
};
use std::future::Future;
use std::sync::Arc;
use std::time::{Duration, SystemTime};

use arc_swap::ArcSwapOption;
use dashmap::DashMap;
use jose_jwk::crypto::KeyInfo;
use serde::{de::Visitor, Deserialize, Deserializer};
use serde::de::Visitor;
use serde::{Deserialize, Deserializer};
use signature::Verifier;
use thiserror::Error;
use tokio::time::Instant;

use crate::{
auth::backend::ComputeCredentialKeys, context::RequestMonitoring,
control_plane::errors::GetEndpointJwksError, http::parse_json_body_with_limit,
intern::RoleNameInt, EndpointId, RoleName,
};
use crate::auth::backend::ComputeCredentialKeys;
use crate::context::RequestMonitoring;
use crate::control_plane::errors::GetEndpointJwksError;
use crate::http::parse_json_body_with_limit;
use crate::intern::RoleNameInt;
use crate::{EndpointId, RoleName};

// TODO(conrad): make these configurable.
const CLOCK_SKEW_LEEWAY: Duration = Duration::from_secs(30);
Expand Down Expand Up @@ -381,10 +381,8 @@ fn verify_rsa_signature(
alg: &jose_jwa::Algorithm,
) -> Result<(), JwtError> {
use jose_jwa::{Algorithm, Signing};
use rsa::{
pkcs1v15::{Signature, VerifyingKey},
RsaPublicKey,
};
use rsa::pkcs1v15::{Signature, VerifyingKey};
use rsa::RsaPublicKey;

let key = RsaPublicKey::try_from(key).map_err(JwtError::InvalidRsaKey)?;

Expand Down Expand Up @@ -655,11 +653,9 @@ impl From<&jose_jwk::Key> for KeyType {

#[cfg(test)]
mod tests {
use crate::RoleName;

use super::*;

use std::{future::IntoFuture, net::SocketAddr, time::SystemTime};
use std::future::IntoFuture;
use std::net::SocketAddr;
use std::time::SystemTime;

use base64::URL_SAFE_NO_PAD;
use bytes::Bytes;
Expand All @@ -672,6 +668,9 @@ mod tests {
use signature::Signer;
use tokio::net::TcpListener;

use super::*;
use crate::RoleName;

fn new_ec_jwk(kid: String) -> (p256::SecretKey, jose_jwk::Jwk) {
let sk = p256::SecretKey::random(&mut OsRng);
let pk = sk.public_key().into();
Expand Down
19 changes: 7 additions & 12 deletions proxy/src/auth/backend/local.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,19 +2,14 @@ use std::net::SocketAddr;

use arc_swap::ArcSwapOption;

use crate::{
auth::backend::jwt::FetchAuthRulesError,
compute::ConnCfg,
context::RequestMonitoring,
control_plane::{
messages::{ColdStartInfo, EndpointJwksResponse, MetricsAuxInfo},
NodeInfo,
},
intern::{BranchIdTag, EndpointIdTag, InternId, ProjectIdTag},
EndpointId,
};

use super::jwt::{AuthRule, FetchAuthRules};
use crate::auth::backend::jwt::FetchAuthRulesError;
use crate::compute::ConnCfg;
use crate::context::RequestMonitoring;
use crate::control_plane::messages::{ColdStartInfo, EndpointJwksResponse, MetricsAuxInfo};
use crate::control_plane::NodeInfo;
use crate::intern::{BranchIdTag, EndpointIdTag, InternId, ProjectIdTag};
use crate::EndpointId;

pub struct LocalBackend {
pub(crate) node_info: NodeInfo,
Expand Down
61 changes: 26 additions & 35 deletions proxy/src/auth/backend/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -17,29 +17,22 @@ use tokio_postgres::config::AuthKeys;
use tracing::{info, warn};

use crate::auth::credentials::check_peer_addr_is_in_list;
use crate::auth::{validate_password_and_exchange, AuthError};
use crate::auth::{self, validate_password_and_exchange, AuthError, ComputeUserInfoMaybeEndpoint};
use crate::cache::Cached;
use crate::config::AuthenticationConfig;
use crate::context::RequestMonitoring;
use crate::control_plane::errors::GetAuthInfoError;
use crate::control_plane::provider::{CachedRoleSecret, ControlPlaneBackend};
use crate::control_plane::AuthSecret;
use crate::control_plane::provider::{
CachedAllowedIps, CachedNodeInfo, CachedRoleSecret, ControlPlaneBackend,
};
use crate::control_plane::{self, Api, AuthSecret};
use crate::intern::EndpointIdInt;
use crate::metrics::Metrics;
use crate::proxy::connect_compute::ComputeConnectBackend;
use crate::proxy::NeonOptions;
use crate::rate_limiter::{BucketRateLimiter, EndpointRateLimiter, RateBucketInfo};
use crate::stream::Stream;
use crate::{
auth::{self, ComputeUserInfoMaybeEndpoint},
config::AuthenticationConfig,
control_plane::{
self,
provider::{CachedAllowedIps, CachedNodeInfo},
Api,
},
stream,
};
use crate::{scram, EndpointCacheKey, EndpointId, RoleName};
use crate::{scram, stream, EndpointCacheKey, EndpointId, RoleName};

/// Alternative to [`std::borrow::Cow`] but doesn't need `T: ToOwned` as we don't need that functionality
pub enum MaybeOwned<'a, T> {
Expand Down Expand Up @@ -500,34 +493,32 @@ impl ComputeConnectBackend for Backend<'_, ComputeCredentials> {

#[cfg(test)]
mod tests {
use std::{net::IpAddr, sync::Arc, time::Duration};
use std::net::IpAddr;
use std::sync::Arc;
use std::time::Duration;

use bytes::BytesMut;
use fallible_iterator::FallibleIterator;
use once_cell::sync::Lazy;
use postgres_protocol::{
authentication::sasl::{ChannelBinding, ScramSha256},
message::{backend::Message as PgMessage, frontend},
};
use postgres_protocol::authentication::sasl::{ChannelBinding, ScramSha256};
use postgres_protocol::message::backend::Message as PgMessage;
use postgres_protocol::message::frontend;
use provider::AuthSecret;
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWriteExt};

use crate::{
auth::{backend::MaskedIp, ComputeUserInfoMaybeEndpoint, IpPattern},
config::AuthenticationConfig,
context::RequestMonitoring,
control_plane::{
self,
provider::{self, CachedAllowedIps, CachedRoleSecret},
CachedNodeInfo,
},
proxy::NeonOptions,
rate_limiter::{EndpointRateLimiter, RateBucketInfo},
scram::{threadpool::ThreadPool, ServerSecret},
stream::{PqStream, Stream},
};

use super::{auth_quirks, jwt::JwkCache, AuthRateLimiter};
use super::jwt::JwkCache;
use super::{auth_quirks, AuthRateLimiter};
use crate::auth::backend::MaskedIp;
use crate::auth::{ComputeUserInfoMaybeEndpoint, IpPattern};
use crate::config::AuthenticationConfig;
use crate::context::RequestMonitoring;
use crate::control_plane::provider::{self, CachedAllowedIps, CachedRoleSecret};
use crate::control_plane::{self, CachedNodeInfo};
use crate::proxy::NeonOptions;
use crate::rate_limiter::{EndpointRateLimiter, RateBucketInfo};
use crate::scram::threadpool::ThreadPool;
use crate::scram::ServerSecret;
use crate::stream::{PqStream, Stream};

struct Auth {
ips: Vec<IpPattern>,
Expand Down
25 changes: 14 additions & 11 deletions proxy/src/auth/credentials.rs
Original file line number Diff line number Diff line change
@@ -1,20 +1,22 @@
//! User credentials used in authentication.
use crate::{
auth::password_hack::parse_endpoint_param,
context::RequestMonitoring,
error::{ReportableError, UserFacingError},
metrics::{Metrics, SniKind},
proxy::NeonOptions,
serverless::SERVERLESS_DRIVER_SNI,
EndpointId, RoleName,
};
use std::collections::HashSet;
use std::net::IpAddr;
use std::str::FromStr;

use itertools::Itertools;
use pq_proto::StartupMessageParams;
use std::{collections::HashSet, net::IpAddr, str::FromStr};
use thiserror::Error;
use tracing::{info, warn};

use crate::auth::password_hack::parse_endpoint_param;
use crate::context::RequestMonitoring;
use crate::error::{ReportableError, UserFacingError};
use crate::metrics::{Metrics, SniKind};
use crate::proxy::NeonOptions;
use crate::serverless::SERVERLESS_DRIVER_SNI;
use crate::{EndpointId, RoleName};

#[derive(Debug, Error, PartialEq, Eq, Clone)]
pub(crate) enum ComputeUserInfoParseError {
#[error("Parameter '{0}' is missing in startup packet.")]
Expand Down Expand Up @@ -249,10 +251,11 @@ fn project_name_valid(name: &str) -> bool {

#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
use ComputeUserInfoParseError::*;

use super::*;

#[test]
fn parse_bare_minimum() -> anyhow::Result<()> {
// According to postgresql, only `user` should be required.
Expand Down
25 changes: 14 additions & 11 deletions proxy/src/auth/flow.rs
Original file line number Diff line number Diff line change
@@ -1,21 +1,24 @@
//! Main authentication flow.
use super::{backend::ComputeCredentialKeys, AuthErrorImpl, PasswordHackPayload};
use crate::{
config::TlsServerEndPoint,
context::RequestMonitoring,
control_plane::AuthSecret,
intern::EndpointIdInt,
sasl,
scram::{self, threadpool::ThreadPool},
stream::{PqStream, Stream},
};
use std::io;
use std::sync::Arc;

use postgres_protocol::authentication::sasl::{SCRAM_SHA_256, SCRAM_SHA_256_PLUS};
use pq_proto::{BeAuthenticationSaslMessage, BeMessage, BeMessage as Be};
use std::{io, sync::Arc};
use tokio::io::{AsyncRead, AsyncWrite};
use tracing::info;

use super::backend::ComputeCredentialKeys;
use super::{AuthErrorImpl, PasswordHackPayload};
use crate::config::TlsServerEndPoint;
use crate::context::RequestMonitoring;
use crate::control_plane::AuthSecret;
use crate::intern::EndpointIdInt;
use crate::sasl;
use crate::scram::threadpool::ThreadPool;
use crate::scram::{self};
use crate::stream::{PqStream, Stream};

/// Every authentication selector is supposed to implement this trait.
pub(crate) trait AuthMethod {
/// Any authentication selector should provide initial backend message
Expand Down
12 changes: 6 additions & 6 deletions proxy/src/auth/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -14,15 +14,15 @@ pub(crate) use password_hack::parse_endpoint_param;
use password_hack::PasswordHackPayload;

mod flow;
use std::io;
use std::net::IpAddr;

pub(crate) use flow::*;
use thiserror::Error;
use tokio::time::error::Elapsed;

use crate::{
control_plane,
error::{ReportableError, UserFacingError},
};
use std::{io, net::IpAddr};
use thiserror::Error;
use crate::control_plane;
use crate::error::{ReportableError, UserFacingError};

/// Convenience wrapper for the authentication error.
pub(crate) type Result<T> = std::result::Result<T, AuthError>;
Expand Down
Loading

1 comment on commit f14e45f

@github-actions
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

5290 tests run: 5073 passed, 0 failed, 217 skipped (full report)


Flaky tests (5)

Postgres 17

Postgres 16

Postgres 15

Postgres 14

Code coverage* (full report)

  • functions: 31.3% (7553 of 24107 functions)
  • lines: 49.2% (60419 of 122920 lines)

* collected from Rust tests only


The comment gets automatically updated with the latest test results
f14e45f at 2024-10-16T14:46:38.802Z :recycle:

Please sign in to comment.