Triple JWT audience/issuer server

auth.go

package main

import (
	"context"
	"fmt"
	"log"
	"net/http"
	"os"

	"github.com/golang-jwt/jwt/v4"
	"github.com/golang-jwt/jwt/v4/request"
)

type PeerType string

const (
	PeerTypeUser    = PeerType("user")    // HTTPS from Voice/Lambda/Action
	PeerTypeDevice  = PeerType("device")  // IoT Device = Bridge
	PeerTypeCluster = PeerType("cluster") // server peers to proxy requests / admin interface
)

type ContextKey int

const (
	JwtTokenKey ContextKey = 0
)

type Auth struct {
	DeviceKeyFn jwt.Keyfunc
	ClientKeyFn jwt.Keyfunc
	PeerKeyFn   jwt.Keyfunc
}

type KeyRoutingfunc = func(*jwt.Token) (key interface{}, handler http.Handler, err error)

func Serve(auth Auth) {
	var h http.Handler = http.DefaultServeMux
	h = JwtMiddleware(jwtRouting(auth, h, h, h))

	err := http.ListenAndServeTLS(":"+os.Getenv("PORT"), os.Getenv("SSL_CERT"), os.Getenv("SSL_KEY"), h)
	if err != http.ErrServerClosed {
		log.Fatal(err)
	}
}

func jwtRouting(auth Auth, delegateWebsocket, delegateClient, delegatePeers http.Handler) KeyRoutingfunc {
	return func(t *jwt.Token) (key interface{}, route http.Handler, err error) {
		claims, isClaims := t.Claims.(jwt.MapClaims)
		if !isClaims {
			return nil, nil, fmt.Errorf("unknown claims")
		}
		if claims.VerifyAudience("ws.example.org", true) {
			// bridges
			key, err = auth.DeviceKeyFn(t)
			route = delegateWebsocket
			return
		}
		if claims.VerifyAudience("api.example.org", true) {
			// clients
			key, err = auth.ClientKeyFn(t)
			route = delegateClient
			return
		}
		if claims.VerifyAudience("peers", true) {
			// peers
			key, err = auth.PeerKeyFn(t)
			route = delegatePeers
			return
		}
		return nil, nil, fmt.Errorf("unknown audience %q", claims["aud"])
	}
}

// JwtMiddleware adds authorization to the context.Context
func JwtMiddleware(keyRoutingFn KeyRoutingfunc) http.Handler {
	return http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
		var delegate http.Handler
		token, err := request.ParseFromRequest(r, request.HeaderExtractor{"Authorization"}, func(t *jwt.Token) (key interface{}, err error) {
			key, delegate, err = keyRoutingFn(t)
			return
		}, request.WithClaims(jwt.MapClaims{}))
		if err != nil {
			log.Println(fmt.Errorf("err in jwt authorization: %w", err).Error())
			http.Error(rw, http.StatusText(http.StatusUnauthorized), http.StatusUnauthorized)
			return
		} else {
			r = r.WithContext(context.WithValue(r.Context(), JwtTokenKey, token))
		}

		if delegate == nil {
			http.NotFound(rw, r)
			return
		}
		delegate.ServeHTTP(rw, r)
	})
}

dev.env

PORT=443
SSL_CERT=devcerts/ssl/server.crt
SSL_KEY=devcerts/ssl/server.key

main.go

package main

import (
	"encoding/json"
	"fmt"
	"log"
	"os"
	"time"

	"github.com/joho/godotenv"
	"gopkg.in/square/go-jose.v2"
	jose_jwt "gopkg.in/square/go-jose.v2/jwt"
)

func main() {
	// Load .env into environment to allow easy development
	err := loadEnv(".env", "dev.env")
	if err != nil {
		log.Fatal("Error loading .env file")
	}

	if len(os.Args) > 2 && os.Args[1] == "sign" {
		keyFile, err := os.OpenFile("./jwk.testec1.private.json", os.O_RDONLY, 0400)
		if err != nil {
			panic(err)
		}
		var key jose.JSONWebKey
		err = json.NewDecoder(keyFile).Decode(&key)
		if err != nil {
			panic(err)
		}
		signer, err := jose.NewSigner(jose.SigningKey{Algorithm: jose.SignatureAlgorithm(key.Algorithm), Key: key.Key}, (&jose.SignerOptions{}).WithType("JWT"))
		if err != nil {
			panic(err)
		}

		raw, err := jose_jwt.Signed(signer).Claims(jose_jwt.Claims{
			Subject:   "device|1",
			Issuer:    "local",
			NotBefore: jose_jwt.NewNumericDate(time.Now()),
			Audience:  jose_jwt.Audience{"ws.example.org"},
		}).CompactSerialize()
		if err != nil {
			panic(err)
		}
		fmt.Println(raw)
		return
	}

	Serve(internal.Auth{})
}

// loadEnv is the optional variant of godotenv.Load: only files that exist are loaded
func loadEnv(files ...string) error {
	var existing []string
	for _, f := range files {
		s, err := os.Stat(f)
		if err == nil && !s.IsDir() {
			existing = append(existing, f)
		}
	}
	return godotenv.Load(existing...)
}