From 70f2f5d750142ee2b24cefea1895f4076feca26f Mon Sep 17 00:00:00 2001
From: Zachary Newell <zlantern@gmail.com>
Date: Wed, 7 Dec 2022 00:08:01 +0000
Subject: [PATCH] Added an OIDC AllowGroups option for authorization.

---
 CHANGELOG.md        |  1 +
 config-example.yaml |  3 +++
 config.go           |  2 ++
 oidc.go             | 38 ++++++++++++++++++++++++++++++++++++++
 4 files changed, 44 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 29702a7a..f7f3394d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,7 @@
 
 ## 0.18.x (2022-xx-xx)
 
+- Added an OIDC AllowGroups Configuration options and authorization check [#1041](https://github.com/juanfont/headscale/pull/1041)
 - Reworked routing and added support for subnet router failover [#1024](https://github.com/juanfont/headscale/pull/1024)
 
 ### Changes
diff --git a/config-example.yaml b/config-example.yaml
index 57866785..4cd33f6c 100644
--- a/config-example.yaml
+++ b/config-example.yaml
@@ -273,6 +273,9 @@ unix_socket_permission: "0770"
 #
 #   allowed_domains:
 #     - example.com
+# Groups from keycloak have a leading '/'
+#   allowed_groups:
+#     - /headscale
 #   allowed_users:
 #     - alice@example.com
 #
diff --git a/config.go b/config.go
index 11883565..233dfe4a 100644
--- a/config.go
+++ b/config.go
@@ -96,6 +96,7 @@ type OIDCConfig struct {
 	ExtraParams                map[string]string
 	AllowedDomains             []string
 	AllowedUsers               []string
+	AllowedGroups              []string
 	StripEmaildomain           bool
 }
 
@@ -568,6 +569,7 @@ func GetHeadscaleConfig() (*Config, error) {
 			ExtraParams:      viper.GetStringMapString("oidc.extra_params"),
 			AllowedDomains:   viper.GetStringSlice("oidc.allowed_domains"),
 			AllowedUsers:     viper.GetStringSlice("oidc.allowed_users"),
+			AllowedGroups:    viper.GetStringSlice("oidc.allowed_groups"),
 			StripEmaildomain: viper.GetBool("oidc.strip_email_domain"),
 		},
 
diff --git a/oidc.go b/oidc.go
index 5ae3f0f0..3eed9187 100644
--- a/oidc.go
+++ b/oidc.go
@@ -25,6 +25,7 @@ const (
 	errEmptyOIDCCallbackParams = Error("empty OIDC callback params")
 	errNoOIDCIDToken           = Error("could not extract ID Token for OIDC callback")
 	errOIDCAllowedDomains      = Error("authenticated principal does not match any allowed domain")
+	errOIDCAllowedGroups       = Error("authenticated principal is not in any allowed group")
 	errOIDCAllowedUsers        = Error("authenticated principal does not match any allowed user")
 	errOIDCInvalidMachineState = Error("requested machine state key expired before authorisation completed")
 	errOIDCNodeKeyMissing      = Error("could not get node key from cache")
@@ -209,6 +210,10 @@ func (h *Headscale) OIDCCallback(
 		return
 	}
 
+	if err := validateOIDCAllowedGroups(writer, h.cfg.OIDC.AllowedGroups, claims); err != nil {
+		return
+	}
+
 	if err := validateOIDCAllowedUsers(writer, h.cfg.OIDC.AllowedUsers, claims); err != nil {
 		return
 	}
@@ -404,6 +409,39 @@ func validateOIDCAllowedDomains(
 	return nil
 }
 
+// validateOIDCAllowedGroups checks if AllowedGroups is provided,
+// and that the user has one group in the list.
+// claims.Groups can be populated by adding a client scope named
+// 'groups' that contains group membership.
+func validateOIDCAllowedGroups(
+	writer http.ResponseWriter,
+	allowedGroups []string,
+	claims *IDTokenClaims,
+) error {
+	if len(allowedGroups) > 0 {
+		for _, group := range allowedGroups {
+			if IsStringInSlice(claims.Groups, group) {
+				return nil
+			}
+		}
+
+		log.Error().Msg("authenticated principal not in any allowed groups")
+		writer.Header().Set("Content-Type", "text/plain; charset=utf-8")
+		writer.WriteHeader(http.StatusBadRequest)
+		_, err := writer.Write([]byte("unauthorized principal (allowed groups)"))
+		if err != nil {
+			log.Error().
+				Caller().
+				Err(err).
+				Msg("Failed to write response")
+		}
+
+		return errOIDCAllowedGroups
+	}
+
+	return nil
+}
+
 // validateOIDCAllowedUsers checks that if AllowedUsers is provided,
 // that the authenticated principal is part of that list.
 func validateOIDCAllowedUsers(