feat: add google oauth (zklogin) behind feature flag

.github/workflows/release.yml

@@ -35,12 +35,26 @@ jobs:
           username: ${{ github.actor }}
           password: ${{ secrets.GITHUB_TOKEN }}
 
-      - name: Build and push Docker image
+      - name: Build and push Local Docker image
         uses: docker/build-push-action@v5
         with:
           context: .
           push: true
           platforms: linux/amd64,linux/arm64
+          build-args:
+            PROFILE=local
+          tags: |
+            ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:local-latest
+            ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:${{ github.event.release.tag_name }}
+
+      - name: Build and push Cloud Docker image
+        uses: docker/build-push-action@v5
+        with:
+          context: .
+          push: true
+          platforms: linux/amd64,linux/arm64
+          build-args:
+            PROFILE=cloud
           tags: |
             ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:latest
             ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:${{ github.event.release.tag_name }}

Dockerfile

@@ -8,6 +8,7 @@ ARG TARGETARCH
 
 # Trace level argument
 ARG TRACE_LEVEL
+ARG PROFILE
 
 # Install build dependencies
 RUN apt-get update && apt-get install -y \
@@ -21,8 +22,13 @@ WORKDIR /usr/src/atoma-proxy
 
 COPY . .
 
+
 # Compile
-RUN RUST_LOG=${TRACE_LEVEL} cargo build --release --bin atoma-proxy
+RUN if [ "$PROFILE" = "local" ]; then \
+    RUST_LOG=${TRACE_LEVEL} cargo build --release --bin atoma-proxy --no-default-features; \
+    else \
+    RUST_LOG=${TRACE_LEVEL} cargo build --release --bin atoma-proxy; \
+    fi
 
 # Final stage
 FROM --platform=$TARGETPLATFORM debian:bullseye-slim

README.md

@@ -114,6 +114,7 @@ service_bind_address = "0.0.0.0:8081"
 secret_key = "secret_key" # Secret key for the tokens generation
 access_token_lifetime = 1 # In minutes
 refresh_token_lifetime = 1 # In days
+google_client_id="" # Google client id for google login (In case google-oauth feature is enabled)
 ```
 
 4. Create required directories
@@ -139,6 +140,10 @@ The deployment consists of two main services:
 - **PostgreSQL**: Manages the database for the Atoma Proxy
 - **Atoma Proxy**: Manages the proxy operations and connects to the Atoma Network
 
+#### Profiles
+- local - this is for targeting the local deployment of the proxy
+- cloud - has the same features as the local, but also enables zklogin using google oauth
+
 #### Service URLs
 
 - Atoma Proxy: `http://localhost:8080` (configured via ATOMA_PROXY_PORT). This is the main service that you will use to interact with the Atoma Network, via an

atoma-auth/Cargo.toml

@@ -32,3 +32,6 @@ sui-sdk-types = { workspace = true, features = ["serde"] }
 thiserror.workspace = true
 tokio.workspace = true
 tracing.workspace = true
+
+[features]
+google-oauth = []

atoma-auth/src/auth.rs

@@ -1,5 +1,10 @@
-use std::{collections::HashMap, str::FromStr, sync::Arc};
+#[cfg(feature = "google-oauth")]
+use std::collections::HashMap;
+use std::{str::FromStr, sync::Arc};
 
+#[cfg(feature = "google-oauth")]
+use crate::google::{self, fetch_google_public_keys};
+use crate::{AtomaAuthConfig, Sui};
 use atoma_state::{types::AtomaAtomaStateManagerEvent, AtomaStateManagerError};
 use atoma_utils::hashing::blake2b_hash;
 use blake2::{
@@ -14,28 +19,29 @@ use fastcrypto::{
     secp256r1::{Secp256r1PublicKey, Secp256r1Signature},
     traits::{ToFromBytes, VerifyingKey},
 };
+#[cfg(feature = "google-oauth")]
 use fastcrypto_zkp::zk_login_utils::Bn254FrElement;
 use flume::Sender;
-use jsonwebtoken::{decode, encode, Algorithm, DecodingKey, EncodingKey, Header, Validation};
+#[cfg(feature = "google-oauth")]
+use jsonwebtoken::Algorithm;
+use jsonwebtoken::{decode, encode, DecodingKey, EncodingKey, Header, Validation};
 use rand::Rng;
 use serde::{Deserialize, Serialize};
 use shared_crypto::intent::{Intent, IntentMessage, PersonalMessage};
+#[cfg(feature = "google-oauth")]
+use sui_sdk::types::crypto::ZkLoginPublicIdentifier;
 use sui_sdk::types::{
     base_types::SuiAddress,
-    crypto::{PublicKey, Signature, SignatureScheme, SuiSignature, ZkLoginPublicIdentifier},
+    crypto::{PublicKey, Signature, SignatureScheme, SuiSignature},
     object::Owner,
     TypeTag,
 };
+#[cfg(feature = "google-oauth")]
 use sui_sdk_types::{SimpleSignature, UserSignature};
 use thiserror::Error;
 use tokio::sync::{oneshot, RwLock};
 use tracing::{error, instrument};
 
-use crate::{
-    google::{self, fetch_google_public_keys},
-    AtomaAuthConfig, Sui,
-};
-
 /// The length of the API token
 const API_TOKEN_LENGTH: usize = 30;
 
@@ -92,8 +98,12 @@ pub enum AuthError {
     SenderOrReceiverNotFound,
     #[error("The payment is not for this user")]
     PaymentNotForThisUser,
+    #[cfg(feature = "google-oauth")]
     #[error("Google error: {0}")]
     GoogleError(#[from] crate::google::GoogleError),
+    #[cfg(not(feature = "google-oauth"))]
+    #[error("ZkLogin not enabled")]
+    ZkLoginNotEnabled,
 }
 
 type Result<T> = std::result::Result<T, AuthError>;
@@ -110,8 +120,10 @@ pub struct Auth {
     state_manager_sender: Sender<AtomaAtomaStateManagerEvent>,
     /// The sui client
     sui: Arc<RwLock<Sui>>,
+    #[cfg(feature = "google-oauth")]
     /// GooglePublicKeys
     google_public_keys: HashMap<String, (DecodingKey, Algorithm)>,
+    #[cfg(feature = "google-oauth")]
     /// Google client id
     google_client_id: String,
 }
@@ -123,14 +135,17 @@ impl Auth {
         state_manager_sender: Sender<AtomaAtomaStateManagerEvent>,
         sui: Arc<RwLock<Sui>>,
     ) -> Result<Self> {
+        #[cfg(feature = "google-oauth")]
         let google_public_keys = fetch_google_public_keys().await?;
         Ok(Self {
             secret_key: config.secret_key,
             access_token_lifetime: config.access_token_lifetime,
             refresh_token_lifetime: config.refresh_token_lifetime,
             state_manager_sender,
             sui,
+            #[cfg(feature = "google-oauth")]
             google_public_keys,
+            #[cfg(feature = "google-oauth")]
             google_client_id: config.google_client_id,
         })
     }
@@ -332,6 +347,7 @@ impl Auth {
     /// This method will check the google oauth token and generate a new refresh and access token
     /// The method will check if the email is present in the claims and store the user in the DB
     /// The method will generate a new refresh and access tokens
+    #[cfg(feature = "google-oauth")]
     #[instrument(level = "info", skip(self))]
     pub async fn check_google_id_token(&self, id_token: &str) -> Result<(String, String)> {
         let claims = google::verify_google_id_token(
@@ -560,6 +576,7 @@ impl Auth {
     ///
     /// * If the signature is not a zk_login signature
     /// * If the signature is not valid
+    #[cfg(feature = "google-oauth")]
     #[instrument(level = "info")]
     async fn get_zk_address(zk_login_signature: &str, tx_digest: &str) -> Result<String> {
         let user_signature = UserSignature::from_base64(zk_login_signature)?;
@@ -706,13 +723,18 @@ impl Auth {
             // The signature is coming from the frontend where the user used his zk credentials to sign the transaction digest as a personal message (digest of the transaction).
             // Now we need to prove that it was the digest he is trying to claim. And if the sui address from the signature is matching the address in the usdc payment transaction.
             match zk_proof_signature {
+                #[cfg(feature = "google-oauth")]
                 Some(signature) => {
                     if sender.to_string()
                         != Self::get_zk_address(&signature, transaction_digest).await?
                     {
                         return Err(AuthError::PaymentNotForThisUser);
                     }
                 }
+                #[cfg(not(feature = "google-oauth"))]
+                Some(_) => {
+                    return Err(AuthError::ZkLoginNotEnabled);
+                }
                 None => {
                     let (result_sender, result_receiver) = oneshot::channel();
                     self.state_manager_sender
@@ -772,15 +794,10 @@ mod test {
 
     use atoma_state::types::AtomaAtomaStateManagerEvent;
     use atoma_sui::AtomaSuiConfig;
-    use chrono::Utc;
     use flume::Receiver;
-    use jsonwebtoken::{decode_header, encode, Algorithm, DecodingKey, EncodingKey, Header};
-    use rand::rngs::OsRng;
-    use rsa::{RsaPrivateKey, RsaPublicKey};
     use tokio::sync::RwLock;
 
-    use crate::google::ISS;
-    use crate::{google, AtomaAuthConfig};
+    use crate::AtomaAuthConfig;
 
     use super::Auth;
     use std::env;
@@ -870,8 +887,13 @@ active_address: "0x939cfcc7fcbc71ce983203bcb36fa498901932ab9293dfa2b271203e71603
     }
 
     async fn setup_test() -> (Auth, Receiver<AtomaAtomaStateManagerEvent>) {
-        let config =
-            AtomaAuthConfig::new("secret".to_string(), 1, 1, "google_client_id".to_string());
+        let config = AtomaAuthConfig::new(
+            "secret".to_string(),
+            1,
+            1,
+            #[cfg(feature = "google-oauth")]
+            "google_client_id".to_string(),
+        );
         let (state_manager_sender, state_manager_receiver) = flume::unbounded();
 
         let sui_config = AtomaSuiConfig::from_file_path(get_config_path());
@@ -997,8 +1019,12 @@ active_address: "0x939cfcc7fcbc71ce983203bcb36fa498901932ab9293dfa2b271203e71603
         }
     }
 
+    #[cfg(feature = "google-oauth")]
     #[tokio::test]
     async fn google_login() {
+        use crate::google::{Claims, ISS};
+        use chrono::Utc;
+        use jsonwebtoken::{encode, Algorithm, DecodingKey, EncodingKey, Header};
         let (mut auth, receiver) = setup_test().await;
         let mock_handle = tokio::task::spawn(async move {
             // First event is for the user to log in to get the tokens
@@ -1031,9 +1057,6 @@ active_address: "0x939cfcc7fcbc71ce983203bcb36fa498901932ab9293dfa2b271203e71603
                 _ => panic!("Unexpected event"),
             }
         });
-        // let private_key = RsaPrivateKey::new(&mut OsRng, 2048).expect("Failed to generate a key");
-        // let public_key = RsaPublicKey::from(&private_key);
-        // (private_key, public_key);
         let encoding_key = EncodingKey::from_secret("fake secret".as_bytes());
         auth.google_public_keys.insert(
             "kid".to_string(),
@@ -1042,12 +1065,14 @@ active_address: "0x939cfcc7fcbc71ce983203bcb36fa498901932ab9293dfa2b271203e71603
                 Algorithm::HS256,
             ),
         );
-        let mut header = Header::default();
-        header.kid = Some("kid".to_string());
-        header.alg = Algorithm::HS256;
+        let header = Header {
+            kid: Some("kid".to_string()),
+            alg: Algorithm::HS256,
+            ..Default::default()
+        };
         let id_token = encode(
             &header,
-            &google::Claims::new(
+            &Claims::new(
                 ISS,
                 "sub",
                 "google_client_id",

atoma-auth/src/config.rs

@@ -12,6 +12,7 @@ pub struct AtomaAuthConfig {
     /// The refresh token lifetime in days.
     pub refresh_token_lifetime: usize,
     /// Google client id.
+    #[cfg(feature = "google-oauth")]
     pub google_client_id: String,
 }
 
@@ -21,12 +22,13 @@ impl AtomaAuthConfig {
         secret_key: String,
         access_token_lifetime: usize,
         refresh_token_lifetime: usize,
-        google_client_id: String,
+        #[cfg(feature = "google-oauth")] google_client_id: String,
     ) -> Self {
         Self {
             secret_key,
             access_token_lifetime,
             refresh_token_lifetime,
+            #[cfg(feature = "google-oauth")]
             google_client_id,
         }
     }

atoma-auth/src/lib.rs

@@ -1,5 +1,6 @@
 mod auth;
 mod config;
+#[cfg(feature = "google-oauth")]
 mod google;
 mod sui;
 

atoma-proxy-service/Cargo.toml

@@ -23,3 +23,6 @@ tracing-subscriber.workspace = true
 tracing.workspace = true
 utoipa = { workspace = true, features = ["axum_extras"] }
 utoipa-swagger-ui = { workspace = true }
+
+[features]
+google-oauth = ["atoma-auth/google-oauth"]

atoma-proxy-service/docs/openapi.yml

@@ -252,38 +252,6 @@ paths:
           description: Failed to get sui address
       security:
       - bearerAuth: []
-  /google_oauth:
-    post:
-      tags:
-      - Auth
-      summary: |-
-        Logs in a user with the proxy service using Google OAuth.
-        This endpoint is used to verify a Google ID token and return an access token.
-      description: |-
-        # Arguments
-
-        * `proxy_service_state` - The shared state containing the state manager
-        * `body` - The request body containing the Google ID token
-
-        # Returns
-
-        * `Result<Json<AuthResponse>>` - A JSON response containing the access and refresh tokens
-      operationId: google_oauth
-      requestBody:
-        content:
-          text/plain:
-            schema:
-              type: string
-        required: true
-      responses:
-        '200':
-          description: Logs in a user with Google OAuth
-          content:
-            text/plain:
-              schema:
-                type: string
-        '500':
-          description: Failed to verify Google ID token
   /current_stacks:
     get:
       tags:
@@ -615,6 +583,38 @@ paths:
               schema: {}
         '500':
           description: Failed to get node distribution
+  /google_oauth:
+    post:
+      tags:
+      - Auth
+      summary: |-
+        Logs in a user with the proxy service using Google OAuth.
+        This endpoint is used to verify a Google ID token and return an access token.
+      description: |-
+        # Arguments
+
+        * `proxy_service_state` - The shared state containing the state manager
+        * `body` - The request body containing the Google ID token
+
+        # Returns
+
+        * `Result<Json<AuthResponse>>` - A JSON response containing the access and refresh tokens
+      operationId: google_oauth
+      requestBody:
+        content:
+          text/plain:
+            schema:
+              type: string
+        required: true
+      responses:
+        '200':
+          description: Logs in a user with Google OAuth
+          content:
+            text/plain:
+              schema:
+                type: string
+        '500':
+          description: Failed to verify Google ID token
 components:
   schemas:
     AuthRequest: