Add ability to override where tokens are read from. (#221)

* Add ability to override where tokens are read from.
* Allows for testing code that reads the token.
* Fix lint errors.

---------

Co-authored-by: Kenny Parnell <[email protected]>

Author: kennyp

dynoid/dynoid.go

@@ -5,6 +5,7 @@ import (
 	"encoding/base64"
 	"encoding/json"
 	"fmt"
+	"io/fs"
 	"log/slog"
 	"os"
 	"path"
@@ -90,7 +91,7 @@ func (s *Subject) MarshalText() ([]byte, error) {
 
 func (s *Subject) UnmarshalText(text []byte) error {
 	if s == nil {
-		*s = Subject{}
+		return fmt.Errorf("cannot unmarshal to a nil pointer")
 	}
 
 	sub := string(text)
@@ -142,11 +143,28 @@ func LocalTokenPath(audience string) string {
 	return fmt.Sprintf("/etc/heroku/dyno-id/%s/token", audience)
 }
 
+type osReader struct{}
+
+func (*osReader) Open(name string) (fs.File, error) {
+	return os.Open(name)
+}
+
+func (*osReader) ReadFile(name string) ([]byte, error) {
+	return os.ReadFile(name)
+}
+
+// DefaultFS is used by [ReadLocal] and [ReadLocalToken] to retrieve tokens.
+//
+// By default they are retrieved via [os.Open] and [os.ReadFile].
+//
+// This is useful when testing code that uses DynoID.
+var DefaultFS fs.ReadFileFS = &osReader{}
+
 // ReadLocal reads the local machines token for the given audience
 //
 // Suitable for passing as a bearer token
 func ReadLocal(audience string) (string, error) {
-	rawToken, err := os.ReadFile(LocalTokenPath(audience))
+	rawToken, err := DefaultFS.ReadFile(LocalTokenPath(audience))
 	if err != nil {
 		return "", err
 	}

dynoid/dynoid_test.go

@@ -1,14 +1,90 @@
-package dynoid
+package dynoid_test
 
 import (
 	"context"
 	"testing"
 
+	"github.com/google/uuid"
+
+	"github.com/heroku/x/dynoid"
 	"github.com/heroku/x/dynoid/dynoidtest"
 )
 
 func TestVerification(t *testing.T) {
-	ctx, iss, err := dynoidtest.NewWithContext(context.Background())
+	ctx, token := GenerateIDToken(t, "heroku")
+
+	verifier := dynoid.NewWithCallback("heroku", dynoid.AllowHerokuHost(dynoidtest.IssuerHost))
+
+	if _, err := verifier.Verify(ctx, token); err != nil {
+		t.Error(err)
+	}
+}
+
+func TestMarshalUnmarshal(t *testing.T) {
+	in := dynoid.Subject{
+		AppID:   "7eeecc9f-b17f-4027-9aa1-ceb8427036c6",
+		AppName: "testing",
+		Dyno:    "web.1",
+	}
+
+	var out dynoid.Subject
+
+	if err := out.UnmarshalText([]byte(in.String())); err != nil {
+		t.Fatalf("failed to unmarshal (%v)", err)
+	}
+
+	if out.AppID != in.AppID {
+		t.Fatalf("AppID missmatch (%q != %q)", out.AppID, in.AppID)
+	}
+
+	if out.AppName != in.AppName {
+		t.Fatalf("AppName missmatch (%q != %q)", out.AppName, in.AppName)
+	}
+
+	if out.Dyno != in.Dyno {
+		t.Fatalf("Dyno missmatch (%q != %q)", out.Dyno, in.Dyno)
+	}
+}
+
+func TestReading(t *testing.T) {
+	oldFS := dynoid.DefaultFS
+	defer func() {
+		dynoid.DefaultFS = oldFS
+	}()
+
+	spaceID := uuid.NewString()
+	appID := uuid.NewString()
+
+	ctx, tk := GenerateIDToken(t, "heroku",
+		dynoidtest.WithSpaceID(spaceID),
+		dynoidtest.WithTokenOpts(dynoidtest.WithSubject(&dynoid.Subject{
+			AppID:   appID,
+			AppName: "testapp",
+			Dyno:    "run.1",
+		})),
+	)
+	dynoid.DefaultFS = dynoidtest.NewFS(map[string]string{
+		"heroku": tk,
+	})
+
+	token, err := dynoid.ReadLocalToken(ctx, "heroku")
+	if err != nil {
+		t.Fatalf("failed to read token (%v)", err)
+	}
+
+	if token.SpaceID != spaceID {
+		t.Fatalf("Unexpected SpaceID %q", token.SpaceID)
+	}
+
+	if token.Subject.AppID != appID {
+		t.Fatalf("Unexpected AppID %q", token.Subject.AppID)
+	}
+}
+
+func GenerateIDToken(t *testing.T, audience string, opts ...dynoidtest.IssuerOpt) (context.Context, string) {
+	t.Helper()
+
+	ctx, iss, err := dynoidtest.NewWithContext(context.Background(), opts...)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -18,9 +94,5 @@ func TestVerification(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	verifier := NewWithCallback("heroku", AllowHerokuHost(dynoidtest.IssuerHost))
-
-	if _, err = verifier.Verify(ctx, token); err != nil {
-		t.Error(err)
-	}
+	return ctx, token
 }

dynoid/dynoidtest/dynoidtest.go

@@ -1,10 +1,12 @@
+// dynoidtest provides helper functions for testing code that uses DynoID
 package dynoidtest
 
 import (
 	"context"
 	"crypto/rand"
 	"crypto/rsa"
 	"encoding/json"
+	"fmt"
 	"net/http"
 	"net/http/httptest"
 	"strings"
@@ -14,43 +16,115 @@ import (
 	"github.com/coreos/go-oidc/v3/oidc"
 	"github.com/golang-jwt/jwt/v4"
 	jose "gopkg.in/square/go-jose.v2"
+
+	"github.com/heroku/x/dynoid"
 )
 
 const (
-	Audience   = "heroku"
+	// IssuerHost is the host used by the dynoidtest.Issuer
 	IssuerHost = "heroku.local"
+
+	DefaultSpaceID = "test"                                 // space id used when one is not provided
+	DefaultAppID   = "00000000-0000-0000-0000-000000000001" // app id used when one is not provided
+	DefaultAppName = "sushi"                                // app name used when one is not provided
+	DefaultDyno    = "web.1"                                // dyno used when one is not provided
 )
 
+// Issuer generates test tokens and provides a client for verifying them.
 type Issuer struct {
-	key *rsa.PrivateKey
+	key       *rsa.PrivateKey
+	spaceID   string
+	tokenOpts []TokenOpt
+}
+
+// IssuerOpt allows the behavior of the issuer to be modified.
+type IssuerOpt interface {
+	apply(*Issuer) error
+}
+
+type issuerOptFunc func(*Issuer) error
+
+func (f issuerOptFunc) apply(i *Issuer) error {
+	return f(i)
 }
 
-func New() (*Issuer, error) {
-	_, iss, err := NewWithContext(context.Background())
+// WithSpaceID allows a spaceID to be supplied instead of using the default
+func WithSpaceID(spaceID string) IssuerOpt {
+	return issuerOptFunc(func(i *Issuer) error {
+		i.spaceID = spaceID
+		return nil
+	})
+}
+
+// WithTokenOpts allows a default set of TokenOpt to be applied to every token
+// generated by the issuer
+func WithTokenOpts(opts ...TokenOpt) IssuerOpt {
+	return issuerOptFunc(func(i *Issuer) error {
+		i.tokenOpts = append(i.tokenOpts, opts...)
+		return nil
+	})
+}
+
+// Create a new Issuer with the supplied opts applied
+func New(opts ...IssuerOpt) (*Issuer, error) {
+	_, iss, err := NewWithContext(context.Background(), opts...)
 	return iss, err
 }
 
-func NewWithContext(ctx context.Context) (context.Context, *Issuer, error) {
+// Create a new Issuer with the supplied opts applied inheriting from the provided context
+func NewWithContext(ctx context.Context, opts ...IssuerOpt) (context.Context, *Issuer, error) {
 	key, err := rsa.GenerateKey(rand.Reader, 2048)
 	if err != nil {
 		return ctx, nil, err
 	}
 
-	iss := &Issuer{key: key}
+	iss := &Issuer{key: key, spaceID: DefaultSpaceID, tokenOpts: []TokenOpt{}}
+	for _, o := range opts {
+		if err := o.apply(iss); err != nil {
+			return ctx, nil, err
+		}
+	}
+
 	ctx = oidc.ClientContext(ctx, iss.HTTPClient())
 
 	return ctx, iss, nil
 }
 
-func (iss *Issuer) GenerateIDToken(clientID string) (string, error) {
+// A TokenOpt modifies the way a token is minted
+type TokenOpt interface {
+	apply(*jwt.RegisteredClaims) error
+}
+
+type tokenOptFunc func(*jwt.RegisteredClaims) error
+
+func (f tokenOptFunc) apply(i *jwt.RegisteredClaims) error {
+	return f(i)
+}
+
+// WithSubject allows the Subject to be different than the default
+func WithSubject(s *dynoid.Subject) TokenOpt {
+	return tokenOptFunc(func(c *jwt.RegisteredClaims) error {
+		c.Subject = s.String()
+		return nil
+	})
+}
+
+// GenerateIDToken returns a new signed token as a string
+func (iss *Issuer) GenerateIDToken(clientID string, opts ...TokenOpt) (string, error) {
 	now := time.Now()
 
 	claims := &jwt.RegisteredClaims{
 		Audience:  jwt.ClaimStrings([]string{clientID}),
 		ExpiresAt: jwt.NewNumericDate(now.Add(5 * time.Minute)),
 		IssuedAt:  jwt.NewNumericDate(now),
-		Issuer:    "https://oidc.heroku.local/issuers/test",
-		Subject:   "app:00000000-0000-0000-0000-000000000001.sushi::dyno:web.1",
+		Issuer:    fmt.Sprintf("https://oidc.heroku.local/issuers/%s", iss.spaceID),
+		Subject:   (&dynoid.Subject{AppID: DefaultAppID, AppName: DefaultAppName, Dyno: DefaultDyno}).String(),
+	}
+
+	for _, o := range append(iss.tokenOpts, opts...) {
+		if err := o.apply(claims); err != nil {
+			return "", err
+		}
 	}
 
 	token := jwt.NewWithClaims(jwt.SigningMethodRS256, claims)
@@ -59,6 +133,7 @@ func (iss *Issuer) GenerateIDToken(clientID string) (string, error) {
 	return token.SignedString(iss.key)
 }
 
+// HTTPClient returns a client that leverages the Issuer to validate tokens.
 func (iss *Issuer) HTTPClient() *http.Client {
 	return &http.Client{Transport: &roundTripper{issuer: iss}}
 }
@@ -72,7 +147,8 @@ type roundTripper struct {
 func (rt *roundTripper) init() {
 	mux := http.NewServeMux()
 
-	mux.HandleFunc("/issuers/test/.well-known/openid-configuration", func(w http.ResponseWriter, r *http.Request) {
+	basePath := fmt.Sprintf("/issuers/%s/.well-known", rt.issuer.spaceID)
+	mux.HandleFunc(basePath+"/openid-configuration", func(w http.ResponseWriter, r *http.Request) {
 		if !strings.EqualFold(r.Method, http.MethodGet) {
 			http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
 			return
@@ -84,17 +160,17 @@ func (rt *roundTripper) init() {
 		w.WriteHeader(http.StatusOK)
 
 		_, _ = w.Write([]byte(`{` +
-			`"issuer":"https://oidc.heroku.local/issuers/test",` +
+			fmt.Sprintf(`"issuer":"https://oidc.heroku.local/issuers/%s",`, rt.issuer.spaceID) +
 			`"authorization_endpoint":"/dummy/authorization",` +
-			`"jwks_uri":"https://oidc.heroku.local/issuers/test/.well-known/jwks.json",` +
+			fmt.Sprintf(`"jwks_uri":"https://oidc.heroku.local/issuers/%s/.well-known/jwks.json",`, rt.issuer.spaceID) +
 			`"response_types_supported":["implicit"],` +
 			`"grant_types_supported":["implicit"],` +
 			`"subject_types_supported":["public"],` +
 			`"id_token_signing_alg_values_supported":["RS256"]` +
 			`}`))
 	})
 
-	mux.HandleFunc("/issuers/test/.well-known/jwks.json", func(w http.ResponseWriter, r *http.Request) {
+	mux.HandleFunc(basePath+"/jwks.json", func(w http.ResponseWriter, r *http.Request) {
 		if !strings.EqualFold(r.Method, http.MethodGet) {
 			http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
 			return