Varnish VCL example for JWT and JWK

jwt_jwk.vcl

#
# JWT - Remote JWK API
#

vcl 4.1;

import edgestash;
import goto;
import http;
import json;
import jwt;
import kvstore;
import std;

acl jwt_localhost
{
    "127.0.0.0"/8;
    "::1"/128;
}

sub vcl_init
{
    # Init JWT and JWK
    new jwt_opts = kvstore.init();

    jwt_opts.set("jwk_scheme", "http");
    jwt_opts.set("jwk_host", "localhost:8008");
    jwt_opts.set("jwk_url", "/jwt/jwk.json");
    jwt_opts.set("alg", "RS256");

    new jwt_reader = jwt.reader();

    call get_jwk_key;
}

# Get the initial JWK key from a remote server
sub get_jwk_key
{
    http.init(0, MEDIUM);
    http.req_set_url(0,
        jwt_opts.get("jwk_scheme") + "://" +
        jwt_opts.get("jwk_host") +
        jwt_opts.get("jwk_url"));

    http.req_send(0);
    http.resp_wait(0);

    if (http.resp_get_status(0) != 200) {
        return (fail("Could not get JWK"));
    }

    jwt_opts.set("jwk_key", http.resp_get_body(0));
}

# Update the JWK key from a remote server
sub update_jwk_key
{
    http.init(0, MEDIUM);
    http.req_set_method(0, "HEAD");
    http.req_set_url(0,http.varnish_url(jwt_opts.get("jwk_url")));
    http.req_set_header(0, "Host", jwt_opts.get("jwk_host"));

    http.req_send(0);
    http.resp_wait(0);
}

sub vcl_recv
{
    # JWK support
    unset req.http.is-jwk;
    if (req.http.Host == jwt_opts.get("jwk_host") &&
            req.url == jwt_opts.get("jwk_url") &&
            client.ip ~ jwt_localhost) {
        set req.http.is-jwk = "true";
        return (hash);
    }

    # Init the JWT key
    jwt_reader.set_jwk(jwt_opts.get("jwk_key"));

    # Parse the JWT
    if (!jwt_reader.parse(req.http.Authorization)) {
        return (synth(401, "MISSING AUTH"));
    }

    # Validate the JWT
    if (!jwt_reader.verify(jwt_opts.get("alg"))) {
        # Attempt to get a new JWK
        call update_jwk_key;
        jwt_reader.set_jwk(jwt_opts.get("jwk_key"));

        # Re-validate with an updated key
        if (!jwt_reader.verify(jwt_opts.get("alg"))) {
            return (synth(401, "VERIFY FAILED"));
        }
    }
}

sub vcl_backend_fetch
{
    # JWK support
    if (bereq.http.is-jwk) {
        set bereq.backend = goto.dns_backend(
                jwt_opts.get("jwk_scheme") + "://" +
                jwt_opts.get("jwk_host"));
        return (fetch);
    }
}

sub vcl_backend_response
{
    # JWK support
    if (bereq.http.is-jwk) {
        set beresp.ttl = 100ms;
        set beresp.grace = 0s;
        edgestash.index_json();
        return (deliver);
    }
}

sub vcl_deliver
{
    # JWK support
    if (req.http.is-jwk) {
        if (obj.hits == 0 && edgestash.is_json()) {
            json.parse_resp_index();
            jwt_opts.set("jwk_key", json.get("."));
        }
        return (deliver);
    }
}