override_test.go

package keyfunc_test

import (
	"crypto/rand"
	"crypto/rsa"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"math/big"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"testing"

	"github.com/dgrijalva/jwt-go"

	"github.com/MicahParks/compatibility-keyfunc"
)

const (
	givenKID  = "givenKID"
	remoteKID = "remoteKID"
)

// pseudoJWKS is a data structure used to JSON marshal a JWKS but is not fully featured.
type pseudoJWKS struct {
	Keys []pseudoJSONKey `json:"keys"`
}

// pseudoJSONKey is a data structure that is used to JSON marshal a JWK that is not fully featured.
type pseudoJSONKey struct {
	KID string `json:"kid"`
	KTY string `json:"kty"`
	E   string `json:"e"`
	N   string `json:"n"`
}

// TestNewGiven tests that given keys will be added to a JWKS with a remote resource.
func TestNewGiven(t *testing.T) {
	tempDir, err := ioutil.TempDir("", "*")
	if err != nil {
		t.Errorf("Failed to create a temporary directory.\nError: %s", err.Error())
		t.FailNow()
	}
	defer func() {
		err = os.RemoveAll(tempDir)
		if err != nil {
			t.Errorf("Failed to remove temporary directory.\nError: %s", err.Error())
			t.FailNow()
		}
	}()

	jwksFile := filepath.Join(tempDir, jwksFilePath)

	givenKeys, givenPrivateKeys, jwksBytes, remotePrivateKeys, err := keysAndJWKS()
	if err != nil {
		t.Errorf("Failed to create cryptographic keys for the test.\nError: %s.", err.Error())
		t.FailNow()
	}

	err = ioutil.WriteFile(jwksFile, jwksBytes, 0600)
	if err != nil {
		t.Errorf("Failed to write JWKS file to temporary directory.\nError: %s", err.Error())
		t.FailNow()
	}

	server := httptest.NewServer(http.FileServer(http.Dir(tempDir)))
	defer server.Close()

	testingRefreshErrorHandler := func(err error) {
		panic(fmt.Sprintf("Unhandled JWKS error.\nError: %s", err.Error()))
	}

	jwksURL := server.URL + jwksFilePath

	options := keyfunc.Options{
		GivenKeys:           givenKeys,
		RefreshErrorHandler: testingRefreshErrorHandler,
	}

	jwks, err := keyfunc.Get(jwksURL, options)
	if err != nil {
		t.Errorf("Failed to get the JWKS the testing URL.\nError: %s", err.Error())
		t.FailNow()
	}

	// Test the given key with a unique key ID.
	createSignParseValidate(t, givenPrivateKeys, jwks, givenKID, true)

	// Test the given key with a non-unique key ID that should be overwritten.
	createSignParseValidate(t, givenPrivateKeys, jwks, remoteKID, false)

	// Test the remote key that should not have been overwritten.
	createSignParseValidate(t, remotePrivateKeys, jwks, remoteKID, true)

	// Change the JWKS options to overwrite remote keys.
	options.GivenKIDOverride = true
	jwks, err = keyfunc.Get(jwksURL, options)
	if err != nil {
		t.Errorf("Failed to recreate JWKS.\nError: %s.", err.Error())
		t.FailNow()
	}

	// Test the given key with a unique key ID.
	createSignParseValidate(t, givenPrivateKeys, jwks, givenKID, true)

	// Test the given key with a non-unique key ID that should overwrite the remote key.
	createSignParseValidate(t, givenPrivateKeys, jwks, remoteKID, true)

	// Test the remote key that should have been overwritten.
	createSignParseValidate(t, remotePrivateKeys, jwks, remoteKID, false)
}

// createSignParseValidate creates, signs, parses, and validates a JWT.
func createSignParseValidate(t *testing.T, keys map[string]*rsa.PrivateKey, jwks *keyfunc.JWKS, kid string, shouldValidate bool) {
	unsignedToken := jwt.New(jwt.SigningMethodRS256)
	unsignedToken.Header[kidAttribute] = kid

	jwtB64, err := unsignedToken.SignedString(keys[kid])
	if err != nil {
		t.Errorf("Failed to sign the JWT.\nError: %s.", err.Error())
		t.FailNow()
	}

	token, err := jwt.Parse(jwtB64, jwks.KeyfuncLegacy)
	if err != nil {
		rsaErr := rsa.ErrVerification
		if !shouldValidate && errors.As(err, &rsaErr) {
			return
		}
		t.Errorf("Failed to parse the JWT.\nError: %s", err.Error())
		t.FailNow()
	}

	if !shouldValidate {
		t.Errorf("The token should not have parsed properly.")
		t.FailNow()
	}

	if !token.Valid {
		t.Errorf("The JWT is not valid.")
		t.FailNow()
	}
}

// keysAndJWKS creates a couple of cryptographic keys and the remote JWKS associated with them.
func keysAndJWKS() (givenKeys map[string]keyfunc.GivenKey, givenPrivateKeys map[string]*rsa.PrivateKey, jwksBytes []byte, remotePrivateKeys map[string]*rsa.PrivateKey, err error) {
	const rsaErrMessage = "failed to create RSA key: %w"
	givenKeys = make(map[string]keyfunc.GivenKey)
	givenPrivateKeys = make(map[string]*rsa.PrivateKey)
	remotePrivateKeys = make(map[string]*rsa.PrivateKey)

	// Create a key not in the remote JWKS.
	key1, err := addRSA(givenKeys, givenKID)
	if err != nil {
		return nil, nil, nil, nil, fmt.Errorf(rsaErrMessage, err)
	}
	givenPrivateKeys[givenKID] = key1

	// Create a key to be overwritten by or override the one with the same key ID in the remote JWKS.
	key2, err := addRSA(givenKeys, remoteKID)
	if err != nil {
		return nil, nil, nil, nil, fmt.Errorf(rsaErrMessage, err)
	}
	givenPrivateKeys[remoteKID] = key2

	// Create a key that exists in the remote JWKS.
	key3, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return nil, nil, nil, nil, fmt.Errorf(rsaErrMessage, err)
	}
	remotePrivateKeys[remoteKID] = key3

	jwks := pseudoJWKS{Keys: []pseudoJSONKey{{
		KID: remoteKID,
		KTY: "RSA",
		E:   base64.RawURLEncoding.EncodeToString(big.NewInt(int64(key3.PublicKey.E)).Bytes()),
		N:   base64.RawURLEncoding.EncodeToString(key3.PublicKey.N.Bytes()),
	}}}

	jwksBytes, err = json.Marshal(jwks)
	if err != nil {
		return nil, nil, nil, nil, fmt.Errorf("failed to marshal the JWKS to JSON: %w", err)
	}

	return givenKeys, givenPrivateKeys, jwksBytes, remotePrivateKeys, nil
}