Refactor cookie generation to use an interface

oauth2/auth.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"net/http"
 	"strings"
+	"time"
 
 	"github.com/influxdata/chronograf"
 )
@@ -60,17 +61,25 @@ func AuthorizedToken(auth Authenticator, te TokenExtractor, logger chronograf.Lo
 			w.WriteHeader(http.StatusUnauthorized)
 			return
 		}
+		ctx := r.Context()
 		// We do not check the validity of the principal.  Those
 		// served further down the chain should do so.
-		principal, err := auth.Authenticate(r.Context(), token)
+		principal, err := auth.Authenticate(ctx, token)
 		if err != nil {
 			log.Error("Invalid token")
 			w.WriteHeader(http.StatusUnauthorized)
 			return
 		}
 
+		token, err = auth.Token(ctx, principal, 10*time.Second)
+		if err != nil {
+			log.Error("Unable to create new token")
+			w.WriteHeader(http.StatusInternalServerError)
+			return
+		}
+
 		// Send the principal to the next handler
-		ctx := context.WithValue(r.Context(), PrincipalKey, principal)
+		ctx = context.WithValue(ctx, PrincipalKey, principal)
 		next.ServeHTTP(w, r.WithContext(ctx))
 		return
 	})

oauth2/mux.go

@@ -1,6 +1,7 @@
 package oauth2
 
 import (
+	"context"
 	"net/http"
 	"time"
 
@@ -130,24 +131,16 @@ func (j *CookieMux) Callback() http.Handler {
 			Subject: id,
 			Issuer:  j.Provider.Name(),
 		}
+		ctx := r.Context()
 		// We create an auth token that will be used by all other endpoints to validate the principal has a claim
-		authToken, err := j.Auth.Token(r.Context(), p, j.cookie.Duration)
+		authToken, err := j.Auth.Token(ctx, p, j.cookie.Duration)
 		if err != nil {
 			log.Error("Unable to create cookie auth token ", err.Error())
 			http.Redirect(w, r, j.FailureURL, http.StatusTemporaryRedirect)
 			return
 		}
 
-		expireCookie := j.Now().UTC().Add(j.cookie.Duration)
+		cookie := j.Generate(ctx, authToken, j.cookie.Duration)
-		cookie := http.Cookie{
-			Name:     j.cookie.Name,
-			Value:    authToken,
-			HttpOnly: true,
-			Path:     "/",
-		}
-		if j.cookie.Duration > 0 {
-			cookie.Expires = expireCookie
-		}
 		log.Info("User ", id, " is authenticated")
 		http.SetCookie(w, &cookie)
 		http.Redirect(w, r, j.SuccessURL, http.StatusTemporaryRedirect)
@@ -157,14 +150,26 @@ func (j *CookieMux) Callback() http.Handler {
 // Logout handler will expire our authentication cookie and redirect to the successURL
 func (j *CookieMux) Logout() http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		deleteCookie := http.Cookie{
+		// To expire a cookie in a browser, you set the expire time to the past.
-			Name:     j.cookie.Name,
+		deleteCookie := j.Generate(context.TODO(), "none", -1*time.Hour)
-			Value:    "none",
-			Expires:  j.Now().UTC().Add(-1 * time.Hour),
-			HttpOnly: true,
-			Path:     "/",
-		}
 		http.SetCookie(w, &deleteCookie)
 		http.Redirect(w, r, j.SuccessURL, http.StatusTemporaryRedirect)
 	})
 }
+
+// Generate will create cookies containing token information
+func (j *CookieMux) Generate(ctx context.Context, token string, expires time.Duration) http.Cookie {
+	cookie := http.Cookie{
+		Name:     j.cookie.Name,
+		Value:    token,
+		HttpOnly: true,
+		Path:     "/",
+	}
+	if expires == 0 {
+		// Sometime in the far future when marty jr gets out of prison.
+		cookie.Expires = time.Date(2030, time.October, 22, 12, 0, 0, 0, time.UTC)
+	} else {
+		cookie.Expires = j.Now().UTC().Add(expires)
+	}
+	return cookie
+}

oauth2/oauth2.go

@@ -73,3 +73,9 @@ type TokenExtractor interface {
 	// Extract will return the token or an error.
 	Extract(r *http.Request) (string, error)
 }
+
+// CookieGenerator takes a token string and returns a cookie to be used in an
+// http response.
+type CookieGenerator interface {
+	Generate(ctx context.Context, token string, expires time.Duration) http.Cookie
+}