Google OAuth with Rust Axum and AsyncGraphQL Backend

18 August, 2023

9 min read

Last updated on 11 February, 2024

In this article I won't put much attention on how web client is written. I will attach some code below, so you'd have rough understanding, but I won't discuss it further. You may also watch video series I did.

import { GoogleLogin } from "@react-oauth/google";
import { loader } from "graphql.macro";
import { useEffect, useState } from "react";
import { Link, useNavigate } from "react-router-dom";
import { OAuthTokenType } from "../../components/authenticated/AuthProvider";
import { useMutation } from "@apollo/client";

const loginMutation = loader("../../queries/login.gql");

export const Login = () => {
  const [login, { error }] = useMutation(loginMutation);
  const navigate = useNavigate();

  const handleSuccess = (response) => {
    const { credential } = response;

    login({
      variables: {
        token: credential,
        oauthTokenType: OAuthTokenType.GOOGLE.toString(),
      },
      onCompleted: (data) => {
        const token = data.login;
        sessionStorage.setItem("token", token);
        const origin = "/home"; // TODO: should be changed later according to user roles
        navigate(origin, { replace: true });
        // eslint-disable-next-line no-restricted-globals
        location.reload();
      },
    });
  };

  return (
    <>
      <div>
        <h1>Login</h1>

        <GoogleLogin
          onSuccess={handleSuccess}
          useOneTap
          auto_select
          theme="filled_black"
        />
      </div>
    </>
  );
};
src/pages/unauthenticated/Login.tsx

For all client major operations I use GraphQL mutations. For example, login mutation.

mutation Login(
  $token: String!,
  $oauthTokenType: String!
) {
  login(input: {
    token: $token,
    oauthTokenType: $oauthTokenType,
  })
}
src/queries/login.gql

There is also rought implementation of registration (just to give you an idea).

import { Route, Routes, useNavigate } from "react-router-dom";
import { Progress } from "../../components/unuathenticated/registration/Progress";
import { FirstStep } from "../../components/unuathenticated/registration/FirstStep";
import { SecondStep } from "../../components/unuathenticated/registration/SecondStep";
import { loader } from "graphql.macro";
import { useEffect, useState } from "react";
import { useMutation } from "@apollo/client";
import { OAuthTokenType } from "../../components/authenticated/AuthProvider";

const registerMutation = loader("../../queries/register.gql");

export const RegistrationPage = () => {
  const [token, setToken] = useState<string | null>(null);
  const [role, setRole] = useState<string | null>(null);
  const [spaceName, setSpaceName] = useState<string | null>(null);
  const [register] = useMutation(registerMutation);
  const navigate = useNavigate();

  useEffect(() => {
    if (token && (spaceName || role)) {
      register({
        variables: {
          token,
          spaceName,
          oauthTokenType: OAuthTokenType.GOOGLE.toString(),
          role,
        },
        onCompleted: (data) => {
          const token = data.register;
          sessionStorage.setItem("token", token);
          const origin = "/home";
          navigate(origin, { replace: true });
          navigate(0)
        },
        onError: (error) => {
          console.log(error);
        },
      });
    }
  // eslint-disable-next-line react-hooks/exhaustive-deps
  }, [token, spaceName, role]);

  return (
    <>
      <Progress />
      <Routes>
        <Route path="/" element={<FirstStep setToken={setToken} />} />
        <Route
          element={<SecondStep setSpaceName={setSpaceName} setRole={setRole} />}
          path="/second"
        />
      </Routes>
    </>
  );
};
src/pages/unauthenticated/Registration.tsx

Registration is divided into two steps, where in the fist step we set token to state after we receive it from Google.

import React, { useEffect } from "react";
import { useAuth } from "../../../hooks/auth";
import { GoogleLogin } from "@react-oauth/google";
import { useNavigate } from "react-router-dom";

export const FirstStep = ({ setToken }) => {
  const navigate = useNavigate();

  return (
    <>
      <GoogleLogin
        onSuccess={(response) => {
          setToken(response.credential);
          navigate("/registration/second");
        }}
        useOneTap
        auto_select
        theme="filled_black"
      />
    </>
  );
};
src/components/unuathenticated/registration/FirstStep.tsx

After we get from user all input (Google Token & custom input I need for my application), we can fire register mutation, which looks like that.

mutation Registration(
  $token: String!
  $spaceName: String
  $oauthTokenType: String!
  $role: String
) {
  register(userInput: {
    token: $token
    spaceName: $spaceName
    oauthTokenType: $oauthTokenType
    role: $role
  })
}
src/queries/register.gql

This mutation would give us token in response. But what token? We have just obtained Google OAuth token, what other token do we need?

Let's review the whole flow.

First phase: getting token from OAuth Provider
Second phase: validating token from provider and getting user data
Third phase: we issue own JWT token for client to give it long-live session
Query Requests with JWT token

Let's check now how backend looks like with that setup.

The whole project structure.

.
├── Cargo.lock
├── Cargo.toml
├── README.md
├── Secrets.toml
├── api
│   ├── Cargo.toml
│   └── src
│       ├── db.rs
│       ├── graphql
│       │   ├── mod.rs
│       │   ├── mutation
│       │   │   ├── bookable_resource.rs
│       │   │   ├── login.rs
│       │   │   ├── mod.rs
│       │   │   ├── plan.rs
│       │   │   ├── registration.rs
│       │   │   └── space.rs
│       │   ├── query
│       │   │   ├── building.rs
│       │   │   ├── mod.rs
│       │   │   ├── plan.rs
│       │   │   ├── space.rs
│       │   │   └── user.rs
│       │   └── schema.rs
│       ├── guard.rs
│       ├── headers.rs
│       └── lib.rs
├── auth
│   ├── Cargo.toml
│   └── src
│       └── lib.rs
├── aws-utils
│   ├── Cargo.toml
│   ├── Readme.md
│   └── src
│       └── lib.rs
├── credentials
├── db
├── entity
│   ├── Cargo.toml
│   └── src
│       ├── bookable_resource.rs
│       ├── building.rs
│       ├── lib.rs
│       ├── plan.rs
│       ├── space.rs
│       ├── space_user_connection.rs
│       └── user.rs
├── migration
│   ├── Cargo.toml
│   ├── README.md
│   └── src
│       ├── lib.rs
│       ├── m20230528_create_table.rs
│       ├── m20230529_create_table.rs
│       ├── m20230626_create_table.rs
│       ├── m20230704_120420_create_table.rs
│       ├── m20230727_113735_space_users_connection.rs
│       ├── main.rs
│       └── statements.rs
├── service
│   ├── Cargo.toml
│   ├── src
│   │   ├── lib.rs
│   │   ├── mutation.rs
│   │   └── query.rs
│   └── tests
│       └── prepare.rs
└── src
    └── main.rs

18 directories, 53 files

The projects contains not only OAuth-related files. But all implementations we'd need are located in auth module, guard, main graphQL Handler in api/src/lib.rs and few GraphQL resolvers, like registration, login, logout and refresh token.

Middleware in guard.rs

Before request reaches GraphQL handler, it hits the middleware. There, token verification is being made. If token is legit, server would inject extension with parsed token into request context.

use crate::StatusCode;
use axum::headers::authorization::Bearer;
use axum::headers::{Authorization, HeaderMapExt};
use axum::http::Request;
use axum::middleware::Next;
use axum::response::Response;

pub async fn guard<T>(mut request: Request<T>, next: Next<T>) -> Result<Response, StatusCode> {
    let token = request
        .headers()
        .typed_get::<Authorization<Bearer>>()
        .map(|header| header.token().to_owned());
    request.extensions_mut().insert(None::<auth::ParsedToken>);

    if let Some(jwt) = token {
        let parsed_token = auth::auth_token(jwt);

        match parsed_token {
            Ok(parsed_token) => {
                request.extensions_mut().insert(Some(parsed_token));
            }
            Err(_) => {}
        }
    }

    Ok(next.run(request).await)
}
api/src/guard.rs

Middleware does validations only on tokens that were issued by Rust server. This is definition of function which does that validation.

...

pub fn auth_token(jwt: String) -> Result<ParsedToken, AuthError> {
    let decoded = decode::<Claims>(
        &jwt,
        &DecodingKey::from_secret(JWT_SECRET),
        &Validation::new(Algorithm::HS512),
    ).map_err(|e| {
        AuthError::JwtError(e)
    })?;

    Ok(ParsedToken {
        role: Role::from_str(&decoded.claims.role).map_err(|_| AuthError::CorruptData)?,
        user_id: decoded.claims.sub,
    })
}
...
auth/src/lib.rs

I use standard Rust crate jsonwebtoken. There are some structs mentioned in the function above, here are their definitions.

#[derive(Debug, Deserialize, Serialize)]
struct Claims {
    sub: i32,
    role: String,
    exp: usize,
}

Errors defined for module with thiserror

#[derive(Debug, Error)]
pub enum AuthError {
    #[error("Missing secret")]
    MissingSecret,

    #[error("google error: {0}")]
    GoogleError(#[from] anyhow::Error),

    #[error("token doesn't contain required user data")]
    OauthTokenDataMissing,

    #[error("Data in token is corrupted")]
    CorruptData,

    #[error("time addition overflow")]
    TimeAdditionOverflow,

    #[error("jwt error: {0}")]
    JwtError(#[from] JwtError),

    #[error("unknown token type")]
    UnknownTokenType,
}

ParsedToken is data server injects in my request context in the middleware.

#[derive(Debug, Clone)]
pub struct ParsedToken {
    pub role: Role,
    pub user_id: i32,
}

After validation has been made, request hits the main GraphQL handler. There is a definition of it below. debug handler pragma is not required, but it's convenient to have it, while you're still working on your application locally.

...

#[axum::debug_handler]
pub async fn graphql_handler(
    schema: Extension<AppSchema>,
    Extension(secret_store): Extension<SecretStore>,
    Extension(token): Extension<Option<auth::ParsedToken>>,
    Extension(headers): Extension<headers::Headers>,
    req: GraphQLRequest,
) -> GraphQLResponse {
    let execute = schema
        .execute(
            req.into_inner()
                .data(token)
                .data(secret_store)
                .data(headers),
        )
        .await;

    execute.into()
}

...
api/src/lib.rs

Main Auth GraphQL Operations

Let's review how login resolver for login mutation works.

...

pub async fn login(&self, ctx: &Context<'_>, input: LoginInput) -> Result<String> {
        let token = input.token;
        let secret_store = ctx
            .data::<SecretStore>()
            .map_err(|_| GraphQLError::ContextExtractionError.extend())?;

        let user_data = auth::validate_oauth_token(
            auth::OauthToken::new(token, input.oauth_token_type)?,
            secret_store,
        )
        .await.map_err(|e| GraphQLError::AuthError(e).extend())?;

        let db = ctx
            .data::<Database>()
            .map_err(|_| GraphQLError::ContextExtractionError.extend())?;

        let (user, role) = Query::get_user_by_email(user_data.email, db.get_connection())
            .await
            .map_err(|e| GraphQLError::DbError(e).extend())?;

        let token =
            auth::create_jwt(&user.id, &role).map_err(|e| GraphQLError::AuthError(e).extend())?;

        let refresh_token = auth::create_refresh_token(&user.id, &role)
            .map_err(|e| GraphQLError::AuthError(e).extend())?;

        let token_max_age = 60 * 60 * 24 * 30;

        ctx.insert_http_header(
            "Set-Cookie",
            format!(
                "refresh_token={}; HttpOnly; Path=/; Max-Age={}",
                refresh_token, token_max_age
            ),
        );

        Ok(token)
    }

...

The most important functions there are validate_oauth_token , create_jwt and create_refresh_token.

Two tokens are being used there: main token (short-lived) and hidden refresh token (long-lived). The refresh token needed to preserve user session, even when main token is expired. It isn't available for JavaScript to get and modify it.

Below are function definitions.

...

pub fn create_jwt(id: &i32, role: &Role) -> Result<String, AuthError> {
    let expiration = Utc::now()
        .checked_add_signed(chrono::Duration::minutes(60))
        .ok_or(AuthError::TimeAdditionOverflow)?
        .timestamp();

    let claims = Claims {
        sub: id.to_owned(),
        role: role.to_string(),
        exp: expiration as usize,
    };
    
    let header = Header::new(Algorithm::HS512);
    encode(&header, &claims, &EncodingKey::from_secret(JWT_SECRET))
        .map_err(|e| AuthError::JwtError(e))
}

...

pub fn create_refresh_token(id: &i32, role: &Role) -> Result<String, AuthError> {
    let expiration = Utc::now()
        .checked_add_signed(chrono::Duration::days(30))
        .ok_or(AuthError::TimeAdditionOverflow)?
        .timestamp();

    let claims = Claims {
        sub: id.to_owned(),
        role: role.to_string(),
        exp: expiration as usize,
    };
    
    let header = Header::new(Algorithm::HS512);
    encode(
        &header,
        &claims,
        &EncodingKey::from_secret(JWT_REFRESH_SECRET),
    )
    .map_err(|e| AuthError::JwtError(e))
}

...

pub async fn validate_oauth_token(
    token: OauthToken,
    secret_store: &SecretStore,
) -> Result<UserDataFromOauth, AuthError> {
    match token.oauth_token_type {
        OAuthTokenType::Google => {
            let google_client_id = secret_store
                .get("GOOGLE_CLIENT_ID")
                .ok_or(AuthError::MissingSecret)?;
            let google_client = AsyncClient::new(google_client_id);
            let data = google_client
                .validate_id_token(&token.token)
                .await
                .map_err(|e| AuthError::GoogleError(e))?;

            let GooglePayload {
                email,
                family_name,
                given_name,
                ..
            } = data;

            let email = email.ok_or(AuthError::OauthTokenDataMissing)?;
            let family_name = family_name.ok_or(AuthError::OauthTokenDataMissing)?;
            let given_name = given_name.ok_or(AuthError::OauthTokenDataMissing)?;

            Ok(UserDataFromOauth {
                email,
                family_name,
                given_name,
                oauth_provider: OauthProvider::Google,
            })
        }
    }
}
auth/src/lib.rs

JWT creation functions for both main and refresh token have a lot of similarities. For validation google oauth token I use client library – google-oauth crate. It abstracts network calls to Google, but you need to know that under the hood, it calls Google Auth API. It means, the less we do those calls, the better.

The next operation is refresh. Here how it looks in GraphQL resolver.

...

pub async fn refresh(&self, ctx: &Context<'_>) -> Result<String> {
        let headers = ctx
            .data::<Headers>()
            .map_err(|_| GraphQLError::ContextExtractionError.extend())?
            .clone();

        let refresh_token = headers
            .get("cookie")
            .ok_or(GraphQLError::Unauthorized {
                reason: "No cookie header",
            }.extend())?
            .to_str()
            .map_err(|_| GraphQLError::Unauthorized {
                reason: "Cookie header is not a string",
            }.extend())?
            .split(";")
            .find(|cookie| cookie.contains("refresh_token"))
            .ok_or(GraphQLError::Unauthorized {
                reason: "No refresh token in cookie",
            }.extend())?
            .split("=")
            .nth(1)
            .ok_or(GraphQLError::Unauthorized {
                reason: "No refresh token in cookie",
            }.extend())?
            .to_owned();

        let new_token = auth::refresh_jwt_token_with_refresh_token(refresh_token)
            .map_err(|e| GraphQLError::AuthError(e).extend())?;

        Ok(new_token)
    }
    
...
api/src/graphql/mutation/login.rs

We see here new auth function refresh_jwt_token_with_refresh_token. Here is its definition.

...

pub fn refresh_jwt_token_with_refresh_token(refresh_token: String) -> Result<String, AuthError> {
    let decoded = decode::<Claims>(
        &refresh_token,
        &DecodingKey::from_secret(JWT_REFRESH_SECRET),
        &Validation::new(Algorithm::HS512),
    )?;

    let role = Role::from_str(&decoded.claims.role).map_err(|_| AuthError::CorruptData)?;

    let expiration = Utc::now()
        .checked_add_signed(chrono::Duration::minutes(60))
        .ok_or(AuthError::TimeAdditionOverflow)?
        .timestamp();

    let claims = Claims {
        sub: decoded.claims.sub,
        role: role.to_string(),
        exp: expiration as usize,
    };
    let header = Header::new(Algorithm::HS512);
    encode(&header, &claims, &EncodingKey::from_secret(JWT_SECRET))
        .map_err(|e| AuthError::JwtError(e))
}

...
auth/src/lib.rs

Based on refresh token in cookie, it issues new main token for requester.

The next operation is registration, which also relies on OAuth and some input from user.

use std::str::FromStr;

use crate::db::Database;
use async_graphql::ErrorExtensions;
use async_graphql::{Context, InputObject, Object, Result};
use auth::{UserDataFromOauth};

use crate::graphql::GraphQLError;
use crate::SecretStore;
use graphql_example_service::Mutation;

#[derive(Default)]
pub struct RegistrationMutation;

#[derive(InputObject, Debug)]
pub struct UserInput {
    pub token: String,
    pub oauth_token_type: String,
    pub space_name: Option<String>,
    pub role: Option<String>,
}

#[Object]
impl RegistrationMutation {
    pub async fn register(&self, ctx: &Context<'_>, user_input: UserInput) -> Result<String> {
        let db = ctx.data::<Database>().unwrap();
        let secret_store = ctx
            .data::<SecretStore>()
            .map_err(|_| GraphQLError::ContextExtractionError.extend())?;

        let user_data = auth::validate_oauth_token(
            auth::OauthToken::new(user_input.token, user_input.oauth_token_type)?,
            secret_store,
        )
        .await
        .map_err(|e| GraphQLError::AuthError(e).extend())?;

        let UserDataFromOauth {
            email,
            family_name,
            given_name,
            ..
        } = user_data;

        let display_name = format!("{} {}", given_name, family_name);
        let role = match user_input.role {
            Some(role) => {
                let role = auth::Role::from_str(&role).map_err(|e| {
                    GraphQLError::InvalidInput {
                        field: "role".to_string(),
                        message: e.to_string(),
                    }
                    .extend()
                })?;
                role
            }
            None => {
                auth::Role::SpaceAdmin
            },
        };
        let (registered_user, role, _) = match role {
            auth::Role::SpaceAdmin => {
                let space_name = user_input.space_name.clone().ok_or(
                    GraphQLError::InvalidInput {
                        field: "space_name".to_string(),
                        message: "space_name is required for SpaceAdmin role".to_string(),
                    }
                    .extend(),
                )?;

                let tuple = Mutation::register_user_space_owner(
                    db.get_connection(),
                    email,
                    display_name,
                    space_name,
                )
                .await
                .map_err(|e| GraphQLError::DbError(e).extend())?;

                tuple
            }
            auth::Role::SpaceUser => {
                let tuple =
                    Mutation::register_user_space_user(db.get_connection(), email, display_name)
                        .await
                        .map_err(|e| GraphQLError::DbError(e).extend())?;

                tuple
            }
            auth::Role::ApplicationAdmin => {
                return Err(GraphQLError::InvalidInput {
                    field: "role".to_string(),
                    message: "ApplicationAdmin role is not allowed for registration".to_string(),
                }
                .extend()
                .into());
            }
        };
        

        let token = auth::create_jwt(&registered_user.id, &role)
            .map_err(|e| GraphQLError::AuthError(e).extend())?;

        let refresh_token = auth::create_refresh_token(&registered_user.id, &role)
            .map_err(|e| GraphQLError::AuthError(e).extend())?;

        let token_max_age = 60 * 60 * 24 * 30;

        ctx.insert_http_header(
            "Set-Cookie",
            format!(
                "refresh_token={}; HttpOnly; Path=/; Max-Age={}",
                refresh_token, token_max_age
            ),
        );

        Ok(token)
    }
}
api/src/graphql/mutation/registration.rs

We receive Google token with user input. It's being done on client side.

And the last operation we need – logout. On server I did nothing, just made some plain client logic for that. This is not real logout and in real applications you shouldn't rely only on that. Server-side logout should be done as well. I just haven't covered it in the article as anyone could do it on their own. Implementing one is pretty straightforward.

...

const handleLogout = () => {
    googleLogout();
    sessionStorage.removeItem("token");
    navigate("/", { replace: true });
    navigate(0) // hard refresh
  };
  
  ...
src/components/authenticated/AuthProvider.tsx

You can subscribe on my newsletters

Let's see if we can become internet friends.

Check also related posts

Troy Köhler

TwitterYouTubeInstagramLinkedin