Merge pull request #6 from Isawan/credential-helper

Implemented authenticated request

.github/workflows/rust.yml

@@ -20,7 +20,7 @@ jobs:
     - name: Build
       run: cargo build --verbose
     - name: Start containers
-      run: docker-compose up -d
+      run: docker-compose up -d --wait
     - name: Install self signed certificate
       run: sudo cp ./resources/test/certs/cert.pem /usr/local/share/ca-certificates/localhost.crt && sudo update-ca-certificates
     - name: Show containers

src/app.rs

@@ -8,17 +8,17 @@ use tower_http::{
 use tracing::Level;
 
 use crate::{
-    config::Args, http::artifacts::artifacts_handler, http::healthcheck::healthcheck_handler,
-    http::index::index_handler, http::version::version_handler, refresh::RefreshRequest,
-    registry::RegistryClient,
+    config::Args, credhelper::database::DatabaseCredentials, http::artifacts::artifacts_handler,
+    http::healthcheck::healthcheck_handler, http::index::index_handler,
+    http::version::version_handler, refresh::RefreshRequest, registry::RegistryClient,
 };
 
 #[derive(Clone)]
 pub(crate) struct AppState {
     pub(crate) s3_client: aws_sdk_s3::Client,
     pub(crate) http_client: reqwest::Client,
     pub(crate) db_client: Pool<Postgres>,
-    pub(crate) registry_client: RegistryClient,
+    pub(crate) registry_client: RegistryClient<DatabaseCredentials>,
     pub(crate) config: Args,
     pub(crate) refresher_tx: mpsc::Sender<RefreshRequest>,
 }
@@ -34,8 +34,12 @@ impl AppState {
         AppState {
             s3_client: s3,
             http_client: http.clone(),
-            db_client: db,
-            registry_client: RegistryClient::new(config.upstream_registry_port, http),
+            db_client: db.clone(),
+            registry_client: RegistryClient::new(
+                config.upstream_registry_port,
+                http,
+                DatabaseCredentials::new(db),
+            ),
             config,
             refresher_tx,
         }

src/credhelper/database.rs

@@ -1,9 +1,12 @@
+use std::marker::Send;
+
 use async_trait::async_trait;
 use sqlx::PgPool;
 
 use super::CredentialHelper;
 
 // Credential helper implementation by storing in the database
+#[derive(Clone)]
 pub struct DatabaseCredentials {
     pool: PgPool,
 }
@@ -14,9 +17,9 @@ impl DatabaseCredentials {
     }
 }
 
-#[async_trait(?Send)]
+#[async_trait]
 impl CredentialHelper for DatabaseCredentials {
-    async fn get(&self, hostname: impl AsRef<str>) -> Result<Option<String>, anyhow::Error> {
+    async fn get(&self, hostname: impl AsRef<str> + Send) -> Result<Option<String>, anyhow::Error> {
         let query = sqlx::query!(
             r#"
             select auth_token from terraform_registry_host
@@ -44,7 +47,7 @@ impl CredentialHelper for DatabaseCredentials {
         Ok(())
     }
 
-    async fn forget(&mut self, hostname: impl AsRef<str>) -> Result<(), anyhow::Error> {
+    async fn forget(&mut self, hostname: impl AsRef<str> + Send) -> Result<(), anyhow::Error> {
         let query = sqlx::query!(
             r#"
             delete from "terraform_registry_host" where "hostname" = $1;

src/credhelper/types.rs

@@ -1,11 +1,12 @@
 use async_trait::async_trait;
+use std::marker::Send;
 
-#[async_trait(?Send)]
+#[async_trait]
 pub trait CredentialHelper {
-    async fn get(&self, hostname: impl AsRef<str>) -> Result<Option<String>, anyhow::Error>;
+    async fn get(&self, hostname: impl AsRef<str> + Send) -> Result<Option<String>, anyhow::Error>;
     async fn store(&mut self, hostname: String, cred: String) -> Result<(), anyhow::Error>;
 
     /// Forget a credential.
     /// NOTE: Deleting a non-existing credential is not an error
-    async fn forget(&mut self, hostname: impl AsRef<str>) -> Result<(), anyhow::Error>;
+    async fn forget(&mut self, hostname: impl AsRef<str> + Send) -> Result<(), anyhow::Error>;
 }

src/http/artifacts.rs

@@ -1,5 +1,6 @@
 use crate::{
     app::AppState,
+    credhelper::CredentialHelper,
     registry::{ProviderResponse, RegistryClient},
 };
 use anyhow::Context;
@@ -192,9 +193,9 @@ async fn get_artifact_from_database(
     Ok(result)
 }
 
-async fn get_upstream(
+async fn get_upstream<T: CredentialHelper>(
     http: Client,
-    registry: RegistryClient,
+    registry: RegistryClient<T>,
     artifact: &ArtifactDetails,
 ) -> Result<Pin<Box<impl Stream<Item = reqwest::Result<Bytes>>>>, anyhow::Error> {
     let provider_path = format!(

src/http/index.rs

@@ -1,5 +1,6 @@
 use crate::{
     app::AppState,
+    credhelper::CredentialHelper,
     error::TerrashineError,
     refresh::{RefreshRequest, RefreshResponse, TerraformProvider},
     registry::{ProviderPlatform, ProviderVersionItem, ProviderVersions, RegistryClient},
@@ -117,9 +118,9 @@ pub(crate) async fn index_handler(
     }
 }
 
-pub(crate) async fn refresh_versions(
+pub(crate) async fn refresh_versions<T: CredentialHelper>(
     db: &PgPool,
-    registry: &RegistryClient,
+    registry: &RegistryClient<T>,
     hostname: &str,
     namespace: &str,
     provider_type: &str,

src/lib.rs

@@ -18,7 +18,9 @@ use tokio::{
 use tokio_util::sync::CancellationToken;
 use tracing::error;
 
-use crate::{refresh::refresher, registry::RegistryClient};
+use crate::{
+    credhelper::database::DatabaseCredentials, refresh::refresher, registry::RegistryClient,
+};
 
 #[derive(Debug)]
 pub struct StartUpNotify {
@@ -68,7 +70,11 @@ pub async fn run(
     };
 
     let refresher_db = db.clone();
-    let refresher_registry = RegistryClient::new(config.upstream_registry_port, http.clone());
+    let refresher_registry = RegistryClient::new(
+        config.upstream_registry_port,
+        http.clone(),
+        DatabaseCredentials::new(db.clone()),
+    );
     let refresher = refresher(
         &refresher_db,
         &refresher_registry,

src/refresh.rs

@@ -1,4 +1,5 @@
 use crate::{
+    credhelper::CredentialHelper,
     error::TerrashineError,
     http::index::refresh_versions,
     registry::{ProviderVersions, RegistryClient},
@@ -42,9 +43,9 @@ pub(crate) enum RefreshResponse {
     ProviderVersionNotStale,
 }
 
-pub(crate) async fn refresher(
+pub(crate) async fn refresher<T: CredentialHelper>(
     db: &PgPool,
-    registry: &RegistryClient,
+    registry: &RegistryClient<T>,
     mut rx: sync::mpsc::Receiver<RefreshRequest>,
     refresh_interval: Duration,
     cancel: CancellationToken,

src/registry/client.rs

@@ -3,20 +3,25 @@ use serde::Deserialize;
 use std::str;
 use url::Url;
 
-use crate::error::TerrashineError;
+use crate::{credhelper::CredentialHelper, error::TerrashineError};
 
 const DISCOVERY_RESPONSE_SIZE_MAX_BYTES: usize = 16384; // 16KB
 const REGISTRY_METADATA_SIZE_MAX_BYTES: usize = 8388608; // 8MB
 
 #[derive(Clone)]
-pub struct RegistryClient {
+pub struct RegistryClient<T> {
     port: u16,
     http: Client,
+    credentials: T,
 }
 
-impl RegistryClient {
-    pub fn new(port: u16, http: Client) -> Self {
-        RegistryClient { port, http }
+impl<T> RegistryClient<T> {
+    pub fn new(port: u16, http: Client, credentials: T) -> Self {
+        RegistryClient {
+            port,
+            http,
+            credentials,
+        }
     }
 }
 
@@ -45,7 +50,7 @@ async fn read_body_limit(
     Ok(())
 }
 
-impl RegistryClient {
+impl<T: CredentialHelper> RegistryClient<T> {
     /// Performs request upstream to handle terraform service discovery protocol
     async fn discover_services(
         &self,
@@ -73,8 +78,9 @@ impl RegistryClient {
         hostname: &str,
         path: &str,
     ) -> Result<A, TerrashineError> {
-        let hostname = hostname;
         let services = self.discover_services(hostname).await?;
+        let auth_token = self.credentials.get(hostname).await?;
+
         if let Some(base_url) = services.providers_v1 {
             let url = Url::parse(&format!("https://{hostname}:{port}", port = self.port))
                 .and_then(|u| u.join(&base_url))
@@ -87,7 +93,12 @@ impl RegistryClient {
                 })?;
             let mut response_buffer = Vec::with_capacity(REGISTRY_METADATA_SIZE_MAX_BYTES);
             tracing::debug!(%url, "GET registry provider");
-            let response = self.http.get(url).send().await?.error_for_status()?;
+            let mut request = self.http.get(url);
+            if let Some(token) = auth_token {
+                request = request.bearer_auth(token);
+            }
+
+            let response = request.send().await?.error_for_status()?;
             read_body_limit(
                 &mut response_buffer,
                 response,