From 2549844b4d4b24bdfddb6cd73140edfebe5ffcfb Mon Sep 17 00:00:00 2001
From: Andy Lo-A-Foe <andy.loafoe@gmail.com>
Date: Tue, 13 Feb 2024 08:18:56 +0100
Subject: [PATCH] FET-118 Implement trust mapping

---
 connector/hsdp/extend_payload.go |  6 ++++
 connector/hsdp/hsdp.go           | 55 ++++++++++++++++++--------------
 connector/hsdp/hsdp_test.go      |  3 ++
 3 files changed, 40 insertions(+), 24 deletions(-)

diff --git a/connector/hsdp/extend_payload.go b/connector/hsdp/extend_payload.go
index 3d99c3cef6..0e1d6a0cf6 100644
--- a/connector/hsdp/extend_payload.go
+++ b/connector/hsdp/extend_payload.go
@@ -21,6 +21,12 @@ func (c *HSDPConnector) ExtendPayload(scopes []string, payload []byte, cdata []b
 
 	c.logger.Info("ExtendPayload called for user: ", cd.Introspect.Username)
 
+	// Check if we have a trusted org mapping
+	aud := originalClaims["aud"].(string)
+	if orgID, ok := c.audienceTrustMap[aud]; ok {
+		trustedOrgID = orgID
+	}
+
 	for _, scope := range scopes {
 		// Experimental fill introspect body into claims
 		if scope == "hsp:iam:introspect" {
diff --git a/connector/hsdp/hsdp.go b/connector/hsdp/hsdp.go
index e5d30516c7..80f2682ab9 100644
--- a/connector/hsdp/hsdp.go
+++ b/connector/hsdp/hsdp.go
@@ -24,15 +24,16 @@ import (
 
 // Config holds configuration options for OpenID Connect logins.
 type Config struct {
-	Issuer         string `json:"issuer"`
-	InsecureIssuer string `json:"insecureIssuer"`
-	ClientID       string `json:"clientID"`
-	ClientSecret   string `json:"clientSecret"`
-	RedirectURI    string `json:"redirectURI"`
-	TrustedOrgID   string `json:"trustedOrgID"`
-	SAML2LoginURL  string `json:"saml2LoginURL"`
-	IAMURL         string `json:"iamURL"`
-	IDMURL         string `json:"idmURL"`
+	Issuer           string           `json:"issuer"`
+	InsecureIssuer   string           `json:"insecureIssuer"`
+	ClientID         string           `json:"clientID"`
+	ClientSecret     string           `json:"clientSecret"`
+	RedirectURI      string           `json:"redirectURI"`
+	TrustedOrgID     string           `json:"trustedOrgID"`
+	AudienceTrustMap AudienceTrustMap `json:"audienceTrustMap"`
+	SAML2LoginURL    string           `json:"saml2LoginURL"`
+	IAMURL           string           `json:"iamURL"`
+	IDMURL           string           `json:"idmURL"`
 
 	// Extensions implemented by HSP IAM
 	Extension
@@ -66,15 +67,18 @@ type Extension struct {
 	IntrospectionEndpoint string `json:"introspection_endpoint"`
 }
 
+type AudienceTrustMap map[string]string
+
 // ConnectorData stores information for sessions authenticated by this connector
 type ConnectorData struct {
-	RefreshToken  []byte
-	AccessToken   []byte
-	Assertion     []byte
-	Groups        []string
-	TrustedIDPOrg string
-	Introspect    iam.IntrospectResponse
-	User          iam.Profile
+	RefreshToken     []byte
+	AccessToken      []byte
+	Assertion        []byte
+	Groups           []string
+	TrustedIDPOrg    string
+	AudienceTrustMap AudienceTrustMap
+	Introspect       iam.IntrospectResponse
+	User             iam.Profile
 }
 
 // Open returns a connector which can be used to log in users through an upstream
@@ -129,14 +133,15 @@ func (c *Config) Open(id string, logger log.Logger) (conn connector.Connector, e
 
 	clientID := c.ClientID
 	return &HSDPConnector{
-		provider:      provider,
-		client:        client,
-		redirectURI:   c.RedirectURI,
-		introspectURI: c.IntrospectionEndpoint,
-		trustedOrgID:  c.TrustedOrgID,
-		samlLoginURL:  c.SAML2LoginURL,
-		clientID:      c.ClientID,
-		clientSecret:  c.ClientSecret,
+		provider:         provider,
+		client:           client,
+		redirectURI:      c.RedirectURI,
+		introspectURI:    c.IntrospectionEndpoint,
+		trustedOrgID:     c.TrustedOrgID,
+		audienceTrustMap: c.AudienceTrustMap,
+		samlLoginURL:     c.SAML2LoginURL,
+		clientID:         c.ClientID,
+		clientSecret:     c.ClientSecret,
 		oauth2Config: &oauth2.Config{
 			ClientID:     clientID,
 			ClientSecret: c.ClientSecret,
@@ -192,6 +197,7 @@ type HSDPConnector struct {
 	tenantGroups              []string
 	insecureSkipEmailVerified bool
 	promptType                string
+	audienceTrustMap          AudienceTrustMap
 }
 
 func (c *HSDPConnector) isSAML() bool {
@@ -413,6 +419,7 @@ func (c *HSDPConnector) createIdentity(ctx context.Context, identity connector.I
 		cd.Groups = identity.Groups
 	}
 	cd.TrustedIDPOrg = trustedOrgID
+	cd.AudienceTrustMap = c.audienceTrustMap
 
 	// Attach connector data
 	connData, err := json.Marshal(&cd)
diff --git a/connector/hsdp/hsdp_test.go b/connector/hsdp/hsdp_test.go
index 9245b84415..0ef7ba0fe3 100644
--- a/connector/hsdp/hsdp_test.go
+++ b/connector/hsdp/hsdp_test.go
@@ -78,6 +78,9 @@ func TestHandleCallback(t *testing.T) {
 				RedirectURI:          fmt.Sprintf("%s/callback", serverURL),
 				BasicAuthUnsupported: &basicAuth,
 				TenantGroups:         []string{"logreaders"},
+				AudienceTrustMap: map[string]string{
+					"clientID": "tenantID",
+				},
 			}
 
 			conn, err := newConnector(config)