From 7f85d95604e734c82de8988df28ab10493617368 Mon Sep 17 00:00:00 2001
From: feedmeapples <aksenov.ro@outlook.com>
Date: Mon, 18 Jul 2022 12:50:31 -0400
Subject: [PATCH] Fix logout that didn't clear ID Token

---
 server/auth/auth.go   | 48 +++++++++++++++++++++++++++++++++++++++++++
 server/routes/auth.go | 13 +++---------
 2 files changed, 51 insertions(+), 10 deletions(-)

diff --git a/server/auth/auth.go b/server/auth/auth.go
index 7a6fd822b9..358a4eb709 100644
--- a/server/auth/auth.go
+++ b/server/auth/auth.go
@@ -91,6 +91,20 @@ func SetUser(c echo.Context, user *User) error {
 	return nil
 }
 
+func ClearUser(c echo.Context) error {
+	err := clearAccessToken(c)
+	if err != nil {
+		return err
+	}
+
+	err = clearIDToken(c)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
 func ValidateAuth(c echo.Context, cfgProvider *config.ConfigProviderWithRefresh) error {
 	cfg, err := cfgProvider.GetConfig()
 	if err != nil {
@@ -162,6 +176,23 @@ func setAccessToken(c echo.Context, token string) error {
 	return nil
 }
 
+func clearAccessToken(c echo.Context) error {
+	sess, _ := session.Get(AuthCookie, c)
+	sess.Options = &sessions.Options{
+		Path:     "/",
+		MaxAge:   -1,
+		HttpOnly: true,
+		SameSite: http.SameSiteStrictMode,
+		Secure:   true,
+	}
+	err := sess.Save(c.Request(), c.Response())
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
 func getAuthorizationExtras(c echo.Context) string {
 	sess, _ := session.Get(AuthExtrasCookie, c)
 	if sess == nil {
@@ -192,3 +223,20 @@ func setIDToken(c echo.Context, idToken *IDToken) error {
 
 	return nil
 }
+
+func clearIDToken(c echo.Context) error {
+	sess, _ := session.Get(AuthExtrasCookie, c)
+	sess.Options = &sessions.Options{
+		Path:     "/",
+		MaxAge:   -1,
+		HttpOnly: true,
+		SameSite: http.SameSiteStrictMode,
+		Secure:   true,
+	}
+	err := sess.Save(c.Request(), c.Response())
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
diff --git a/server/routes/auth.go b/server/routes/auth.go
index 1afde9e786..b2fe50b01d 100644
--- a/server/routes/auth.go
+++ b/server/routes/auth.go
@@ -33,8 +33,6 @@ import (
 
 	"github.com/coreos/go-oidc/v3/oidc"
 	"github.com/gorilla/securecookie"
-	"github.com/gorilla/sessions"
-	"github.com/labstack/echo-contrib/session"
 	"github.com/labstack/echo/v4"
 	"github.com/temporalio/ui-server/v2/server/auth"
 	"github.com/temporalio/ui-server/v2/server/config"
@@ -150,15 +148,10 @@ func authenticateCb(ctx context.Context, oauthCfg *oauth2.Config, provider *oidc
 }
 
 func logout(c echo.Context) error {
-	sess, _ := session.Get(auth.AuthCookie, c)
-	sess.Options = &sessions.Options{
-		Path:     "/",
-		MaxAge:   -1,
-		HttpOnly: true,
-		SameSite: http.SameSiteStrictMode,
-		Secure:   true,
+	err := auth.ClearUser(c)
+	if err != nil {
+		return echo.NewHTTPError(http.StatusInternalServerError, "unable to clear user: "+err.Error())
 	}
-	sess.Save(c.Request(), c.Response())
 
 	returnUrl := c.Request().Header.Get("Referer")
 	if returnUrl == "" {