Customize your GraphQL Federation authentication and authorization with Grafbase Extensions

We recently launched Grafbase extensions, a powerful way to cutomize the Grafbase Gateway with your own logic using WebAssembly. While our previous blog posts focused on extensions that resolve fields without requiring subgraph infrastructure, such as: rest and nats, today we'll explore how to implement custom authentication and authorization with extensions.

Authentication extensions run before any GraphQL processing, restricting access at the request level. Authorization extensions, on the other hand, provide fine-grained access control over which fields, objects and other elements can be accessed.

A complete end to end example is available in our Grafbase repository. And the Grafbase SDK documentation covers everything you can do with extensions.

Before building your own extension, consider using existing ones! Currently available extensions include:

• JWT: authenticate clients with a JWT token
• Authenticated: Restrict access to unauthenticated clients.
• Requires Scopes: Grant access only to clients with appropriate OAuth scopes

Find more extensions in the Extensions Marketplace.

Installing an extension is straightforward, add them to your configuration:

# grafbase.toml
[extensions]
authenticated = "1"

[extensions.jwt]
version = "1"

[extensions.jwt.config]
url = "https://example.com/.well-known/jwks.json"

And install them with the Grafbase CLI:

grafbase extension install

All extension code is freely available under Apache-2 license on GitHub. Feel free to open an issue for improvements or you can fork them to better suit your needs. Extensions don't need to be in the Marketplace to be used, you can host them yourself.

Let's start by creating a simple authentication extension. First, bootstrap the project with the Grafbase CLI:

grafbase extension init --type authentication guard

This creates a Rust project using the Grafbase SDK. The extension.toml file defines the extension's metadata such as its name and type. We also a dummy integration test prepared that starts the Grafbase Gateway.

 guard
 ├─ src
 │  └─ lib.rs
 ├─ tests
 │  └─ integration_tests.rs
 ├─ Cargo.toml
 ├─ Cargo.lock
 └─ extension.toml

The src/lib.rs containing the business logic. The extension must implement the AuthenticationExtension trait. The new() method is called when the extension is loaded, and the authenticate() method is called for each incoming request.

use grafbase_sdk::{
    AuthenticationExtension,
    types::{Configuration, Error, ErrorResponse, GatewayHeaders, Token},
};

#[derive(AuthenticationExtension)]
struct Guard;

impl AuthenticationExtension for Guard {
    fn new(config: Configuration) -> Result<Self, Error> {
        Ok(Self)
    }

    fn authenticate(&mut self, headers: &GatewayHeaders) -> Result<Token, ErrorResponse> {
        Err(ErrorResponse::unauthorized())
    }
}

First, let's load any configuration that you might have:

#[derive(AuthenticationExtension)]
struct Guard {
    config: Config,
}

#[derive(serde::Deserialize)]
struct Config {
    header_name: String,
    secret: String
}

impl AuthenticationExtension for Guard {
    fn new(config: Configuration) -> Result<Self, Error> {
        Ok(Self {
            config: config.deserialize()?,
        })
    }
    // ...
}

The provided Configuration holds settings from grafbase.toml which can include environment variables:

# grafbase.toml
[extensions.guard.config]
header_name = "Authorization"
secret = "{{ env.SECRET_KEY }}"

With the help of serde you can parse it and validate simultaneously. You can initialize your extension however you want. If an error is returned, all authentication calls will fail.

Now let's see how authenticate() works. It receives the read-only GatewayHeaders and returns either a Token or an ErrorResponse. A Token can be either anonymous or hold bytes. There's no constraint on the nature of the bytes. For example, the JWT extension provides the JWT claims as JSON. An ErrorResponse is special kind of error that prevents any further execution in the gateway and allows returning a custom HTTP status code and one or multiple GraphQL errors.

Let's implement simple logic that validates whether we have a header with the appropriate secret:

fn authenticate(&mut self, headers: &GatewayHeaders) -> Result<Token, ErrorResponse> {
    let Some(value) = headers.get(&self.config.header_name) else {
        return Err(ErrorResponse::unauthorized().with_error("Missing authorization header"));
    };

    if value != self.config.secret {
        return Err(ErrorResponse::unauthorized().with_error("Invalid token"));
    };

    Ok(Token::from_bytes(vec![]))
}

The error response we return will have a 401 status code in this example, but you can also customize it. The error message is optional, and multiple errors can be added. Here we choose to return a Token with empty bytes as there is no data to forward but still want to distinguish it from anonymous tokens if we want to use Authenticated extension later on. The Grafbase Gateway can be configured to provide an anonymous token if authentication fails instead of returning an error:

# grafbase.toml
authentication.default = "anonymous"

To finalize our extension, we need to build it with:

grafbase extension build

This will create a build directory containing everything needed by the Grafbase Gateway. You can then use the extension with:

# grafbase.toml
[extensions.guard]
path = "<..>/guard/build"

[extensions.guard.config]
header_name = "Authorization"
secret = "{{ env.SECRET_KEY }}"

For more a more advanced example see our JWT (GitHub) extension, our Grafbase SDK and our end to end authentication example.

Since WebAssembly is single-threaded, the Grafbase Gateway maintains a pool of WebAssembly instances. The maximum pool size and other settings can be controlled in the configuration. To share data across instances, the Grafbase SDK provides a simple in-memory cache. The JWT extension uses this to share the downloaded JWKS for example.

An authorization extension provides fine-grained control over client access to resources. Similar to our authentication extension, we start by initializing our project with the Grafbase CLI:

cargo run -p grafbase extension init --type authorization warden

The structure resembles the other extension but includes a key difference, the definitions.graphql file:

 warden
 ├─ tests
 │  └─ integration_tests.rs
 ├─ src
 │  └─ lib.rs
 ├─ Cargo.toml
 ├─ Cargo.lock
 ├─ extension.toml
 └─ definitions.graphql

Unlike authentication extensions, authorization dependeds on what the client is requesting. Let's examine how this works with the @authenticated directive from the Authenticated extension:

# authenticated's definitions.graphql
directive @authenticated on FIELD_DEFINITION | OBJECT | INTERFACE | ENUM | SCALAR

This directive definition allows subgraph owners to protect specific schema elements:

# subgraph schema
extend schema
  @link(
    url: "https://grafbase.com/extensions/authenticated/1.0.0"
    import: ["@authenticated"]
  )

type Query {
  me: User! @authenticated
  user(id: ID!): User!
}

type User {
  name: String
}

When a client requests query { me { name} }, the gateway calls the Authenticated. However, it won't be called for query { user(id: 1) { name } }. The @link imports extension directive, and the URL can point to any location containing the extension's build directory content.

You can create powerful authorization patterns, the Grafbase Gateway will validate every use of your directives:

• Define multiple directives
• Use custom input types, scalars and enums for your directive arguments.
• Inject query and/or response data into directive arguments

For example, using the InputValueSet scalar from the Grafbase spec allows injecting field arguments:

# extension defintions.graphql
extend schema
  @link(url: "https://specs.grafbase.com/grafbase", import: ["InputValueSet"])

# '*' means all arguments are injected
directive @protect(arguments: InputValueSet = "*") on FIELD_DEFINITION

Subgraph owners can then protect endpoints with logic that depends on arguments:

# subgraph schema
extend schema @link(url: "file:///<path>/warden/build", import: ["@protect"])

type Query {
  user(id: ID!): User! @protect
}

type User {
  name: String
}

Let's examine the AuthorizationExtension trait. Similarly to before we can inject custom configuration in the new() method and the authorize_query() method has the business logic:

use grafbase_sdk::{
    AuthorizationExtension, IntoQueryAuthorization,
    types::{AuthorizationDecisions, Configuration, Error, ErrorResponse, QueryElements, SubgraphHeaders, Token},
};

#[derive(AuthorizationExtension)]
struct Warden;

impl AuthorizationExtension for Warden {
    fn new(config: Configuration) -> Result<Self, Error> {
        Ok(Self)
    }

    fn authorize_query(
        &mut self,
        headers: &mut SubgraphHeaders,
        token: Token,
        elements: QueryElements<'_>,
    ) -> Result<impl IntoQueryAuthorization, ErrorResponse> {
        Ok(AuthorizationDecisions::deny_all("Not authorized"))
    }
}

It receives:

• headers: Mutable subgraph headers that will be sent to subgraphs. They're the result of applying the global header rules to the gateway headers, but subgraph-specific ones aren't applied yet.
• token: The token from an authentication extension.
• elements: Query elements, such as fields or objects, requiring authorization.

Each QueryElement provides:

• directive_site(): Location metadata such as the field name.
• directive_arguments(): The directive arguments after data injection. In our example with @protect, given the query query { user(id: 1) { name} } this would return {"arguments": {"id": 1}}.

The method returns AuthorizationDecisions determining access for each element.

Here's an implementation that only allows access to the user with ID 1:

fn authorize_query(
    &mut self,
    headers: &mut SubgraphHeaders,
    token: Token,
    elements: QueryElements<'_>,
) -> Result<impl IntoQueryAuthorization, ErrorResponse> {
    #[derive(serde::Deserialize)]
    struct Protect<T> {
        arguments: T,
    }

    #[derive(serde::Deserialize)]
    struct UserArguments {
        id: String,
    }

    let mut builder = AuthorizationDecisions::deny_some_builder();
    for element in elements {
        let DirectiveSite::FieldDefinition(field) = element.directive_site() else {
            unreachable!()
        };
        match (field.parent_type_name(), field.name()) {
            ("Query", "user") => {
                let protect: Protect<UserArguments> = element.directive_arguments()?;
                let user_id = protect.arguments.id;
                if user_id != "1" {
                    builder.deny(element, "Access denied");
                }
            }
            _ => unreachable!(),
        }
    }

    Ok(builder.build())
}

When access is denied, the field won't be queried from the subgraph. Authorization extensions can also prepare authorization information for subgraphs based on accessed fields:

# extension defintions.graphql
directive @auth(
  scopes: [String!]!
) on FIELD_DEFINITION | OBJECT | INTERFACE | ENUM | SCALAR

And generate a new token for subgraphs based on accessed fields:

fn authorize_query(
    &mut self,
    headers: &mut SubgraphHeaders,
    token: Token,
    elements: QueryElements<'_>,
) -> Result<impl IntoQueryAuthorization, ErrorResponse> {
    let mut merged_scopes = HashSet::new();

    #[derive(serde::Deserialize)]
    struct Arguments {
        scopes: Vec<String>,
    }
    for element in elements {
        let Arguments { scopes } = element.directive_arguments()?;
        merged_scopes.extend(scopes);
    }

    // build_token is a custom function with your logic.
    let subgraph_token = build_token(merged_scopes);
    headers.set(http::header::AUTHORIZATION, [subgraph_token]);

    Ok(AuthorizationDecisions::grant_all())
}

To minimize performance impact, Authorization is applied after query planning. It removes unnecessary fields and requirements from subgraph queries, similar to how @include and @skip directives work in the Grafbase Gateway.

Beyond controlling field access, authorization extensions can modify GraphQL responses before they reach the client when authorization depends on response data

# extension defintions.graphql
extend schema
  @link(url: "https://specs.grafbase.com/grafbase", import: ["FieldSet"])

# FieldSet defines which fields from the object should be injected.
directive @sensitive(fields: FieldSet!) on OBJECT

# subgraph schema
extend schema @link(url: "file:///<path>/warden/build", import: ["@sensitive"])

type Query {
  users: [User!]
}

type User @sensitive(fields: "id") {
  id: ID!
  name: String
}

Here we implement with the same business logic as our previous example, but this time relying on response data

impl AuthorizationExtension for Warden {
    // ...

    fn authorize_query(
        &mut self,
        headers: &mut SubgraphHeaders,
        token: Token,
        elements: QueryElements<'_>,
    ) -> Result<impl IntoQueryAuthorization, ErrorResponse> {
        let state: Vec<u8> = token.into_bytes().unwrap_or_default();
        Ok((AuthorizationDecisions::grant_all(), state))
    }

    fn authorize_response(
        &mut self,
        state: Vec<u8>,
        elements: ResponseElements<'_>,
    ) -> Result<AuthorizationDecisions, Error> {
        #[derive(serde::Deserialize)]
        struct Arguments<T> {
            fields: T,
        }

        #[derive(serde::Deserialize)]
        struct User {
            id: String,
        }

        let mut builder = AuthorizationDecisions::deny_some_builder();
        for element in elements {
            for item in element.items() {
                let Arguments { fields: User { id } } = item.directive_arguments()?;
                if id != "1" {
                    builder.deny(item, "Access denied");
                }
            }
        }

        Ok(builder.build())
    }
}

authorize_response() method is called with only two arguments: state which are custom bytes provided by authorize_query() and all the elements requiring authorization. It is only called for directive with arguments depending on response data. Unlike query authorization, this approach doesn't prevent execution of queries like query { users { name} }, but instead modifies the response afterward. The Grafbase Gateway ensures that any fields needed for authorization (via FieldSet) are requested from the subgraph.

Congratulations on making it to the end! Authentication and authorization are complex topics, especially in GraphQL with its highly dynamic nature. Our extensions offer what we believe is an efficient solution to it giving you fine-grained control over the logic applied by the Grafbase Gateway.

We're eager to hear your thoughts and experiences! Extensions doc / Grafbase Gateway GitHub / Discord