Move context functions from middleware package to dynoid package.

Consistent naming of parameters.
heroku · Oct 16, 2024 · b9c682f · b9c682f
1 parent cf37430
commit b9c682f
Show file tree

Hide file tree

Showing 4 changed files with 107 additions and 67 deletions.
diff --git a/dynoid/README.md b/dynoid/README.md
@@ -27,11 +27,13 @@ import "github.com/heroku/x/dynoid"
 ## Index
 
 - [Variables](<#variables>)
+- [func ContextWithError\(ctx context.Context, err error\) context.Context](<#ContextWithError>)
+- [func ContextWithToken\(ctx context.Context, t \*Token\) context.Context](<#ContextWithToken>)
 - [func LocalTokenPath\(audience string\) string](<#LocalTokenPath>)
 - [func ReadLocal\(audience string\) \(string, error\)](<#ReadLocal>)
 - [type IssuerCallback](<#IssuerCallback>)
-  - [func AllowHerokuHost\(host string\) IssuerCallback](<#AllowHerokuHost>)
-  - [func AllowHerokuSpace\(host string, spaceIDs ...string\) IssuerCallback](<#AllowHerokuSpace>)
+  - [func AllowHerokuHost\(issuerHost string\) IssuerCallback](<#AllowHerokuHost>)
+  - [func AllowHerokuSpace\(issuerHost string, spaceIDs ...string\) IssuerCallback](<#AllowHerokuSpace>)
 - [type MalformedTokenError](<#MalformedTokenError>)
   - [func \(e \*MalformedTokenError\) Error\(\) string](<#MalformedTokenError.Error>)
   - [func \(e \*MalformedTokenError\) Unwrap\(\) error](<#MalformedTokenError.Unwrap>)
@@ -41,13 +43,14 @@ import "github.com/heroku/x/dynoid"
   - [func \(s \*Subject\) String\(\) string](<#Subject.String>)
   - [func \(s \*Subject\) UnmarshalText\(text \[\]byte\) error](<#Subject.UnmarshalText>)
 - [type Token](<#Token>)
+  - [func FromContext\(ctx context.Context\) \(\*Token, error\)](<#FromContext>)
   - [func ReadLocalToken\(ctx context.Context, audience string\) \(\*Token, error\)](<#ReadLocalToken>)
   - [func \(t \*Token\) LogValue\(\) slog.Value](<#Token.LogValue>)
 - [type UntrustedIssuerError](<#UntrustedIssuerError>)
   - [func \(e \*UntrustedIssuerError\) Error\(\) string](<#UntrustedIssuerError.Error>)
 - [type Verifier](<#Verifier>)
-  - [func New\(clientID string\) \*Verifier](<#New>)
-  - [func NewWithCallback\(clientID string, callback IssuerCallback\) \*Verifier](<#NewWithCallback>)
+  - [func New\(audience string\) \*Verifier](<#New>)
+  - [func NewWithCallback\(audience string, callback IssuerCallback\) \*Verifier](<#NewWithCallback>)
   - [func \(v \*Verifier\) Verify\(ctx context.Context, rawIDToken string\) \(\*Token, error\)](<#Verifier.Verify>)
 
 
@@ -63,6 +66,32 @@ This is useful when testing code that uses DynoID.
 var DefaultFS fs.ReadFileFS = &osReader{}
 ```
 
+<a name="ErrTokenNotSet"></a>
+
+```go
+var (
+    ErrTokenNotSet = errors.New("token not set") // returned when neither a token nor an error is set
+)
+```
+
+<a name="ContextWithError"></a>
+## func [ContextWithError](<https://github.com/heroku/x/blob/master/dynoid/context.go#L31>)
+
+```go
+func ContextWithError(ctx context.Context, err error) context.Context
+```
+
+ContextWithError adds the given error to the Context to be retrieved later by calling FromContext
+
+<a name="ContextWithToken"></a>
+## func [ContextWithToken](<https://github.com/heroku/x/blob/master/dynoid/context.go#L25>)
+
+```go
+func ContextWithToken(ctx context.Context, t *Token) context.Context
+```
+
+ContextWithToken adds the given Token to the Context to be retrieved later by calling FromContext
+
 <a name="LocalTokenPath"></a>
 ## func [LocalTokenPath](<https://github.com/heroku/x/blob/master/dynoid/dynoid.go#L146>)
 
@@ -96,7 +125,7 @@ type IssuerCallback func(issuer string) error
 ### func [AllowHerokuHost](<https://github.com/heroku/x/blob/master/dynoid/dynoid.go#L63>)
 
 ```go
-func AllowHerokuHost(host string) IssuerCallback
+func AllowHerokuHost(issuerHost string) IssuerCallback
 ```
 
 AllowHerokuHost verifies that the issuer is from Heroku for the given host domain
@@ -105,7 +134,7 @@ AllowHerokuHost verifies that the issuer is from Heroku for the given host domai
 ### func [AllowHerokuSpace](<https://github.com/heroku/x/blob/master/dynoid/dynoid.go#L198>)
 
 ```go
-func AllowHerokuSpace(host string, spaceIDs ...string) IssuerCallback
+func AllowHerokuSpace(issuerHost string, spaceIDs ...string) IssuerCallback
 ```
 
 AllowHerokuSpace verifies that the issuer is from Heroku for the given host and space id.
@@ -201,6 +230,15 @@ type Token struct {
 }
 ```
 
+<a name="FromContext"></a>
+### func [FromContext](<https://github.com/heroku/x/blob/master/dynoid/context.go#L36>)
+
+```go
+func FromContext(ctx context.Context) (*Token, error)
+```
+
+FromContext returns the Token or error associated with the given Context
+
 <a name="ReadLocalToken"></a>
 ### func [ReadLocalToken](<https://github.com/heroku/x/blob/master/dynoid/dynoid.go#L185>)
 
@@ -255,7 +293,7 @@ type Verifier struct {
 ### func [New](<https://github.com/heroku/x/blob/master/dynoid/dynoid.go#L225>)
 
 ```go
-func New(clientID string) *Verifier
+func New(audience string) *Verifier
 ```
 
 Instantiate a new Verifier without an IssuerCallback set.
@@ -266,7 +304,7 @@ The IssuerCallback must be set before calling Verify or an error will be returne
 ### func [NewWithCallback](<https://github.com/heroku/x/blob/master/dynoid/dynoid.go#L234>)
 
 ```go
-func NewWithCallback(clientID string, callback IssuerCallback) *Verifier
+func NewWithCallback(audience string, callback IssuerCallback) *Verifier
 ```
 
 Instantiate a new Verifier with the IssuerCallback set.
@@ -582,12 +620,10 @@ import "github.com/heroku/x/dynoid/middleware"
 ## Index
 
 - [Variables](<#variables>)
-- [func AddToContext\(ctx context.Context, token \*dynoid.Token, err error\) context.Context](<#AddToContext>)
 - [func Authorize\(audience string, callback dynoid.IssuerCallback\) func\(http.Handler\) http.Handler](<#Authorize>)
 - [func AuthorizeSameSpace\(audience string\) func\(http.Handler\) http.Handler](<#AuthorizeSameSpace>)
 - [func AuthorizeSpaces\(audience string, spaces ...string\) func\(http.Handler\) http.Handler](<#AuthorizeSpaces>)
 - [func AuthorizeSpacesWithIssuer\(audience, issuer string, spaces ...string\) func\(http.Handler\) http.Handler](<#AuthorizeSpacesWithIssuer>)
-- [func FromContext\(ctx context.Context\) \(\*dynoid.Token, error\)](<#FromContext>)
 - [func Populate\(audience string, callback dynoid.IssuerCallback\) func\(http.Handler\) http.Handler](<#Populate>)
 
 
@@ -601,17 +637,8 @@ var (
 )
 ```
 
-<a name="AddToContext"></a>
-## func [AddToContext](<https://github.com/heroku/x/blob/master/dynoid/middleware/dynoid.go#L88>)
-
-```go
-func AddToContext(ctx context.Context, token *dynoid.Token, err error) context.Context
-```
-
-AddToContext adds the Token to the given context
-
 <a name="Authorize"></a>
-## func [Authorize](<https://github.com/heroku/x/blob/master/dynoid/middleware/dynoid.go#L38>)
+## func [Authorize](<https://github.com/heroku/x/blob/master/dynoid/middleware/dynoid.go#L30>)
 
 ```go
 func Authorize(audience string, callback dynoid.IssuerCallback) func(http.Handler) http.Handler
@@ -620,7 +647,7 @@ func Authorize(audience string, callback dynoid.IssuerCallback) func(http.Handle
 Authorize populates the dyno identity blocks requests where the callback fails.
 
 <a name="AuthorizeSameSpace"></a>
-## func [AuthorizeSameSpace](<https://github.com/heroku/x/blob/master/dynoid/middleware/dynoid.go#L56>)
+## func [AuthorizeSameSpace](<https://github.com/heroku/x/blob/master/dynoid/middleware/dynoid.go#L48>)
 
 ```go
 func AuthorizeSameSpace(audience string) func(http.Handler) http.Handler
@@ -629,7 +656,7 @@ func AuthorizeSameSpace(audience string) func(http.Handler) http.Handler
 AuthorizeSameSpace restricts access to tokens from the same space/issuer for the given audience.
 
 <a name="AuthorizeSpaces"></a>
-## func [AuthorizeSpaces](<https://github.com/heroku/x/blob/master/dynoid/middleware/dynoid.go#L70>)
+## func [AuthorizeSpaces](<https://github.com/heroku/x/blob/master/dynoid/middleware/dynoid.go#L62>)
 
 ```go
 func AuthorizeSpaces(audience string, spaces ...string) func(http.Handler) http.Handler
@@ -638,25 +665,16 @@ func AuthorizeSpaces(audience string, spaces ...string) func(http.Handler) http.
 AuthorizeSpaces populates the dyno identity and blocks any requests that aren't from one of the given spaces.
 
 <a name="AuthorizeSpacesWithIssuer"></a>
-## func [AuthorizeSpacesWithIssuer](<https://github.com/heroku/x/blob/master/dynoid/middleware/dynoid.go#L83>)
+## func [AuthorizeSpacesWithIssuer](<https://github.com/heroku/x/blob/master/dynoid/middleware/dynoid.go#L75>)
 
 ```go
 func AuthorizeSpacesWithIssuer(audience, issuer string, spaces ...string) func(http.Handler) http.Handler
 ```
 
 AuthorizeSpacesWithIssuer populates the dyno identity and blocks any requests that aren't from one of the given spaces and issuer.
 
-<a name="FromContext"></a>
-## func [FromContext](<https://github.com/heroku/x/blob/master/dynoid/middleware/dynoid.go#L95>)
-
-```go
-func FromContext(ctx context.Context) (*dynoid.Token, error)
-```
-
-FromContext fetches the Token from the context
-
 <a name="Populate"></a>
-## func [Populate](<https://github.com/heroku/x/blob/master/dynoid/middleware/dynoid.go#L27>)
+## func [Populate](<https://github.com/heroku/x/blob/master/dynoid/middleware/dynoid.go#L19>)
 
 ```go
 func Populate(audience string, callback dynoid.IssuerCallback) func(http.Handler) http.Handler

diff --git a/dynoid/context.go b/dynoid/context.go
@@ -0,0 +1,42 @@
+package dynoid
+
+import (
+	"context"
+	"errors"
+)
+
+var (
+	ErrTokenNotSet = errors.New("token not set") // returned when neither a token nor an error is set
+)
+
+type dynoidCtxKey uint
+
+const (
+	tokenCtxKey dynoidCtxKey = iota
+)
+
+type maybeToken struct {
+	*Token
+	error
+}
+
+// ContextWithToken adds the given Token to the Context to be retrieved later
+// by calling FromContext
+func ContextWithToken(ctx context.Context, t *Token) context.Context {
+	return context.WithValue(ctx, tokenCtxKey, &maybeToken{Token: t})
+}
+
+// ContextWithError adds the given error to the Context to be retrieved later
+// by calling FromContext
+func ContextWithError(ctx context.Context, err error) context.Context {
+	return context.WithValue(ctx, tokenCtxKey, &maybeToken{error: err})
+}
+
+// FromContext returns the Token or error associated with the given Context
+func FromContext(ctx context.Context) (*Token, error) {
+	if v, ok := ctx.Value(tokenCtxKey).(*maybeToken); ok {
+		return v.Token, v.error
+	}
+
+	return nil, ErrTokenNotSet
+}
diff --git a/dynoid/dynoid.go b/dynoid/dynoid.go
@@ -60,9 +60,9 @@ type IssuerCallback func(issuer string) error
 
 // AllowHerokuHost verifies that the issuer is from Heroku for the given host
 // domain
-func AllowHerokuHost(host string) IssuerCallback {
+func AllowHerokuHost(issuerHost string) IssuerCallback {
 	return func(issuer string) error {
-		if !strings.HasPrefix(issuer, fmt.Sprintf("https://oidc.%v/", host)) {
+		if !strings.HasPrefix(issuer, fmt.Sprintf("https://oidc.%v/", issuerHost)) {
 			return &UntrustedIssuerError{Issuer: issuer}
 		}
 
@@ -195,10 +195,10 @@ func ReadLocalToken(ctx context.Context, audience string) (*Token, error) {
 
 // AllowHerokuSpace verifies that the issuer is from Heroku for the given host
 // and space id.
-func AllowHerokuSpace(host string, spaceIDs ...string) IssuerCallback {
+func AllowHerokuSpace(issuerHost string, spaceIDs ...string) IssuerCallback {
 	return func(issuer string) error {
 		for _, id := range spaceIDs {
-			if iss := fmt.Sprintf("https://oidc.%s/spaces/%s", host, id); iss == issuer {
+			if iss := fmt.Sprintf("https://oidc.%s/spaces/%s", issuerHost, id); iss == issuer {
 				return nil
 			}
 		}
@@ -222,17 +222,17 @@ type Verifier struct {
 //
 // The IssuerCallback must be set before calling Verify or an error will be
 // returned.
-func New(clientID string) *Verifier {
+func New(audience string) *Verifier {
 	return &Verifier{
-		config:    &oidc.Config{ClientID: clientID},
+		config:    &oidc.Config{ClientID: audience},
 		mu:        &sync.RWMutex{},
 		providers: make(map[string]*oidc.Provider),
 	}
 }
 
 // Instantiate a new Verifier with the IssuerCallback set.
-func NewWithCallback(clientID string, callback IssuerCallback) *Verifier {
-	v := New(clientID)
+func NewWithCallback(audience string, callback IssuerCallback) *Verifier {
+	v := New(audience)
 	v.IssuerCallback = callback
 	return v
 }

diff --git a/dynoid/middleware/dynoid.go b/dynoid/middleware/dynoid.go
@@ -1,7 +1,6 @@
 package middleware
 
 import (
-	"context"
 	"errors"
 	"fmt"
 	"net/http"
@@ -11,13 +10,6 @@ import (
 	"github.com/heroku/x/dynoid"
 )
 
-type ctxKeyDynoID int
-
-const (
-	DynoIDKey ctxKeyDynoID = iota
-	DynoIDErrKey
-)
-
 var (
 	ErrTokenMissing = errors.New("token not found")
 )
@@ -40,7 +32,7 @@ func Authorize(audience string, callback dynoid.IssuerCallback) func(http.Handle
 
 	return func(next http.Handler) http.Handler {
 		return populate(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-			if _, err := FromContext(r.Context()); err != nil {
+			if _, err := dynoid.FromContext(r.Context()); err != nil {
 				w.WriteHeader(http.StatusForbidden)
 				fmt.Fprint(w, http.StatusText(http.StatusForbidden))
 				return
@@ -84,21 +76,6 @@ func AuthorizeSpacesWithIssuer(audience, issuer string, spaces ...string) func(h
 	return Authorize(audience, dynoid.AllowHerokuSpace(issuer, spaces...))
 }
 
-// AddToContext adds the Token to the given context
-func AddToContext(ctx context.Context, token *dynoid.Token, err error) context.Context {
-	ctx = context.WithValue(ctx, DynoIDKey, token)
-	ctx = context.WithValue(ctx, DynoIDErrKey, err)
-	return ctx
-}
-
-// FromContext fetches the Token from the context
-func FromContext(ctx context.Context) (*dynoid.Token, error) {
-	token, _ := ctx.Value(DynoIDKey).(*dynoid.Token)
-	err, _ := ctx.Value(DynoIDErrKey).(error)
-
-	return token, err
-}
-
 func populateDynoID(audience string, callback dynoid.IssuerCallback) func(*http.Request) *http.Request {
 	verifier := dynoid.NewWithCallback(audience, callback)
 
@@ -107,12 +84,15 @@ func populateDynoID(audience string, callback dynoid.IssuerCallback) func(*http.
 
 		rawToken := tokenFromHeader(r)
 		if rawToken == "" {
-			return r.WithContext(AddToContext(ctx, nil, ErrTokenMissing))
+			return r.WithContext(dynoid.ContextWithError(ctx, ErrTokenMissing))
 		}
 
 		token, err := verifier.Verify(r.Context(), rawToken)
+		if err != nil {
+			return r.WithContext(dynoid.ContextWithError(ctx, err))
+		}
 
-		return r.WithContext(AddToContext(ctx, token, err))
+		return r.WithContext(dynoid.ContextWithToken(ctx, token))
 	}
 }