-
Notifications
You must be signed in to change notification settings - Fork 0
refactor: make token a watch channel #39
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
Changes from all commits
b7a2eb1
0def8bb
7b5900c
270f8b4
d0077f3
5f96be5
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -4,13 +4,17 @@ use crate::{ | |
deps::tracing::{error, info}, | ||
utils::from_env::FromEnv, | ||
}; | ||
use core::fmt; | ||
use oauth2::{ | ||
basic::{BasicClient, BasicTokenType}, | ||
AuthUrl, ClientId, ClientSecret, EmptyExtraTokenFields, EndpointNotSet, EndpointSet, | ||
HttpClientError, RequestTokenError, StandardErrorResponse, StandardTokenResponse, TokenUrl, | ||
AccessToken, AuthUrl, ClientId, ClientSecret, EmptyExtraTokenFields, EndpointNotSet, | ||
EndpointSet, HttpClientError, RefreshToken, RequestTokenError, Scope, StandardErrorResponse, | ||
StandardTokenResponse, TokenResponse, TokenUrl, | ||
}; | ||
use tokio::{ | ||
sync::watch::{self, Ref}, | ||
task::JoinHandle, | ||
}; | ||
use std::sync::{Arc, Mutex}; | ||
use tokio::task::JoinHandle; | ||
|
||
type Token = StandardTokenResponse<EmptyExtraTokenFields, BasicTokenType>; | ||
|
||
|
@@ -57,38 +61,17 @@ impl OAuthConfig { | |
} | ||
} | ||
|
||
/// A shared token that can be read and written to by multiple threads. | ||
#[derive(Debug, Clone, Default)] | ||
pub struct SharedToken(Arc<Mutex<Option<Token>>>); | ||
|
||
impl SharedToken { | ||
/// Read the token from the shared token. | ||
pub fn read(&self) -> Option<Token> { | ||
self.0.lock().unwrap().clone() | ||
} | ||
|
||
/// Write a new token to the shared token. | ||
pub fn write(&self, token: Token) { | ||
let mut lock = self.0.lock().unwrap(); | ||
*lock = Some(token); | ||
} | ||
|
||
/// Check if the token is authenticated. | ||
pub fn is_authenticated(&self) -> bool { | ||
self.0.lock().unwrap().is_some() | ||
} | ||
} | ||
|
||
/// A self-refreshing, periodically fetching authenticator for the block | ||
/// builder. This task periodically fetches a new token, and stores it in a | ||
/// [`SharedToken`]. | ||
/// builder. This task periodically fetches a new token, and sends it to all | ||
/// active [`SharedToken`]s via a [`tokio::sync::watch`] channel.. | ||
#[derive(Debug)] | ||
pub struct Authenticator { | ||
/// Configuration | ||
pub config: OAuthConfig, | ||
config: OAuthConfig, | ||
client: MyOAuthClient, | ||
token: SharedToken, | ||
reqwest: reqwest::Client, | ||
|
||
token: watch::Sender<Option<Token>>, | ||
} | ||
|
||
impl Authenticator { | ||
|
@@ -99,6 +82,8 @@ impl Authenticator { | |
.set_auth_uri(AuthUrl::from_url(config.oauth_authenticate_url.clone())) | ||
.set_token_uri(TokenUrl::from_url(config.oauth_token_url.clone())); | ||
|
||
// NB: this is MANDATORY | ||
// https://docs.rs/oauth2/latest/oauth2/#security-warning | ||
let rq_client = reqwest::Client::builder() | ||
.redirect(reqwest::redirect::Policy::none()) | ||
.build() | ||
|
@@ -107,8 +92,8 @@ impl Authenticator { | |
Self { | ||
config: config.clone(), | ||
client, | ||
token: Default::default(), | ||
reqwest: rq_client, | ||
token: watch::channel(None).0, | ||
} | ||
} | ||
|
||
|
@@ -129,20 +114,20 @@ impl Authenticator { | |
|
||
/// Returns true if there is Some token set | ||
pub fn is_authenticated(&self) -> bool { | ||
self.token.is_authenticated() | ||
self.token.borrow().is_some() | ||
} | ||
|
||
/// Sets the Authenticator's token to the provided value | ||
fn set_token(&self, token: StandardTokenResponse<EmptyExtraTokenFields, BasicTokenType>) { | ||
self.token.write(token); | ||
self.token.send_replace(Some(token)); | ||
} | ||
|
||
/// Returns the currently set token | ||
pub fn token(&self) -> SharedToken { | ||
self.token.clone() | ||
self.token.subscribe().into() | ||
} | ||
|
||
/// Fetches an oauth token | ||
/// Fetches an oauth token. | ||
pub async fn fetch_oauth_token( | ||
&self, | ||
) -> Result< | ||
|
@@ -161,25 +146,184 @@ impl Authenticator { | |
Ok(token_result) | ||
} | ||
|
||
/// Spawns a task that periodically fetches a new token every 300 seconds. | ||
pub fn spawn(self) -> JoinHandle<()> { | ||
/// Get a reference to the OAuth configuration. | ||
pub const fn config(&self) -> &OAuthConfig { | ||
&self.config | ||
} | ||
|
||
/// Create a future that contains the periodic refresh loop. | ||
async fn task_future(self) { | ||
let interval = self.config.oauth_token_refresh_interval; | ||
|
||
let handle: JoinHandle<()> = tokio::spawn(async move { | ||
loop { | ||
info!("Refreshing oauth token"); | ||
match self.authenticate().await { | ||
Ok(_) => { | ||
info!("Successfully refreshed oauth token"); | ||
} | ||
Err(e) => { | ||
error!(%e, "Failed to refresh oauth token"); | ||
} | ||
}; | ||
let _sleep = tokio::time::sleep(tokio::time::Duration::from_secs(interval)).await; | ||
} | ||
}); | ||
|
||
handle | ||
loop { | ||
info!("Refreshing oauth token"); | ||
match self.authenticate().await { | ||
Ok(_) => { | ||
info!("Successfully refreshed oauth token"); | ||
} | ||
Err(e) => { | ||
error!(%e, "Failed to refresh oauth token"); | ||
} | ||
}; | ||
let _sleep = tokio::time::sleep(tokio::time::Duration::from_secs(interval)).await; | ||
} | ||
} | ||
|
||
/// Spawns a task that periodically fetches a new token. The refresh | ||
/// interval may be configured via the | ||
/// [`OAuthConfig::oauth_token_refresh_interval`] property. | ||
pub fn spawn(self) -> JoinHandle<()> { | ||
tokio::spawn(self.task_future()) | ||
} | ||
} | ||
|
||
/// A shared token, wrapped in a [`tokio::sync::watch`] Receiver. The token is | ||
/// periodically refreshed by an [`Authenticator`] task, and can be awaited | ||
/// for when it becomes available. | ||
/// | ||
/// This allows multiple tasks to wait for the token to be available, and | ||
/// provides a way to check if the token is authenticated without blocking. | ||
/// Please consult the [`Receiver`] documentation for caveats regarding | ||
/// usage. | ||
/// | ||
/// [`Receiver`]: tokio::sync::watch::Receiver | ||
#[derive(Debug, Clone)] | ||
pub struct SharedToken(watch::Receiver<Option<Token>>); | ||
|
||
impl From<watch::Receiver<Option<Token>>> for SharedToken { | ||
fn from(inner: watch::Receiver<Option<Token>>) -> Self { | ||
Self(inner) | ||
} | ||
} | ||
|
||
impl SharedToken { | ||
/// Wait for the token to be available, and get a reference to the secret. | ||
/// | ||
/// This is implemented using [`Receiver::wait_for`], and has the same | ||
/// blocking, panics, errors, and cancel safety. However, it uses a clone | ||
/// of the [`watch::Receiver`] and will not update the local view of the | ||
/// channel. | ||
Comment on lines
+203
to
+205
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. The actual plus for this fn is that this:
no? There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
this is not really relevant, as we don't have any situations where we wait for the next updated token. We always care only about the immediate value. If we hypothetically wanted to take some action on token refresh, we would care about updating There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I see the actual + for this function compared to the previous version with the rwlock that you can There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. the other thing here is not requiring |
||
/// | ||
/// [`Receiver::wait_for`]: tokio::sync::watch::Receiver::wait_for | ||
pub async fn secret(&self) -> Result<String, watch::error::RecvError> { | ||
Ok(self | ||
.clone() | ||
.token() | ||
.await? | ||
.access_token() | ||
.secret() | ||
.to_owned()) | ||
} | ||
|
||
/// Wait for the token to be available, then get a reference to it. | ||
/// | ||
/// Holding this reference will block the background task from updating | ||
/// the token until it is dropped, so it is recommended to drop this | ||
/// reference as soon as possible. | ||
/// | ||
/// This is implemented using [`Receiver::wait_for`], and has the same | ||
/// blocking, panics, errors, and cancel safety. Unlike [`Self::secret`] | ||
/// it is NOT implemented using a clone, and will update the local view of | ||
/// the channel. | ||
prestwich marked this conversation as resolved.
Show resolved
Hide resolved
|
||
/// | ||
/// Generally, prefer using [`Self::secret`] for simple use cases, and | ||
/// this when deeper inspection of the token is required. | ||
/// | ||
/// [`Receiver::wait_for`]: tokio::sync::watch::Receiver::wait_for | ||
pub async fn token(&mut self) -> Result<TokenRef<'_>, watch::error::RecvError> { | ||
self.0.wait_for(Option::is_some).await.map(Into::into) | ||
} | ||
|
||
/// Create a future that will resolve when the token is ready. | ||
/// | ||
/// This is implemented using [`Receiver::wait_for`], and has the same | ||
/// blocking, panics, errors, and cancel safety. | ||
/// | ||
/// [`Receiver::wait_for`]: tokio::sync::watch::Receiver::wait_for | ||
pub async fn wait(&self) -> Result<(), watch::error::RecvError> { | ||
self.clone().0.wait_for(Option::is_some).await.map(drop) | ||
} | ||
|
||
/// Borrow the current token, if available. If called before the token is | ||
/// set by the authentication task, this will return `None`. | ||
/// | ||
/// Holding this reference will block the background task from updating | ||
/// the token until it is dropped, so it is recommended to drop this | ||
/// reference as soon as possible. | ||
/// | ||
/// This is implemented using [`Receiver::borrow`]. | ||
/// | ||
/// [`Receiver::borrow`]: tokio::sync::watch::Receiver::borrow | ||
pub fn borrow(&mut self) -> Ref<'_, Option<Token>> { | ||
self.0.borrow() | ||
} | ||
|
||
/// Check if the background task has produced an authentication token. | ||
/// | ||
/// Holding this reference will block the background task from updating | ||
/// the token until it is dropped, so it is recommended to drop this | ||
/// reference as soon as possible. | ||
/// | ||
/// This is implemented using [`Receiver::borrow`]. | ||
/// | ||
/// [`Receiver::borrow`]: tokio::sync::watch::Receiver::borrow | ||
pub fn is_authenticated(&self) -> bool { | ||
self.0.borrow().is_some() | ||
} | ||
} | ||
|
||
/// A reference to token data, contained in a [`SharedToken`]. | ||
/// | ||
/// This is implemented using [`watch::Ref`], and as a result holds a lock on | ||
/// the token data. Holding this reference will block the background task | ||
/// from updating the token until it is dropped, so it is recommended to drop | ||
/// this reference as soon as possible. | ||
pub struct TokenRef<'a> { | ||
inner: Ref<'a, Option<Token>>, | ||
} | ||
|
||
impl<'a> From<Ref<'a, Option<Token>>> for TokenRef<'a> { | ||
fn from(inner: Ref<'a, Option<Token>>) -> Self { | ||
Self { inner } | ||
} | ||
} | ||
|
||
impl fmt::Debug for TokenRef<'_> { | ||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { | ||
f.debug_struct("TokenRef").finish_non_exhaustive() | ||
} | ||
} | ||
|
||
impl<'a> TokenRef<'a> { | ||
/// Get a reference to the inner token. | ||
pub fn inner(&'a self) -> &'a Token { | ||
self.inner.as_ref().unwrap() | ||
} | ||
|
||
/// Get a reference to the [`AccessToken`] contained in the token. | ||
pub fn access_token(&self) -> &AccessToken { | ||
self.inner().access_token() | ||
} | ||
|
||
/// Get a reference to the [`TokenType`] instance contained in the token. | ||
/// | ||
/// [`TokenType`]: oauth2::TokenType | ||
pub fn token_type(&self) -> &<Token as TokenResponse>::TokenType { | ||
self.inner().token_type() | ||
} | ||
|
||
/// Get a reference to the current token's expiration time, if it has one. | ||
pub fn expires_in(&self) -> Option<std::time::Duration> { | ||
self.inner().expires_in() | ||
} | ||
|
||
/// Get a reference to the refresh token, if it exists. | ||
pub fn refresh_token(&self) -> Option<&RefreshToken> { | ||
self.inner().refresh_token() | ||
} | ||
|
||
/// Get a reference to the scopes associated with the token, if any. | ||
pub fn scopes(&self) -> Option<&Vec<Scope>> { | ||
self.inner().scopes() | ||
} | ||
} |
Uh oh!
There was an error while loading. Please reload this page.