From bb45bfcdecac35b6afa09673bd17a780b258a541 Mon Sep 17 00:00:00 2001
From: Dave Grijalva <grijalva@gmail.com>
Date: Mon, 6 Jun 2016 15:27:44 -0700
Subject: [PATCH] add new interface Extractor for making token extraction
 pluggable

---
 request/extractor.go    | 75 +++++++++++++++++++++++++++++++++++++++++
 request/oauth2.go       | 25 ++++++++++++++
 request/request.go      | 31 +++++------------
 request/request_test.go |  2 +-
 4 files changed, 109 insertions(+), 24 deletions(-)
 create mode 100644 request/extractor.go
 create mode 100644 request/oauth2.go

diff --git a/request/extractor.go b/request/extractor.go
new file mode 100644
index 0000000..57ea2dd
--- /dev/null
+++ b/request/extractor.go
@@ -0,0 +1,75 @@
+package request
+
+import (
+	"errors"
+	"net/http"
+)
+
+// Errors
+var (
+	ErrNoTokenInRequest = errors.New("no token present in request")
+)
+
+// Interface for extracting a token from an HTTP request
+type Extractor interface {
+	ExtractToken(*http.Request) (string, error)
+}
+
+// Extract token from headers
+type HeaderExtractor []string
+
+func (e HeaderExtractor) ExtractToken(req *http.Request) (string, error) {
+	// loop over header names and return the first one that contains data
+	for _, header := range e {
+		if ah := req.Header.Get(header); ah != "" {
+			return ah, nil
+		}
+	}
+	return "", ErrNoTokenInRequest
+}
+
+// Extract token from request args
+type ArgumentExtractor []string
+
+func (e ArgumentExtractor) ExtractToken(req *http.Request) (string, error) {
+	// Make sure form is parsed
+	req.ParseMultipartForm(10e6)
+
+	// loop over arg names and return the first one that contains data
+	for _, arg := range e {
+		if ah := req.Form.Get(arg); ah != "" {
+			return ah, nil
+		}
+	}
+
+	return "", ErrNoTokenInRequest
+}
+
+// Tries extractors in order until one works or an error occurs
+type MultiExtractor []Extractor
+
+func (e MultiExtractor) ExtractToken(req *http.Request) (string, error) {
+	// loop over header names and return the first one that contains data
+	for _, extractor := range e {
+		if tok, err := extractor.ExtractToken(req); tok != "" {
+			return tok, nil
+		} else if err != ErrNoTokenInRequest {
+			return "", err
+		}
+	}
+	return "", ErrNoTokenInRequest
+}
+
+// Wrap an Extractor in this to post-process the value before it's handed off
+type PostExtractionFilter struct {
+	Extractor
+	Filter func(string) (string, error)
+}
+
+func (e *PostExtractionFilter) ExtractToken(req *http.Request) (string, error) {
+	if tok, err := e.Extractor.ExtractToken(req); tok != "" {
+		return e.Filter(tok)
+	} else {
+		return "", err
+	}
+}
diff --git a/request/oauth2.go b/request/oauth2.go
new file mode 100644
index 0000000..bbb6ac2
--- /dev/null
+++ b/request/oauth2.go
@@ -0,0 +1,25 @@
+package request
+
+import (
+	"strings"
+)
+
+// Extract Authorization header and strip 'Bearer ' from it
+var AuthorizationHeaderExtractor = &PostExtractionFilter{
+	HeaderExtractor{"Authorization"},
+	func(tok string) (string, error) {
+		// Should be a bearer token
+		if len(tok) > 6 && strings.ToUpper(tok[0:7]) == "BEARER " {
+			return tok[7:], nil
+		}
+		return tok, nil
+	},
+}
+
+// Extractor for OAuth2 access tokens
+var OAuth2Extractor = &MultiExtractor{
+	// Look for authorization token first
+	AuthorizationHeaderExtractor,
+	// Extract access_token from form or GET argument
+	&ArgumentExtractor{"access_token"},
+}
diff --git a/request/request.go b/request/request.go
index 98e1be2..9daaa48 100644
--- a/request/request.go
+++ b/request/request.go
@@ -1,39 +1,24 @@
 package request
 
 import (
-	"errors"
 	"github.com/dgrijalva/jwt-go"
 	"net/http"
-	"strings"
-)
-
-// Errors
-var (
-	ErrNoTokenInRequest = errors.New("no token present in request")
 )
 
 // Try to find the token in an http.Request.
 // This method will call ParseMultipartForm if there's no token in the header.
 // Currently, it looks in the Authorization header as well as
 // looking for an 'access_token' request parameter in req.Form.
-func ParseFromRequest(req *http.Request, keyFunc jwt.Keyfunc) (token *jwt.Token, err error) {
-	return ParseFromRequestWithClaims(req, jwt.MapClaims{}, keyFunc)
+func ParseFromRequest(req *http.Request, extractor Extractor, keyFunc jwt.Keyfunc) (token *jwt.Token, err error) {
+	return ParseFromRequestWithClaims(req, extractor, jwt.MapClaims{}, keyFunc)
 }
 
-func ParseFromRequestWithClaims(req *http.Request, claims jwt.Claims, keyFunc jwt.Keyfunc) (token *jwt.Token, err error) {
-	// Look for an Authorization header
-	if ah := req.Header.Get("Authorization"); ah != "" {
-		// Should be a bearer token
-		if len(ah) > 6 && strings.ToUpper(ah[0:7]) == "BEARER " {
-			return jwt.ParseWithClaims(ah[7:], claims, keyFunc)
-		}
+func ParseFromRequestWithClaims(req *http.Request, extractor Extractor, claims jwt.Claims, keyFunc jwt.Keyfunc) (token *jwt.Token, err error) {
+	// Extract token from request
+	tokStr, err := extractor.ExtractToken(req)
+	if err != nil {
+		return nil, err
 	}
 
-	// Look for "access_token" parameter
-	req.ParseMultipartForm(10e6)
-	if tokStr := req.Form.Get("access_token"); tokStr != "" {
-		return jwt.ParseWithClaims(tokStr, claims, keyFunc)
-	}
-
-	return nil, ErrNoTokenInRequest
+	return jwt.ParseWithClaims(tokStr, claims, keyFunc)
 }
diff --git a/request/request_test.go b/request/request_test.go
index 6b09521..504e1bc 100644
--- a/request/request_test.go
+++ b/request/request_test.go
@@ -65,7 +65,7 @@ func TestParseRequest(t *testing.T) {
 				r.Header.Set(k, tokenString)
 			}
 		}
-		token, err := ParseFromRequestWithClaims(r, jwt.MapClaims{}, keyfunc)
+		token, err := ParseFromRequestWithClaims(r, OAuth2Extractor, jwt.MapClaims{}, keyfunc)
 
 		if token == nil {
 			t.Errorf("[%v] Token was not found: %v", data.name, err)