Make HTTP middleware easier to use

appengine/cmd/tokenserver/frontend/main.go

Index: appengine/cmd/tokenserver/frontend/main.go
diff --git a/appengine/cmd/tokenserver/frontend/main.go b/appengine/cmd/tokenserver/frontend/main.go
index 8c397a9855d0db98ae243e794adff0bca50d96f2..5301a4ff11845781ea66399f85ac627e0bbf6f9a 100644
--- a/appengine/cmd/tokenserver/frontend/main.go
+++ b/appengine/cmd/tokenserver/frontend/main.go
@@ -6,34 +6,34 @@
 // module.
 //
 // It stitches together all the code.
 package frontend
 
 import (
 	"net/http"
 	"sync"
 
 	"github.com/golang/protobuf/proto"
-	"github.com/julienschmidt/httprouter"
 	"golang.org/x/net/context"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 
 	"github.com/luci/gae/service/info"
 	"github.com/luci/luci-go/appengine/gaeauth/server"
 	"github.com/luci/luci-go/appengine/gaemiddleware"
 	"github.com/luci/luci-go/appengine/tsmon"
 	"github.com/luci/luci-go/common/logging"
 	"github.com/luci/luci-go/server/auth"
 	"github.com/luci/luci-go/server/auth/machine"
 	"github.com/luci/luci-go/server/discovery"
 	"github.com/luci/luci-go/server/prpc"
+	"github.com/luci/luci-go/server/router"
 
 	"github.com/luci/luci-go/common/api/tokenserver/admin/v1"
 	"github.com/luci/luci-go/common/api/tokenserver/identity/v1"
 	"github.com/luci/luci-go/common/api/tokenserver/minter/v1"
 
 	"github.com/luci/luci-go/appengine/cmd/tokenserver/services/admin/certauthorities"
 	"github.com/luci/luci-go/appengine/cmd/tokenserver/services/admin/serviceaccounts"
 	"github.com/luci/luci-go/appengine/cmd/tokenserver/services/identity/identityfetcher"
 	"github.com/luci/luci-go/appengine/cmd/tokenserver/services/minter/tokenminter"
 )
@@ -75,106 +75,106 @@ func adminPrelude(serviceName string) func(context.Context, string, proto.Messag
 		case err != nil:
 			return nil, grpc.Errorf(codes.Internal, "can't check ACL - %s", err)
 		case !admin:
 			return nil, grpc.Errorf(codes.PermissionDenied, "not an admin")
 		}
 		return c, nil
 	}
 }
 
 func init() {
-	router := httprouter.New()
-	base := gaemiddleware.BaseProd
+	r := router.New()
+	basemw := gaemiddleware.BaseProd()
 
 	// Install auth, config and tsmon handlers.
-	gaemiddleware.InstallHandlers(router, base)
+	gaemiddleware.InstallHandlers(r, basemw)
 
 	// The service has no UI, so just redirect to stock RPC explorer.
-	router.GET("/", func(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
-		http.Redirect(w, r, "/rpcexplorer/", http.StatusFound)
+	r.GET("/", nil, func(c *router.Context) {
+		http.Redirect(c.Writer, c.Request, "/rpcexplorer/", http.StatusFound)
 	})
 
 	// Optional warmup routes.
-	router.GET("/_ah/warmup", base(warmupHandler))
-	router.GET("/_ah/start", base(warmupHandler))
+	r.GET("/_ah/warmup", basemw, warmupHandler)
+	r.GET("/_ah/start", basemw, warmupHandler)
 
 	// Backend routes used for cron and task queues.
-	router.GET("/internal/cron/read-config", base(gaemiddleware.RequireCron(readConfigCron)))
-	router.GET("/internal/cron/fetch-crl", base(gaemiddleware.RequireCron(fetchCRLCron)))
+	r.GET("/internal/cron/read-config", append(basemw, gaemiddleware.RequireCron), readConfigCron)
+	r.GET("/internal/cron/fetch-crl", append(basemw, gaemiddleware.RequireCron), fetchCRLCron)
 
 	// Install all RPC servers.
 	api := prpc.Server{
 		Authenticator: auth.Authenticator{
 			&server.OAuth2Method{Scopes: []string{server.EmailScope}},
 			&machine.MachineTokenAuthMethod{},
 		},
 		UnaryServerInterceptor: tsmon.NewGrpcUnaryInterceptor(nil),
 	}
 	admin.RegisterCertificateAuthoritiesServer(&api, caServerWithAuth)
 	admin.RegisterServiceAccountsServer(&api, serviceAccountsServerWithAuth)
 	identity.RegisterIdentityFetcherServer(&api, identityFetcher)
 	minter.RegisterTokenMinterServer(&api, tokenMinterServerWithoutAuth) // auth inside
 	discovery.Enable(&api)
-	api.InstallHandlers(router, base)
+	api.InstallHandlers(r, basemw)
 
 	// Expose all this stuff.
-	http.DefaultServeMux.Handle("/", router)
+	http.DefaultServeMux.Handle("/", r)
 }
 
 /// Routes.
 
 // warmupHandler warms in-memory caches.
-func warmupHandler(c context.Context, w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
-	if err := server.Warmup(c); err != nil {
+func warmupHandler(c *router.Context) {
+	if err := server.Warmup(c.Context); err != nil {
 		panic(err) // let panic catcher deal with it
 	}
-	w.WriteHeader(http.StatusOK)
+	c.Writer.WriteHeader(http.StatusOK)
 }
 
 // readConfigCron is handler for /internal/cron/read-config GAE cron task.
-func readConfigCron(c context.Context, w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
+func readConfigCron(c *router.Context) {
 	// Don't override manually imported configs with 'nil' on devserver.
-	if info.Get(c).IsDevAppServer() {
-		w.WriteHeader(http.StatusOK)
+	if info.Get(c.Context).IsDevAppServer() {
+		c.Writer.WriteHeader(http.StatusOK)
 		return
 	}
-	if _, err := caServerWithoutAuth.ImportConfig(c, nil); err != nil {
+	if _, err := caServerWithoutAuth.ImportConfig(c.Context, nil); err != nil {
 		panic(err) // let panic catcher deal with it
 	}
-	w.WriteHeader(http.StatusOK)
+	c.Writer.WriteHeader(http.StatusOK)
 }
 
 // fetchCRLCron is handler for /internal/cron/fetch-crl GAE cron task.
-func fetchCRLCron(c context.Context, w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
-	list, err := caServerWithoutAuth.ListCAs(c, nil)
+func fetchCRLCron(c *router.Context) {
+	list, err := caServerWithoutAuth.ListCAs(c.Context, nil)
 	if err != nil {
 		panic(err) // let panic catcher deal with it
 	}
 
 	// Fetch CRL of each active CA in parallel. In practice there are very few
 	// CAs there (~= 1), so the risk of OOM is small.
 	wg := sync.WaitGroup{}
 	errs := make([]error, len(list.Cn))
 	for i, cn := range list.Cn {
 		wg.Add(1)
 		go func(i int, cn string) {
 			defer wg.Done()
-			_, err := caServerWithoutAuth.FetchCRL(c, &admin.FetchCRLRequest{Cn: cn})
+			_, err := caServerWithoutAuth.FetchCRL(c.Context, &admin.FetchCRLRequest{Cn: cn})
 			if err != nil {
-				logging.Errorf(c, "FetchCRL(%q) failed - %s", cn, err)
+				logging.Errorf(c.Context, "FetchCRL(%q) failed - %s", cn, err)
 				errs[i] = err
 			}
 		}(i, cn)
 	}
 	wg.Wait()
 
 	// Retry cron job only on transient errors. On fatal errors let it rerun one
 	// minute later, as usual, to avoid spamming logs with errors.
 	status := http.StatusOK
 	for _, err = range errs {
 		if grpc.Code(err) == codes.Internal {
 			status = http.StatusInternalServerError
 			break
 		}
 	}
-	w.WriteHeader(status)
+	c.Writer.WriteHeader(status)
 }