From e3d04f279f689a702b9e0939cb0d8dabbe309593 Mon Sep 17 00:00:00 2001
From: Connor Peet <connor@peet.io>
Date: Mon, 6 May 2024 09:47:37 -0700
Subject: [PATCH] cli: support refresh token in tunnel user login (#212106)

---
 cli/src/auth.rs             |  9 +++++++--
 cli/src/commands/args.rs    |  7 +++++--
 cli/src/commands/tunnels.rs | 12 +++++++++---
 3 files changed, 21 insertions(+), 7 deletions(-)

diff --git a/cli/src/auth.rs b/cli/src/auth.rs
index 9d5c9b73fdb..67f1bfa6bc7 100644
--- a/cli/src/auth.rs
+++ b/cli/src/auth.rs
@@ -480,6 +480,7 @@ impl Auth {
 		&self,
 		provider: Option<AuthProvider>,
 		access_token: Option<String>,
+		refresh_token: Option<String>,
 	) -> Result<StoredCredential, AnyError> {
 		let provider = match provider {
 			Some(p) => p,
@@ -490,8 +491,12 @@ impl Auth {
 			Some(t) => StoredCredential {
 				provider,
 				access_token: t,
-				refresh_token: None,
-				expires_at: None,
+				// if a refresh token is given, assume it's valid now but refresh it
+				// soon in order to get the real expiry time.
+				expires_at: refresh_token
+					.as_ref()
+					.map(|_| Utc::now() + chrono::Duration::minutes(5)),
+				refresh_token,
 			},
 			None => self.do_device_code_flow_with_provider(provider).await?,
 		};
diff --git a/cli/src/commands/args.rs b/cli/src/commands/args.rs
index 1eaaa57353e..8a943826615 100644
--- a/cli/src/commands/args.rs
+++ b/cli/src/commands/args.rs
@@ -788,11 +788,14 @@ pub enum TunnelUserSubCommands {
 
 #[derive(Args, Debug, Clone)]
 pub struct LoginArgs {
-	/// An access token to store for authentication. Note: this will not be
-	/// refreshed if it expires!
+	/// An access token to store for authentication.
 	#[clap(long, requires = "provider")]
 	pub access_token: Option<String>,
 
+	/// An access token to store for authentication.
+	#[clap(long, requires = "access_token")]
+	pub refresh_token: Option<String>,
+
 	/// The auth provider to use. If not provided, a prompt will be shown.
 	#[clap(value_enum, long)]
 	pub provider: Option<AuthProvider>,
diff --git a/cli/src/commands/tunnels.rs b/cli/src/commands/tunnels.rs
index f06cd9a1e2a..1755dbbfaef 100644
--- a/cli/src/commands/tunnels.rs
+++ b/cli/src/commands/tunnels.rs
@@ -274,10 +274,11 @@ pub async fn service(
 pub async fn user(ctx: CommandContext, user_args: TunnelUserSubCommands) -> Result<i32, AnyError> {
 	let auth = Auth::new(&ctx.paths, ctx.log.clone());
 	match user_args {
-		TunnelUserSubCommands::Login(login_args) => {
+		TunnelUserSubCommands::Login(mut login_args) => {
 			auth.login(
 				login_args.provider.map(|p| p.into()),
-				login_args.access_token.to_owned(),
+				login_args.access_token.take(),
+				login_args.refresh_token.take(),
 			)
 			.await?;
 		}
@@ -488,7 +489,12 @@ pub async fn forward(
 		forward_args.login.provider.take(),
 		forward_args.login.access_token.take(),
 	) {
-		auth.login(Some(p.into()), Some(at)).await?;
+		auth.login(
+			Some(p.into()),
+			Some(at),
+			forward_args.login.refresh_token.take(),
+		)
+		.await?;
 	}
 
 	let mut tunnels = DevTunnels::new_port_forwarding(&ctx.log, auth, &ctx.paths);