Support user defined jwt auth and sdk functions

CHANGELOG.md

@@ -6,6 +6,7 @@ _NOTE: This is the first fully open-source release, using the name "Modus" for t
 "Hypermode" still refers to the company and the commercial hosting platform - but not the framework.
 In previous releases, the name "Hypermode" was used for all three._
 
+- Support user defined jwt auth and sdk functions [#405](https://github.com/hypermodeinc/modus/pull/405)
 - Migrate from Hypermode to Modus [#412](https://github.com/hypermodeinc/modus/pull/412)
 - Import WasmExtractor code [#415](https://github.com/hypermodeinc/modus/pull/415)
 - Import Manifest code [#416](https://github.com/hypermodeinc/modus/pull/416)

cspell.json

@@ -1,7 +1,10 @@
 {
   "version": "0.2",
   "language": "en",
-  "flagWords": ["hte", "teh"],
+  "flagWords": [
+    "hte",
+    "teh"
+  ],
   "words": [
     "abstractlogger",
     "Albus",
@@ -153,6 +156,7 @@
     "typedarray",
     "uids",
     "uncategorized",
+    "Unmarshalled",
     "unmarshalling",
     "unnest",
     "upsert",

runtime/go.mod

@@ -23,6 +23,7 @@ require (
 	github.com/getsentry/sentry-go v0.29.0
 	github.com/go-viper/mapstructure/v2 v2.2.1
 	github.com/goccy/go-json v0.10.3
+	github.com/golang-jwt/jwt/v5 v5.2.0
 	github.com/google/renameio v1.0.1
 	github.com/google/uuid v1.6.0
 	github.com/jackc/pgx/v5 v5.7.1

runtime/go.sum

@@ -111,6 +111,8 @@ github.com/goccy/go-json v0.10.3/go.mod h1:oq7eo15ShAhp70Anwd5lgX2pLfOS3QCiwU/PU
 github.com/godbus/dbus/v5 v5.0.4/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA=
 github.com/gogo/protobuf v1.3.2 h1:Ov1cvc58UF3b5XjBnZv7+opcTcQFZebYjWzi34vdm4Q=
 github.com/gogo/protobuf v1.3.2/go.mod h1:P1XiOD3dCwIKUDQYPy72D8LYyHL2YPYrpS2s69NZV8Q=
+github.com/golang-jwt/jwt/v5 v5.2.0 h1:d/ix8ftRUorsN+5eMIlF4T6J8CAt9rch3My2winC1Jw=
+github.com/golang-jwt/jwt/v5 v5.2.0/go.mod h1:pqrtFR0X4osieyHYxtmOUWsAWrfe1Q5UVIyoH402zdk=
 github.com/golang/mock v1.6.0 h1:ErTB+efbowRARo13NNdxyJji2egdxLGQhRaY+DUumQc=
 github.com/golang/mock v1.6.0/go.mod h1:p6yTPP+5HYm5mzsMV8JkE6ZKdX+/wYM6Hr+LicevLPs=
 github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI=

runtime/httpserver/server.go

@@ -24,6 +24,7 @@ import (
 	"github.com/hypermodeinc/modus/runtime/graphql"
 	"github.com/hypermodeinc/modus/runtime/logger"
 	"github.com/hypermodeinc/modus/runtime/metrics"
+	"github.com/hypermodeinc/modus/runtime/middleware"
 	"github.com/hypermodeinc/modus/runtime/utils"
 
 	"github.com/rs/cors"
@@ -111,8 +112,8 @@ func startHttpServer(ctx context.Context, addresses ...string) {
 func GetHandlerMux() http.Handler {
 	mux := http.NewServeMux()
 
-	// Register our main endpoint with instrumentation.
-	mux.Handle("/graphql", metrics.InstrumentHandler(graphql.GraphQLRequestHandler, "graphql"))
+	// Register our main endpoints with instrumentation.
+	mux.Handle("/graphql", metrics.InstrumentHandler(middleware.HandleJWT(graphql.GraphQLRequestHandler), "graphql"))
 
 	// Register metrics endpoint which uses the Prometheus scraping protocol.
 	// We do not instrument it with the InstrumentHandler so that any scraper (eg. OTel)

runtime/main.go

@@ -20,6 +20,7 @@ import (
 	"github.com/hypermodeinc/modus/runtime/config"
 	"github.com/hypermodeinc/modus/runtime/httpserver"
 	"github.com/hypermodeinc/modus/runtime/logger"
+	"github.com/hypermodeinc/modus/runtime/middleware"
 	"github.com/hypermodeinc/modus/runtime/services"
 	"github.com/hypermodeinc/modus/runtime/utils"
 
@@ -63,6 +64,9 @@ func main() {
 	// Set local mode if debugging is enabled
 	local := utils.DebugModeEnabled()
 
+	// Retrieve auth private keys
+	middleware.Init()
+
 	// Start the HTTP server to listen for requests.
 	// Note, this function blocks, and handles shutdown gracefully.
 	httpserver.Start(ctx, local)

runtime/middleware/jwt.go

@@ -0,0 +1,133 @@
+/*
+ * Copyright 2024 Hypermode Inc.
+ * Licensed under the terms of the Apache License, Version 2.0
+ * See the LICENSE file that accompanied this code for further details.
+ *
+ * SPDX-FileCopyrightText: 2024 Hypermode Inc. <[email protected]>
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package middleware
+
+import (
+	"context"
+	"crypto/rsa"
+	"encoding/json"
+	"net/http"
+	"os"
+	"strings"
+
+	"github.com/hypermodeinc/modus/runtime/config"
+	"github.com/hypermodeinc/modus/runtime/logger"
+	"github.com/hypermodeinc/modus/runtime/utils"
+
+	"github.com/golang-jwt/jwt/v5"
+)
+
+type jwtClaimsKey string
+
+const jwtClaims jwtClaimsKey = "jwt_claims"
+
+var rsaPublicKeys map[string]*rsa.PublicKey
+
+func Init() {
+	privKeysStr := os.Getenv("MODUS_RSA_PEMS")
+	if privKeysStr == "" {
+		return
+	}
+	var publicKeysStr map[string]string
+	err := json.Unmarshal([]byte(privKeysStr), &publicKeysStr)
+	if err != nil {
+		logger.Error(context.Background()).Err(err).Msg("JWT private keys unmarshalling error")
+		return
+	}
+	rsaPublicKeys = make(map[string]*rsa.PublicKey)
+	for key, value := range publicKeysStr {
+		publicKey, err := jwt.ParseRSAPublicKeyFromPEM([]byte(value))
+		if err != nil {
+			logger.Error(context.Background()).Err(err).Msg("JWT public key parsing error")
+			return
+		}
+		rsaPublicKeys[key] = publicKey
+	}
+}
+
+var jwtParser = new(jwt.Parser)
+
+func HandleJWT(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		var ctx context.Context = r.Context()
+		tokenStr := r.Header.Get("Authorization")
+		if tokenStr != "" {
+			if s, found := strings.CutPrefix(tokenStr, "Bearer "); found {
+				tokenStr = s
+			} else {
+				logger.Error(ctx).Msg("Invalid JWT token format, Bearer required")
+				http.Error(w, "Invalid JWT token format, Bearer required", http.StatusUnauthorized)
+				return
+			}
+		}
+
+		if len(rsaPublicKeys) == 0 {
+			if !config.IsDevEnvironment() || tokenStr == "" {
+				next.ServeHTTP(w, r)
+				return
+			}
+			token, _, err := jwtParser.ParseUnverified(tokenStr, jwt.MapClaims{})
+			if err != nil {
+				logger.Warn(ctx).Err(err).Msg("Error parsing JWT token. Continuing since running in development")
+				next.ServeHTTP(w, r)
+				return
+			}
+			if claims, ok := token.Claims.(jwt.MapClaims); ok {
+				ctx = addClaimsToContext(ctx, claims)
+			}
+			next.ServeHTTP(w, r.WithContext(ctx))
+			return
+
+		}
+
+		var token *jwt.Token
+		var err error
+
+		for _, publicKey := range rsaPublicKeys {
+			token, err = jwtParser.Parse(tokenStr, func(token *jwt.Token) (interface{}, error) {
+				return publicKey, nil
+			})
+			if err == nil {
+				logger.Info(ctx).Msg("JWT token parsed successfully")
+				break
+			}
+		}
+		if err != nil {
+			if config.IsDevEnvironment() {
+				logger.Debug(r.Context()).Err(err).Msg("JWT parse error")
+				next.ServeHTTP(w, r)
+				return
+			}
+			logger.Error(r.Context()).Err(err).Msg("JWT parse error")
+			http.Error(w, "Invalid JWT token", http.StatusUnauthorized)
+		}
+
+		if claims, ok := token.Claims.(jwt.MapClaims); ok {
+			ctx = addClaimsToContext(ctx, claims)
+		}
+		next.ServeHTTP(w, r.WithContext(ctx))
+	})
+}
+
+func addClaimsToContext(ctx context.Context, claims jwt.MapClaims) context.Context {
+	claimsJson, err := utils.JsonSerialize(claims)
+	if err != nil {
+		logger.Error(ctx).Err(err).Msg("JWT claims serialization error")
+		return ctx
+	}
+	return context.WithValue(ctx, jwtClaims, string(claimsJson))
+}
+
+func GetJWTClaims(ctx context.Context) string {
+	if claims, ok := ctx.Value(jwtClaims).(string); ok {
+		return claims
+	}
+	return ""
+}

runtime/wasmhost/wasmhost.go

@@ -17,6 +17,7 @@ import (
 
 	"github.com/hypermodeinc/modus/runtime/functions"
 	"github.com/hypermodeinc/modus/runtime/logger"
+	"github.com/hypermodeinc/modus/runtime/middleware"
 	"github.com/hypermodeinc/modus/runtime/plugins"
 	"github.com/hypermodeinc/modus/runtime/utils"
 
@@ -115,11 +116,13 @@ func (host *wasmHost) GetModuleInstance(ctx context.Context, plugin *plugins.Plu
 	// for concurrency and performance reasons.
 	// See https://github.com/tetratelabs/wazero/pull/2275
 	// And https://gophers.slack.com/archives/C040AKTNTE0/p1719587772724619?thread_ts=1719522663.531579&cid=C040AKTNTE0
+	jwtClaims := middleware.GetJWTClaims(ctx)
 	cfg := wazero.NewModuleConfig().
 		WithName("").
 		WithSysWalltime().WithSysNanotime().
 		WithRandSource(rand.Reader).
-		WithStdout(wOut).WithStderr(wErr)
+		WithStdout(wOut).WithStderr(wErr).
+		WithEnv("JWT_CLAIMS", jwtClaims)
 
 	// Instantiate the plugin as a module.
 	// NOTE: This will also invoke the plugin's `_start` function,

sdk/assemblyscript/examples/auth/.prettierrc

@@ -0,0 +1,3 @@
+{
+  "plugins": ["assemblyscript-prettier"]
+}

sdk/assemblyscript/examples/auth/README.md

@@ -0,0 +1,3 @@
+# Modus Auth Example
+
+This is an example showing how to access auth claims from an incoming request.

sdk/assemblyscript/examples/auth/asconfig.json

@@ -0,0 +1,6 @@
+{
+  "extends": "./node_modules/@hypermode/modus-sdk-as/plugin.asconfig.json",
+  "options": {
+    "transform": ["@hypermode/modus-sdk-as/transform", "json-as/transform"]
+  }
+}

sdk/assemblyscript/examples/auth/assembly/index.ts

@@ -0,0 +1,27 @@
+/*
+ * This example is part of the Modus project, licensed under the Apache License 2.0.
+ * You may modify and use this example in accordance with the license.
+ * See the LICENSE file that accompanied this code for further details.
+ */
+
+import { auth } from "@hypermode/modus-sdk-as";
+
+// This is a simple example of a claims class that can be used to parse the JWT claims.
+@json
+export class ExampleClaims {
+  public exp!: i64;
+  public iat!: i64;
+  public iss!: string;
+  public jti!: string;
+  public nbf!: i64;
+  public sub!: string;
+
+  // This is an example of a custom claim that can be used to parse the user ID.
+  @alias("user-id")
+  public userId!: string;
+}
+
+// This function can be used to get the JWT claims, and parse them into the Claims class.
+export function getJWTClaims(): ExampleClaims {
+  return auth.getJWTClaims<ExampleClaims>();
+}

sdk/assemblyscript/examples/auth/assembly/tsconfig.json

@@ -0,0 +1,4 @@
+{
+  "extends": "assemblyscript/std/assembly.json",
+  "include": ["./**/*.ts"]
+}

sdk/assemblyscript/examples/auth/eslint.config.js

@@ -0,0 +1,11 @@
+// @ts-check
+
+import eslint from "@eslint/js";
+import tseslint from "typescript-eslint";
+import aseslint from "@hypermode/modus-sdk-as/tools/assemblyscript-eslint";
+
+export default tseslint.config(
+  eslint.configs.recommended,
+  ...tseslint.configs.recommended,
+  aseslint.config,
+);

sdk/assemblyscript/examples/auth/hypermode.json

@@ -0,0 +1,12 @@
+{
+  "$schema": "https://manifest.hypermode.com/hypermode.json",
+  "models": {
+    // No models are used by this example, but if you add any, they would go here.
+  },
+  "hosts": {
+    // No hosts are used by this example, but if you add any, they would go here.
+  },
+  "collections": {
+    // No collections are used by this example, but if you add any, they would go here.
+  }
+}