refactor router to reduce duplication and remove an external dependency

.dockerignore

@@ -2,9 +2,8 @@
 *
 
 # Whitelist required files
-!.env.encrypted
-!scripts/*
 !lambda/*
+!router/*
 !server/*
 !u2fsimulator/*
 !u2fserver/*

go.mod

@@ -11,7 +11,6 @@ require (
 	github.com/fxamacker/cbor/v2 v2.8.0
 	github.com/go-webauthn/webauthn v0.11.2
 	github.com/google/uuid v1.6.0
-	github.com/gorilla/mux v1.8.1
 	github.com/kelseyhightower/envconfig v1.4.0
 	github.com/pquerna/otp v1.5.0
 	github.com/satori/go.uuid v1.2.0

go.sum

@@ -52,8 +52,6 @@ github.com/google/go-tpm v0.9.3 h1:+yx0/anQuGzi+ssRqeD6WpXjW2L/V0dItUayO0i9sRc=
 github.com/google/go-tpm v0.9.3/go.mod h1:h9jEsEECg7gtLis0upRBQU+GhYVH6jMjrFxI8u6bVUY=
 github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
 github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
-github.com/gorilla/mux v1.8.1 h1:TuBL49tXwgrFYWhqrNgrUNEY92u81SPhu7sTdzQEiWY=
-github.com/gorilla/mux v1.8.1/go.mod h1:AKf9I4AEqPTmMytcMc0KkNouC66V3BtZ4qD5fmWSiMQ=
 github.com/kelseyhightower/envconfig v1.4.0 h1:Im6hONhd3pLkfDFsbRgu68RDNkGF1r3dvMUtDTo2cv8=
 github.com/kelseyhightower/envconfig v1.4.0/go.mod h1:cccZRl6mQpaq41TPp5QxidR+Sa3axMbJDNb//FQX6Gg=
 github.com/kr/pretty v0.2.1/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI=

lambda/main.go

@@ -2,8 +2,6 @@ package main
 
 import (
 	"context"
-	"encoding/json"
-	"fmt"
 	"io"
 	"log"
 	"net/http"
@@ -13,10 +11,10 @@ import (
 
 	"github.com/aws/aws-lambda-go/events"
 	"github.com/aws/aws-lambda-go/lambda"
-	"github.com/gorilla/mux"
 	"github.com/kelseyhightower/envconfig"
 
 	mfa "github.com/silinternational/serverless-mfa-api-go"
+	"github.com/silinternational/serverless-mfa-api-go/router"
 )
 
 var envConfig mfa.EnvConfig
@@ -63,65 +61,15 @@ func credentialToDelete(req events.APIGatewayProxyRequest) (string, bool) {
 	return credID, true
 }
 
-func addDeleteCredentialParamForMux(r *http.Request, key, value string) *http.Request {
-	vars := map[string]string{key: value}
-	return mux.SetURLVars(r, vars)
-}
-
 func handler(ctx context.Context, req events.APIGatewayProxyRequest) (events.APIGatewayProxyResponse, error) {
 	r := httpRequestFromProxyRequest(ctx, req)
+	w := newLambdaResponseWriter()
 
 	app := mfa.NewApp(envConfig)
+	mux := router.NewMux(app)
 
-	if !strings.HasPrefix(r.URL.Path, "/api-key") {
-		user, err := mfa.AuthenticateRequest(r)
-		if err != nil {
-			return clientError(http.StatusUnauthorized, fmt.Sprintf("unable to authenticate request: %s", err))
-		}
-		newCtx := context.WithValue(r.Context(), mfa.UserContextKey, user)
-		r = r.WithContext(newCtx)
-	}
-
-	// Use custom lambda http.ResponseWriter
-	w := newLambdaResponseWriter()
+	mux.ServeHTTP(w, r)
 
-	// This (delete /webauthn/credential/abc123) expects a path param that is
-	//   the id that was previously returned as
-	//   the key_handle_hash from the FinishRegistration call.
-	// Alternatively, if the id param indicates that a legacy U2F key should be removed
-	//   (e.g. by matching the string "u2f")
-	//   then that user is saved with all of its legacy u2f fields blanked out.
-	if credIdToDelete, ok := credentialToDelete(req); ok {
-		// add the id to the context in order for mux to retrieve it
-		r = addDeleteCredentialParamForMux(r, mfa.IDParam, credIdToDelete)
-		app.DeleteCredential(w, r)
-		// Routes other than /webauthn/delete/credential/abc213
-	} else {
-		route := strings.ToLower(fmt.Sprintf("%s %s", req.HTTPMethod, req.Path))
-
-		switch route {
-		case "post /api-key":
-			app.CreateApiKey(w, r)
-		case "post /api-key/activate":
-			app.ActivateApiKey(w, r)
-		case "post /api-key/rotate":
-			app.RotateApiKey(w, r)
-		case "post /totp":
-			app.CreateTOTP(w, r)
-		case "post /webauthn/login":
-			app.BeginLogin(w, r)
-		case "put /webauthn/login":
-			app.FinishLogin(w, r)
-		case "post /webauthn/register":
-			app.BeginRegistration(w, r)
-		case "put /webauthn/register":
-			app.FinishRegistration(w, r)
-		case "delete /webauthn/user":
-			app.DeleteUser(w, r)
-		default:
-			return clientError(http.StatusNotFound, fmt.Sprintf("The requested route is not supported: %s", route))
-		}
-	}
 	headers := map[string]string{}
 	for k, v := range w.Header() {
 		headers[k] = v[0]
@@ -134,20 +82,6 @@ func handler(ctx context.Context, req events.APIGatewayProxyRequest) (events.API
 	}, nil
 }
 
-// clientError helper for send responses relating to client errors.
-func clientError(status int, body string) (events.APIGatewayProxyResponse, error) {
-	type cError struct {
-		Error string
-	}
-
-	js, _ := json.Marshal(cError{Error: body})
-
-	return events.APIGatewayProxyResponse{
-		StatusCode: status,
-		Body:       string(js),
-	}, nil
-}
-
 func httpRequestFromProxyRequest(ctx context.Context, req events.APIGatewayProxyRequest) *http.Request {
 	headers := http.Header{}
 	for k, v := range req.Headers {
@@ -164,5 +98,6 @@ func httpRequestFromProxyRequest(ctx context.Context, req events.APIGatewayProxy
 		RequestURI:    req.Path,
 		URL:           requestURL,
 	}
+
 	return r.WithContext(ctx)
 }

lambda/main_test.go

@@ -1,14 +1,10 @@
 package main
 
 import (
-	"net/http/httptest"
 	"testing"
 
 	"github.com/aws/aws-lambda-go/events"
-	"github.com/gorilla/mux"
 	"github.com/stretchr/testify/require"
-
-	mfa "github.com/silinternational/serverless-mfa-api-go"
 )
 
 func TestCredentialToDelete(t *testing.T) {
@@ -69,15 +65,3 @@ func TestCredentialToDelete(t *testing.T) {
 		})
 	}
 }
-
-func TestAddDeleteCredentialParamForMux(t *testing.T) {
-	assert := require.New(t)
-	r := httptest.NewRequest("DELETE", "/webauthn/credential/abc123", nil)
-
-	credId := "abc123"
-	r = addDeleteCredentialParamForMux(r, mfa.IDParam, credId)
-	params := mux.Vars(r)
-	got, ok := params[mfa.IDParam]
-	assert.True(ok, "didn't find key in mux vars: %v", params)
-	assert.Equal(credId, got, "incorrect param value")
-}

router/main.go

@@ -0,0 +1,73 @@
+package router
+
+import (
+	"net/http"
+
+	mfa "github.com/silinternational/serverless-mfa-api-go"
+)
+
+// NewMux forms a new ServeMux router, see https://pkg.go.dev/net/http#ServeMux.
+func NewMux(app *mfa.App) *http.ServeMux {
+	mux := http.NewServeMux()
+
+	for _, r := range getRoutes(app) {
+		mux.Handle(r.Pattern, authenticationMiddleware(r.HandlerFunc))
+	}
+	return mux
+}
+
+// route is used to pass information about a particular route.
+type route struct {
+	Pattern     string
+	HandlerFunc http.HandlerFunc
+}
+
+// getRoutes returns a list of routes for the server
+func getRoutes(app *mfa.App) []route {
+	return []route{
+		{
+			Pattern:     "POST /api-key/activate",
+			HandlerFunc: app.ActivateApiKey,
+		},
+		{
+			Pattern:     "POST /api-key/rotate",
+			HandlerFunc: app.RotateApiKey,
+		},
+		{
+			Pattern:     "POST /api-key",
+			HandlerFunc: app.CreateApiKey,
+		},
+		{
+			Pattern:     "POST /webauthn/register",
+			HandlerFunc: app.BeginRegistration,
+		},
+		{
+			Pattern:     "PUT /webauthn/register",
+			HandlerFunc: app.FinishRegistration,
+		},
+		{
+			Pattern:     "POST /webauthn/login",
+			HandlerFunc: app.BeginLogin,
+		},
+		{
+			Pattern:     "PUT /webauthn/login",
+			HandlerFunc: app.FinishLogin,
+		},
+		{
+			Pattern:     "DELETE /webauthn/user",
+			HandlerFunc: app.DeleteUser,
+		},
+		{ // This expects a path param that is the id that was previously returned
+			// as the key_handle_hash from the FinishRegistration call.
+			// Alternatively, if the id param indicates that a legacy U2F key should be removed
+			//	 (e.g. by matching the string "u2f")
+			//   then that user is saved with all of its legacy u2f fields blanked out.
+			Pattern:     "DELETE /webauthn/credential/{" + mfa.IDParam + "}/",
+			HandlerFunc: app.DeleteCredential,
+		},
+		{
+			Pattern:     "POST /totp",
+			HandlerFunc: app.CreateTOTP,
+		},
+	}
+}

server/authentication_middleware.go → router/middleware.go

@@ -1,4 +1,4 @@
-package main
+package router
 
 import (
 	"context"

server/main.go

@@ -1,16 +1,14 @@
 package main
 
 import (
-	"encoding/json"
-	"fmt"
 	"log"
 	"net/http"
 	"os"
 
-	"github.com/gorilla/mux"
 	"github.com/kelseyhightower/envconfig"
 
 	mfa "github.com/silinternational/serverless-mfa-api-go"
+	"github.com/silinternational/serverless-mfa-api-go/router"
 )
 
 var envConfig mfa.EnvConfig
@@ -29,110 +27,7 @@ func main() {
 	// ListenAndServe starts an HTTP server with a given address and
 	// handler defined in NewRouter.
 	log.Println("Starting service on port 8080")
-	router := newRouter(mfa.NewApp(envConfig))
-	log.Fatal(http.ListenAndServe(":8080", router))
-}
-
-// route is used to pass information about a particular route.
-type route struct {
-	Name        string
-	Method      string
-	Pattern     string
-	HandlerFunc http.HandlerFunc
-}
-
-// getRoutes returns a list of routes for the server
-func getRoutes(app *mfa.App) []route {
-	return []route{
-		{
-			Name:        "ActivateApiKey",
-			Method:      "POST",
-			Pattern:     "/api-key/activate",
-			HandlerFunc: app.ActivateApiKey,
-		},
-		{
-			Name:        "RotateApiKey",
-			Method:      "POST",
-			Pattern:     "/api-key/rotate",
-			HandlerFunc: app.RotateApiKey,
-		},
-		{
-			Name:        "CreateApiKey",
-			Method:      "POST",
-			Pattern:     "/api-key",
-			HandlerFunc: app.CreateApiKey,
-		},
-		{
-			Name:        "FinishRegistration",
-			Method:      "PUT",
-			Pattern:     "/webauthn/register",
-			HandlerFunc: app.FinishRegistration,
-		},
-		{
-			Name:        "BeginLogin",
-			Method:      "POST",
-			Pattern:     "/webauthn/login",
-			HandlerFunc: app.BeginLogin,
-		},
-		{
-			Name:        "FinishLogin",
-			Method:      "PUT",
-			Pattern:     "/webauthn/login",
-			HandlerFunc: app.FinishLogin,
-		},
-		{
-			Name:        "DeleteUser",
-			Method:      "DELETE",
-			Pattern:     "/webauthn/user",
-			HandlerFunc: app.DeleteUser,
-		},
-		{ // This expects a path param that is the id that was previously returned
-			// as the key_handle_hash from the FinishRegistration call.
-			// Alternatively, if the id param indicates that a legacy U2F key should be removed
-			//	 (e.g. by matching the string "u2f")
-			//   then that user is saved with all of its legacy u2f fields blanked out.
-			Name:        "DeleteCredential",
-			Method:      "DELETE",
-			Pattern:     fmt.Sprintf("/webauthn/credential/{%s}", mfa.IDParam),
-			HandlerFunc: app.DeleteCredential,
-		},
-		{
-			Name:        "CreateTOTP",
-			Method:      "POST",
-			Pattern:     "/totp",
-			HandlerFunc: app.CreateTOTP,
-		},
-	}
-}
-
-// newRouter forms a new mux router, see https://github.com/gorilla/mux.
-func newRouter(app *mfa.App) *mux.Router {
-	// Create a basic router.
-	router := mux.NewRouter().StrictSlash(true)
-
-	// authenticate request based on api key and secret in headers
-	// also adds user to context
-	router.Use(authenticationMiddleware)
-
-	// Assign the handlers to run when endpoints are called.
-	for _, route := range getRoutes(app) {
-		router.Methods(route.Method).Path(route.Pattern).Name(route.Name).Handler(route.HandlerFunc)
-	}
-
-	router.NotFoundHandler = router.NewRoute().HandlerFunc(notFound).GetHandler()
-	return router
-}
-
-func notFound(w http.ResponseWriter, r *http.Request) {
-	w.Header().Set("Content-Type", "application/json; charset=UTF-8")
-	w.WriteHeader(http.StatusNotFound)
-
-	notFound := map[string]string{
-		"Method":     r.Method,
-		"URL":        r.URL.String(),
-		"RequestURI": r.RequestURI,
-	}
-	if err := json.NewEncoder(w).Encode(notFound); err != nil {
-		log.Printf("ERROR could not marshal not found message to JSON: %s", err)
-	}
+	app := mfa.NewApp(envConfig)
+	mux := router.NewMux(app)
+	log.Fatal(http.ListenAndServe(":8080", mux))
 }