SERVER-70703: Implement OIDC refresh flow in legacy shell

7 files changed, 174 insertions, 43 deletions

diff --git a/src/mongo/client/sasl_oidc_client_conversation.cpp b/src/mongo/client/sasl_oidc_client_conversation.cpp
index 762fef2d324..c97d9330814 100644
--- a/src/mongo/client/sasl_oidc_client_conversation.cpp
+++ b/src/mongo/client/sasl_oidc_client_conversation.cpp
@@ -32,11 +32,12 @@
 #include "mongo/client/sasl_oidc_client_conversation.h"
 
 #include "mongo/base/data_range.h"
-#include "mongo/base/data_type_validated.h"
 #include "mongo/bson/json.h"
 #include "mongo/client/mongo_uri.h"
 #include "mongo/client/sasl_client_session.h"
+#include "mongo/client/sasl_oidc_client_params_gen.h"
 #include "mongo/db/auth/oidc_protocol_gen.h"
+#include "mongo/rpc/object_check.h"
 #include "mongo/shell/program_runner.h"
 #include "mongo/util/net/http_client.h"
 
@@ -48,37 +49,40 @@ constexpr auto kRequestScopesParameterName = "scope"_sd;
 constexpr auto kGrantTypeParameterName = "grant_type"_sd;
 constexpr auto kGrantTypeParameterDeviceCodeValue =
     "urn:ietf:params:oauth:grant-type:device_code"_sd;
+constexpr auto kGrantTypeParameterRefreshTokenValue = "refresh_token"_sd;
 constexpr auto kDeviceCodeParameterName = "device_code"_sd;
-
-std::string buildPostBody(StringData clientId,
-                          const boost::optional<StringData>& clientSecret,
-                          const boost::optional<std::vector<StringData>>& requestScopes,
-                          const boost::optional<std::string>& deviceCode) {
-    StringBuilder sb;
-    sb << kClientIdParameterName << "=" << uriEncode(clientId);
-
-    if (clientSecret && !clientSecret->empty()) {
-        sb << "&" << kClientSecretParameterName << "=" << uriEncode(clientSecret.get());
+constexpr auto kRefreshTokenParameterName = kGrantTypeParameterRefreshTokenValue;
+
+inline void appendPostBodyRequiredParams(StringBuilder* sb,
+                                         StringData clientId,
+                                         const boost::optional<StringData>& clientSecret) {
+    *sb << kClientIdParameterName << "=" << uriEncode(clientId);
+    if (clientSecret) {
+        *sb << "&" << kClientSecretParameterName << "=" << uriEncode(clientSecret->toString());
     }
+}
 
-    if (requestScopes && requestScopes->size() > 0) {
-        sb << "&" << kRequestScopesParameterName << "=";
+inline void appendPostBodyDeviceCodeRequestParams(
+    StringBuilder* sb, const boost::optional<std::vector<StringData>>& requestScopes) {
+    if (requestScopes) {
+        *sb << "&" << kRequestScopesParameterName << "=";
         for (std::size_t i = 0; i < requestScopes->size(); i++) {
-            sb << uriEncode(requestScopes.get()[i]);
+            *sb << uriEncode(requestScopes.get()[i]);
             if (i < requestScopes->size() - 1) {
-                sb << uriEncode(" ");
+                *sb << "%20";
             }
         }
     }
+}
 
-    if (deviceCode && !deviceCode->empty()) {
-        // If the device code is provided, the request must explicitly specify the grant type as
-        // device code.
-        sb << "&" << kGrantTypeParameterName << "=" << kGrantTypeParameterDeviceCodeValue;
-        sb << "&" << kDeviceCodeParameterName << "=" << uriEncode(deviceCode.get());
-    }
+inline void appendPostBodyTokenRequestParams(StringBuilder* sb, StringData deviceCode) {
+    *sb << "&" << kGrantTypeParameterName << "=" << kGrantTypeParameterDeviceCodeValue << "&"
+        << kDeviceCodeParameterName << "=" << uriEncode(deviceCode);
+}
 
-    return sb.str();
+inline void appendPostBodyRefreshFlowParams(StringBuilder* sb, StringData refreshToken) {
+    *sb << "&" << kGrantTypeParameterName << "=" << kGrantTypeParameterRefreshTokenValue << "&"
+        << kRefreshTokenParameterName << "=" << uriEncode(refreshToken);
 }
 
 BSONObj doPostRequest(HttpClient* httpClient, StringData endPoint, const std::string& requestBody) {
@@ -101,12 +105,15 @@ std::pair<std::string, std::string> doDeviceAuthorizationGrantFlow(
     auto clientId = serverReply.getClientId();
     uassert(ErrorCodes::BadValue, "Encountered empty client ID in server reply", !clientId.empty());
 
+    // Cache clientId for potential refresh flow uses in the future.
+    oidcClientGlobalParams.oidcClientId = clientId.toString();
+
     // Construct body of POST request to device authorization endpoint based on provided
     // parameters.
-    auto deviceCodeRequest = buildPostBody(clientId,
-                                           serverReply.getClientSecret(),
-                                           serverReply.getRequestScopes(),
-                                           boost::none /* deviceCode */);
+    StringBuilder deviceCodeRequestSb;
+    appendPostBodyRequiredParams(&deviceCodeRequestSb, clientId, serverReply.getClientSecret());
+    appendPostBodyDeviceCodeRequestParams(&deviceCodeRequestSb, serverReply.getRequestScopes());
+    auto deviceCodeRequest = deviceCodeRequestSb.str();
 
     // Retrieve device code and user verification URI from IdP.
     auto httpClient = HttpClient::createWithoutConnectionPool();
@@ -116,35 +123,39 @@ std::pair<std::string, std::string> doDeviceAuthorizationGrantFlow(
         doPostRequest(httpClient.get(), deviceAuthorizationEndpoint, deviceCodeRequest);
 
     // Simulate end user login via user verification URI.
-    auto deviceCode = deviceAuthorizationResponseObj["device_code"_sd].String();
-    auto activationEndpoint =
-        deviceAuthorizationResponseObj["verification_uri_complete"_sd].String();
-    oidcClientGlobalParams.oidcIdPAuthCallback(principalName, activationEndpoint);
+    auto deviceAuthorizationResponse = OIDCDeviceAuthorizationResponse::parse(
+        IDLParserContext{"oidcDeviceAuthorizationResponse"}, deviceAuthorizationResponseObj);
+    oidcClientGlobalParams.oidcIdPAuthCallback(
+        principalName, deviceAuthorizationResponse.getVerificationUriComplete());
 
     // Poll token endpoint for access and refresh tokens. It should return immediately since
     // the shell blocks on the authenticationSimulator until it completes, but poll anyway.
-    auto tokenRequest = buildPostBody(
-        clientId, serverReply.getClientSecret(), boost::none /* requestScopes */, deviceCode);
+    StringBuilder tokenRequestSb;
+    appendPostBodyRequiredParams(&tokenRequestSb, clientId, serverReply.getClientSecret());
+    appendPostBodyTokenRequestParams(&tokenRequestSb, deviceAuthorizationResponse.getDeviceCode());
+    auto tokenRequest = tokenRequestSb.str();
 
     while (true) {
         BSONObj tokenResponseObj =
             doPostRequest(httpClient.get(), serverReply.getTokenEndpoint(), tokenRequest);
+        auto tokenResponse =
+            OIDCTokenResponse::parse(IDLParserContext{"oidcTokenResponse"}, tokenResponseObj);
 
         // The token endpoint will either respond with the tokens or {"error":
         // "authorization pending"}.
-        bool hasAccessToken = tokenResponseObj.hasField("access_token"_sd);
-        bool hasError = tokenResponseObj.hasField("error"_sd);
+        bool hasAccessToken = tokenResponse.getAccessToken().has_value();
+        bool hasError = tokenResponse.getError().has_value();
         uassert(ErrorCodes::UnknownError,
                 fmt::format("Received unrecognized reply from token endpoint: {}",
                             tokenResponseObj.toString()),
                 hasAccessToken || hasError);
 
         if (hasAccessToken) {
-            auto accessToken = tokenResponseObj["access_token"_sd].String();
+            auto accessToken = tokenResponse.getAccessToken()->toString();
 
             // If a refresh token was also provided, cache that as well.
-            if (tokenResponseObj.hasField("refresh_token"_sd)) {
-                return {accessToken, tokenResponseObj["refresh_token"_sd].String()};
+            if (tokenResponse.getRefreshToken()) {
+                return {accessToken, tokenResponse.getRefreshToken()->toString()};
             }
 
             return {accessToken, ""};
@@ -153,10 +164,10 @@ std::pair<std::string, std::string> doDeviceAuthorizationGrantFlow(
         // Assert that the error returned with "authorization pending", which indicates that
         // the token endpoint has not perceived end-user authentication yet and we should
         // poll again.
+        auto error = tokenResponse.getError()->toString();
         uassert(ErrorCodes::UnknownError,
-                fmt::format("Received unexpected error from token endpoint: {}",
-                            tokenResponseObj["error"_sd].String()),
-                tokenResponseObj["error"_sd].String() != "authorization pending");
+                fmt::format("Received unexpected error from token endpoint: {}", error),
+                error == "authorization pending");
     }
 
     MONGO_UNREACHABLE
@@ -184,6 +195,44 @@ StatusWith<bool> SaslOIDCClientConversation::step(StringData inputData, std::str
     }
 }
 
+StatusWith<std::string> SaslOIDCClientConversation::doRefreshFlow() try {
+    // The refresh flow can only be performed if a successful auth attempt has already occurred.
+    uassert(ErrorCodes::IllegalOperation,
+            "Cannot perform refresh flow without previously-successful auth attempt",
+            !oidcClientGlobalParams.oidcRefreshToken.empty() &&
+                !oidcClientGlobalParams.oidcClientId.empty() &&
+                !oidcClientGlobalParams.oidcTokenEndpoint.empty());
+
+    StringBuilder refreshFlowRequestBuilder;
+    appendPostBodyRequiredParams(&refreshFlowRequestBuilder,
+                                 oidcClientGlobalParams.oidcClientId,
+                                 StringData(oidcClientGlobalParams.oidcClientSecret));
+    appendPostBodyRefreshFlowParams(&refreshFlowRequestBuilder,
+                                    oidcClientGlobalParams.oidcRefreshToken);
+
+    auto refreshFlowRequestBody = refreshFlowRequestBuilder.str();
+
+    auto httpClient = HttpClient::createWithoutConnectionPool();
+    httpClient->setHeaders(
+        {"Accept: application/json", "Content-Type: application/x-www-form-urlencoded"});
+    BSONObj refreshFlowResponseObj = doPostRequest(
+        httpClient.get(), oidcClientGlobalParams.oidcTokenEndpoint, refreshFlowRequestBody);
+    auto refreshResponse =
+        OIDCTokenResponse::parse(IDLParserContext{"oidcRefreshResponse"}, refreshFlowResponseObj);
+
+    // New tokens should be supplied immediately.
+    uassert(ErrorCodes::UnknownError,
+            "Failed to retrieve refreshed access token",
+            refreshResponse.getAccessToken());
+    if (refreshResponse.getRefreshToken()) {
+        oidcClientGlobalParams.oidcRefreshToken = refreshResponse.getRefreshToken()->toString();
+    }
+
+    return refreshResponse.getAccessToken()->toString();
+} catch (const DBException& ex) {
+    return ex.toStatus();
+}
+
 StatusWith<bool> SaslOIDCClientConversation::_firstStep(std::string* outputData) {
     // If an access token was provided without a username, proceed to the second step and send it
     // directly to the server.
@@ -230,6 +279,9 @@ StatusWith<bool> SaslOIDCClientConversation::_secondStep(StringData input,
                     (tokenEndpoint.startsWith("https://"_sd) ||
                      tokenEndpoint.startsWith("http://localhost"_sd)));
 
+        // Cache the token endpoint for potential reuse during the refresh flow.
+        oidcClientGlobalParams.oidcTokenEndpoint = tokenEndpoint.toString();
+
         // Try device authorization grant flow first if provided, falling back to authorization code
         // flow.
         if (serverReply.getDeviceAuthorizationEndpoint()) {
diff --git a/src/mongo/client/sasl_oidc_client_conversation.h b/src/mongo/client/sasl_oidc_client_conversation.h
index f44674e4dff..da234b43391 100644
--- a/src/mongo/client/sasl_oidc_client_conversation.h
+++ b/src/mongo/client/sasl_oidc_client_conversation.h
@@ -53,6 +53,10 @@ public:
 
     StatusWith<bool> step(StringData inputData, std::string* outputData) override;
 
+    // Refreshes oidcClientGlobalParams.accessToken using oidcClientGlobalParams.refreshToken,
+    // returning the acquired access token if successful.
+    static StatusWith<std::string> doRefreshFlow();
+
 private:
     // Step of the conversation - can be 1, 2, or 3.
     int _step{0};
diff --git a/src/mongo/client/sasl_oidc_client_params.h b/src/mongo/client/sasl_oidc_client_params.h
index fe28b8934c1..994b60944d6 100644
--- a/src/mongo/client/sasl_oidc_client_params.h
+++ b/src/mongo/client/sasl_oidc_client_params.h
@@ -31,18 +31,22 @@
 
 #include "mongo/base/string_data.h"
 
+#include <boost/optional/optional.hpp>
+
+#include "mongo/base/string_data.h"
+
 namespace mongo {
 /**
  * OIDC Client parameters
  */
 struct OIDCClientGlobalParams {
     /**
-     * Access Token.
+     * Access Token. Populated either by configuration or token acquisition flow.
      */
     std::string oidcAccessToken;
 
     /**
-     * Refresh Token.
+     * Refresh Token. Populated during token acquisition flow.
      */
     std::string oidcRefreshToken;
 
@@ -51,6 +55,20 @@ struct OIDCClientGlobalParams {
      * authentication. This should be provided by tests, presumably as a JS function.
      */
     std::function<void(StringData, StringData)> oidcIdPAuthCallback;
+    /**
+     * Client ID. Populated via server SASL reply.
+     */
+    std::string oidcClientId;
+
+    /**
+     * Client Secret. Populated via server SASL reply.
+     */
+    std::string oidcClientSecret;
+
+    /**
+     * Token endpoint. Populated via server SASL reply.
+     */
+    std::string oidcTokenEndpoint;
 };
 
 extern OIDCClientGlobalParams oidcClientGlobalParams;
diff --git a/src/mongo/client/sasl_oidc_client_params.idl b/src/mongo/client/sasl_oidc_client_params.idl
index af073911e2d..0ea16c025ca 100644
--- a/src/mongo/client/sasl_oidc_client_params.idl
+++ b/src/mongo/client/sasl_oidc_client_params.idl
@@ -19,3 +19,36 @@ configs:
             grant flow.
         arg_vartype: String
         cpp_varname: oidcClientGlobalParams.oidcAccessToken
+
+structs:
+    OIDCDeviceAuthorizationResponse:
+        description: "IdP response from the deviceAuthorization endpoint."
+        strict: false
+        fields:
+            device_code:
+                description: "Device code to use in token request"
+                cpp_name: deviceCode
+                type: string
+            verification_uri_complete:
+                description: "URI for end user authentication"
+                cpp_name: verificationUriComplete
+                type: string
+
+    OIDCTokenResponse:
+        description: IdP response from the token endpoint.
+        strict: false
+        fields:
+            access_token:
+                description: "Access token returned to be sent to the server."
+                cpp_name: accessToken
+                type: string
+                optional: true
+            refresh_token:
+                description: "Refresh token returned to be used for token reacquisition."
+                cpp_name: refreshToken
+                type: string
+                optional: true
+            error:
+                description: "Error message returned by the token endpoint."
+                type: string
+                optional: true
diff --git a/src/mongo/scripting/mozjs/mongo.cpp b/src/mongo/scripting/mozjs/mongo.cpp
index 4dd98cfa891..876f977e2e2 100644
--- a/src/mongo/scripting/mozjs/mongo.cpp
+++ b/src/mongo/scripting/mozjs/mongo.cpp
@@ -82,6 +82,7 @@ const JSFunctionSpec MongoBase::methods[] = {
     MONGO_ATTACH_JS_CONSTRAINED_METHOD_NO_PROTO(_runCommandImpl, MongoExternalInfo),
     MONGO_ATTACH_JS_CONSTRAINED_METHOD_NO_PROTO(_startSession, MongoExternalInfo),
     MONGO_ATTACH_JS_CONSTRAINED_METHOD_NO_PROTO(_setOIDCIdPAuthCallback, MongoExternalInfo),
+    MONGO_ATTACH_JS_CONSTRAINED_METHOD_NO_PROTO(_refreshAccessToken, MongoExternalInfo),
     JS_FS_END,
 };
 
@@ -659,6 +660,11 @@ void MongoBase::Functions::_setOIDCIdPAuthCallback::call(JSContext* cx, JS::Call
     args.rval().setUndefined();
 }
 
+void MongoBase::Functions::_refreshAccessToken::call(JSContext* cx, JS::CallArgs args) {
+    auto accessToken = uassertStatusOK(SaslOIDCClientConversation::doRefreshFlow());
+    ValueReader(cx, args.rval()).fromStringData(accessToken);
+}
+
 void MongoExternalInfo::Functions::load::call(JSContext* cx, JS::CallArgs args) {
     auto scope = getScope(cx);
 
diff --git a/src/mongo/scripting/mozjs/mongo.h b/src/mongo/scripting/mozjs/mongo.h
index e9c47a0de84..8a73734ac6e 100644
--- a/src/mongo/scripting/mozjs/mongo.h
+++ b/src/mongo/scripting/mozjs/mongo.h
@@ -80,9 +80,10 @@ struct MongoBase : public BaseInfo {
         MONGO_DECLARE_JS_FUNCTION(_runCommandImpl);
         MONGO_DECLARE_JS_FUNCTION(_startSession);
         MONGO_DECLARE_JS_FUNCTION(_setOIDCIdPAuthCallback);
+        MONGO_DECLARE_JS_FUNCTION(_refreshAccessToken);
     };
 
-    static const JSFunctionSpec methods[22];
+    static const JSFunctionSpec methods[23];
 
     static const char* const className;
     static const unsigned classFlags = JSCLASS_HAS_PRIVATE;
diff --git a/src/mongo/shell/session.js b/src/mongo/shell/session.js
index 9aefeca7885..b077301f0cd 100644
--- a/src/mongo/shell/session.js
+++ b/src/mongo/shell/session.js
@@ -362,6 +362,23 @@ var {
                     }
                 }
 
+                // Handle ErrorCodes.Reauthentication first.
+                if (res !== undefined && res.code === ErrorCodes.ReauthenticationRequired) {
+                    try {
+                        const accessToken = client._refreshAccessToken();
+                        assert(client.getDB('$external').auth({
+                            oidcAccessToken: accessToken,
+                            mechanism: 'MONGODB-OIDC'
+                        }));
+                        continue;
+                    } catch (e) {
+                        // Could not automatically reauthenticate, return the error response
+                        // as-is.
+                        jsTest.log('Assertion thrown when performing refresh flow: ' + e);
+                        return res;
+                    }
+                }
+
                 if (numRetries > 0) {
                     --numRetries;