Merge pull request #72 from Workiva/cors_whitelist

Require origin whitelist for CORS middleware

Author: tylertreat-wf

rest/api.go

@@ -89,10 +89,17 @@ func NewConfiguration() *Configuration {
 	}
 }
 
+// MiddlewareError is returned by Middleware to indicate that a request should
+// not be served.
+type MiddlewareError struct {
+	Code     int
+	Response []byte
+}
+
 // Middleware can be passed in to API#Start and API#StartTLS and will be
-// invoked on every request to a route handled by the API. Returns true if the
-// request should be terminated, false if it should continue.
-type Middleware func(w http.ResponseWriter, r *http.Request) bool
+// invoked on every request to a route handled by the API. Returns a
+// MiddlewareError if the request should be terminated.
+type Middleware func(w http.ResponseWriter, r *http.Request) *MiddlewareError
 
 // middlewareProxy proxies an http.Handler by invoking middleware before
 // passing the request to the Handler. It implements the http.Handler
@@ -106,7 +113,9 @@ type middlewareProxy struct {
 // proxied http.Handler.
 func (m *middlewareProxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	for _, middleware := range m.middleware {
-		if middleware(w, r) {
+		if err := middleware(w, r); err != nil {
+			w.WriteHeader(err.Code)
+			w.Write(err.Response)
 			return
 		}
 	}

rest/api_test.go

@@ -1082,14 +1082,14 @@ func (h *httpHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 }
 
 // Ensures that middlewareProxy invokes middleware and doesn't delegate to the
-// wrapped http.Handler if any middleware return true.
+// wrapped http.Handler if any middleware return error.
 func TestMiddlewareProxyTerminate(t *testing.T) {
 	assert := assert.New(t)
 	handler := &httpHandler{}
 	called := false
-	middleware := func(w http.ResponseWriter, r *http.Request) bool {
+	middleware := func(w http.ResponseWriter, r *http.Request) *MiddlewareError {
 		called = true
-		return true
+		return &MiddlewareError{}
 	}
 	req, _ := http.NewRequest("GET", "http://example.com/foo", nil)
 	w := httptest.NewRecorder()
@@ -1102,14 +1102,14 @@ func TestMiddlewareProxyTerminate(t *testing.T) {
 }
 
 // Ensures that middlewareProxy invokes middleware and delegates to the wrapped
-// http.Handler if all middleware return false.
+// http.Handler if all middleware return nil.
 func TestMiddlewareProxyDelegate(t *testing.T) {
 	assert := assert.New(t)
 	handler := &httpHandler{}
 	called := false
-	middleware := func(w http.ResponseWriter, r *http.Request) bool {
+	middleware := func(w http.ResponseWriter, r *http.Request) *MiddlewareError {
 		called = true
-		return false
+		return nil
 	}
 	req, _ := http.NewRequest("GET", "http://example.com/foo", nil)
 	w := httptest.NewRecorder()

rest/example_middleware_test.go

@@ -78,11 +78,11 @@ func HandlerMiddleware(next http.Handler) http.Handler {
 }
 
 // Global API middleware is implemented as a function which takes an
-// http.ResponseWriter and http.Request and returns a bool indicating if the
-// request should terminate or not.
-func GlobalMiddleware(w http.ResponseWriter, r *http.Request) bool {
+// http.ResponseWriter and http.Request and returns a MiddlewareError if the
+// request should terminate.
+func GlobalMiddleware(w http.ResponseWriter, r *http.Request) *MiddlewareError {
 	log.Println(r)
-	return false
+	return nil
 }
 
 // This example shows how to implement request middleware. ResourceHandlers

rest/middleware/cors.go

@@ -1,16 +1,79 @@
 package middleware
 
-import "net/http"
-
-// CORSMiddleware enables cross-origin requests. It implements the Middleware
-// interface.
-func CORSMiddleware(w http.ResponseWriter, r *http.Request) bool {
-	if origin := r.Header.Get("Origin"); origin != "" {
-		w.Header().Set("Access-Control-Allow-Origin", origin)
-		w.Header().Set("Access-Control-Allow-Methods", "POST, GET, OPTIONS, PUT, DELETE")
-		w.Header()["Access-Control-Allow-Headers"] = r.Header["Access-Control-Request-Headers"]
-		w.Header().Set("Access-Control-Allow-Credentials", "true")
+import (
+	"net/http"
+	"net/url"
+	"strings"
+
+	"github.com/Workiva/go-rest/rest"
+)
+
+// NewCORSMiddleware returns a Middleware which enables cross-origin requests.
+// Origin must match the supplied whitelist (which supports wildcards). Returns
+// a MiddlewareError if the request should be terminated.
+func NewCORSMiddleware(originWhitelist []string) rest.Middleware {
+	return func(w http.ResponseWriter, r *http.Request) *rest.MiddlewareError {
+		originMatch := false
+		if origin := r.Header.Get("Origin"); checkOrigin(origin, originWhitelist) {
+			w.Header().Set("Access-Control-Allow-Origin", origin)
+			w.Header().Set("Access-Control-Allow-Methods", "POST, GET, OPTIONS, PUT, DELETE")
+			w.Header()["Access-Control-Allow-Headers"] = r.Header["Access-Control-Request-Headers"]
+			w.Header().Set("Access-Control-Allow-Credentials", "true")
+			originMatch = true
+		}
+
+		var err *rest.MiddlewareError
+		if r.Method == "OPTIONS" {
+			err = &rest.MiddlewareError{Code: http.StatusOK}
+		} else if !originMatch {
+			err = &rest.MiddlewareError{
+				Code:     http.StatusBadRequest,
+				Response: []byte("Origin does not match whitelist"),
+			}
+		}
+		return err
+	}
+}
+
+// checkOrigin checks if the given origin is contained in the origin whitelist.
+// Returns true if the origin is in the whitelist, false if not.
+func checkOrigin(origin string, whitelist []string) bool {
+	url, err := url.Parse(origin)
+	if err != nil {
+		return false
+	}
+	originComponents := strings.Split(url.Host, ".")
+
+checkWhitelist:
+	for _, whitelisted := range whitelist {
+		if whitelisted == "*" {
+			return true
+		}
+
+		whitelistedComponents := strings.Split(whitelisted, ".")
+
+		if len(originComponents) != len(whitelistedComponents) {
+			// Do not match, try next host in whitelist.
+			continue
+		}
+
+		for i, originComponent := range originComponents {
+			whitelistedComponent := whitelistedComponents[i]
+			if whitelistedComponent == "*" {
+				// Wildcard, check next component.
+				continue
+			}
+
+			if originComponent != whitelistedComponent {
+				// Mismatch, try next host in whitelist.
+				continue checkWhitelist
+			}
+		}
+
+		// Origin matches whitelisted domain.
+		return true
 	}
 
-	return r.Method == "OPTIONS"
+	// No matches.
+	return false
 }

rest/middleware/cors_test.go

@@ -8,26 +8,59 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
-// Ensures that CORSMiddleware applies the headers needed for CORS and returns
-// true for non-OPTIONS requests.
-func TestCORSMiddleware(t *testing.T) {
+// Ensures that CORSMiddleware applies the headers needed for CORS when * is
+// present in the whitelist.
+func TestCORSMiddlewareAll(t *testing.T) {
 	assert := assert.New(t)
 	req, _ := http.NewRequest("GET", "http://example.com/foo", nil)
-	req.Header.Set("Origin", "abc")
+	req.Header.Set("Origin", "http://foo.com")
 	req.Header.Set("Access-Control-Request-Headers", "def")
 	w := httptest.NewRecorder()
-	assert.False(CORSMiddleware(w, req))
+	assert.Nil(NewCORSMiddleware([]string{"*"})(w, req))
 
-	assert.Equal("abc", w.Header().Get("Access-Control-Allow-Origin"))
+	assert.Equal("http://foo.com", w.Header().Get("Access-Control-Allow-Origin"))
 	assert.Equal("POST, GET, OPTIONS, PUT, DELETE", w.Header().Get("Access-Control-Allow-Methods"))
 	assert.Equal([]string{"def"}, w.Header()["Access-Control-Allow-Headers"])
 	assert.Equal([]string{"true"}, w.Header()["Access-Control-Allow-Credentials"])
 }
 
-// Ensures that CORSMiddleware returns true for OPTIONS requests.
+// Ensures that CORSMiddleware applies the headers needed for CORS and respects
+// the origin whitelist.
+func TestCORSMiddlewareWhitelist(t *testing.T) {
+	assert := assert.New(t)
+	req, _ := http.NewRequest("GET", "http://example.com/foo", nil)
+	req.Header.Set("Origin", "http://foo.wdesk.com")
+	req.Header.Set("Access-Control-Request-Headers", "def")
+	w := httptest.NewRecorder()
+	middleware := NewCORSMiddleware([]string{"blah.wdesk.org", "*.wdesk.com"})
+	assert.Nil(middleware(w, req))
+
+	assert.Equal("http://foo.wdesk.com", w.Header().Get("Access-Control-Allow-Origin"))
+	assert.Equal("POST, GET, OPTIONS, PUT, DELETE", w.Header().Get("Access-Control-Allow-Methods"))
+	assert.Equal([]string{"def"}, w.Header()["Access-Control-Allow-Headers"])
+	assert.Equal([]string{"true"}, w.Header()["Access-Control-Allow-Credentials"])
+
+	// Mismatched origin
+	req, _ = http.NewRequest("GET", "http://example.com/foo", nil)
+	req.Header.Set("Origin", "http://baz.wdesk.org")
+	req.Header.Set("Access-Control-Request-Headers", "def")
+	w = httptest.NewRecorder()
+	err := middleware(w, req)
+	assert.NotNil(err)
+	assert.Equal(http.StatusBadRequest, err.Code)
+
+	assert.Equal("", w.Header().Get("Access-Control-Allow-Origin"))
+	assert.Equal("", w.Header().Get("Access-Control-Allow-Methods"))
+	assert.Nil(w.Header()["Access-Control-Allow-Headers"])
+	assert.Nil(w.Header()["Access-Control-Allow-Credentials"])
+}
+
+// Ensures that CORSMiddleware returns a MiddlewareError with a 200 response
+// code.
 func TestCORSMiddlewareOptionsRequest(t *testing.T) {
 	req, _ := http.NewRequest("OPTIONS", "http://example.com/foo", nil)
-	req.Header.Set("Origin", "abc")
+	req.Header.Set("Origin", "http://foo.com")
 	w := httptest.NewRecorder()
-	assert.True(t, CORSMiddleware(w, req))
+	err := NewCORSMiddleware([]string{"foo.com"})(w, req)
+	assert.Equal(t, http.StatusOK, err.Code)
 }