From 8838f289cae3eead6f5b0d5dab2a63c4d3b1031e Mon Sep 17 00:00:00 2001
From: Jackson Argo <jackson@startree.ai>
Date: Fri, 4 Aug 2023 11:34:35 -0400
Subject: [PATCH] add regex for oidc group matching

Signed-off-by: Jackson Argo <jackson@startree.ai>
---
 connector/oidc/oidc.go      | 17 +++++++++++++++++
 connector/oidc/oidc_test.go | 19 +++++++++++++++++++
 2 files changed, 36 insertions(+)

diff --git a/connector/oidc/oidc.go b/connector/oidc/oidc.go
index ff4713c2705..8bbd321e172 100644
--- a/connector/oidc/oidc.go
+++ b/connector/oidc/oidc.go
@@ -8,6 +8,7 @@ import (
 	"fmt"
 	"net/http"
 	"net/url"
+	"regexp"
 	"strings"
 	"time"
 
@@ -87,6 +88,9 @@ type Config struct {
 		// Configurable key which contains the groups claims
 		GroupsKey string `json:"groups"` // defaults to "groups"
 	} `json:"claimMapping"`
+
+	// Regex filter applied to the groups
+	GroupsRegex string `json:"groupsRegex"`
 }
 
 // Domains that don't support basic auth. golang.org/x/oauth2 has an internal
@@ -161,6 +165,14 @@ func (c *Config) Open(id string, logger log.Logger) (conn connector.Connector, e
 		c.PromptType = "consent"
 	}
 
+	var groupsRegex *regexp.Regexp
+	if c.GroupsRegex != "" {
+		groupsRegex, err = regexp.Compile(c.GroupsRegex)
+		if err != nil {
+			logger.Warnf("ignoring invalid regex `%s`", c.GroupsRegex)
+		}
+	}
+
 	clientID := c.ClientID
 	return &oidcConnector{
 		provider:    provider,
@@ -189,6 +201,7 @@ func (c *Config) Open(id string, logger log.Logger) (conn connector.Connector, e
 		preferredUsernameKey:      c.ClaimMapping.PreferredUsernameKey,
 		emailKey:                  c.ClaimMapping.EmailKey,
 		groupsKey:                 c.ClaimMapping.GroupsKey,
+		groupsRegex:               groupsRegex,
 	}, nil
 }
 
@@ -216,6 +229,7 @@ type oidcConnector struct {
 	preferredUsernameKey      string
 	emailKey                  string
 	groupsKey                 string
+	groupsRegex               *regexp.Regexp
 }
 
 func (c *oidcConnector) Close() error {
@@ -419,6 +433,9 @@ func (c *oidcConnector) createIdentity(ctx context.Context, identity connector.I
 		if found {
 			for _, v := range vs {
 				if s, ok := v.(string); ok {
+					if c.groupsRegex != nil && !c.groupsRegex.MatchString(s) {
+						continue
+					}
 					groups = append(groups, s)
 				} else {
 					return identity, fmt.Errorf("malformed \"%v\" claim", groupsKey)
diff --git a/connector/oidc/oidc_test.go b/connector/oidc/oidc_test.go
index 29e8875ea76..58ad54c2ece 100644
--- a/connector/oidc/oidc_test.go
+++ b/connector/oidc/oidc_test.go
@@ -62,6 +62,7 @@ func TestHandleCallback(t *testing.T) {
 		expectPreferredUsername   string
 		expectedEmailField        string
 		token                     map[string]interface{}
+		groupsRegex               string
 	}{
 		{
 			name:               "simpleCase",
@@ -288,6 +289,23 @@ func TestHandleCallback(t *testing.T) {
 				"email_verified": true,
 			},
 		},
+		{
+			name:               "groupsRegex",
+			userIDKey:          "", // not configured
+			userNameKey:        "", // not configured
+			groupsRegex:        `^.*\d$`,
+			expectUserID:       "subvalue",
+			expectUserName:     "namevalue",
+			expectGroups:       []string{"group1", "group2"},
+			expectedEmailField: "emailvalue",
+			token: map[string]interface{}{
+				"sub":            "subvalue",
+				"name":           "namevalue",
+				"groups":         []string{"group1", "group2", "groupA", "groupB"},
+				"email":          "emailvalue",
+				"email_verified": true,
+			},
+		},
 	}
 
 	for _, tc := range tests {
@@ -319,6 +337,7 @@ func TestHandleCallback(t *testing.T) {
 				InsecureEnableGroups:      true,
 				BasicAuthUnsupported:      &basicAuth,
 				OverrideClaimMapping:      tc.overrideClaimMapping,
+				GroupsRegex:               tc.groupsRegex,
 			}
 			config.ClaimMapping.PreferredUsernameKey = tc.preferredUsernameKey
 			config.ClaimMapping.EmailKey = tc.emailKey