Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

refactor: Improve AuthenticationStorage #1026

Merged
merged 24 commits into from
Jan 17, 2025
Merged
7 changes: 2 additions & 5 deletions crates/rattler-bin/src/commands/create.rs
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@ use rattler_conda_types::{
Channel, ChannelConfig, GenericVirtualPackage, MatchSpec, ParseStrictness, Platform,
PrefixRecord, RepoDataRecord, Version,
};
use rattler_networking::{AuthenticationMiddleware, AuthenticationStorage};
use rattler_networking::AuthenticationMiddleware;
use rattler_repodata_gateway::{Gateway, RepoData, SourceConfig};
use rattler_solve::{
libsolv_c::{self},
Expand Down Expand Up @@ -147,11 +147,8 @@ pub async fn create(opt: Opt) -> anyhow::Result<()> {
.build()
.expect("failed to create client");

let authentication_storage = AuthenticationStorage::default();
let download_client = reqwest_middleware::ClientBuilder::new(download_client)
.with_arc(Arc::new(AuthenticationMiddleware::new(
authentication_storage,
)))
.with_arc(Arc::new(AuthenticationMiddleware::default()?))
.with(rattler_networking::OciMiddleware)
.with(rattler_networking::GCSMiddleware)
.build();
Expand Down
1 change: 0 additions & 1 deletion crates/rattler_networking/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,6 @@ async-trait = { workspace = true }
base64 = { workspace = true }
chrono = { workspace = true }
dirs = { workspace = true }
fslock = { workspace = true }
google-cloud-auth = { workspace = true, optional = true }
google-cloud-token = { workspace = true, optional = true }
http = { workspace = true }
Expand Down
22 changes: 15 additions & 7 deletions crates/rattler_networking/src/authentication_middleware.rs
Original file line number Diff line number Diff line change
@@ -1,16 +1,17 @@
//! `reqwest` middleware that authenticates requests with data from the `AuthenticationStorage`

Check warning on line 1 in crates/rattler_networking/src/authentication_middleware.rs

View workflow job for this annotation

GitHub Actions / Format and Lint

Diff in /home/runner/work/rattler/rattler/crates/rattler_networking/src/authentication_middleware.rs
use crate::{Authentication, AuthenticationStorage};
use async_trait::async_trait;
use base64::prelude::BASE64_STANDARD;
use base64::Engine;
use reqwest::{Request, Response};
use reqwest_middleware::{Middleware, Next};
use std::path::{Path, PathBuf};

Check warning on line 8 in crates/rattler_networking/src/authentication_middleware.rs

View workflow job for this annotation

GitHub Actions / Format and Lint

Diff in /home/runner/work/rattler/rattler/crates/rattler_networking/src/authentication_middleware.rs
use std::sync::OnceLock;
use url::Url;
use anyhow::Result;

/// `reqwest` middleware to authenticate requests
#[derive(Clone, Default)]
#[derive(Clone)]
pub struct AuthenticationMiddleware {
auth_storage: AuthenticationStorage,
}
Expand Down Expand Up @@ -53,6 +54,13 @@
Self { auth_storage }
}

/// Create a new authentication middleware with the default authentication storage
pub fn default() -> Result<Self> {
Ok(Self {
auth_storage: AuthenticationStorage::default()?,
})
}

/// Authenticate the given URL with the given authentication information
fn authenticate_url(url: Url, auth: &Option<Authentication>) -> Url {
if let Some(credentials) = auth {
Expand Down Expand Up @@ -176,7 +184,7 @@
#[test]
fn test_store_fallback() -> anyhow::Result<()> {
let tdir = tempdir()?;
let mut storage = AuthenticationStorage::new();
let mut storage = AuthenticationStorage::empty();
storage.add_backend(Arc::from(FileStorage::new(
tdir.path().to_path_buf().join("auth.json"),
)?));
Expand All @@ -191,7 +199,7 @@
#[tokio::test]
async fn test_conda_token_storage() -> anyhow::Result<()> {
let tdir = tempdir()?;
let mut storage = AuthenticationStorage::new();
let mut storage = AuthenticationStorage::empty();
storage.add_backend(Arc::from(FileStorage::new(
tdir.path().to_path_buf().join("auth.json"),
)?));
Expand Down Expand Up @@ -245,7 +253,7 @@
#[tokio::test]
async fn test_bearer_storage() -> anyhow::Result<()> {
let tdir = tempdir()?;
let mut storage = AuthenticationStorage::new();
let mut storage = AuthenticationStorage::empty();
storage.add_backend(Arc::from(FileStorage::new(
tdir.path().to_path_buf().join("auth.json"),
)?));
Expand Down Expand Up @@ -305,7 +313,7 @@
#[tokio::test]
async fn test_basic_auth_storage() -> anyhow::Result<()> {
let tdir = tempdir()?;
let mut storage = AuthenticationStorage::new();
let mut storage = AuthenticationStorage::empty();
storage.add_backend(Arc::from(FileStorage::new(
tdir.path().to_path_buf().join("auth.json"),
)?));
Expand Down Expand Up @@ -383,7 +391,7 @@
("*.com", false),
] {
let tdir = tempdir()?;
let mut storage = AuthenticationStorage::new();
let mut storage = AuthenticationStorage::empty();
storage.add_backend(Arc::from(FileStorage::new(
tdir.path().to_path_buf().join("auth.json"),
)?));
Expand Down Expand Up @@ -418,7 +426,7 @@
.to_str()
.unwrap(),
),
|| AuthenticationStorage::from_env().unwrap(),
|| AuthenticationStorage::default().unwrap(),
);

let host = "test.example.com";
Expand Down
Original file line number Diff line number Diff line change
@@ -1,17 +1,16 @@
//! file storage for passwords.
use anyhow::Result;

Check warning on line 2 in crates/rattler_networking/src/authentication_storage/backends/file.rs

View workflow job for this annotation

GitHub Actions / Format and Lint

Diff in /home/runner/work/rattler/rattler/crates/rattler_networking/src/authentication_storage/backends/file.rs
use fslock::LockFile;
use std::collections::BTreeMap;
use std::path::Path;
use std::sync::Arc;
use std::{path::PathBuf, sync::Mutex};
use std::sync::{Arc, RwLock};
use std::path::PathBuf;

use crate::authentication_storage::StorageBackend;
use crate::Authentication;

#[derive(Clone, Debug)]
struct FileStorageCache {
cache: BTreeMap<String, Authentication>,
content: BTreeMap<String, Authentication>,
file_exists: bool,
}

Expand All @@ -25,7 +24,7 @@
/// The cache of the file storage
/// This is used to avoid reading the file from disk every time
/// a credential is accessed
cache: Arc<Mutex<FileStorageCache>>,
cache: Arc<RwLock<FileStorageCache>>,
pavelzw marked this conversation as resolved.
Show resolved Hide resolved
}

/// An error that can occur when accessing the file storage
Expand All @@ -35,87 +34,54 @@
#[error("IO error: {0}")]
IOError(#[from] std::io::Error),

/// Failed to lock the file storage file
#[error("failed to lock file storage file {0}")]
FailedToLock(String, #[source] std::io::Error),

/// An error occurred when (de)serializing the credentials
#[error("JSON error: {0}")]
JSONError(#[from] serde_json::Error),
}

/// Lock the file storage file for reading and writing. This will block until the lock is
/// acquired.
fn lock_file_storage(path: &Path, write: bool) -> Result<Option<LockFile>, FileStorageError> {
if !write && !path.exists() {
return Ok(None);
}

std::fs::create_dir_all(path.parent().unwrap())?;
let path = path.with_extension("lock");
let mut lock = fslock::LockFile::open(&path)
.map_err(|e| FileStorageError::FailedToLock(path.to_string_lossy().into_owned(), e))?;

// First try to lock the file without block. If we can't immediately get the lock we block and issue a debug message.
if !lock
.try_lock_with_pid()
.map_err(|e| FileStorageError::FailedToLock(path.to_string_lossy().into_owned(), e))?
{
tracing::debug!("waiting for lock on {}", path.to_string_lossy());
lock.lock_with_pid()
.map_err(|e| FileStorageError::FailedToLock(path.to_string_lossy().into_owned(), e))?;
}

Ok(Some(lock))
}

impl FileStorageCache {
pub fn from_path(path: &Path) -> Result<Self, FileStorageError> {
let file_exists = path.exists();
let cache = if file_exists {
lock_file_storage(path, false)?;
let content = if file_exists {
let file = std::fs::File::open(path)?;
let reader = std::io::BufReader::new(file);
serde_json::from_reader(reader)?
} else {
BTreeMap::new()

Check warning on line 50 in crates/rattler_networking/src/authentication_storage/backends/file.rs

View workflow job for this annotation

GitHub Actions / Format and Lint

Diff in /home/runner/work/rattler/rattler/crates/rattler_networking/src/authentication_storage/backends/file.rs
};

Ok(Self { cache, file_exists })
Ok(Self { content, file_exists })
}
}

impl FileStorage {
/// Create a new file storage with the given path
pub fn new(path: PathBuf) -> Result<Self, FileStorageError> {
// read the JSON file if it exists, and store it in the cache
let cache = Arc::new(Mutex::new(FileStorageCache::from_path(&path)?));
let cache = Arc::new(RwLock::new(FileStorageCache::from_path(&path)?));

Ok(Self { path, cache })
}

/// Read the JSON file and deserialize it into a `BTreeMap`, or return an empty `BTreeMap` if the
/// Updates the cache by reading the JSON file and deserializing it into a `BTreeMap`, or return an empty `BTreeMap` if the
/// file does not exist
fn read_json(&self) -> Result<BTreeMap<String, Authentication>, FileStorageError> {
let new_cache = FileStorageCache::from_path(&self.path)?;
let mut cache = self.cache.lock().unwrap();
cache.cache = new_cache.cache;
let mut cache = self.cache.write().unwrap();
cache.content = new_cache.content;
cache.file_exists = new_cache.file_exists;

Ok(cache.cache.clone())
Ok(cache.content.clone())
}

/// Serialize the given `BTreeMap` and write it to the JSON file
fn write_json(&self, dict: &BTreeMap<String, Authentication>) -> Result<(), FileStorageError> {
let _lock = lock_file_storage(&self.path, true)?;

let file = std::fs::File::create(&self.path)?;
let writer = std::io::BufWriter::new(file);
serde_json::to_writer(writer, dict)?;

// Store the new data in the cache
let mut cache = self.cache.lock().unwrap();
cache.cache = dict.clone();
let mut cache = self.cache.write().unwrap();
cache.content = dict.clone();
cache.file_exists = true;

Ok(())
Expand All @@ -130,8 +96,8 @@
}

fn get(&self, host: &str) -> Result<Option<crate::Authentication>> {
let cache = self.cache.lock().unwrap();
Ok(cache.cache.get(host).cloned())
let cache = self.cache.read().unwrap();
Ok(cache.content.get(host).cloned())
}

fn delete(&self, host: &str) -> Result<()> {
Expand All @@ -151,8 +117,8 @@
path.push("credentials.json");
Self::new(path.clone()).unwrap_or(Self {
path,
cache: Arc::new(Mutex::new(FileStorageCache {
cache: BTreeMap::new(),
cache: Arc::new(RwLock::new(FileStorageCache {
content: BTreeMap::new(),
file_exists: false,
})),
})
Expand Down
56 changes: 28 additions & 28 deletions crates/rattler_networking/src/authentication_storage/storage.rs
Original file line number Diff line number Diff line change
Expand Up @@ -25,54 +25,48 @@ pub struct AuthenticationStorage {
cache: Arc<Mutex<HashMap<String, Option<Authentication>>>>,
}

impl Default for AuthenticationStorage {
fn default() -> Self {
let mut storage = Self::new();

storage.add_backend(Arc::from(KeyringAuthenticationStorage::default()));
storage.add_backend(Arc::from(FileStorage::default()));
storage.add_backend(Arc::from(NetRcStorage::from_env().unwrap_or_else(
|(path, err)| {
tracing::warn!("error reading netrc file from {}: {}", path.display(), err);
NetRcStorage::default()
},
)));

storage
}
}

impl AuthenticationStorage {
/// Create a new authentication storage with no backends
pub fn new() -> Self {
pub fn empty() -> Self {
Self {
backends: vec![],
cache: Arc::new(Mutex::new(HashMap::new())),
}
}

/// Create a new authentication storage with the default backends
/// respecting the `RATTLER_AUTH_FILE` environment variable.
/// If the variable is set, the file storage backend will be used
/// with the path specified in the variable
pub fn from_env() -> Result<Self> {
/// Following order:
/// - file storage from $RATTLER_AUTH_FILE (if set)
pavelzw marked this conversation as resolved.
Show resolved Hide resolved
/// - file storage from the default location
/// - keyring storage
/// - netrc storage
pub fn default() -> Result<Self> {
let mut storage = Self::empty();

if let Ok(auth_file) = std::env::var("RATTLER_AUTH_FILE") {
let path = std::path::Path::new(&auth_file);

tracing::info!(
"\"RATTLER_AUTH_FILE\" environment variable set, using file storage at {}",
auth_file
);

Ok(Self::from_file(path)?)
} else {
Ok(Self::default())
storage.add_backend(Arc::from(FileStorage::new(path.into())?));
}
storage.add_backend(Arc::from(FileStorage::default()));
storage.add_backend(Arc::from(KeyringAuthenticationStorage::default()));
storage.add_backend(Arc::from(NetRcStorage::from_env().unwrap_or_else(
|(path, err)| {
tracing::warn!("error reading netrc file from {}: {}", path.display(), err);
NetRcStorage::default()
},
)));

Ok(storage)
}

/// Create a new authentication storage with just a file storage backend
/// Returns an error if the file is not found.
pub fn from_file(path: &std::path::Path) -> Result<Self> {
Copy link
Contributor Author

Choose a reason for hiding this comment

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

this was introduced in #645 but i don't think this is used anywhere so i don't think this brings us much benefit, especially now since RATTLER_AUTH_FILE is handled directly in AuthenticationStorage::default()

let mut storage = Self::new();
let mut storage = Self::empty();
let backend = FileStorage::new(path.to_path_buf()).map_err(|e| {
anyhow!(
"Error creating file storage backend from file ({}): {}",
Expand All @@ -91,6 +85,12 @@ impl AuthenticationStorage {
self.backends.push(backend);
}

/// Add a new storage backend to the authentication storage at the given index
/// (backends are tried in the order they are added)
pub fn insert_backend(&mut self, index: usize, backend: Arc<dyn StorageBackend + Send + Sync>) {
Copy link
Contributor Author

Choose a reason for hiding this comment

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

This will be needed for pixi's authentication-file-override in its global config.

My suggestion would be the following order:

  • file storage from $RATTLER_AUTH_FILE (if set)
  • file specified in authentication-file-override in pixi global config (if set) (pixi-only that's why we need to insert_backend(1, ...) in pixi)
  • file storage from the default location
  • keyring storage
  • netrc storage

self.backends.insert(index, backend);
}

/// Store the given authentication information for the given host
pub fn store(&self, host: &str, authentication: &Authentication) -> Result<()> {
{
Expand Down
2 changes: 1 addition & 1 deletion crates/rattler_repodata_gateway/src/fetch/jlap/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@
//! pub async fn main() {
//! let subdir_url = Url::parse("https://conda.anaconda.org/conda-forge/osx-64/").unwrap();
//! let client = reqwest_middleware::ClientBuilder::new(reqwest::Client::new())
//! .with_arc(Arc::new(AuthenticationMiddleware::default()))
//! .with_arc(Arc::new(AuthenticationMiddleware::default().unwrap()))
//! .build();
//! let cache = Path::new("./cache");
//! let current_repo_data = cache.join("c93ef9c9.json");
Expand Down
2 changes: 1 addition & 1 deletion crates/rattler_repodata_gateway/src/fetch/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1472,7 +1472,7 @@ mod test {

let client = Client::builder().no_gzip().build().unwrap();
let authenticated_client = reqwest_middleware::ClientBuilder::new(client)
.with_arc(Arc::new(AuthenticationMiddleware::default()))
.with_arc(Arc::new(AuthenticationMiddleware::default().unwrap()))
.build();

let result = fetch_repo_data(
Expand Down
Loading