Make HTTP middleware easier to use

server/auth/openid/method.go

Index: server/auth/openid/method.go
diff --git a/server/auth/openid/method.go b/server/auth/openid/method.go
index f3f3d1207161b7ec6c709e809a77a298e87755cb..74abaaa94124c79ae3c925e23d21c4131a44f580 100644
--- a/server/auth/openid/method.go
+++ b/server/auth/openid/method.go
@@ -10,14 +10,13 @@ import (
 	"net/url"
 	"time"
 
-	"github.com/julienschmidt/httprouter"
 	"golang.org/x/net/context"
 
 	"github.com/luci/luci-go/common/clock"
 	"github.com/luci/luci-go/common/errors"
 	"github.com/luci/luci-go/common/logging"
 	"github.com/luci/luci-go/server/auth"
-	"github.com/luci/luci-go/server/middleware"
+	"github.com/luci/luci-go/server/router"
 )
 
 // These are installed into a HTTP router by AuthMethod.InstallHandlers(...).
@@ -52,10 +51,10 @@ type AuthMethod struct {
 
 // InstallHandlers installs HTTP handlers used in OpenID protocol. Must be
 // installed in server HTTP router for OpenID authentication flow to work.
-func (m *AuthMethod) InstallHandlers(r *httprouter.Router, base middleware.Base) {
-	r.GET(loginURL, base(m.loginHandler))
-	r.GET(logoutURL, base(m.logoutHandler))
-	r.GET(callbackURL, base(m.callbackHandler))
+func (m *AuthMethod) InstallHandlers(r *router.Router, handlers []router.Handler) {
+	r.GET(loginURL, append(handlers, m.loginHandler)...)
+	r.GET(logoutURL, append(handlers, m.logoutHandler)...)
+	r.GET(callbackURL, append(handlers, m.callbackHandler)...)
 }
 
 // Warmup prepares local caches. It's optional.
@@ -118,16 +117,18 @@ func (m *AuthMethod) LogoutURL(c context.Context, dest string) (string, error) {
 ////
 
 // loginHandler initiates login flow by redirecting user to OpenID login page.
-func (m *AuthMethod) loginHandler(c context.Context, rw http.ResponseWriter, r *http.Request, p httprouter.Params) {
-	dest, err := normalizeURL(r.URL.Query().Get("r"))
+func (m *AuthMethod) loginHandler(c *router.Context) {
+	dest, err := normalizeURL(c.Request.URL.Query().Get("r"))
 	if err != nil {
-		replyError(c, rw, err, "Bad redirect URI (%q) - %s", dest, err)
+		replyError(c.Context, c.Writer, err, "Bad redirect URI (%q) - %s", dest, err)
+		c.Abort()
 		return
 	}
 
-	cfg, err := fetchCachedSettings(c)
+	cfg, err := fetchCachedSettings(c.Context)
 	if err != nil {
-		replyError(c, rw, err, "Can't load OpenID settings - %s", err)
+		replyError(c.Context, c.Writer, err, "Can't load OpenID settings - %s", err)
+		c.Abort()
 		return
 	}
 
@@ -135,53 +136,59 @@ func (m *AuthMethod) loginHandler(c context.Context, rw http.ResponseWriter, r *
 	// in callback URI handler. See callbackHandler.
 	state := map[string]string{
 		"dest_url": dest,
-		"host_url": r.Host,
+		"host_url": c.Request.Host,
 	}
-	authURI, err := authenticationURI(c, cfg, state)
+	authURI, err := authenticationURI(c.Context, cfg, state)
 	if err != nil {
-		replyError(c, rw, err, "Can't generate authentication URI - %s", err)
+		replyError(c.Context, c.Writer, err, "Can't generate authentication URI - %s", err)
+		c.Abort()
 		return
 	}
-	http.Redirect(rw, r, authURI, http.StatusFound)
+	http.Redirect(c.Writer, c.Request, authURI, http.StatusFound)
 }
 
 // logoutHandler nukes active session and redirect back to destination URL.
-func (m *AuthMethod) logoutHandler(c context.Context, rw http.ResponseWriter, r *http.Request, p httprouter.Params) {
-	dest, err := normalizeURL(r.URL.Query().Get("r"))
+func (m *AuthMethod) logoutHandler(c *router.Context) {
+	dest, err := normalizeURL(c.Request.URL.Query().Get("r"))
 	if err != nil {
-		replyError(c, rw, err, "Bad redirect URI (%q) - %s", dest, err)
+		replyError(c.Context, c.Writer, err, "Bad redirect URI (%q) - %s", dest, err)
 		return
 	}
 
 	// Close a session if there's one.
-	sid, err := decodeSessionCookie(c, r)
+	sid, err := decodeSessionCookie(c.Context, c.Request)
 	if err != nil {
-		replyError(c, rw, err, "Error when decoding session cookie - %s", err)
+		replyError(c.Context, c.Writer, err, "Error when decoding session cookie - %s", err)
+		c.Abort()
 		return
 	}
 	if sid != "" {
-		if err = m.SessionStore.CloseSession(c, sid); err != nil {
-			replyError(c, rw, err, "Error when closing the session - %s", err)
+		if err = m.SessionStore.CloseSession(c.Context, sid); err != nil {
+			replyError(c.Context, c.Writer, err, "Error when closing the session - %s", err)
+			c.Abort()
 			return
 		}
 	}
 
 	// Nuke all session cookies to get to a completely clean state.
-	removeCookie(rw, r, sessionCookieName)
-	m.removeIncompatibleCookies(rw, r)
+	removeCookie(c.Writer, c.Request, sessionCookieName)
+	m.removeIncompatibleCookies(c.Writer, c.Request)
 
 	// Redirect to the final destination.
-	http.Redirect(rw, r, dest, http.StatusFound)
+	http.Redirect(c.Writer, c.Request, dest, http.StatusFound)
 }
 
 // callbackHandler handles redirect from OpenID backend. Parameters contain
 // authorization code that can be exchanged for user profile.
-func (m *AuthMethod) callbackHandler(c context.Context, rw http.ResponseWriter, r *http.Request, p httprouter.Params) {
+func (m *AuthMethod) callbackHandler(ctx *router.Context) {
+	c, rw, r := ctx.Context, ctx.Writer, ctx.Request
+
 	// This code path is hit when user clicks "Deny" on consent page.
 	q := r.URL.Query()
 	errorMsg := q.Get("error")
 	if errorMsg != "" {
 		replyError(c, rw, errors.New("login error"), "OpenID login error: %s", errorMsg)
+		ctx.Abort()
 		return
 	}
 
@@ -189,16 +196,19 @@ func (m *AuthMethod) callbackHandler(c context.Context, rw http.ResponseWriter,
 	code := q.Get("code")
 	if code == "" {
 		replyError(c, rw, errors.New("login error"), "Missing 'code' parameter")
+		ctx.Abort()
 		return
 	}
 	stateTok := q.Get("state")
 	if stateTok == "" {
 		replyError(c, rw, errors.New("login error"), "Missing 'state' parameter")
+		ctx.Abort()
 		return
 	}
 	state, err := validateStateToken(c, stateTok)
 	if err != nil {
 		replyError(c, rw, err, "Failed to validate 'state' token")
+		ctx.Abort()
 		return
 	}
 
@@ -207,6 +217,7 @@ func (m *AuthMethod) callbackHandler(c context.Context, rw http.ResponseWriter,
 	dest, err := normalizeURL(state["dest_url"])
 	if err != nil {
 		replyError(c, rw, err, "Bad redirect URI (%q) - %s", dest, err)
+		ctx.Abort()
 		return
 	}
 
@@ -229,6 +240,7 @@ func (m *AuthMethod) callbackHandler(c context.Context, rw http.ResponseWriter,
 		url.Host = state["host_url"]
 		logging.Warningf(c, "Redirecting to callback URI on another host %q", url.Host)
 		http.Redirect(rw, r, url.String(), http.StatusFound)
+		ctx.Abort()
 		return
 	}
 
@@ -236,11 +248,13 @@ func (m *AuthMethod) callbackHandler(c context.Context, rw http.ResponseWriter,
 	cfg, err := fetchCachedSettings(c)
 	if err != nil {
 		replyError(c, rw, err, "Can't load OpenID settings - %s", err)
+		ctx.Abort()
 		return
 	}
 	uid, user, err := handleAuthorizationCode(c, cfg, code)
 	if err != nil {
 		replyError(c, rw, err, "Error when fetching user profile - %s", err)
+		ctx.Abort()
 		return
 	}
 
@@ -248,6 +262,7 @@ func (m *AuthMethod) callbackHandler(c context.Context, rw http.ResponseWriter,
 	prevSid, err := decodeSessionCookie(c, r)
 	if err != nil {
 		replyError(c, rw, err, "Error when decoding session cookie - %s", err)
+		ctx.Abort()
 		return
 	}
 
@@ -256,6 +271,7 @@ func (m *AuthMethod) callbackHandler(c context.Context, rw http.ResponseWriter,
 	sid, err := m.SessionStore.OpenSession(c, uid, user, expTime)
 	if err != nil {
 		replyError(c, rw, err, "Error when creating the session - %s", err)
+		ctx.Abort()
 		return
 	}
 
@@ -263,6 +279,7 @@ func (m *AuthMethod) callbackHandler(c context.Context, rw http.ResponseWriter,
 	if prevSid != "" {
 		if err = m.SessionStore.CloseSession(c, sid); err != nil {
 			replyError(c, rw, err, "Error when closing the session - %s", err)
+			ctx.Abort()
 			return
 		}
 	}
@@ -271,6 +288,7 @@ func (m *AuthMethod) callbackHandler(c context.Context, rw http.ResponseWriter,
 	cookie, err := makeSessionCookie(c, sid, !m.Insecure)
 	if err != nil {
 		replyError(c, rw, err, "Can't make session cookie - %s", err)
+		ctx.Abort()
 		return
 	}
 	http.SetCookie(rw, cookie)