From e14ddf08c3b7356b4ea23ff641ba21e6b24d7e00 Mon Sep 17 00:00:00 2001
From: Steven Silvester <steven.silvester@ieee.org>
Date: Wed, 11 Sep 2024 12:09:26 -0500
Subject: [PATCH 1/2] GODRIVER-3331 Fix default authSource for SRV connections
 (#1795)

(cherry picked from commit c5b97054ba037f040d9a788395d09d47597e23b8)
---
 .evergreen/config.yml                        |  9 +++++---
 mongo/options/clientoptions_test.go          | 14 ++++++++++++-
 x/mongo/driver/connstring/connstring.go      |  4 ++++
 x/mongo/driver/connstring/connstring_test.go | 22 ++++++++++++++++++++
 4 files changed, 45 insertions(+), 4 deletions(-)

diff --git a/.evergreen/config.yml b/.evergreen/config.yml
index ed5b5ae242..944ab2153d 100644
--- a/.evergreen/config.yml
+++ b/.evergreen/config.yml
@@ -1715,7 +1715,7 @@ tasks:
   - name: "testgcpkms-task"
     commands:
     - command: shell.exec
-      type: setup
+      type: test
       params:
         shell: "bash"
         working_dir: src/go.mongodb.org/mongo-driver
@@ -1796,7 +1796,7 @@ tasks:
   - name: "testazurekms-task"
     commands:
     - command: shell.exec
-      type: setup
+      type: test
       params:
         shell: "bash"
         working_dir: src/go.mongodb.org/mongo-driver
@@ -1862,6 +1862,7 @@ tasks:
           role_arn: ${LAMBDA_AWS_ROLE_ARN}
           duration_seconds: 3600
       - command: shell.exec
+        type: test
         params:
           working_dir: src/go.mongodb.org/mongo-driver
           shell: bash
@@ -1884,6 +1885,7 @@ tasks:
   - name: "oidc-auth-test-azure"
     commands:
     - command: shell.exec
+      type: test
       params:
         working_dir: src/go.mongodb.org/mongo-driver
         shell: bash
@@ -1909,6 +1911,7 @@ tasks:
   - name: "oidc-auth-test-gcp"
     commands:
     - command: shell.exec
+      type: test
       params:
         working_dir: src/go.mongodb.org/mongo-driver
         shell: bash
@@ -2604,7 +2607,7 @@ buildvariants:
   - name: testoidc-variant
     display_name: "OIDC"
     run_on:
-      - ubuntu2204-large
+      - ubuntu2204-small
     expansions:
       GO_DIST: "/opt/golang/go1.22"
     tasks:
diff --git a/mongo/options/clientoptions_test.go b/mongo/options/clientoptions_test.go
index 527b95468f..cc4f8e3e61 100644
--- a/mongo/options/clientoptions_test.go
+++ b/mongo/options/clientoptions_test.go
@@ -1285,7 +1285,7 @@ func TestSetURIopts(t *testing.T) {
 			wantErrs: nil,
 		},
 		{
-			name: "tmp",
+			name: "oidc azure",
 			uri:  "mongodb://example.com/?authMechanism=MONGODB-OIDC&authMechanismProperties=TOKEN_RESOURCE:mongodb://test-cluster,ENVIRONMENT:azureManagedIdentities",
 			wantopts: &ClientOptions{
 				Hosts: []string{"example.com"},
@@ -1296,6 +1296,18 @@ func TestSetURIopts(t *testing.T) {
 			},
 			wantErrs: nil,
 		},
+		{
+			name: "oidc gcp",
+			uri:  "mongodb://test.mongodb.net/?authMechanism=MONGODB-OIDC&authMechanismProperties=ENVIRONMENT:gcp,TOKEN_RESOURCE:mongodb://test-cluster",
+			wantopts: &ClientOptions{
+				Hosts: []string{"test.mongodb.net"},
+				Auth: &Credential{AuthMechanism: "MONGODB-OIDC", AuthSource: "$external", AuthMechanismProperties: map[string]string{
+					"ENVIRONMENT":    "gcp",
+					"TOKEN_RESOURCE": "mongodb://test-cluster"}},
+				HTTPClient: httputil.DefaultHTTPClient,
+			},
+			wantErrs: nil
+		},
 		{
 			name: "comma in key:value pair causes error",
 			uri:  "mongodb://example.com/?authMechanismProperties=TOKEN_RESOURCE:mongodb://host1%2Chost2",
diff --git a/x/mongo/driver/connstring/connstring.go b/x/mongo/driver/connstring/connstring.go
index ece143a1e4..d45087e1ba 100644
--- a/x/mongo/driver/connstring/connstring.go
+++ b/x/mongo/driver/connstring/connstring.go
@@ -292,6 +292,10 @@ func (u *ConnString) setDefaultAuthParams(dbName string) error {
 		}
 		fallthrough
 	case "mongodb-aws", "mongodb-x509", "mongodb-oidc":
+		// dns.LookupTXT will get "authSource=admin" from Atlas hosts.
+		if u.AuthSource == "admin" {
+			u.AuthSource = "$external"
+		}
 		if u.AuthSource == "" {
 			u.AuthSource = "$external"
 		} else if u.AuthSource != "$external" {
diff --git a/x/mongo/driver/connstring/connstring_test.go b/x/mongo/driver/connstring/connstring_test.go
index 3a3e68a636..b203114f3e 100644
--- a/x/mongo/driver/connstring/connstring_test.go
+++ b/x/mongo/driver/connstring/connstring_test.go
@@ -90,6 +90,28 @@ func TestAuthSource(t *testing.T) {
 			}
 		})
 	}
+
+	tests = []struct {
+		s        string
+		expected string
+		err      bool
+	}{
+		{s: "authMechanismProperties=ENVIRONMENT:gcp,TOKEN_RESOURCE:mongodb://test-cluster", expected: "$external"},
+	}
+
+	for _, test := range tests {
+		s := fmt.Sprintf("mongodb://test.mongodb.net/?authMechanism=MONGODB-OIDC&/%s", test.s)
+		t.Run(s, func(t *testing.T) {
+			cs, err := connstring.ParseAndValidate(s)
+			if test.err {
+				require.Error(t, err)
+			} else {
+				require.NoError(t, err)
+				require.Equal(t, test.expected, cs.AuthSource)
+			}
+		})
+	}
+
 }
 
 func TestConnect(t *testing.T) {

From 740aef01c229ba7256db75d76ae90a78f3da9a1b Mon Sep 17 00:00:00 2001
From: Steven Silvester <steve.silvester@mongodb.com>
Date: Wed, 11 Sep 2024 15:12:27 -0500
Subject: [PATCH 2/2] lint

---
 mongo/options/clientoptions_test.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mongo/options/clientoptions_test.go b/mongo/options/clientoptions_test.go
index cc4f8e3e61..c235b0e6e6 100644
--- a/mongo/options/clientoptions_test.go
+++ b/mongo/options/clientoptions_test.go
@@ -1306,7 +1306,7 @@ func TestSetURIopts(t *testing.T) {
 					"TOKEN_RESOURCE": "mongodb://test-cluster"}},
 				HTTPClient: httputil.DefaultHTTPClient,
 			},
-			wantErrs: nil
+			wantErrs: nil,
 		},
 		{
 			name: "comma in key:value pair causes error",