[proxy] Refactoring

This patch attempts to fix some of the technical debt
we had to introduce in previous patches.

proxy/src/auth/backend/console.rs

@@ -0,0 +1,225 @@
+//! Cloud API V2.
+
+use crate::{
+    auth::{self, AuthFlow, ClientCredentials, DatabaseInfo},
+    compute,
+    error::UserFacingError,
+    scram,
+    stream::PqStream,
+    url::ApiUrl,
+};
+use serde::{Deserialize, Serialize};
+use std::{future::Future, io};
+use thiserror::Error;
+use tokio::io::{AsyncRead, AsyncWrite};
+use utils::pq_proto::{BeMessage as Be, BeParameterStatusMessage};
+
+pub type Result<T> = std::result::Result<T, ConsoleAuthError>;
+
+#[derive(Debug, Error)]
+pub enum ConsoleAuthError {
+    #[error(transparent)]
+    BadProjectName(#[from] auth::credentials::ProjectNameError),
+
+    // We shouldn't include the actual secret here.
+    #[error("Bad authentication secret")]
+    BadSecret,
+
+    #[error("Console responded with a malformed compute address: '{0}'")]
+    BadComputeAddress(String),
+
+    #[error("Console responded with a malformed JSON: '{0}'")]
+    BadResponse(#[from] serde_json::Error),
+
+    /// HTTP status (other than 200) returned by the console.
+    #[error("Console responded with an HTTP status: {0}")]
+    HttpStatus(reqwest::StatusCode),
+
+    #[error(transparent)]
+    Io(#[from] std::io::Error),
+}
+
+impl UserFacingError for ConsoleAuthError {
+    fn to_string_client(&self) -> String {
+        use ConsoleAuthError::*;
+        match self {
+            BadProjectName(e) => e.to_string_client(),
+            _ => "Internal error".to_string(),
+        }
+    }
+}
+
+// TODO: convert into an enum with "error"
+#[derive(Serialize, Deserialize, Debug)]
+struct GetRoleSecretResponse {
+    role_secret: String,
+}
+
+// TODO: convert into an enum with "error"
+#[derive(Serialize, Deserialize, Debug)]
+struct GetWakeComputeResponse {
+    address: String,
+}
+
+/// Auth secret which is managed by the cloud.
+pub enum AuthInfo {
+    /// Md5 hash of user's password.
+    Md5([u8; 16]),
+
+    /// [SCRAM](crate::scram) authentication info.
+    Scram(scram::ServerSecret),
+}
+
+#[must_use]
+pub(super) struct Api<'a> {
+    endpoint: &'a ApiUrl,
+    creds: &'a ClientCredentials,
+    /// Cache project name, since we'll need it several times.
+    project: &'a str,
+}
+
+impl<'a> Api<'a> {
+    /// Construct an API object containing the auth parameters.
+    pub(super) fn new(endpoint: &'a ApiUrl, creds: &'a ClientCredentials) -> Result<Self> {
+        Ok(Self {
+            endpoint,
+            creds,
+            project: creds.project_name()?,
+        })
+    }
+
+    /// Authenticate the existing user or throw an error.
+    pub(super) async fn handle_user(
+        self,
+        client: &mut PqStream<impl AsyncRead + AsyncWrite + Unpin + Send>,
+    ) -> auth::Result<compute::NodeInfo> {
+        handle_user(client, &self, Self::get_auth_info, Self::wake_compute).await
+    }
+
+    async fn get_auth_info(&self) -> Result<AuthInfo> {
+        let mut url = self.endpoint.clone();
+        url.path_segments_mut().push("proxy_get_role_secret");
+        url.query_pairs_mut()
+            .append_pair("project", self.project)
+            .append_pair("role", &self.creds.user);
+
+        // TODO: use a proper logger
+        println!("cplane request: {url}");
+
+        let resp = reqwest::get(url.into_inner()).await.map_err(io_error)?;
+        if !resp.status().is_success() {
+            return Err(ConsoleAuthError::HttpStatus(resp.status()));
+        }
+
+        let response: GetRoleSecretResponse =
+            serde_json::from_str(&resp.text().await.map_err(io_error)?)?;
+
+        scram::ServerSecret::parse(response.role_secret.as_str())
+            .map(AuthInfo::Scram)
+            .ok_or(ConsoleAuthError::BadSecret)
+    }
+
+    /// Wake up the compute node and return the corresponding connection info.
+    async fn wake_compute(&self) -> Result<DatabaseInfo> {
+        let mut url = self.endpoint.clone();
+        url.path_segments_mut().push("proxy_wake_compute");
+        url.query_pairs_mut().append_pair("project", self.project);
+
+        // TODO: use a proper logger
+        println!("cplane request: {url}");
+
+        let resp = reqwest::get(url.into_inner()).await.map_err(io_error)?;
+        if !resp.status().is_success() {
+            return Err(ConsoleAuthError::HttpStatus(resp.status()));
+        }
+
+        let response: GetWakeComputeResponse =
+            serde_json::from_str(&resp.text().await.map_err(io_error)?)?;
+
+        let (host, port) = parse_host_port(&response.address)
+            .ok_or(ConsoleAuthError::BadComputeAddress(response.address))?;
+
+        Ok(DatabaseInfo {
+            host,
+            port,
+            dbname: self.creds.dbname.to_owned(),
+            user: self.creds.user.to_owned(),
+            password: None,
+        })
+    }
+}
+
+/// Common logic for user handling in API V2.
+/// We reuse this for a mock API implementation in [`super::postgres`].
+pub(super) async fn handle_user<'a, Endpoint, GetAuthInfo, WakeCompute>(
+    client: &mut PqStream<impl AsyncRead + AsyncWrite + Unpin>,
+    endpoint: &'a Endpoint,
+    get_auth_info: impl FnOnce(&'a Endpoint) -> GetAuthInfo,
+    wake_compute: impl FnOnce(&'a Endpoint) -> WakeCompute,
+) -> auth::Result<compute::NodeInfo>
+where
+    GetAuthInfo: Future<Output = Result<AuthInfo>>,
+    WakeCompute: Future<Output = Result<DatabaseInfo>>,
+{
+    let auth_info = get_auth_info(endpoint).await?;
+
+    let flow = AuthFlow::new(client);
+    let scram_keys = match auth_info {
+        AuthInfo::Md5(_) => {
+            // TODO: decide if we should support MD5 in api v2
+            return Err(auth::AuthErrorImpl::auth_failed("MD5 is not supported").into());
+        }
+        AuthInfo::Scram(secret) => {
+            let scram = auth::Scram(&secret);
+            Some(compute::ScramKeys {
+                client_key: flow.begin(scram).await?.authenticate().await?.as_bytes(),
+                server_key: secret.server_key.as_bytes(),
+            })
+        }
+    };
+
+    client
+        .write_message_noflush(&Be::AuthenticationOk)?
+        .write_message_noflush(&BeParameterStatusMessage::encoding())?;
+
+    Ok(compute::NodeInfo {
+        db_info: wake_compute(endpoint).await?,
+        scram_keys,
+    })
+}
+
+/// Upcast (almost) any error into an opaque [`io::Error`].
+pub(super) fn io_error(e: impl Into<Box<dyn std::error::Error + Send + Sync>>) -> io::Error {
+    io::Error::new(io::ErrorKind::Other, e)
+}
+
+fn parse_host_port(input: &str) -> Option<(String, u16)> {
+    let (host, port) = input.split_once(':')?;
+    Some((host.to_owned(), port.parse().ok()?))
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use serde_json::json;
+
+    #[test]
+    fn parse_db_info() -> anyhow::Result<()> {
+        let _: DatabaseInfo = serde_json::from_value(json!({
+            "host": "localhost",
+            "port": 5432,
+            "dbname": "postgres",
+            "user": "john_doe",
+            "password": "password",
+        }))?;
+
+        let _: DatabaseInfo = serde_json::from_value(json!({
+            "host": "localhost",
+            "port": 5432,
+            "dbname": "postgres",
+            "user": "john_doe",
+        }))?;
+
+        Ok(())
+    }
+}

proxy/src/auth/backend/legacy_console.rs

@@ -0,0 +1,204 @@
+//! Cloud API V1.
+
+use super::DatabaseInfo;
+use crate::{
+    auth::{self, ClientCredentials},
+    compute,
+    error::UserFacingError,
+    stream::PqStream,
+    waiters,
+};
+use serde::{Deserialize, Serialize};
+use thiserror::Error;
+use tokio::io::{AsyncRead, AsyncWrite};
+use utils::pq_proto::{BeMessage as Be, BeParameterStatusMessage};
+
+#[derive(Debug, Error)]
+pub enum AuthErrorImpl {
+    /// Authentication error reported by the console.
+    #[error("Authentication failed: {0}")]
+    AuthFailed(String),
+
+    /// HTTP status (other than 200) returned by the console.
+    #[error("Console responded with an HTTP status: {0}")]
+    HttpStatus(reqwest::StatusCode),
+
+    #[error("Console responded with a malformed JSON: {0}")]
+    MalformedResponse(#[from] serde_json::Error),
+
+    #[error(transparent)]
+    Transport(#[from] reqwest::Error),
+
+    #[error(transparent)]
+    WaiterRegister(#[from] waiters::RegisterError),
+
+    #[error(transparent)]
+    WaiterWait(#[from] waiters::WaitError),
+}
+
+#[derive(Debug, Error)]
+#[error(transparent)]
+pub struct AuthError(Box<AuthErrorImpl>);
+
+impl AuthError {
+    /// Smart constructor for authentication error reported by `mgmt`.
+    pub fn auth_failed(msg: impl Into<String>) -> Self {
+        Self(Box::new(AuthErrorImpl::AuthFailed(msg.into())))
+    }
+}
+
+impl<T> From<T> for AuthError
+where
+    AuthErrorImpl: From<T>,
+{
+    fn from(e: T) -> Self {
+        Self(Box::new(e.into()))
+    }
+}
+
+impl UserFacingError for AuthError {
+    fn to_string_client(&self) -> String {
+        use AuthErrorImpl::*;
+        match self.0.as_ref() {
+            AuthFailed(_) | HttpStatus(_) => self.to_string(),
+            _ => "Internal error".to_string(),
+        }
+    }
+}
+
+// NOTE: the order of constructors is important.
+// https://serde.rs/enum-representations.html#untagged
+#[derive(Serialize, Deserialize, Debug)]
+#[serde(untagged)]
+enum ProxyAuthResponse {
+    Ready { conn_info: DatabaseInfo },
+    Error { error: String },
+    NotReady { ready: bool }, // TODO: get rid of `ready`
+}
+
+async fn authenticate_proxy_client(
+    auth_endpoint: &reqwest::Url,
+    creds: &ClientCredentials,
+    md5_response: &str,
+    salt: &[u8; 4],
+    psql_session_id: &str,
+) -> Result<DatabaseInfo, AuthError> {
+    let mut url = auth_endpoint.clone();
+    url.query_pairs_mut()
+        .append_pair("login", &creds.user)
+        .append_pair("database", &creds.dbname)
+        .append_pair("md5response", md5_response)
+        .append_pair("salt", &hex::encode(salt))
+        .append_pair("psql_session_id", psql_session_id);
+
+    super::with_waiter(psql_session_id, |waiter| async {
+        println!("cloud request: {}", url);
+        // TODO: leverage `reqwest::Client` to reuse connections
+        let resp = reqwest::get(url).await?;
+        if !resp.status().is_success() {
+            return Err(AuthErrorImpl::HttpStatus(resp.status()).into());
+        }
+
+        let auth_info: ProxyAuthResponse = serde_json::from_str(resp.text().await?.as_str())?;
+        println!("got auth info: #{:?}", auth_info);
+
+        use ProxyAuthResponse::*;
+        let db_info = match auth_info {
+            Ready { conn_info } => conn_info,
+            Error { error } => return Err(AuthErrorImpl::AuthFailed(error).into()),
+            NotReady { .. } => waiter.await?.map_err(AuthErrorImpl::AuthFailed)?,
+        };
+
+        Ok(db_info)
+    })
+    .await
+}
+
+async fn handle_existing_user(
+    auth_endpoint: &reqwest::Url,
+    client: &mut PqStream<impl AsyncRead + AsyncWrite + Unpin + Send>,
+    creds: &ClientCredentials,
+) -> Result<compute::NodeInfo, auth::AuthError> {
+    let psql_session_id = super::link::new_psql_session_id();
+    let md5_salt = rand::random();
+
+    client
+        .write_message(&Be::AuthenticationMD5Password(md5_salt))
+        .await?;
+
+    // Read client's password hash
+    let msg = client.read_password_message().await?;
+    let md5_response = parse_password(&msg).ok_or(auth::AuthErrorImpl::MalformedPassword)?;
+
+    let db_info = authenticate_proxy_client(
+        auth_endpoint,
+        creds,
+        md5_response,
+        &md5_salt,
+        &psql_session_id,
+    )
+    .await?;
+
+    client
+        .write_message_noflush(&Be::AuthenticationOk)?
+        .write_message_noflush(&BeParameterStatusMessage::encoding())?;
+
+    Ok(compute::NodeInfo {
+        db_info,
+        scram_keys: None,
+    })
+}
+
+pub async fn handle_user(
+    auth_endpoint: &reqwest::Url,
+    auth_link_uri: &reqwest::Url,
+    client: &mut PqStream<impl AsyncRead + AsyncWrite + Unpin + Send>,
+    creds: &ClientCredentials,
+) -> auth::Result<compute::NodeInfo> {
+    if creds.is_existing_user() {
+        handle_existing_user(auth_endpoint, client, creds).await
+    } else {
+        super::link::handle_user(auth_link_uri, client).await
+    }
+}
+
+fn parse_password(bytes: &[u8]) -> Option<&str> {
+    std::str::from_utf8(bytes).ok()?.strip_suffix('\0')
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use serde_json::json;
+
+    #[test]
+    fn test_proxy_auth_response() {
+        // Ready
+        let auth: ProxyAuthResponse = serde_json::from_value(json!({
+            "ready": true,
+            "conn_info": DatabaseInfo::default(),
+        }))
+        .unwrap();
+        assert!(matches!(
+            auth,
+            ProxyAuthResponse::Ready {
+                conn_info: DatabaseInfo { .. }
+            }
+        ));
+
+        // Error
+        let auth: ProxyAuthResponse = serde_json::from_value(json!({
+            "ready": false,
+            "error": "too bad, so sad",
+        }))
+        .unwrap();
+        assert!(matches!(auth, ProxyAuthResponse::Error { .. }));
+
+        // NotReady
+        let auth: ProxyAuthResponse = serde_json::from_value(json!({
+            "ready": false,
+        }))
+        .unwrap();
+        assert!(matches!(auth, ProxyAuthResponse::NotReady { .. }));
+    }
+}

proxy/src/auth/backend/link.rs

@@ -0,0 +1,47 @@
+use crate::{auth, compute, stream::PqStream};
+use tokio::io::{AsyncRead, AsyncWrite};
+use utils::pq_proto::{BeMessage as Be, BeParameterStatusMessage};
+
+fn hello_message(redirect_uri: &str, session_id: &str) -> String {
+    format!(
+        concat![
+            "Welcome to Neon!\n",
+            "Authenticate by visiting:\n",
+            "    {redirect_uri}{session_id}\n\n",
+        ],
+        redirect_uri = redirect_uri,
+        session_id = session_id,
+    )
+}
+
+pub fn new_psql_session_id() -> String {
+    hex::encode(rand::random::<[u8; 8]>())
+}
+
+pub async fn handle_user(
+    redirect_uri: &reqwest::Url,
+    client: &mut PqStream<impl AsyncRead + AsyncWrite + Unpin>,
+) -> auth::Result<compute::NodeInfo> {
+    let psql_session_id = new_psql_session_id();
+    let greeting = hello_message(redirect_uri.as_str(), &psql_session_id);
+
+    let db_info = super::with_waiter(psql_session_id, |waiter| async {
+        // Give user a URL to spawn a new database
+        client
+            .write_message_noflush(&Be::AuthenticationOk)?
+            .write_message_noflush(&BeParameterStatusMessage::encoding())?
+            .write_message(&Be::NoticeResponse(&greeting))
+            .await?;
+
+        // Wait for web console response (see `mgmt`)
+        waiter.await?.map_err(auth::AuthErrorImpl::auth_failed)
+    })
+    .await?;
+
+    client.write_message_noflush(&Be::NoticeResponse("Connecting to database."))?;
+
+    Ok(compute::NodeInfo {
+        db_info,
+        scram_keys: None,
+    })
+}

proxy/src/auth/backend/postgres.rs

@@ -0,0 +1,88 @@
+//! Local mock of Cloud API V2.
+
+use crate::{
+    auth::{
+        self,
+        backend::console::{self, io_error, AuthInfo, Result},
+        ClientCredentials, DatabaseInfo,
+    },
+    compute, scram,
+    stream::PqStream,
+    url::ApiUrl,
+};
+use tokio::io::{AsyncRead, AsyncWrite};
+
+#[must_use]
+pub(super) struct Api<'a> {
+    endpoint: &'a ApiUrl,
+    creds: &'a ClientCredentials,
+}
+
+impl<'a> Api<'a> {
+    /// Construct an API object containing the auth parameters.
+    pub(super) fn new(endpoint: &'a ApiUrl, creds: &'a ClientCredentials) -> Result<Self> {
+        Ok(Self { endpoint, creds })
+    }
+
+    /// Authenticate the existing user or throw an error.
+    pub(super) async fn handle_user(
+        self,
+        client: &mut PqStream<impl AsyncRead + AsyncWrite + Unpin + Send>,
+    ) -> auth::Result<compute::NodeInfo> {
+        // We reuse user handling logic from a production module.
+        console::handle_user(client, &self, Self::get_auth_info, Self::wake_compute).await
+    }
+
+    /// This implementation fetches the auth info from a local postgres instance.
+    async fn get_auth_info(&self) -> Result<AuthInfo> {
+        // Perhaps we could persist this connection, but then we'd have to
+        // write more code for reopening it if it got closed, which doesn't
+        // seem worth it.
+        let (client, connection) =
+            tokio_postgres::connect(self.endpoint.as_str(), tokio_postgres::NoTls)
+                .await
+                .map_err(io_error)?;
+
+        tokio::spawn(connection);
+        let query = "select rolpassword from pg_catalog.pg_authid where rolname = $1";
+        let rows = client
+            .query(query, &[&self.creds.user])
+            .await
+            .map_err(io_error)?;
+
+        match &rows[..] {
+            // We can't get a secret if there's no such user.
+            [] => Err(io_error(format!("unknown user '{}'", self.creds.user)).into()),
+
+            // We shouldn't get more than one row anyway.
+            [row, ..] => {
+                let entry = row.try_get(0).map_err(io_error)?;
+                scram::ServerSecret::parse(entry)
+                    .map(AuthInfo::Scram)
+                    .or_else(|| {
+                        // It could be an md5 hash if it's not a SCRAM secret.
+                        let text = entry.strip_prefix("md5")?;
+                        Some(AuthInfo::Md5({
+                            let mut bytes = [0u8; 16];
+                            hex::decode_to_slice(text, &mut bytes).ok()?;
+                            bytes
+                        }))
+                    })
+                    // Putting the secret into this message is a security hazard!
+                    .ok_or(console::ConsoleAuthError::BadSecret)
+            }
+        }
+    }
+
+    /// We don't need to wake anything locally, so we just return the connection info.
+    async fn wake_compute(&self) -> Result<DatabaseInfo> {
+        Ok(DatabaseInfo {
+            // TODO: handle that near CLI params parsing
+            host: self.endpoint.host_str().unwrap_or("localhost").to_owned(),
+            port: self.endpoint.port().unwrap_or(5432),
+            dbname: self.creds.dbname.to_owned(),
+            user: self.creds.user.to_owned(),
+            password: None,
+        })
+    }
+}