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

Micro-optimize GCS token access #1353

Merged
merged 2 commits into from
Jan 30, 2024
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 4 additions & 4 deletions crates/symbolicator-service/src/download/gcs.rs
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@ use crate::utils::gcs::{self, GcsToken};
use crate::utils::http::DownloadTimeouts;

/// An LRU cache for GCS OAuth tokens.
type GcsTokenCache = moka::future::Cache<Arc<GcsSourceKey>, CacheEntry<Arc<GcsToken>>>;
type GcsTokenCache = moka::future::Cache<Arc<GcsSourceKey>, CacheEntry<GcsToken>>;

/// Downloader implementation that supports the GCS source.
#[derive(Debug)]
Expand All @@ -35,16 +35,16 @@ impl GcsDownloader {
///
/// If the cache contains a valid token, then this token is returned. Otherwise, a new token is
/// requested from GCS and stored in the cache.
async fn get_token(&self, source_key: &Arc<GcsSourceKey>) -> CacheEntry<Arc<GcsToken>> {
async fn get_token(&self, source_key: &Arc<GcsSourceKey>) -> CacheEntry<GcsToken> {
metric!(counter("source.gcs.token.access") += 1);

let init = Box::pin(async {
metric!(counter("source.gcs.token.computation") += 1);
let token = gcs::request_new_token(&self.client, source_key).await;
token.map(Arc::new).map_err(CacheError::from)
token.map_err(CacheError::from)
});
let replace_if =
|entry: &CacheEntry<Arc<GcsToken>>| entry.as_ref().map_or(true, |t| t.is_expired());
|entry: &CacheEntry<GcsToken>| entry.as_ref().map_or(true, |t| t.is_expired());

self.token_cache
.entry_by_ref(source_key)
Expand Down
42 changes: 23 additions & 19 deletions crates/symbolicator-service/src/utils/gcs.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,9 @@
//! Access to Google Cloud Storeage

use std::sync::Arc;

use chrono::{DateTime, Duration, Utc};
use jsonwebtoken::errors::Error as JwtError;
use jsonwebtoken::EncodingKey;
use reqwest::Client;
use serde::{Deserialize, Serialize};
Expand All @@ -10,9 +13,9 @@ use url::Url;
use symbolicator_sources::GcsSourceKey;

/// A JWT token usable for GCS.
#[derive(Debug)]
#[derive(Debug, Clone)]
pub struct GcsToken {
access_token: String,
bearer_token: Arc<str>,
expires_at: DateTime<Utc>,
}

Expand All @@ -23,8 +26,8 @@ impl GcsToken {
}

/// The token in the HTTP Bearer-header format, header value only.
pub fn bearer_token(&self) -> String {
format!("Bearer {}", self.access_token)
pub fn bearer_token(&self) -> &str {
&self.bearer_token
}
}

Expand Down Expand Up @@ -57,22 +60,11 @@ pub enum GcsError {
#[error("failed to construct URL")]
InvalidUrl,
#[error("failed encoding JWT")]
Jwt(#[from] jsonwebtoken::errors::Error),
Jwt(#[from] JwtError),
#[error("failed to send authentication request")]
Auth(#[source] reqwest::Error),
}

/// Returns the JWT key parsed from a string.
///
/// Because Google provides this key in JSON format a lot of users just copy-paste this key
/// directly, leaving the escaped newlines from the JSON-encoding in place. In normal
/// base64 this should not occur so we pre-process the key to convert these back to real
/// newlines, ensuring they are in the correct PEM format.
fn key_from_string(key: &str) -> Result<EncodingKey, jsonwebtoken::errors::Error> {
let buffer = key.replace("\\n", "\n");
EncodingKey::from_rsa_pem(buffer.as_bytes())
}

/// Returns the URL for an object.
///
/// This can be used to e.g. fetch the objects's metadata.
Expand All @@ -92,8 +84,19 @@ pub fn download_url(bucket: &str, object: &str) -> Result<Url, GcsError> {
Ok(url)
}

/// Returns the JWT key parsed from a string.
///
/// Because Google provides this key in JSON format a lot of users just copy-paste this key
/// directly, leaving the escaped newlines from the JSON-encoding in place. In normal
/// base64 this should not occur so we pre-process the key to convert these back to real
/// newlines, ensuring they are in the correct PEM format.
fn key_from_string(key: &str) -> Result<EncodingKey, JwtError> {
let buffer = key.replace("\\n", "\n");
EncodingKey::from_rsa_pem(buffer.as_bytes())
}

/// Computes a JWT authentication assertion for the given GCS bucket.
fn get_auth_jwt(source_key: &GcsSourceKey, expiration: i64) -> Result<String, GcsError> {
fn get_auth_jwt(source_key: &GcsSourceKey, expiration: i64) -> Result<String, JwtError> {
let header = jsonwebtoken::Header::new(jsonwebtoken::Algorithm::RS256);

let jwt_claims = JwtClaims {
Expand All @@ -106,7 +109,7 @@ fn get_auth_jwt(source_key: &GcsSourceKey, expiration: i64) -> Result<String, Gc

let key = key_from_string(&source_key.private_key)?;

Ok(jsonwebtoken::encode(&header, &jwt_claims, &key)?)
jsonwebtoken::encode(&header, &jwt_claims, &key)
}

/// Requests a new GCS OAuth token.
Expand All @@ -133,9 +136,10 @@ pub async fn request_new_token(
.json::<GcsTokenResponse>()
.await
.map_err(GcsError::Auth)?;
let bearer_token = format!("Bearer {}", token.access_token).into();

Ok(GcsToken {
access_token: token.access_token,
bearer_token,
expires_at,
})
}
Expand Down
Loading