From 477b3caf415d1f65f71dd366d9ebc5c04c64c099 Mon Sep 17 00:00:00 2001
From: Basanth Jenu H B <basanthjenuhb@gmail.com>
Date: Sun, 10 Sep 2023 11:12:26 +0530
Subject: [PATCH] feat: filter sso groups based on regex (#11774)

Signed-off-by: bjenuhb <Basanth_JenuHB@intuit.com>
Co-authored-by: bjenuhb <Basanth_JenuHB@intuit.com>
---
 config/sso.go          |  7 ++--
 server/auth/sso/sso.go | 80 +++++++++++++++++++++++++++++-------------
 2 files changed, 59 insertions(+), 28 deletions(-)

diff --git a/config/sso.go b/config/sso.go
index 4c1a18254e2f..6e288d2d3cf2 100644
--- a/config/sso.go
+++ b/config/sso.go
@@ -18,9 +18,10 @@ type SSOConfig struct {
 	Scopes        []string        `json:"scopes,omitempty"`
 	SessionExpiry metav1.Duration `json:"sessionExpiry,omitempty"`
 	// customGroupClaimName will override the groups claim name
-	CustomGroupClaimName string `json:"customGroupClaimName,omitempty"`
-	UserInfoPath         string `json:"userInfoPath,omitempty"`
-	InsecureSkipVerify   bool   `json:"insecureSkipVerify,omitempty"`
+	CustomGroupClaimName string   `json:"customGroupClaimName,omitempty"`
+	UserInfoPath         string   `json:"userInfoPath,omitempty"`
+	InsecureSkipVerify   bool     `json:"insecureSkipVerify,omitempty"`
+	FilterGroupsRegex    []string `json:"filterGroupsRegex,omitempty"`
 }
 
 func (c SSOConfig) GetSessionExpiry() time.Duration {
diff --git a/server/auth/sso/sso.go b/server/auth/sso/sso.go
index bd9603e55637..2d238bba0ab5 100644
--- a/server/auth/sso/sso.go
+++ b/server/auth/sso/sso.go
@@ -9,6 +9,7 @@ import (
 	"crypto/x509"
 	"fmt"
 	"net/http"
+	"regexp"
 	"strings"
 	"time"
 
@@ -49,18 +50,19 @@ var _ Interface = &sso{}
 type Config = config.SSOConfig
 
 type sso struct {
-	config          *oauth2.Config
-	issuer          string
-	idTokenVerifier *oidc.IDTokenVerifier
-	httpClient      *http.Client
-	baseHRef        string
-	secure          bool
-	privateKey      crypto.PrivateKey
-	encrypter       jose.Encrypter
-	rbacConfig      *config.RBACConfig
-	expiry          time.Duration
-	customClaimName string
-	userInfoPath    string
+	config            *oauth2.Config
+	issuer            string
+	idTokenVerifier   *oidc.IDTokenVerifier
+	httpClient        *http.Client
+	baseHRef          string
+	secure            bool
+	privateKey        crypto.PrivateKey
+	encrypter         jose.Encrypter
+	rbacConfig        *config.RBACConfig
+	expiry            time.Duration
+	customClaimName   string
+	userInfoPath      string
+	filterGroupsRegex []*regexp.Regexp
 }
 
 func (s *sso) IsRBACEnabled() bool {
@@ -181,25 +183,38 @@ func newSso(
 	if err != nil {
 		return nil, fmt.Errorf("failed to create JWT encrpytor: %w", err)
 	}
-	lf := log.Fields{"redirectUrl": config.RedirectURL, "issuer": c.Issuer, "issuerAlias": "DISABLED", "clientId": c.ClientID, "scopes": config.Scopes, "insecureSkipVerify": c.InsecureSkipVerify}
+
+	var filterGroupsRegex []*regexp.Regexp
+	if c.FilterGroupsRegex != nil && len(c.FilterGroupsRegex) > 0 {
+		for _, regex := range c.FilterGroupsRegex {
+			compiledRegex, err := regexp.Compile(regex)
+			if err != nil {
+				return nil, fmt.Errorf("failed to compile sso.filterGroupRegex: %s %w", regex, err)
+			}
+			filterGroupsRegex = append(filterGroupsRegex, compiledRegex)
+		}
+	}
+
+	lf := log.Fields{"redirectUrl": config.RedirectURL, "issuer": c.Issuer, "issuerAlias": "DISABLED", "clientId": c.ClientID, "scopes": config.Scopes, "insecureSkipVerify": c.InsecureSkipVerify, "filterGroupsRegex": c.FilterGroupsRegex}
 	if c.IssuerAlias != "" {
 		lf["issuerAlias"] = c.IssuerAlias
 	}
 	log.WithFields(lf).Info("SSO configuration")
 
 	return &sso{
-		config:          config,
-		idTokenVerifier: idTokenVerifier,
-		baseHRef:        baseHRef,
-		httpClient:      httpClient,
-		secure:          secure,
-		privateKey:      privateKey,
-		encrypter:       encrypter,
-		rbacConfig:      c.RBAC,
-		expiry:          c.GetSessionExpiry(),
-		customClaimName: c.CustomGroupClaimName,
-		userInfoPath:    c.UserInfoPath,
-		issuer:          c.Issuer,
+		config:            config,
+		idTokenVerifier:   idTokenVerifier,
+		baseHRef:          baseHRef,
+		httpClient:        httpClient,
+		secure:            secure,
+		privateKey:        privateKey,
+		encrypter:         encrypter,
+		rbacConfig:        c.RBAC,
+		expiry:            c.GetSessionExpiry(),
+		customClaimName:   c.CustomGroupClaimName,
+		userInfoPath:      c.UserInfoPath,
+		issuer:            c.Issuer,
+		filterGroupsRegex: filterGroupsRegex,
 	}, nil
 }
 
@@ -280,6 +295,21 @@ func (s *sso) HandleCallback(w http.ResponseWriter, r *http.Request) {
 			return
 		}
 	}
+
+	// only return groups that match at least one of the regexes
+	if s.filterGroupsRegex != nil && len(s.filterGroupsRegex) > 0 {
+		var filteredGroups []string
+		for _, group := range groups {
+			for _, regex := range s.filterGroupsRegex {
+				if regex.MatchString(group) {
+					filteredGroups = append(filteredGroups, group)
+					break
+				}
+			}
+		}
+		groups = filteredGroups
+	}
+
 	argoClaims := &types.Claims{
 		Claims: jwt.Claims{
 			Issuer:  issuer,