implement logging in with the oauth tokens on the server side

2 files changed, 83 insertions, 10 deletions

diff --git a/teleterm/src/oauth.rs b/teleterm/src/oauth.rs
index 4e4ad72..0c6627c 100644
--- a/teleterm/src/oauth.rs
+++ b/teleterm/src/oauth.rs
@@ -43,8 +43,15 @@ pub trait Oauth {
                 Error::ExchangeCode { msg }
             })
             .and_then(|token| {
-                cache_refresh_token(token_cache_file, &token)
-                    .map(move |_| token.access_token().secret().to_string())
+                let access_token = token.access_token().secret().to_string();
+                let refresh_token =
+                    token.refresh_token().unwrap().secret().to_string();
+                cache_refresh_token(
+                    token_cache_file,
+                    &access_token,
+                    &refresh_token,
+                )
+                .map(move |_| access_token)
             });
         Box::new(fut)
     }
@@ -65,12 +72,28 @@ pub trait Oauth {
                 Error::ExchangeCode { msg }
             })
             .and_then(|token| {
-                cache_refresh_token(token_cache_file, &token)
-                    .map(move |_| token.access_token().secret().to_string())
+                let access_token = token.access_token().secret().to_string();
+                let refresh_token =
+                    token.refresh_token().unwrap().secret().to_string();
+                cache_refresh_token(
+                    token_cache_file,
+                    &access_token,
+                    &refresh_token,
+                )
+                .map(move |_| access_token)
             });
         Box::new(fut)
     }
 
+    fn save_tokens(
+        &self,
+        access_token: &str,
+        refresh_token: &str,
+    ) -> Box<dyn futures::Future<Item = (), Error = Error> + Send> {
+        let token_cache_file = self.server_token_file(false).unwrap();
+        cache_refresh_token(token_cache_file, access_token, refresh_token)
+    }
+
     fn get_username_from_access_token(
         self: Box<Self>,
         token: &str,
@@ -116,13 +139,10 @@ fn client_id_file(
 
 fn cache_refresh_token(
     token_cache_file: std::path::PathBuf,
-    token: &oauth2::basic::BasicTokenResponse,
+    access_token: &str,
+    refresh_token: &str,
 ) -> Box<dyn futures::Future<Item = (), Error = Error> + Send> {
-    let token_data = format!(
-        "{}\n{}\n",
-        token.refresh_token().unwrap().secret(),
-        token.access_token().secret(),
-    );
+    let token_data = format!("{}\n{}\n", refresh_token, access_token);
     let fut = tokio::fs::File::create(token_cache_file.clone())
         .with_context(move || crate::error::CreateFile {
             filename: token_cache_file.to_string_lossy().to_string(),
diff --git a/teleterm/src/server.rs b/teleterm/src/server.rs
index f23fac2..866edac 100644
--- a/teleterm/src/server.rs
+++ b/teleterm/src/server.rs
@@ -633,6 +633,51 @@ impl<S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Send + 'static>
         Ok(Some(Box::new(fut)))
     }
 
+    fn handle_message_oauth_response_token(
+        &mut self,
+        conn: &mut Connection<S>,
+        access_token: &str,
+        refresh_token: &str,
+    ) -> Result<
+        Option<
+            Box<
+                dyn futures::Future<
+                        Item = (ConnectionState, crate::protocol::Message),
+                        Error = Error,
+                    > + Send,
+            >,
+        >,
+    > {
+        let client = conn.oauth_client.take().ok_or_else(|| {
+            Error::UnexpectedMessage {
+                message: crate::protocol::Message::oauth_response_token(
+                    access_token,
+                    refresh_token,
+                ),
+            }
+        })?;
+
+        let term_info = conn.state.term_info().unwrap().clone();
+        let access_token = access_token.to_string();
+        let fut = client.save_tokens(&access_token, refresh_token).and_then(
+            move |_| {
+                client.get_username_from_access_token(&access_token).map(
+                    |username| {
+                        (
+                            ConnectionState::LoggedIn {
+                                term_info,
+                                username: username.clone(),
+                            },
+                            crate::protocol::Message::logged_in(&username),
+                        )
+                    },
+                )
+            },
+        );
+
+        Ok(Some(Box::new(fut)))
+    }
+
     fn handle_accepted_message(
         &mut self,
         conn: &mut Connection<S>,
@@ -676,6 +721,14 @@ impl<S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Send + 'static>
             crate::protocol::Message::OauthResponseCode { code } => {
                 self.handle_message_oauth_response_code(conn, &code)
             }
+            crate::protocol::Message::OauthResponseToken {
+                access_token,
+                refresh_token,
+            } => self.handle_message_oauth_response_token(
+                conn,
+                &access_token,
+                &refresh_token,
+            ),
             m => Err(Error::UnauthenticatedMessage { message: m }),
         }
     }