From a412ee55ae812be1e97c61c506309d13c1f5bda0 Mon Sep 17 00:00:00 2001
From: "Gerasimos (Makis) Maropoulos" <kataras2006@hotmail.com>
Date: Sat, 17 Oct 2020 15:22:42 +0300
Subject: [PATCH] jwt: add VerifyJSON and ReadJSON helpers

---
 middleware/jwt/jwt.go      | 57 ++++++++++++++++++++++++++++++++++++++
 middleware/jwt/jwt_test.go | 36 ++++++++++++++++++++++--
 2 files changed, 90 insertions(+), 3 deletions(-)

diff --git a/middleware/jwt/jwt.go b/middleware/jwt/jwt.go
index 37df4622..6cc8c167 100644
--- a/middleware/jwt/jwt.go
+++ b/middleware/jwt/jwt.go
@@ -2,6 +2,7 @@ package jwt
 
 import (
 	"crypto"
+	"encoding/json"
 	"os"
 	"strings"
 	"time"
@@ -449,6 +450,17 @@ func Get(ctx *context.Context) interface{} {
 		switch v := tok.Value.(type) {
 		case *context.Map:
 			return *v
+		case *json.RawMessage:
+			// This is useful when we can accept more than one
+			// type of JWT token in the same request path,
+			// but we also want to keep type safety.
+			// Usage:
+			// type myClaims struct { Roles []string `json:"roles"`}
+			// v := jwt.Get(ctx)
+			// var claims myClaims
+			// jwt.Unmarshal(v, &claims)
+			// [...claims.Roles]
+			return *v
 		default:
 			return v
 		}
@@ -621,7 +633,19 @@ func (j *JWT) TokenPair(refreshMaxAge time.Duration, refreshClaims interface{},
 // - The Context Logout method is set if Blocklist was initialized
 // Any error is captured to the Context,
 // which can be retrieved by a `ctx.GetErr()` call.
+//
+// See `VerifyJSON` too.
 func (j *JWT) Verify(newPtr func() interface{}, expections ...Expectation) context.Handler {
+	if newPtr == nil {
+		newPtr = func() interface{} {
+			// Return a map here as the default type one,
+			// as it does allow .Get callers to access its fields with ease
+			// (although, I always recommend using structs for type-safety and
+			// also they can accept a required tag option too).
+			return &context.Map{}
+		}
+	}
+
 	expections = append(expections, MeetRequirements(newPtr()))
 
 	return func(ctx *context.Context) {
@@ -647,6 +671,39 @@ func (j *JWT) Verify(newPtr func() interface{}, expections ...Expectation) conte
 	}
 }
 
+// VerifyJSON works like `Verify` but instead it
+// binds its "newPtr" function to return a raw JSON message.
+// This allows the caller to bind this JSON message to any Go structure (or map).
+// This is useful when we can accept more than one
+// type of JWT token in the same request path,
+// but we also want to keep type safety.
+// Usage:
+// app.Use(jwt.VerifyJSON())
+// Inside a route Handler:
+// claims := struct { Roles []string `json:"roles"`}{}
+// jwt.ReadJSON(ctx, &claims)
+// ...access to claims.Roles as []string
+func (j *JWT) VerifyJSON(expections ...Expectation) context.Handler {
+	return j.Verify(func() interface{} {
+		return new(json.RawMessage)
+	})
+}
+
+// ReadJSON is a helper which binds "claimsPtr" to the
+// raw JSON token claims.
+// Use inside the handlers when `VerifyJSON()` middleware was registered.
+func ReadJSON(ctx *context.Context, claimsPtr interface{}) error {
+	v := Get(ctx)
+	if v == nil {
+		return ErrMissing
+	}
+	data, ok := v.(json.RawMessage)
+	if !ok {
+		return ErrMissing
+	}
+	return Unmarshal(data, claimsPtr)
+}
+
 // NewUser returns a new User based on the given "opts".
 // The caller can modify the User until its `GetToken` is called.
 func (j *JWT) NewUser(opts ...UserOption) *User {
diff --git a/middleware/jwt/jwt_test.go b/middleware/jwt/jwt_test.go
index f42e80c1..b3e8c422 100644
--- a/middleware/jwt/jwt_test.go
+++ b/middleware/jwt/jwt_test.go
@@ -192,9 +192,7 @@ func TestVerifyMap(t *testing.T) {
 	})
 
 	// Test map + Verify middleware.
-	userAPI.Post("/middleware", j.Verify(func() interface{} {
-		return &iris.Map{} // or &map[string]interface{}{}
-	}), func(ctx iris.Context) {
+	userAPI.Post("/middleware", j.Verify(nil), func(ctx iris.Context) {
 		claims := jwt.Get(ctx)
 		ctx.JSON(claims)
 	})
@@ -265,6 +263,38 @@ func TestVerifyStruct(t *testing.T) {
 	e.POST("/user").WithHeader("Authorization", "Bearer "+token).Expect().Status(httptest.StatusUnauthorized)
 }
 
+func TestVerifyJSON(t *testing.T) {
+	j := jwt.HMAC(testMaxAge, "secret", "itsa16bytesecret")
+
+	app := iris.New()
+	app.Get("/user/auth", func(ctx iris.Context) {
+		err := j.WriteToken(ctx, iris.Map{"roles": []string{"admin"}})
+		if err != nil {
+			ctx.StopWithError(iris.StatusUnauthorized, err)
+			return
+		}
+	})
+
+	app.Post("/", j.VerifyJSON(), func(ctx iris.Context) {
+		claims := struct {
+			Roles []string `json:"roles"`
+		}{}
+		jwt.ReadJSON(ctx, &claims)
+		ctx.JSON(claims)
+	})
+
+	e := httptest.New(t, app, httptest.LogLevel("error"))
+	token := e.GET("/user/auth").Expect().Status(httptest.StatusOK).Body().Raw()
+	if token == "" {
+		t.Fatalf("empty token")
+	}
+
+	e.POST("/").WithHeader("Authorization", "Bearer "+token).Expect().
+		Status(httptest.StatusOK).JSON().Equal(iris.Map{"roles": []string{"admin"}})
+
+	e.POST("/").Expect().Status(httptest.StatusUnauthorized)
+}
+
 func TestVerifyUserAndExpected(t *testing.T) { // Tests the jwt.User struct + context validator + expected.
 	maxAge := testMaxAge / 2
 	j := jwt.HMAC(maxAge, "secret", "itsa16bytesecret")