Repository: MystenLabs/sui
File: crates/sui-pg-db/src/tls.rs (Rust)

Code before:
            .supported_verify_schemes()
    }
}

/// Establish a PostgreSQL connection with custom TLS configuration using tokio-postgres. The
/// returned connection is compatible with diesel-async. This is needed because diesel-async does
/// not expose TLS configuration. Even with a TLS connector, the actual usage is negotiated with the
/// postgres server:
/// - Client sends SSLRequest, server responds 'S' (TLS supported) or 'N' (TLS not supported)
/// - If server supports TLS: encrypted connection using the provided config
/// - If server doesn't support TLS: falls back to plaintext (when using sslmode=prefer default)
/// - If sslmode=disable in URL: no TLS attempted regardless of connector
pub(crate) async fn establish_tls_connection(
    database_url: &str,
    tls_config: ClientConfig,
) -> ConnectionResult<AsyncPgConnectionWithId> {
    let tls = MakeRustlsConnect::new(tls_config);
    let (client, conn) = tokio_postgres::connect(database_url, tls)
        .await
        .map_err(|e| ConnectionError::BadConnection(e.to_string()))?;

    // The `conn` object performs actual IO with the database, and tokio-postgres suggests spawning
    // it off to run in the background. This will resolve only when the connection is closed, either
    // because of a fatal error or because its associated Client has dropped and all outstanding
    // work has completed.
    let conn_id = Uuid::new_v4();
    tokio::spawn(async move {
        if let Err(e) = conn.await {
            error!(%conn_id, "Database connection terminated: {e}");
        }
    });

    // Users interact with the database through the client object. We convert it into an
    // AsyncPgConnection so it can be compatible with diesel.
    match AsyncPgConnection::try_from(client).await {
        Ok(mut c) => {
            #[derive(QueryableByName)]
            struct Pid {
                #[diesel(sql_type = diesel::sql_types::Integer)]
                pg_backend_pid: i32,
            }

            let pid = diesel::sql_query("SELECT pg_backend_pid()")
                .get_result::<Pid>(&mut c)
                .await
                .map_err(ConnectionError::CouldntSetupConfiguration)?;

            info!(%conn_id, pid = ?pid.pg_backend_pid, "Database connection established");
            Ok(AsyncPgConnectionWithId {
                id: conn_id,
                conn: c,

Review comment (bug):
Instead of creating a new `conn_id` and associating this to `pg_backend_pid()`, can we use the `pg_backend_pid()` directly to represent the conn? If we fail to establish a connection then there isn't anything for us to trace through anyways right

Code after:
        _dss: &DigitallySignedStruct,
    ) -> Result<HandshakeSignatureValid, rustls::Error> {
        Ok(HandshakeSignatureValid::assertion())
    }

    fn supported_verify_schemes(&self) -> Vec<rustls::SignatureScheme> {
        rustls::client::WebPkiServerVerifier::builder(Arc::new(root_certs()))
            .build()
            .unwrap()
            .supported_verify_schemes()
    }
}

/// Establish a PostgreSQL connection with custom TLS configuration using tokio-postgres. The
/// returned connection is compatible with diesel-async. This is needed because diesel-async does
/// not expose TLS configuration. Even with a TLS connector, the actual usage is negotiated with the
/// postgres server:
/// - Client sends SSLRequest, server responds 'S' (TLS supported) or 'N' (TLS not supported)
/// - If server supports TLS: encrypted connection using the provided config
/// - If server doesn't support TLS: falls back to plaintext (when using sslmode=prefer default)
/// - If sslmode=disable in URL: no TLS attempted regardless of connector
pub(crate) async fn establish_tls_connection(
    database_url: &str,
    tls_config: ClientConfig,
) -> ConnectionResult<AsyncPgConnectionWithId> {
    let tls = MakeRustlsConnect::new(tls_config);
    let (client, tokio_conn) = tokio_postgres::connect(database_url, tls)
        .await
        .map_err(|e| ConnectionError::BadConnection(e.to_string()))?;

    // The `conn` object performs actual IO with the database, and tokio-postgres suggests spawning
    // it off to run in the background. This will resolve only when the connection is closed,
    // either because of a fatal error or because its associated Client has dropped and all
    // outstanding work has completed.
    let h_conn = tokio::spawn(tokio_conn);

    // Users interact with the database through the client object. We convert it into an
    // AsyncPgConnection so it can be compatible with diesel.
    let mut conn = AsyncPgConnection::try_from(client).await?;

    #[derive(QueryableByName)]
    struct Pid {
        #[diesel(sql_type = Integer)]
        pid: i32,
    }

    let Pid { pid } = diesel::sql_query("SELECT pg_backend_pid() AS pid")
        .get_result(&mut conn)
        .await
        .map_err(ConnectionError::CouldntSetupConfiguration)?;