Authenticating WebSocket connections in GraphQL with JWT, Asp.Net Core & HotChocolate

UPDATE

The latest versions of HotChocolate change how the socket session override classes work. Here is a quick example of some code that is working for 13.8.1.

JWT auth code is the same, follow the rest of the guide for more information and use this code as a base for your DefaultSocketSessionInterceptor implementation.

public class SocketConnectPayload
{
    public string Authorization { get; set; }
}

public class SubscriptionAuthMiddleware : DefaultSocketSessionInterceptor
{
    public override async ValueTask<ConnectionStatus> OnConnectAsync(ISocketSession connection, IOperationMessagePayload message, CancellationToken cancellationToken)
    {
        try
        {
            using var scope = connection.Connection.HttpContext.RequestServices.CreateScope();
            var payload = message.As<SocketConnectPayload>();
            var jwtHeader = payload.Authorization;

            if (!jwtHeader.StartsWith("Bearer "))
                return ConnectionStatus.Reject("Unauthorized");

            var token = jwtHeader.Replace("Bearer ", "");

            var opts = scope.ServiceProvider.GetRequiredService<JwtBearerOptions>();

            var claims = new JwtBearerBacker(opts).IsJwtValid(token);
            if (claims == null)
                return ConnectionStatus.Reject("Unauthoized(invalid token)");

            var userId = claims.FindFirst(ClaimTypes.NameIdentifier)?.Value ?? "";
            connection.Connection.HttpContext.Items["userId"] = userId;

            // Find the user
            using var ctx = scope.ServiceProvider.GetRequiredService<IDbContextFactory<DataContext>>().CreateDbContext();
            var user = await ctx.Users.FirstOrDefaultAsync(p => p.UserId == userId);
            if (user == null) return ConnectionStatus.Reject("Profile not created");

            connection.Connection.HttpContext.Items["userObj"] = user;

            return ConnectionStatus.Accept();
        }
        catch (Exception ex)
        {
            return ConnectionStatus.Reject(ex.Message);
        }
    }
}

Also, as mentioned by @marcinjaniak in the comments, this code does not handle JWT expirations on already established websocket connections (see their comment for an example on handling this).

The Problem

Whilst building a deployment tool for our backend services, I spent quite some time trying to work out how to correctly authenticate GraphQL Subscription(websocket) connections in HotChocolate.

My first thought was headers, but after quite some research I found that I cannot inject an authorization header into a websocket as the spec doesn't allow for headers and web browser just don't support it.

But after quite a bit of digging and figuring out how we'd managed this in node projects in the past, I found that the authorization header actually is still sent when the subscription websocket is connected, but rather than a header, it is sent with a payload that arrives on the initial connected state.

The Solution

So, now that I worked out how the authorization header can be sent, the next step was figuring out how to get this payload in HotChocolate.

After some digging, I found there is an extension method to services.AddGraphQLServer() entitled .AddSocketSessionInterceptor<T>().

Some additional digging found that T should be a class that extends from ISocketSessionInterceptor, a class with 3 methods for handling websocket events: OnCloseAsync, OnRequestAsync and OnConnectAsync. In this case, OnConnectAsync was the method I needed to parse our connection payload.

Building SubscriptionAuthMiddleware

So first off, we should setup our class, extend from ISocketSessionInterceptor and add our extension methods. Here is a skeleton class to start:

public class SubscriptionAuthMiddleware : ISocketSessionInterceptor 
{
    public async ValueTask OnCloseAsync(ISocketConnection connection, CancellationToken cancellationToken) { }
    public async ValueTask OnRequestAsync(ISocketConnection connection, IQueryRequestBuilder requestBuilder, CancellationToken cancellationToken) { }

    /* We don't need the above two methods, just this one */
    public async ValueTask<ConnectionStatus> OnConnectAsync(ISocketConnection connection, InitializeConnectionMessage message, CancellationToken cancellationToken) 
    {
        try 
        {

        }
        catch(Exception ex)
        {
            // If we catch any exceptions, reject the connection.
            // This is probably not ideal, there is likely a way to return a message
            // but I didn't look that far into it.
            return ConnectionStatus.Reject(ex.Message);
        }
    }
}

So now we have our class ready, we can start adding to OnConnectAsync. The payload we recieve can be found under message.Payload.
So lets grab our JWT, and fail if the JWT is not present. Add this code inside our try statement:

var jwtHeader = message.Payload["Authorization"] as string;

if (string.IsNullOrEmpty(jwtHeader) || !jwtHeader.StartsWith("Bearer "))
    return ConnectionStatus.Reject("Unauthorized");

var token = jwtHeader.Replace("Bearer ", "");

This will grab "Authorization" from payload if it exists, then check it contains the 'Bearer' part of the header. If it doesn't, reject the connection.
After that, we parse out the "Bearer " from the header and are left with our JWT.

JWT Config Shenanigans

Before we continue, we're going to have to take a brief adventure over to our JWT setup in Startup.cs as we need to add our JWT config to the service provider.
Simply put, we're going to extract the JWT Config out of .AddJwtBearer, give it its own variable and add it to our service provider. I won't provide much details on how to do this as it's fairly simple. This is what I have afterwards:

var opts = new JwtBearerOptions()
{
    Authority = Configuration["Auth:Authority"],
    TokenValidationParameters = new TokenValidationParameters
    {
        ValidateIssuer = true,
        ValidIssuer = Configuration["Auth:Authority"],
        ValidateAudience = true,
        ValidAudience = Configuration["Auth:Audience"],
        ValidateLifetime = true
    }
};
services.AddSingleton(opts);

services
    .AddAuthentication(JwtBearerDefaults.AuthenticationScheme)
    .AddJwtBearer(options =>
    {
        options.Authority = opts.Authority;
        options.TokenValidationParameters = opts.TokenValidationParameters;
    });

This should get us by, we'll come back to why we did that later when we get to configuring our JWT validation.

Building our JWT Validator

Before we can parse our JWT, we are going to need to build a class to validate it outside the usual ASP.Net Core Authentication pipeline.
Before we continue, I'd like to give credit to Christopher J. where I got the JWTBearerBacker that I built off of.

Lets start by creating a skeleton class:

public class JwtBearerBacker
{
    public JwtBearerOptions Options { get; private set; }

    public JwtBearerBacker(JwtBearerOptions options)
    {
        this.Options = options;
    }

    public ClaimsPrincipal IsJwtValid(string token)
    {

    }
}

Here we have a simple class that accepts the JWTBearerOptions we added to our ServiceProvider earlier. This class will be built within our OnConnectAsync method later.
Now, I use Firebase for authentication and I ran into an issue with the original code from Christopher J.. It did not work with Google's JWT signing keys, so I had to add some code to get it to pull that information from our Firebase project.

In our IsJwtValid method, start by adding this code, and if you're using Firebase, replace 'MYPROJECTID' with your project ID. Otherwise, replace the entire URL with your link to openid-configuration. You may be able to use the code from Christopher J. directly, but it will depend on your Auth provider.

var configManager = new ConfigurationManager<OpenIdConnectConfiguration>(
    $"https://securetoken.google.com/MYPROJECTID/.well-known/openid-configuration",
    new OpenIdConnectConfigurationRetriever()
);

var openidconfig = configManager.GetConfigurationAsync().Result;

This code will initialise a ConfigurationManager that will load our openID config and keys from Firebase for us!
Next, we will recreate our TokenValidationParamaters with some additional values included(such as our signing keys):

Options.TokenValidationParameters = new TokenValidationParameters()
{
    ValidateIssuer = true,
    ValidIssuer = "https://securetoken.google.com/MYPROJECTID",
    ValidateAudience = true,
    ValidAudience = "MYPROJECTID",
    ValidateLifetime = true,

    ValidateIssuerSigningKey = true,
    IssuerSigningKeys = openidconfig.SigningKeys,
};

You may want to find a way to inject the configuration into your JWTBearerBacker class if you wish to use config for your Firebase credentials.
Next on our list is a for-loop! This one is from the original code mostly, it loops through the JWT validators and attempts to validate our token with them. If an exception is caught due to key errors, it may attempt to refresh the configuration for up-to-date keys.

Add this code after the above code in your IsJwtValid method:

List<Exception> validationFailures = null;
SecurityToken validatedToken;
foreach (var validator in Options.SecurityTokenValidators)
{
    // Ensure we can even read the token at all
    if (validator.CanReadToken(token))
    {
        try
        {
            // Try to return a ClaimsPrincipal if we can
            // Otherwise an exception is thrown, caught and we continue on.
            return validator
                .ValidateToken(token, Options.TokenValidationParameters, out validatedToken);
        }
        catch (Exception ex)
        {
            // If the keys are invalid, refresh config
            if (Options.RefreshOnIssuerKeyNotFound && Options.ConfigurationManager != null
                && ex is SecurityTokenSignatureKeyNotFoundException)
            {
                Options.ConfigurationManager.RequestRefresh();
            }

            // Add to our list of failures. This was from the OG code
            // Not sure what we need it for.
            if (validationFailures == null)
                validationFailures = new List<Exception>(1);

            validationFailures.Add(ex);
            continue;
        }
    }
}

// No user could be found
return null;

There we have it, a complete JWT validation class that'll automatically pull the JWT keys as needed!
Now, back to our websocket code..

Returning to SubscriptionAuthMiddleware

Next thing we need to do is pull our JWT config from the ServiceProvider, which we can get from our socket context,
then attempt to validate the JWT we got from the payload with our JWT class:

var opts = connection.HttpContext.RequestServices.GetRequiredService<JwtBearerOptions>();

var claims = new JwtBearerBacker(opts).IsJwtValid(token);
    if (claims == null)
        return ConnectionStatus.Reject("Unauthoized(invalid token)");

If the could not be found, that means the JWT couldn't be validated, so we reject the connection.

Last thing we need to do is grab the User ID an add it to our HttpContext so we can access it later in our Subscription handlers!

// Grab our User ID
var userId = claims.FindFirst(ClaimTypes.NameIdentifier)?.Value ?? "";

// Add it to our HttpContext
connection.HttpContext.Items["userId"] = userId;

// Accept the websocket connection
return ConnectionStatus.Accept();

Woo-hoo! We built an authorization handler for Subscriptions! Now, lets head back to Startup and add it to our GraphQl setup.
In Startup.cs, under ConfigureServices, find services.AddGraphQLServer() and add the end of all your methods(ensure you have .AddAuthorization in there),
add the following line:

.AddSocketSessionInterceptor<SubscriptionAuthMiddleware>();

This will tell HotChocolate to proxy our subscription requests through our middleware class.

Setting up a subscription

Before you can access the HttpContext in your subscription, you'll have to head back to ConfigureServices and add this line:

services.AddHttpContextAccessor();

I'll assume you know how to add your subscription class to HotChocolate and set it up.
Once you have one, you can inject the accessor into your subscription handler method and access the UserID like so:

[Subscribe]
[Topic("MyTopic")]
public async Task<MyReturnType> OnMyTopic([EventMessage] MyType e, [Service] IHttpContextAccessor hctx)
{
    var userId = hctx.HttpContext.Items["userId"] as string;
    if (userId == null)
        throw new Exception("Unauthorized");

    // ...
    return myReturnType;
}

And that's it, you can now authorize your Subscription requests!

Extra: How to add the header to URQL in Typescript

You can initialise your URQL subscription like so to inject the Authorize header(Credit to my friend Wheatley for this code):

const subscriptionClient = new SubscriptionClient("WEBSOCKETURL", {
  reconnect: true,
  connectionParams: async() => {
    const token = ...;
    return {
        Authorization: `Bearer ${token}`,
    }    
  },
})

The End!

Well that's it! For reference, here are our completed code:

SubscriptionAuthMiddleware

public class SubscriptionAuthMiddleware : ISocketSessionInterceptor 
{
    public async ValueTask OnCloseAsync(ISocketConnection connection, CancellationToken cancellationToken) { }
    public async ValueTask OnRequestAsync(ISocketConnection connection, IQueryRequestBuilder requestBuilder, CancellationToken cancellationToken) { }

    /* We don't need the above two methods, just this one */
    public async ValueTask<ConnectionStatus> OnConnectAsync(ISocketConnection connection, InitializeConnectionMessage message, CancellationToken cancellationToken) 
    {
        try 
        {
            var jwtHeader = message.Payload["Authorization"] as string;

            if (string.IsNullOrEmpty(jwtHeader) || !jwtHeader.StartsWith("Bearer "))
                return ConnectionStatus.Reject("Unauthorized");

            var token = jwtHeader.Replace("Bearer ", "");
            var opts = connection.HttpContext.RequestServices.GetRequiredService<JwtBearerOptions>();

            var claims = new JwtBearerBacker(opts).IsJwtValid(token);
                if (claims == null)
                    return ConnectionStatus.Reject("Unauthoized(invalid token)");

            // Grab our User ID
            var userId = claims.FindFirst(ClaimTypes.NameIdentifier)?.Value ?? "";

            // Add it to our HttpContext
            connection.HttpContext.Items["userId"] = userId;

            // Accept the websocket connection
            return ConnectionStatus.Accept();
        }
        catch(Exception ex)
        {
            // If we catch any exceptions, reject the connection.
            // This is probably not ideal, there is likely a way to return a message
            // but I didn't look that far into it.
            return ConnectionStatus.Reject(ex.Message);
        }
    }
}

JwtBearerBacker

public class JwtBearerBacker
{
    public JwtBearerOptions Options { get; private set; }

    public JwtBearerBacker(JwtBearerOptions options)
    {
        this.Options = options;
    }

    public ClaimsPrincipal IsJwtValid(string token)
    {
        var configManager = new ConfigurationManager<OpenIdConnectConfiguration>(
            $"https://securetoken.google.com/MYPROJECTID/.well-known/openid-configuration",
            new OpenIdConnectConfigurationRetriever()
        );

        var openidconfig = configManager.GetConfigurationAsync().Result;
        Options.TokenValidationParameters = new TokenValidationParameters()
        {
            ValidateIssuer = true,
            ValidIssuer = "https://securetoken.google.com/MYPROJECTID",
            ValidateAudience = true,
            ValidAudience = "MYPROJECTID",
            ValidateLifetime = true,

            ValidateIssuerSigningKey = true,
            IssuerSigningKeys = openidconfig.SigningKeys,
        };

        List<Exception> validationFailures = null;
        SecurityToken validatedToken;
        foreach (var validator in Options.SecurityTokenValidators)
        {
            // Ensure we can even read the token at all
            if (validator.CanReadToken(token))
            {
                try
                {
                    // Try to return a ClaimsPrincipal if we can
                    // Otherwise an exception is thrown, caught and we continue on.
                    return validator
                        .ValidateToken(token, Options.TokenValidationParameters, out validatedToken);
                }
                catch (Exception ex)
                {
                    // If the keys are invalid, refresh config
                    if (Options.RefreshOnIssuerKeyNotFound && Options.ConfigurationManager != null
                        && ex is SecurityTokenSignatureKeyNotFoundException)
                    {
                        Options.ConfigurationManager.RequestRefresh();
                    }

                    // Add to our list of failures. This was from the OG code
                    // Not sure what we need it for.
                    if (validationFailures == null)
                        validationFailures = new List<Exception>(1);

                    validationFailures.Add(ex);
                    continue;
                }
            }
        }

        // No user could be found
        return null;
    }
}

Blog