From 9c05d2f38ab726d30191a0bafd8afd489ba5a9f2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Torjus=20H=C3=A5kestad?= <torjus@usit.uio.no>
Date: Fri, 27 Oct 2023 21:46:54 +0200
Subject: [PATCH] Improve authmw

---
 authmw/token.go      | 18 +++++++++++++++++-
 authmw/token_test.go | 31 +++++++++++++++++++++++++++++++
 go.mod               |  2 +-
 go.sum               |  4 ++--
 4 files changed, 51 insertions(+), 4 deletions(-)
 create mode 100644 authmw/token_test.go

diff --git a/authmw/token.go b/authmw/token.go
index 304ee56..add588a 100644
--- a/authmw/token.go
+++ b/authmw/token.go
@@ -16,6 +16,12 @@ import (
 	"go.opentelemetry.io/otel"
 )
 
+type ctxType string
+
+var ctxKeyClaims ctxType = "claims"
+
+var ErrNoClaimsInRequest = fmt.Errorf("no claims in request")
+
 func VerifyToken(authURL string, permittedRoles []string) func(http.Handler) http.Handler {
 	fn := func(next http.Handler) http.Handler {
 		fn := func(w http.ResponseWriter, r *http.Request) {
@@ -128,7 +134,7 @@ func VerifyToken(authURL string, permittedRoles []string) func(http.Handler) htt
 
 			// Add claims to request context
 			if claims, ok := token.Claims.(*auth.MicrofilmClaims); ok && token.Valid {
-				ctx := context.WithValue(r.Context(), "claims", claims)
+				ctx := context.WithValue(r.Context(), ctxKeyClaims, claims)
 				next.ServeHTTP(w, r.WithContext(ctx))
 				return
 			}
@@ -140,3 +146,13 @@ func VerifyToken(authURL string, permittedRoles []string) func(http.Handler) htt
 
 	return fn
 }
+
+func ClaimsFromCtx(ctx context.Context) (*auth.MicrofilmClaims, error) {
+	rawValue := ctx.Value(ctxKeyClaims)
+	value, ok := rawValue.(*auth.MicrofilmClaims)
+	if ok {
+		return value, nil
+	}
+
+	return nil, ErrNoClaimsInRequest
+}
diff --git a/authmw/token_test.go b/authmw/token_test.go
new file mode 100644
index 0000000..a0dd8b2
--- /dev/null
+++ b/authmw/token_test.go
@@ -0,0 +1,31 @@
+package authmw
+
+import (
+	"context"
+	"testing"
+
+	"git.t-juice.club/microfilm/auth"
+	"github.com/golang-jwt/jwt/v5"
+	"github.com/google/go-cmp/cmp"
+)
+
+func TestClaimsFromContext(t *testing.T) {
+	claims := &auth.MicrofilmClaims{
+		Role: "admin",
+		RegisteredClaims: jwt.RegisteredClaims{
+			Issuer:  "test",
+			Subject: "subject",
+		},
+	}
+	ctx := context.WithValue(context.Background(), ctxKeyClaims, claims)
+
+	retrieved, err := ClaimsFromCtx(ctx)
+	if err != nil {
+		t.Fatalf("Unable to retrieve claims")
+	}
+
+	if diff := cmp.Diff(claims, retrieved); diff != "" {
+		t.Fatalf("Claims diff: %s", diff)
+	}
+	return
+}
diff --git a/go.mod b/go.mod
index 87fff6f..aeefa5b 100644
--- a/go.mod
+++ b/go.mod
@@ -6,6 +6,7 @@ require (
 	git.t-juice.club/microfilm/users v0.1.2
 	github.com/go-chi/chi/v5 v5.0.10
 	github.com/golang-jwt/jwt/v5 v5.0.0
+	github.com/google/go-cmp v0.6.0
 	github.com/google/uuid v1.3.1
 	github.com/nats-io/nats.go v1.31.0
 	github.com/nats-io/nkeys v0.4.5
@@ -27,7 +28,6 @@ require (
 	github.com/golang/protobuf v1.5.3 // indirect
 	github.com/grpc-ecosystem/grpc-gateway/v2 v2.16.0 // indirect
 	github.com/klauspost/compress v1.17.0 // indirect
-	github.com/nats-io/nkeys v0.4.5 // indirect
 	github.com/nats-io/nuid v1.0.1 // indirect
 	github.com/russross/blackfriday/v2 v2.1.0 // indirect
 	github.com/xrash/smetrics v0.0.0-20201216005158-039620a65673 // indirect
diff --git a/go.sum b/go.sum
index 933cb4e..5223f9a 100644
--- a/go.sum
+++ b/go.sum
@@ -24,8 +24,8 @@ github.com/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaS
 github.com/golang/protobuf v1.5.3 h1:KhyjKVUg7Usr/dYsdSqoFveMYd5ko72D+zANwlG1mmg=
 github.com/golang/protobuf v1.5.3/go.mod h1:XVQd3VNwM+JqD3oG2Ue2ip4fOMUkwXdXDdiuN0vRsmY=
 github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
-github.com/google/go-cmp v0.5.9 h1:O2Tfq5qg4qc4AmwVlvv0oLiVAGB7enBSJ2x2DqQFi38=
-github.com/google/go-cmp v0.5.9/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
+github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI=
+github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
 github.com/google/uuid v1.3.1 h1:KjJaJ9iWZ3jOFZIf1Lqf4laDRCasjl0BCmnEGxkdLb4=
 github.com/google/uuid v1.3.1/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/grpc-ecosystem/grpc-gateway/v2 v2.16.0 h1:YBftPWNWd4WwGqtY2yeZL2ef8rHAxPBD8KFhJpmcqms=