From 2466e65f43f61fa87fce0069cc06defe94f1e856 Mon Sep 17 00:00:00 2001
From: RobKenis <rob.kenis@hotmail.com>
Date: Sat, 12 Oct 2024 12:52:47 +0200
Subject: [PATCH] support multiple subjects in oidc ping (#4475)

Resolves: #4466
---
 pkg/auth/auth.go      |  3 +-
 pkg/auth/oidc.go      | 27 ++++++++++++------
 pkg/auth/oidc_test.go | 64 +++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 85 insertions(+), 9 deletions(-)
 create mode 100644 pkg/auth/oidc_test.go

diff --git a/pkg/auth/auth.go b/pkg/auth/auth.go
index 9d8db7b6..ae706986 100644
--- a/pkg/auth/auth.go
+++ b/pkg/auth/auth.go
@@ -50,7 +50,8 @@ func NewAuthVerifier(cfg v1.AuthServerConfig) (authVerifier Verifier) {
 	case v1.AuthMethodToken:
 		authVerifier = NewTokenAuth(cfg.AdditionalScopes, cfg.Token)
 	case v1.AuthMethodOIDC:
-		authVerifier = NewOidcAuthVerifier(cfg.AdditionalScopes, cfg.OIDC)
+		tokenVerifier := NewTokenVerifier(cfg.OIDC)
+		authVerifier = NewOidcAuthVerifier(cfg.AdditionalScopes, tokenVerifier)
 	}
 	return authVerifier
 }
diff --git a/pkg/auth/oidc.go b/pkg/auth/oidc.go
index d87420ff..40ce060f 100644
--- a/pkg/auth/oidc.go
+++ b/pkg/auth/oidc.go
@@ -87,14 +87,18 @@ func (auth *OidcAuthProvider) SetNewWorkConn(newWorkConnMsg *msg.NewWorkConn) (e
 	return err
 }
 
+type TokenVerifier interface {
+	Verify(context.Context, string) (*oidc.IDToken, error)
+}
+
 type OidcAuthConsumer struct {
 	additionalAuthScopes []v1.AuthScope
 
-	verifier         *oidc.IDTokenVerifier
-	subjectFromLogin string
+	verifier          TokenVerifier
+	subjectsFromLogin []string
 }
 
-func NewOidcAuthVerifier(additionalAuthScopes []v1.AuthScope, cfg v1.AuthOIDCServerConfig) *OidcAuthConsumer {
+func NewTokenVerifier(cfg v1.AuthOIDCServerConfig) TokenVerifier {
 	provider, err := oidc.NewProvider(context.Background(), cfg.Issuer)
 	if err != nil {
 		panic(err)
@@ -105,9 +109,14 @@ func NewOidcAuthVerifier(additionalAuthScopes []v1.AuthScope, cfg v1.AuthOIDCSer
 		SkipExpiryCheck:   cfg.SkipExpiryCheck,
 		SkipIssuerCheck:   cfg.SkipIssuerCheck,
 	}
+	return provider.Verifier(&verifierConf)
+}
+
+func NewOidcAuthVerifier(additionalAuthScopes []v1.AuthScope, verifier TokenVerifier) *OidcAuthConsumer {
 	return &OidcAuthConsumer{
 		additionalAuthScopes: additionalAuthScopes,
-		verifier:             provider.Verifier(&verifierConf),
+		verifier:             verifier,
+		subjectsFromLogin:    []string{},
 	}
 }
 
@@ -116,7 +125,9 @@ func (auth *OidcAuthConsumer) VerifyLogin(loginMsg *msg.Login) (err error) {
 	if err != nil {
 		return fmt.Errorf("invalid OIDC token in login: %v", err)
 	}
-	auth.subjectFromLogin = token.Subject
+	if !slices.Contains(auth.subjectsFromLogin, token.Subject) {
+		auth.subjectsFromLogin = append(auth.subjectsFromLogin, token.Subject)
+	}
 	return nil
 }
 
@@ -125,11 +136,11 @@ func (auth *OidcAuthConsumer) verifyPostLoginToken(privilegeKey string) (err err
 	if err != nil {
 		return fmt.Errorf("invalid OIDC token in ping: %v", err)
 	}
-	if token.Subject != auth.subjectFromLogin {
+	if !slices.Contains(auth.subjectsFromLogin, token.Subject) {
 		return fmt.Errorf("received different OIDC subject in login and ping. "+
-			"original subject: %s, "+
+			"original subjects: %s, "+
 			"new subject: %s",
-			auth.subjectFromLogin, token.Subject)
+			auth.subjectsFromLogin, token.Subject)
 	}
 	return nil
 }
diff --git a/pkg/auth/oidc_test.go b/pkg/auth/oidc_test.go
new file mode 100644
index 00000000..58054186
--- /dev/null
+++ b/pkg/auth/oidc_test.go
@@ -0,0 +1,64 @@
+package auth_test
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/coreos/go-oidc/v3/oidc"
+	"github.com/stretchr/testify/require"
+
+	"github.com/fatedier/frp/pkg/auth"
+	v1 "github.com/fatedier/frp/pkg/config/v1"
+	"github.com/fatedier/frp/pkg/msg"
+)
+
+type mockTokenVerifier struct{}
+
+func (m *mockTokenVerifier) Verify(ctx context.Context, subject string) (*oidc.IDToken, error) {
+	return &oidc.IDToken{
+		Subject: subject,
+	}, nil
+}
+
+func TestPingWithEmptySubjectFromLoginFails(t *testing.T) {
+	r := require.New(t)
+	consumer := auth.NewOidcAuthVerifier([]v1.AuthScope{v1.AuthScopeHeartBeats}, &mockTokenVerifier{})
+	err := consumer.VerifyPing(&msg.Ping{
+		PrivilegeKey: "ping-without-login",
+		Timestamp:    time.Now().UnixMilli(),
+	})
+	r.Error(err)
+	r.Contains(err.Error(), "received different OIDC subject in login and ping")
+}
+
+func TestPingAfterLoginWithNewSubjectSucceeds(t *testing.T) {
+	r := require.New(t)
+	consumer := auth.NewOidcAuthVerifier([]v1.AuthScope{v1.AuthScopeHeartBeats}, &mockTokenVerifier{})
+	err := consumer.VerifyLogin(&msg.Login{
+		PrivilegeKey: "ping-after-login",
+	})
+	r.NoError(err)
+
+	err = consumer.VerifyPing(&msg.Ping{
+		PrivilegeKey: "ping-after-login",
+		Timestamp:    time.Now().UnixMilli(),
+	})
+	r.NoError(err)
+}
+
+func TestPingAfterLoginWithDifferentSubjectFails(t *testing.T) {
+	r := require.New(t)
+	consumer := auth.NewOidcAuthVerifier([]v1.AuthScope{v1.AuthScopeHeartBeats}, &mockTokenVerifier{})
+	err := consumer.VerifyLogin(&msg.Login{
+		PrivilegeKey: "login-with-first-subject",
+	})
+	r.NoError(err)
+
+	err = consumer.VerifyPing(&msg.Ping{
+		PrivilegeKey: "ping-with-different-subject",
+		Timestamp:    time.Now().UnixMilli(),
+	})
+	r.Error(err)
+	r.Contains(err.Error(), "received different OIDC subject in login and ping")
+}