Save server specific login tokens in server config

Author: aceysmith

crates/cli/src/config.rs

@@ -24,6 +24,7 @@ pub struct ServerConfig {
     pub host: String,
     pub protocol: String,
     pub ecdsa_public_key: Option<String>,
+    pub spacetimedb_token: Option<String>,
 }
 
 impl ServerConfig {
@@ -40,6 +41,7 @@ impl ServerConfig {
         set_table_opt_value(edit, HOST_KEY, Some(&from.host));
         set_table_opt_value(edit, PROTOCOL_KEY, Some(&from.protocol));
         set_table_opt_value(edit, ECDSA_PUBLIC_KEY, from.ecdsa_public_key.as_deref());
+        set_table_opt_value(edit, SPACETIMEDB_TOKEN_KEY, from.spacetimedb_token.as_deref());
     }
 
     fn nick_or_host(&self) -> &str {
@@ -105,11 +107,13 @@ impl TryFrom<&toml_edit::Table> for ServerConfig {
         let host = read_str(table, HOST_KEY)?;
         let protocol = read_str(table, PROTOCOL_KEY)?;
         let ecdsa_public_key = read_opt_str(table, ECDSA_PUBLIC_KEY)?;
+        let spacetimedb_token = read_opt_str(table, SPACETIMEDB_TOKEN_KEY)?;
         Ok(ServerConfig {
             nickname,
             host,
             protocol,
             ecdsa_public_key,
+            spacetimedb_token,
         })
     }
 }
@@ -161,12 +165,14 @@ impl RawConfig {
             protocol: "http".to_string(),
             nickname: Some("local".to_string()),
             ecdsa_public_key: None,
+            spacetimedb_token: None,
         };
         let maincloud = ServerConfig {
             host: "maincloud.spacetimedb.com".to_string(),
             protocol: "https".to_string(),
             nickname: Some("maincloud".to_string()),
             ecdsa_public_key: None,
+            spacetimedb_token: None,
         };
         RawConfig {
             default_server: local.nickname.clone(),
@@ -245,6 +251,7 @@ impl RawConfig {
             host,
             protocol,
             ecdsa_public_key,
+            spacetimedb_token: None
         });
         Ok(())
     }
@@ -444,13 +451,57 @@ Fetch the server's fingerprint with:
         self.web_session_token = Some(token);
     }
 
+
+    /// Return the spacetimedb_token for the server named by `server`.
+    ///
+    /// Returns an `Err` if there is no server configuration.
+    /// Returns `None` if the server configuration exists, but does not have a token saved.
+    pub fn server_spacetimedb_token(&self, server: &str) -> anyhow::Result<Option<&str>> {
+        self.find_server(server)
+            .map(|cfg| cfg.spacetimedb_token.as_deref())
+            .with_context(|| format!("Cannot find spacetimedb_token for unknown server: {server}"))
+    }
+
+    /// Return the spacetimedb_token for the default server.
+    ///
+    /// Returns an `Err` if there is no default server configuration.
+    /// Returns `None` if the server configuration exists, but does not have a token saved.
+    pub fn default_server_spacetimedb_token(&self) -> anyhow::Result<Option<&str>> {
+        if let Some(server) = &self.default_server {
+            self.server_spacetimedb_token(server)
+        } else {
+            Err(anyhow::anyhow!(NO_DEFAULT_SERVER_ERROR_MESSAGE))
+        }
+    }
+
     pub fn set_spacetimedb_token(&mut self, token: String) {
         self.spacetimedb_token = Some(token);
     }
 
+    /// Store the spacetimedb_token for the server named by `server`.
+    ///
+    /// Returns an `Err` if no such server configuration exists.
+    pub fn set_server_spacetimedb_token(&mut self, server: &str, spacetimedb_token: String) -> anyhow::Result<()> {
+        let cfg = self.find_server_mut(server)?;
+        cfg.spacetimedb_token = Some(spacetimedb_token);
+        Ok(())
+    }
+
+    /// Store the spacetimedb_token for the default server.
+    ///
+    /// Returns an `Err` if there is no default server configuration.
+    pub fn set_default_server_spacetimedb_token(&mut self, spacetimedb_token: String) -> anyhow::Result<()> {
+        let cfg = self.default_server_mut()?;
+        cfg.spacetimedb_token = Some(spacetimedb_token);
+        Ok(())
+    }
+
     pub fn clear_login_tokens(&mut self) {
         self.web_session_token = None;
         self.spacetimedb_token = None;
+        for cfg in &mut self.server_configs {
+            cfg.spacetimedb_token = None;
+        }
     }
 }
 
@@ -808,6 +859,14 @@ Update the server's fingerprint with:
         self.home.set_spacetimedb_token(token);
     }
 
+    pub fn set_server_spacetimedb_token(&mut self, server: Option<&str>, token: String) -> anyhow::Result<()> {
+        if let Some(server) = server {
+            self.home.set_server_spacetimedb_token(server, token)
+        } else {
+            self.home.set_default_server_spacetimedb_token(token)
+        }
+    }
+
     pub fn clear_login_tokens(&mut self) {
         self.home.clear_login_tokens();
     }
@@ -816,8 +875,19 @@ Update the server's fingerprint with:
         self.home.web_session_token.as_ref()
     }
 
-    pub fn spacetimedb_token(&self) -> Option<&String> {
-        self.home.spacetimedb_token.as_ref()
+    pub fn spacetimedb_token(&self, server: Option<&str>, allow_fallback: bool) -> Option<String> {
+        let token = if let Some(server) = server {
+            self.home.server_spacetimedb_token(server)
+        } else {
+            self.home.default_server_spacetimedb_token()
+        };
+        if let Ok(Some(token)) = token {
+            Some(token.to_string())
+        } else if allow_fallback {
+            self.home.spacetimedb_token.clone()
+        } else {
+            None
+        }
     }
 }
 
@@ -1001,12 +1071,14 @@ default_server = "local"
                     host: "127.0.0.1:3000".to_string(),
                     protocol: "http".to_string(),
                     ecdsa_public_key: None,
+                    spacetimedb_token: None,
                 },
                 ServerConfig {
                     nickname: Some("testnet".to_string()),
                     host: "testnet.spacetimedb.com".to_string(),
                     protocol: "https".to_string(),
                     ecdsa_public_key: None,
+                    spacetimedb_token: None,
                 },
             ];
             config.home.spacetimedb_token =

crates/cli/src/subcommands/login.rs

@@ -1,4 +1,4 @@
-use crate::util::decode_identity;
+use crate::{common_args, util::decode_identity};
 use crate::Config;
 use clap::{Arg, ArgAction, ArgGroup, ArgMatches, Command};
 use reqwest::Url;
@@ -42,6 +42,9 @@ fn get_subcommands() -> Vec<Command> {
                 .action(ArgAction::SetTrue)
                 .help("Also show the auth token"),
         )
+        .arg(
+            common_args::server().help("The server used to log in to a SpacetimeDB server directly, without going through a global auth server")
+        )
         .about("Show the current login info")]
 }
 
@@ -53,7 +56,7 @@ pub async fn exec(mut config: Config, args: &ArgMatches) -> Result<(), anyhow::E
     let spacetimedb_token: Option<&String> = args.get_one("spacetimedb-token");
     let host: &String = args.get_one("auth-host").unwrap();
     let host = Url::parse(host)?;
-    let server_issued_login: Option<&String> = args.get_one("server");
+    let server_issued_login: Option<&str> = args.get_one::<String>("server").map(|s| s.as_ref());
 
     if let Some(token) = spacetimedb_token {
         config.set_spacetimedb_token(token.clone());
@@ -63,9 +66,9 @@ pub async fn exec(mut config: Config, args: &ArgMatches) -> Result<(), anyhow::E
 
     if let Some(server) = server_issued_login {
         let host = Url::parse(&config.get_host_url(Some(server))?)?;
-        spacetimedb_token_cached(&mut config, &host, true).await?;
+        spacetimedb_token_cached(&mut config, &host, server_issued_login, true).await?;
     } else {
-        spacetimedb_token_cached(&mut config, &host, false).await?;
+        spacetimedb_token_cached(&mut config, &host, None, false).await?;
     }
 
     Ok(())
@@ -80,16 +83,21 @@ async fn exec_subcommand(config: Config, cmd: &str, args: &ArgMatches) -> Result
 
 async fn exec_show(config: Config, args: &ArgMatches) -> Result<(), anyhow::Error> {
     let include_token = args.get_flag("token");
+    let server_issued_login: Option<&str> = args.get_one::<String>("server").map(|s| s.as_ref());
 
-    let token = if let Some(token) = config.spacetimedb_token() {
+    let token = if let Some(token) = config.spacetimedb_token(server_issued_login, false) {
         token
     } else {
         println!("You are not logged in. Run `spacetime login` to log in.");
         return Ok(());
     };
 
-    let identity = decode_identity(token)?;
-    println!("You are logged in as {}", identity);
+    let identity = decode_identity(&token)?;
+    if let Some(server) = server_issued_login {
+        println!("You are logged in to server \"{}\" as {}", server, identity);
+    } else {
+        println!("You are logged in as {}", identity);
+    }
 
     if include_token {
         println!("Your auth token (don't share this!) is {}", token);
@@ -98,19 +106,20 @@ async fn exec_show(config: Config, args: &ArgMatches) -> Result<(), anyhow::Erro
     Ok(())
 }
 
-async fn spacetimedb_token_cached(config: &mut Config, host: &Url, direct_login: bool) -> anyhow::Result<String> {
+async fn spacetimedb_token_cached(config: &mut Config, host: &Url, server: Option<&str>, direct_login: bool) -> anyhow::Result<String> {
     // Currently, this token does not expire. However, it will at some point in the future. When that happens,
     // this code will need to happen before any request to a spacetimedb server, rather than at the end of the login flow here.
-    if let Some(token) = config.spacetimedb_token() {
+    // If we are logging in to a specific server, we should not accept the root token as verification that we are logged in.
+    if let Some(token) = config.spacetimedb_token(server, !direct_login) {
         println!("You are already logged in.");
         println!("If you want to log out, use spacetime logout.");
         Ok(token.clone())
     } else {
-        spacetimedb_login_force(config, host, direct_login).await
+        spacetimedb_login_force(config, host, server, direct_login).await
     }
 }
 
-pub async fn spacetimedb_login_force(config: &mut Config, host: &Url, direct_login: bool) -> anyhow::Result<String> {
+pub async fn spacetimedb_login_force(config: &mut Config, host: &Url, server: Option<&str>, direct_login: bool) -> anyhow::Result<String> {
     let token = if direct_login {
         let token = spacetimedb_direct_login(host).await?;
         println!("We have logged in directly to your target server.");
@@ -120,7 +129,11 @@ pub async fn spacetimedb_login_force(config: &mut Config, host: &Url, direct_log
         let session_token = web_login_cached(config, host).await?;
         spacetimedb_login(host, &session_token).await?
     };
-    config.set_spacetimedb_token(token.clone());
+    if direct_login {
+        config.set_server_spacetimedb_token(server, token.clone())?;
+    } else {
+        config.set_spacetimedb_token(token.clone());
+    }
     config.save();
 
     Ok(token)

crates/cli/src/subcommands/publish.rs

@@ -160,8 +160,8 @@ pub async fn exec(mut config: Config, args: &ArgMatches) -> Result<(), anyhow::E
     let res = builder.body(program_bytes).send().await?;
     if res.status() == StatusCode::UNAUTHORIZED && !anon_identity {
         // If we're not in the `anon_identity` case, then we have already forced the user to log in above (using `get_auth_header`), so this should be safe to unwrap.
-        let token = config.spacetimedb_token().unwrap();
-        let identity = decode_identity(token)?;
+        let token = config.spacetimedb_token(server, true).unwrap();
+        let identity = decode_identity(&token)?;
         let err = res.text().await?;
         return unauth_error_context(
             Err(anyhow::anyhow!(err)),
@@ -192,8 +192,8 @@ pub async fn exec(mut config: Config, args: &ArgMatches) -> Result<(), anyhow::E
                 anyhow::bail!("You need to be logged in as the owner of {name} to publish to {name}",);
             }
             // If we're not in the `anon_identity` case, then we have already forced the user to log in above (using `get_auth_header`), so this should be safe to unwrap.
-            let token = config.spacetimedb_token().unwrap();
-            let identity = decode_identity(token)?;
+            let token = config.spacetimedb_token(server, true).unwrap();
+            let identity = decode_identity(&token)?;
             //TODO(jdetter): Have a nice name generator here, instead of using some abstract characters
             // we should perhaps generate fun names like 'green-fire-dragon' instead
             let suggested_tld: String = identity.chars().take(12).collect();

crates/cli/src/util.rs

@@ -299,7 +299,7 @@ pub async fn get_login_token_or_log_in(
     target_server: Option<&str>,
     interactive: bool,
 ) -> anyhow::Result<String> {
-    if let Some(token) = config.spacetimedb_token() {
+    if let Some(token) = config.spacetimedb_token(target_server, true) {
         return Ok(token.clone());
     }
 
@@ -314,10 +314,10 @@ pub async fn get_login_token_or_log_in(
 
     if full_login {
         let host = Url::parse(DEFAULT_AUTH_HOST)?;
-        spacetimedb_login_force(config, &host, false).await
+        spacetimedb_login_force(config, &host, None, false).await
     } else {
         let host = Url::parse(&config.get_host_url(target_server)?)?;
-        spacetimedb_login_force(config, &host, true).await
+        spacetimedb_login_force(config, &host, target_server, true).await
     }
 }