From 852ca2502a650abb125c41169dcd0673ec08af02 Mon Sep 17 00:00:00 2001
From: Jordan Liggitt <jliggitt@redhat.com>
Date: Fri, 8 Apr 2016 17:36:21 -0400
Subject: [PATCH] Wrap login requests to clear in-memory session

---
 pkg/cmd/server/origin/auth.go            | 17 ++++++++++++++--
 pkg/cmd/server/origin/auth_config.go     | 14 +++++++++----
 pkg/cmd/server/origin/handler_wrapper.go | 26 ++++++++++++++++++++++++
 3 files changed, 51 insertions(+), 6 deletions(-)
 create mode 100644 pkg/cmd/server/origin/handler_wrapper.go

diff --git a/pkg/cmd/server/origin/auth.go b/pkg/cmd/server/origin/auth.go
index f4b86bc3e188..5dc7d4b1d443 100644
--- a/pkg/cmd/server/origin/auth.go
+++ b/pkg/cmd/server/origin/auth.go
@@ -81,8 +81,7 @@ const (
 // then returns an array of strings indicating what endpoints were started
 // (these are format strings that will expect to be sent a single string value).
 func (c *AuthConfig) InstallAPI(container *restful.Container) ([]string, error) {
-	// TODO: register into container
-	mux := container.ServeMux
+	mux := c.getMux(container)
 
 	accessTokenStorage := accesstokenetcd.NewREST(c.EtcdHelper, c.EtcdBackends...)
 	accessTokenRegistry := accesstokenregistry.NewRegistry(accessTokenStorage)
@@ -174,6 +173,20 @@ func (c *AuthConfig) InstallAPI(container *restful.Container) ([]string, error)
 	}, nil
 }
 
+func (c *AuthConfig) getMux(container *restful.Container) cmdutil.Mux {
+	// Register directly into the container's mux
+	if c.HandlerWrapper == nil {
+		return container.ServeMux
+	}
+
+	// Wrap all handlers before registering into the container's mux
+	// This lets us do things like defer session clearing to the end of a request
+	return &handlerWrapperMux{
+		mux:     container.ServeMux,
+		wrapper: c.HandlerWrapper,
+	}
+}
+
 func (c *AuthConfig) getErrorHandler() (*errorpage.ErrorPage, error) {
 	errorTemplate := ""
 	if c.Options.Templates != nil {
diff --git a/pkg/cmd/server/origin/auth_config.go b/pkg/cmd/server/origin/auth_config.go
index 4a24f43c22f5..889cfcc38c72 100644
--- a/pkg/cmd/server/origin/auth_config.go
+++ b/pkg/cmd/server/origin/auth_config.go
@@ -38,6 +38,8 @@ type AuthConfig struct {
 	IdentityRegistry identityregistry.Registry
 
 	SessionAuth *session.Authenticator
+
+	HandlerWrapper handlerWrapper
 }
 
 func BuildAuthConfig(options configapi.MasterConfig) (*AuthConfig, error) {
@@ -68,13 +70,15 @@ func BuildAuthConfig(options configapi.MasterConfig) (*AuthConfig, error) {
 	}
 
 	var sessionAuth *session.Authenticator
+	var sessionHandlerWrapper handlerWrapper
 	if options.OAuthConfig.SessionConfig != nil {
 		secure := isHTTPS(options.OAuthConfig.MasterPublicURL)
-		auth, err := BuildSessionAuth(secure, options.OAuthConfig.SessionConfig)
+		auth, wrapper, err := buildSessionAuth(secure, options.OAuthConfig.SessionConfig)
 		if err != nil {
 			return nil, err
 		}
 		sessionAuth = auth
+		sessionHandlerWrapper = wrapper
 	}
 
 	// Build the list of valid redirect_uri prefixes for a login using the openshift-web-console client to redirect to
@@ -101,18 +105,20 @@ func BuildAuthConfig(options configapi.MasterConfig) (*AuthConfig, error) {
 		UserRegistry:     userRegistry,
 
 		SessionAuth: sessionAuth,
+
+		HandlerWrapper: sessionHandlerWrapper,
 	}
 
 	return ret, nil
 }
 
-func BuildSessionAuth(secure bool, config *configapi.SessionConfig) (*session.Authenticator, error) {
+func buildSessionAuth(secure bool, config *configapi.SessionConfig) (*session.Authenticator, handlerWrapper, error) {
 	secrets, err := getSessionSecrets(config.SessionSecretsFile)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 	sessionStore := session.NewStore(secure, int(config.SessionMaxAgeSeconds), secrets...)
-	return session.NewAuthenticator(sessionStore, config.SessionName), nil
+	return session.NewAuthenticator(sessionStore, config.SessionName), sessionStore, nil
 }
 
 func getSessionSecrets(filename string) ([]string, error) {
diff --git a/pkg/cmd/server/origin/handler_wrapper.go b/pkg/cmd/server/origin/handler_wrapper.go
new file mode 100644
index 000000000000..526fc0d9b58e
--- /dev/null
+++ b/pkg/cmd/server/origin/handler_wrapper.go
@@ -0,0 +1,26 @@
+package origin
+
+import (
+	"net/http"
+
+	cmdutil "github.com/openshift/origin/pkg/cmd/util"
+)
+
+type handlerWrapper interface {
+	Wrap(http.Handler) http.Handler
+}
+
+// handlerWrapperMux wraps all handlers before registering them in the contained mux
+type handlerWrapperMux struct {
+	mux     cmdutil.Mux
+	wrapper handlerWrapper
+}
+
+var _ = cmdutil.Mux(&handlerWrapperMux{})
+
+func (m *handlerWrapperMux) Handle(pattern string, handler http.Handler) {
+	m.mux.Handle(pattern, m.wrapper.Wrap(handler))
+}
+func (m *handlerWrapperMux) HandleFunc(pattern string, handler func(http.ResponseWriter, *http.Request)) {
+	m.mux.Handle(pattern, m.wrapper.Wrap(http.HandlerFunc(handler)))
+}