From 3003a4f99ae886463232aa5dbc1ce64633483b30 Mon Sep 17 00:00:00 2001
From: Eric <illuminating.me@gmail.com>
Date: Sun, 7 Jan 2024 23:25:44 +0800
Subject: [PATCH] feat(webserver): complete `github_auth` method (#1169)

* feat(webserver): register user when oauth by github

* [autofix.ci] apply automated fixes

* add comment

---------

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
---
 ee/tabby-db/src/invitations.rs         | 18 +++++++++++++++
 ee/tabby-webserver/src/oauth/mod.rs    |  1 +
 ee/tabby-webserver/src/schema/auth.rs  |  3 +++
 ee/tabby-webserver/src/service/auth.rs | 32 ++++++++++++++++++++++----
 4 files changed, 50 insertions(+), 4 deletions(-)

diff --git a/ee/tabby-db/src/invitations.rs b/ee/tabby-db/src/invitations.rs
index 24cc09e3be32..abfe2d07b6b7 100644
--- a/ee/tabby-db/src/invitations.rs
+++ b/ee/tabby-db/src/invitations.rs
@@ -69,6 +69,24 @@ impl DbConn {
         Ok(token?)
     }
 
+    pub async fn get_invitation_by_email(&self, email: &str) -> Result<Option<InvitationDAO>> {
+        let email = email.to_owned();
+        let token = self
+            .conn
+            .call(|conn| {
+                Ok(conn
+                    .query_row(
+                        r#"SELECT id, email, code, created_at FROM invitations WHERE email = ?"#,
+                        [email],
+                        InvitationDAO::from_row,
+                    )
+                    .optional())
+            })
+            .await?;
+
+        Ok(token?)
+    }
+
     pub async fn create_invitation(&self, email: String) -> Result<i32> {
         if self.get_user_by_email(&email).await?.is_some() {
             return Err(anyhow!("User already registered"));
diff --git a/ee/tabby-webserver/src/oauth/mod.rs b/ee/tabby-webserver/src/oauth/mod.rs
index 1f554120af54..6da96abe2ade 100644
--- a/ee/tabby-webserver/src/oauth/mod.rs
+++ b/ee/tabby-webserver/src/oauth/mod.rs
@@ -62,6 +62,7 @@ async fn github_callback(
         }
         Err(GithubAuthError::InvalidVerificationCode) => Err(StatusCode::BAD_REQUEST),
         Err(GithubAuthError::CredentialNotActive) => Err(StatusCode::NOT_FOUND),
+        Err(GithubAuthError::UserNotInvited) => Err(StatusCode::UNAUTHORIZED),
         Err(e) => {
             error!("Failed to authenticate with Github: {:?}", e);
             Err(StatusCode::INTERNAL_SERVER_ERROR)
diff --git a/ee/tabby-webserver/src/schema/auth.rs b/ee/tabby-webserver/src/schema/auth.rs
index b7ecffff9a6f..5cc22085970e 100644
--- a/ee/tabby-webserver/src/schema/auth.rs
+++ b/ee/tabby-webserver/src/schema/auth.rs
@@ -159,6 +159,9 @@ pub enum GithubAuthError {
     #[error("The Github credential is not active")]
     CredentialNotActive,
 
+    #[error("The user is not invited to access the system")]
+    UserNotInvited,
+
     #[error(transparent)]
     Other(#[from] anyhow::Error),
 
diff --git a/ee/tabby-webserver/src/service/auth.rs b/ee/tabby-webserver/src/service/auth.rs
index 6cf859124458..d60a271c671e 100644
--- a/ee/tabby-webserver/src/service/auth.rs
+++ b/ee/tabby-webserver/src/service/auth.rs
@@ -357,11 +357,35 @@ impl AuthenticationService for DbConn {
             return Err(GithubAuthError::CredentialNotActive);
         }
 
-        let _email = client.fetch_user_email(code, credential).await?;
+        let email = client.fetch_user_email(code, credential).await?;
 
-        // TODO: auto register & generate token
+        let user = if let Some(user) = self.get_user_by_email(&email).await? {
+            user
+        } else {
+            let Some(invitation) = self.get_invitation_by_email(&email).await? else {
+                return Err(GithubAuthError::UserNotInvited);
+            };
+            // it's ok to set password to empty string here, because
+            // 1. both `register` & `token_auth` mutation will do input validation, so empty password won't be accepted
+            // 2. `password_verify` will always return false for empty password hash read from user table
+            // so user created here is only able to login by github oauth, normal login won't work
+            let id = self
+                .create_user_with_invitation(email, "".to_owned(), false, invitation.id)
+                .await?;
+            self.get_user(id).await?.unwrap()
+        };
+
+        let refresh_token = generate_refresh_token();
+        self.create_refresh_token(user.id, &refresh_token).await?;
+
+        let access_token = generate_jwt(JWTPayload::new(user.email.clone(), user.is_admin))
+            .map_err(|_| GithubAuthError::Unknown)?;
 
-        Ok(GithubAuthResponse::default())
+        let resp = GithubAuthResponse {
+            access_token,
+            refresh_token,
+        };
+        Ok(resp)
     }
 }
 
@@ -492,7 +516,7 @@ mod tests {
                 email.to_owned(),
                 password.to_owned(),
                 password.to_owned(),
-                Some(invitation.code.clone())
+                Some(invitation.code.clone()),
             )
             .await
             .is_ok());