Redirect WebSockets to different backends in Azure API Management

Azure API Management supports proxying WebSocket connections between clients and multiple backends. Adding a WebSocket API to API Management configures the handshake operation automatically. But what if we want to chose a backend based on the WebSockets URL? Let's explore our options.

First of all big thanks to my colleague Yann Duval, who co-authored the solutions blow!

Before diving into the several options, we need to understand, how establishing a WebSocket connection works. When a client tries to connect to ws://foo, the first request that is actually made behind the curtain is a GET request to http://foo, which then (if successful) upgrades the connection to a WebSocket connection using the ws:// protocol from that moment on. This first GET request is called the handshake. We will use this knowledge to intercept the handshake request and eventually redirect that request before upgrading to an actual WebSocket connection.

Routing based on query parameters

The first and most straight-forward way to configure routing is based on query parameters. This is also one that is most compliant with the WebSocket specification, as the handshake endpoint should be available at the same URL as the WebSocket itself, so query parameters are a way to keep the URL consistent and still add additional properties to the request.

When we use query parameter for choosing a backend, we can use the auto-generated handshake operation that API Management adds when creating a WebSocket API and add a custom policy to its Inbound policies.

<policies>
    <inbound>
        <base />
        <set-backend-service base-url="@{
            string myParam = context.Request.Url.Query.GetValueOrDefault("foo", "lorem");
            if (myParam.Equals("bar"))
            {
               return "wss://backend-a";
            }
            else
            {
                return "wss://backend-b";
            }
        }" />
    </inbound>
    <backend>
        <base />
    </backend>
    <outbound>
        <base />
    </outbound>
    <on-error>
        <base />
    </on-error>
</policies>

When a client now tries to open a WebSocket connection against wss://my-apim-url/my-websocket?foo=bar , the API Management will proxy that request to backend-a, otherwise to backend-b.

Routing based on path segments

If (for some reason) you can't use query parameters but need to check URL path segments to decide which backend to choose, you can not use the auto-generated onHandshake operation anymore, because its URL is immutable and will always listen to the same URL as the WebSocket.

In this scenario, we fist need to expose each WebSocket backend individually at the API Management level. We then later add one additional fake WebSocket Handshake API, which returns a redirect to one of the actual WebSocket APIs based on the URL path segments.

After adding each WebSocket backend individually, we add an HTTP-API to our API Management instance, for which you manually need to create a GET operation, that listens on the /* path an acts as our fake handshake. For this operation, we add an Inbound policy, which instead of setting the backend service returns a 302 redirect response.

<policies>
    <inbound>
        <return-response>
            <base />
            <set-status code="302" reason="Redirect" />
            <set-header name="Location" exists-action="override">
                <value>@{
                    string foo = context.Request.OriginalUrl.Path.Split('/')[2];
                    if (myParam.Equals("bar"))
                    {
                        return "wss://my-apim-url/websocket-a";
                    }
                    else
                    {
                        return "wss://my-apim-url/websocket-b";
                    }
                }</value>
            </set-header>
        </return-response>
    </inbound>
    <backend>
        <base />
    </backend>
    <outbound>
        <base />
    </outbound>
    <on-error>
        <base />
    </on-error>
</policies>

When a client tries to open a WebSocket connection against wss://my-apim-url/my-fake-websocket/bar , the API Management will return a redirect to wss://my-apim-url/websocket-a, which should be configured to point to the internal backend wss://backend-a. Otherwise it will redirect to wss://my-apim-url/websocket-b, which should be configured to point to the internal backend wss://backend-b.

Routing with custom logic

If the logic for deciding which backend to use becomes more complex or requires asynchronous tasks like database lookups, we can also deploy a custom service somewhere, which performs a similar logic than above and then returns a 302 redirect response.

Azure offers several hosting options for such a service. Unfortunately, at the time of writing this blog post, Azure Functions refused to answer on requests coming from the wss:// protocol, so I would recommend deploying the service as an Azure Container App to achieve a serverless solution.

This redirect-server should then also be added to API Management as HTTP-API to fake the WebSocket handshake just as in the option above.

A simple Node.js server using Express.js could look like this:

const express = require('express');

const app = express();
const port = process.env.PORT || 3000;

app.get('/*', (req, res) => {
  const foo = req.url.split('/')[2];    
  if (foo === 'bar') {
    res.setHeader("Location", "wss://my-apim-url/websocket-a");
  } else {
    res.setHeader("Location", "wss://my-apim-url/websocket-b");
  }
  
  res.writeHead(302);
  res.end();
});

app.listen(port, () => {
  console.log(`Server listening on port ${port}`);
});

Following redirects

As per RFC, clients are not required to follow redirects. So during during our tests, we had to explicitly enable the followRedirects flag in the ws WebSocket library we used. This might vary between client library, but keep in mind, that redirects need to be enabled for the above scenarios to work.

var Ws = require('ws')
var ws = new Ws('wss://my-apim-url/my-fake-websocket/bar', {
  followRedirects: true,
  headers: {
    "Ocp-Apim-Subscription-Key": "..." 
  }
});

ws.on('message', function message(data) {
  console.log('received: %s', data);
});

ws.on('error', function message(data) {
  console.log('error: %s', data);
});