oak-team/oak-gpui
3
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

WIP: Clear cached credentials if authentication fails

Browse Source 
Still need to actually handle an HTTP response from the server indicating there was an invalid token.

Co-Authored-By: Max Brunsfeld <maxbrunsfeld@gmail.com>
This commit is contained in:
Nathan Sobo
2021-09-14 19:19:11 -06:00
co-authored by Max Brunsfeld
parent 77a4a36eb3
commit 4a9918979e
8 changed files with 149 additions and 67 deletions
Download Patch File Download Diff File Expand all files Collapse all files
zed/src/rpc.rs
+75 -16
View File
@@ -15,10 +15,11 @@ use std::{
time::{Duration, Instant},
};
use surf::Url;
use thiserror::Error;
pub use zrpc::{proto, ConnectionId, PeerId, TypedEnvelope};
use zrpc::{
proto::{AnyTypedEnvelope, EntityMessage, EnvelopedMessage, RequestMessage},
Conn, Peer, Receipt,
Connection, Peer, Receipt,
};
lazy_static! {
@@ -32,10 +33,32 @@ pub struct Client {
authenticate:
Option<Box<dyn 'static + Send + Sync + Fn(&AsyncAppContext) -> Task<Result<Credentials>>>>,
establish_connection: Option<
Box<dyn 'static + Send + Sync + Fn(&Credentials, &AsyncAppContext) -> Task<Result<Conn>>>,
Box<
dyn 'static
+ Send
+ Sync
+ Fn(
&Credentials,
&AsyncAppContext,
) -> Task<Result<Connection, EstablishConnectionError>>,
>,
>,
}
#[derive(Error, Debug)]
pub enum EstablishConnectionError {
#[error("invalid access token")]
InvalidAccessToken,
#[error("{0}")]
Other(anyhow::Error),
}
impl EstablishConnectionError {
pub fn other(error: impl Into<anyhow::Error> + Send + Sync) -> Self {
Self::Other(error.into())
}
}
#[derive(Copy, Clone, Debug)]
pub enum Status {
SignedOut,
@@ -122,7 +145,10 @@ impl Client {
#[cfg(any(test, feature = "test-support"))]
pub fn override_establish_connection<F>(&mut self, connect: F) -> &mut Self
where
F: 'static + Send + Sync + Fn(&Credentials, &AsyncAppContext) -> Task<Result<Conn>>,
F: 'static
+ Send
+ Sync
+ Fn(&Credentials, &AsyncAppContext) -> Task<Result<Connection, EstablishConnectionError>>,
{
self.establish_connection = Some(Box::new(connect));
self
@@ -288,13 +314,18 @@ impl Client {
Ok(())
}
Err(err) => {
eprintln!("error in authenticate and connect {}", err);
if matches!(err, EstablishConnectionError::InvalidAccessToken) {
eprintln!("nuking credentials");
self.state.write().credentials.take();
}
self.set_status(Status::ConnectionError, cx);
Err(err)
Err(err)?
}
}
}
async fn set_connection(self: &Arc<Self>, conn: Conn, cx: &AsyncAppContext) {
async fn set_connection(self: &Arc<Self>, conn: Connection, cx: &AsyncAppContext) {
let (connection_id, handle_io, mut incoming) = self.peer.add_connection(conn).await;
cx.foreground()
.spawn({
@@ -359,7 +390,7 @@ impl Client {
self: &Arc<Self>,
credentials: &Credentials,
cx: &AsyncAppContext,
) -> Task<Result<Conn>> {
) -> Task<Result<Connection, EstablishConnectionError>> {
if let Some(callback) = self.establish_connection.as_ref() {
callback(credentials, cx)
} else {
@@ -371,28 +402,43 @@ impl Client {
self: &Arc<Self>,
credentials: &Credentials,
cx: &AsyncAppContext,
) -> Task<Result<Conn>> {
) -> Task<Result<Connection, EstablishConnectionError>> {
let request = Request::builder().header(
"Authorization",
format!("{} {}", credentials.user_id, credentials.access_token),
);
cx.background().spawn(async move {
if let Some(host) = ZED_SERVER_URL.strip_prefix("https://") {
let stream = smol::net::TcpStream::connect(host).await?;
let request = request.uri(format!("wss://{}/rpc", host)).body(())?;
let stream = smol::net::TcpStream::connect(host)
.await
.map_err(EstablishConnectionError::other)?;
let request = request
.uri(format!("wss://{}/rpc", host))
.body(())
.map_err(EstablishConnectionError::other)?;
let (stream, _) = async_tungstenite::async_tls::client_async_tls(request, stream)
.await
.context("websocket handshake")?;
Ok(Conn::new(stream))
.context("websocket handshake")
.map_err(EstablishConnectionError::other)?;
Ok(Connection::new(stream))
} else if let Some(host) = ZED_SERVER_URL.strip_prefix("http://") {
let stream = smol::net::TcpStream::connect(host).await?;
let request = request.uri(format!("ws://{}/rpc", host)).body(())?;
let stream = smol::net::TcpStream::connect(host)
.await
.map_err(EstablishConnectionError::other)?;
let request = request
.uri(format!("ws://{}/rpc", host))
.body(())
.map_err(EstablishConnectionError::other)?;
let (stream, _) = async_tungstenite::client_async(request, stream)
.await
.context("websocket handshake")?;
Ok(Conn::new(stream))
.context("websocket handshake")
.map_err(EstablishConnectionError::other)?;
Ok(Connection::new(stream))
} else {
Err(anyhow!("invalid server url: {}", *ZED_SERVER_URL))
Err(EstablishConnectionError::other(anyhow!(
"invalid server url: {}",
*ZED_SERVER_URL
)))
}
})
}
@@ -591,6 +637,19 @@ mod tests {
cx.foreground().advance_clock(Duration::from_secs(10));
while !matches!(status.recv().await, Some(Status::Connected { .. })) {}
assert_eq!(server.auth_count(), 1); // Client reused the cached credentials when reconnecting
server.forbid_connections();
server.disconnect().await;
while !matches!(status.recv().await, Some(Status::ReconnectionError { .. })) {}
// Clear cached credentials after authentication fails
server.roll_access_token();
server.allow_connections();
cx.foreground().advance_clock(Duration::from_secs(10));
assert_eq!(server.auth_count(), 1);
cx.foreground().advance_clock(Duration::from_secs(10));
while !matches!(status.recv().await, Some(Status::Connected { .. })) {}
assert_eq!(server.auth_count(), 2); // Client re-authenticated due to an invalid token
}
#[test]
zed/src/test.rs
+34 -15
View File
@@ -4,7 +4,7 @@ use crate::{
fs::RealFs,
http::{HttpClient, Request, Response, ServerResponse},
language::LanguageRegistry,
rpc::{self, Client, Credentials},
rpc::{self, Client, Credentials, EstablishConnectionError},
settings::{self, ThemeRegistry},
time::ReplicaId,
user::UserStore,
@@ -26,7 +26,7 @@ use std::{
},
};
use tempdir::TempDir;
use zrpc::{proto, Conn, ConnectionId, Peer, Receipt, TypedEnvelope};
use zrpc::{proto, Connection, ConnectionId, Peer, Receipt, TypedEnvelope};
#[cfg(test)]
#[ctor::ctor]
@@ -210,6 +210,8 @@ pub struct FakeServer {
connection_id: Mutex<Option<ConnectionId>>,
forbid_connections: AtomicBool,
auth_count: AtomicUsize,
access_token: AtomicUsize,
user_id: u64,
}
impl FakeServer {
@@ -224,6 +226,8 @@ impl FakeServer {
connection_id: Default::default(),
forbid_connections: Default::default(),
auth_count: Default::default(),
access_token: Default::default(),
user_id: client_user_id,
});
Arc::get_mut(client)
@@ -232,8 +236,8 @@ impl FakeServer {
let server = server.clone();
move |cx| {
server.auth_count.fetch_add(1, SeqCst);
let access_token = server.access_token.load(SeqCst).to_string();
cx.spawn(move |_| async move {
let access_token = "the-token".to_string();
Ok(Credentials {
user_id: client_user_id,
access_token,
@@ -244,11 +248,10 @@ impl FakeServer {
.override_establish_connection({
let server = server.clone();
move |credentials, cx| {
assert_eq!(credentials.user_id, client_user_id);
assert_eq!(credentials.access_token, "the-token");
let credentials = credentials.clone();
cx.spawn({
let server = server.clone();
move |cx| async move { server.connect(&cx).await }
move |cx| async move { server.establish_connection(&credentials, &cx).await }
})
}
});
@@ -266,23 +269,39 @@ impl FakeServer {
self.incoming.lock().take();
}
async fn connect(&self, cx: &AsyncAppContext) -> Result<Conn> {
async fn establish_connection(
&self,
credentials: &Credentials,
cx: &AsyncAppContext,
) -> Result<Connection, EstablishConnectionError> {
assert_eq!(credentials.user_id, self.user_id);
if self.forbid_connections.load(SeqCst) {
Err(anyhow!("server is forbidding connections"))
} else {
let (client_conn, server_conn, _) = Conn::in_memory();
let (connection_id, io, incoming) = self.peer.add_connection(server_conn).await;
cx.background().spawn(io).detach();
*self.incoming.lock() = Some(incoming);
*self.connection_id.lock() = Some(connection_id);
Ok(client_conn)
Err(EstablishConnectionError::Other(anyhow!(
"server is forbidding connections"
)))?
}
if credentials.access_token != self.access_token.load(SeqCst).to_string() {
Err(EstablishConnectionError::InvalidAccessToken)?
}
let (client_conn, server_conn, _) = Connection::in_memory();
let (connection_id, io, incoming) = self.peer.add_connection(server_conn).await;
cx.background().spawn(io).detach();
*self.incoming.lock() = Some(incoming);
*self.connection_id.lock() = Some(connection_id);
Ok(client_conn)
}
pub fn auth_count(&self) -> usize {
self.auth_count.load(SeqCst)
}
pub fn roll_access_token(&self) {
self.access_token.fetch_add(1, SeqCst);
}
pub fn forbid_connections(&self) {
self.forbid_connections.store(true, SeqCst);
}