Merge pull request #50 from MicahParks/response_extractor

Add response extractor

README.md

@@ -70,7 +70,7 @@ jwksURL := os.Getenv("JWKS_URL")
 
 // Confirm the environment variable is not empty.
 if jwksURL == "" {
-	log.Fatalln("JWKS_URL environment variable must be populated.")
+log.Fatalln("JWKS_URL environment variable must be populated.")
 }
 ```
 
@@ -79,9 +79,9 @@ if jwksURL == "" {
 Via HTTP:
 ```go
 // Create the JWKS from the resource at the given URL.
-jwks, err := keyfunc.Get(jwksURL, keyfunc.Options{})
+jwks, err := keyfunc.Get(jwksURL, keyfunc.Options{}) // See recommended options in the examples directory.
 if err != nil {
-	log.Fatalf("Failed to get the JWKS from the given URL.\nError: %s", err)
+log.Fatalf("Failed to get the JWKS from the given URL.\nError: %s", err)
 }
 ```
 Via JSON:
@@ -92,7 +92,7 @@ var jwksJSON = json.RawMessage(`{"keys":[{"kid":"zXew0UJ1h6Q4CCcd_9wxMzvcp5cEBif
 // Create the JWKS from the resource at the given URL.
 jwks, err := keyfunc.NewJSON(jwksJSON)
 if err != nil {
-	log.Fatalf("Failed to create JWKS from JSON.\nError: %s", err)
+log.Fatalf("Failed to create JWKS from JSON.\nError: %s", err)
 }
 ```
 Via a given key:
@@ -103,7 +103,7 @@ uniqueKeyID := "myKeyID"
 
 // Create the JWKS from the HMAC key.
 jwks := keyfunc.NewGiven(map[string]keyfunc.GivenKey{
-	uniqueKeyID: keyfunc.NewGivenHMAC(key),
+uniqueKeyID: keyfunc.NewGivenHMAC(key),
 })
 ```
 
@@ -117,7 +117,7 @@ features mentioned at the bottom of this `README.md`.
 // Parse the JWT.
 token, err := jwt.Parse(jwtB64, jwks.Keyfunc)
 if err != nil {
-	return nil, fmt.Errorf("failed to parse token: %w", err)
+return nil, fmt.Errorf("failed to parse token: %w", err)
 }
 ```
 
@@ -151,6 +151,10 @@ These features can be configured by populating fields in the
 * A custom HTTP client can be used.
 * A custom HTTP request factory can be provided to create HTTP requests for the remote JWKS resource. For example, an
   HTTP header can be added to indicate a User-Agent.
+* A custom HTTP response extractor can be provided to get the raw JWKS JSON from the `*http.Response`. For example, the
+  HTTP response code could be checked. Implementations are responsible for closing the response body. By default, the
+  response body is read and closed, the status code is ignored. The default behavior is likely to be changed soon.
+  See https://github.com/MicahParks/keyfunc/issues/48.
 * A map of JWT key IDs (`kid`) to keys can be given and used for the `jwt.Keyfunc`. For an example, see
   the `examples/given` directory.
 * A copy of the latest raw JWKS `[]byte` can be returned.

examples/aws_cognito/main.go

@@ -30,6 +30,7 @@ func main() {
 		RefreshRateLimit:  time.Minute * 5,
 		RefreshTimeout:    time.Second * 10,
 		RefreshUnknownKID: true,
+		ResponseExtractor: keyfunc.ResponseExtractorStatusOK,
 	}
 
 	// Create the JWKS from the resource at the given URL.

examples/ctx/main.go

@@ -22,6 +22,7 @@ func main() {
 		RefreshErrorHandler: func(err error) {
 			log.Printf("There was an error with the jwt.Keyfunc\nError: %s", err.Error())
 		},
+		ResponseExtractor: keyfunc.ResponseExtractorStatusOK,
 	}
 
 	// Create the JWKS from the resource at the given URL.

examples/given/main.go

@@ -45,6 +45,7 @@ func main() {
 		RefreshRateLimit:  time.Minute * 5,
 		RefreshTimeout:    time.Second * 10,
 		RefreshUnknownKID: true,
+		ResponseExtractor: keyfunc.ResponseExtractorStatusOK,
 	}
 
 	// Create the JWKS from the resource at the given URL.

examples/interval/main.go

@@ -23,6 +23,7 @@ func main() {
 		RefreshErrorHandler: func(err error) {
 			log.Printf("There was an error with the jwt.Keyfunc\nError: %s", err.Error())
 		},
+		ResponseExtractor: keyfunc.ResponseExtractorStatusOK,
 	}
 
 	// Create the JWKS from the resource at the given URL.

examples/keycloak/main.go

@@ -26,6 +26,7 @@ func main() {
 		RefreshRateLimit:  time.Minute * 5,
 		RefreshTimeout:    time.Second * 10,
 		RefreshUnknownKID: true,
+		ResponseExtractor: keyfunc.ResponseExtractorStatusOK,
 	}
 
 	// Create the JWKS from the resource at the given URL.

examples/recommended_options/main.go

@@ -31,6 +31,7 @@ func main() {
 		RefreshRateLimit:  time.Minute * 5,
 		RefreshTimeout:    time.Second * 10,
 		RefreshUnknownKID: true,
+		ResponseExtractor: keyfunc.ResponseExtractorStatusOK,
 	}
 
 	// Create the JWKS from the resource at the given URL.

get.go

@@ -3,6 +3,8 @@ package keyfunc
 import (
 	"bytes"
 	"context"
+	"encoding/json"
+	"fmt"
 	"io"
 	"net/http"
 	"sync"
@@ -29,6 +31,16 @@ func Get(jwksURL string, options Options) (jwks *JWKS, err error) {
 	if jwks.requestFactory == nil {
 		jwks.requestFactory = defaultRequestFactory
 	}
+	if jwks.responseExtractor == nil {
+		jwks.responseExtractor = func(ctx context.Context, resp *http.Response) (json.RawMessage, error) {
+			// This behavior is likely going to change in favor of checking the response code.
+			// See https://github.com/MicahParks/keyfunc/issues/48
+
+			//goland:noinspection GoUnhandledErrorResult
+			defer resp.Body.Close()
+			return io.ReadAll(resp.Body)
+		}
+	}
 	if jwks.refreshTimeout == 0 {
 		jwks.refreshTimeout = defaultRefreshTimeout
 	}
@@ -141,19 +153,17 @@ func (j *JWKS) refresh() (err error) {
 
 	req, err := j.requestFactory(ctx, j.jwksURL)
 	if err != nil {
-		return err
+		return fmt.Errorf("failed to create request via factory function: %w", err)
 	}
 
 	resp, err := j.client.Do(req)
 	if err != nil {
 		return err
 	}
-	//goland:noinspection GoUnhandledErrorResult
-	defer resp.Body.Close()
 
-	jwksBytes, err := io.ReadAll(resp.Body)
+	jwksBytes, err := j.responseExtractor(ctx, resp)
 	if err != nil {
-		return err
+		return fmt.Errorf("failed to extract response via extractor function: %w", err)
 	}
 
 	// Only reprocess if the JWKS has changed.

jwks.go

@@ -50,6 +50,7 @@ type JWKS struct {
 	refreshTimeout      time.Duration
 	refreshUnknownKID   bool
 	requestFactory      func(ctx context.Context, url string) (*http.Request, error)
+	responseExtractor   func(ctx context.Context, resp *http.Response) (json.RawMessage, error)
 }
 
 // rawJWKS represents a JWKS in JSON format.

jwks_test.go

@@ -358,7 +358,7 @@ func TestRawJWKS(t *testing.T) {
 	}
 
 	raw := jwks.RawJWKS()
-	if bytes.Compare(raw, []byte(jwksJSON)) != 0 {
+	if !bytes.Equal(raw, []byte(jwksJSON)) {
 		t.Fatalf("Raw JWKS does not match remote JWKS resource.")
 	}
 
@@ -367,7 +367,7 @@ func TestRawJWKS(t *testing.T) {
 	copy(raw, emptySlice)
 
 	nextRaw := jwks.RawJWKS()
-	if bytes.Compare(nextRaw, emptySlice) == 0 {
+	if bytes.Equal(nextRaw, emptySlice) {
 		t.Fatalf("Raw JWKS is not a copy.")
 	}
 }

options.go

@@ -2,10 +2,17 @@ package keyfunc
 
 import (
 	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"io"
 	"net/http"
 	"time"
 )
 
+// ErrInvalidHTTPStatusCode indicates that the HTTP status code is invalid.
+var ErrInvalidHTTPStatusCode = errors.New("invalid HTTP status code")
+
 // Options represents the configuration options for a JWKS.
 //
 // If RefreshInterval and or RefreshUnknownKID is not nil, then a background goroutine will be launched to refresh the
@@ -58,6 +65,22 @@ type Options struct {
 	// RequestFactory creates HTTP requests for the remote JWKS resource located at the given url. For example, an
 	// HTTP header could be added to indicate a User-Agent.
 	RequestFactory func(ctx context.Context, url string) (*http.Request, error)
+
+	// ResponseExtractor consumes a *http.Response and produces the raw JSON for the JWKS. By default, the raw JSON is
+	// expected in the response body and the response's status code is not checked.
+	//
+	// The default behavior is likely to change soon. See this relevant GitHub issue:
+	// https://github.com/MicahParks/keyfunc/issues/48
+	ResponseExtractor func(ctx context.Context, resp *http.Response) (json.RawMessage, error)
+}
+
+// ResponseExtractorStatusOK is meant to be used as the ResponseExtractor field for Options. It confirms that response
+// status code is 200 OK and returns the raw JSON from the response body.
+func ResponseExtractorStatusOK(ctx context.Context, resp *http.Response) (json.RawMessage, error) {
+	if resp.StatusCode != http.StatusOK {
+		return nil, fmt.Errorf("%w: %d", ErrInvalidHTTPStatusCode, resp.StatusCode)
+	}
+	return io.ReadAll(resp.Body)
 }
 
 // applyOptions applies the given options to the given JWKS.
@@ -81,4 +104,5 @@ func applyOptions(jwks *JWKS, options Options) {
 	jwks.refreshTimeout = options.RefreshTimeout
 	jwks.refreshUnknownKID = options.RefreshUnknownKID
 	jwks.requestFactory = options.RequestFactory
+	jwks.responseExtractor = options.ResponseExtractor
 }

options_test.go

@@ -0,0 +1,45 @@
+package keyfunc_test
+
+import (
+	"errors"
+	"net/http"
+	"net/http/httptest"
+	"sync"
+	"testing"
+
+	"github.com/MicahParks/keyfunc"
+)
+
+func TestResponseExtractorStatusOK(t *testing.T) {
+	var mux sync.Mutex
+	statusCode := http.StatusOK
+
+	server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
+		mux.Lock()
+		writer.WriteHeader(statusCode)
+		mux.Unlock()
+		_, _ = writer.Write([]byte(jwksJSON))
+	}))
+	defer server.Close()
+
+	options := keyfunc.Options{
+		ResponseExtractor: keyfunc.ResponseExtractorStatusOK,
+	}
+	jwks, err := keyfunc.Get(server.URL, options)
+	if err != nil {
+		t.Fatalf("Failed to get JWK Set from server.\nError: %s", err)
+	}
+
+	if len(jwks.ReadOnlyKeys()) == 0 {
+		t.Fatalf("Expected JWK Set to have keys.")
+	}
+
+	mux.Lock()
+	statusCode = http.StatusInternalServerError
+	mux.Unlock()
+
+	_, err = keyfunc.Get(server.URL, options)
+	if !errors.Is(err, keyfunc.ErrInvalidHTTPStatusCode) {
+		t.Fatalf("Expected error to be ErrInvalidHTTPStatusCode.\nError: %s", err)
+	}
+}