MB-65972: Add support for auto-renewal JWTs for MemcachedConnection

Allow MemcachedConnection to keep a JWT builder around and update
the JWT in use when the server returns AuthStale.

Change-Id: Ia386459210e555dca51ce90e0a2993bc41729b22
Reviewed-on: https://review.couchbase.org/c/kv_engine/+/225359
Reviewed-by: Vesko Karaganev <vesko.karaganev@couchbase.com>
Tested-by: Build Bot <build@couchbase.com>

Author: trondn

json_web_token/builder.cc

@@ -18,9 +18,12 @@
 namespace cb::jwt {
 class BuilderImpl : public Builder {
 public:
-    BuilderImpl(nlohmann::json initial_header, nlohmann::json initial_payload)
+    BuilderImpl(nlohmann::json initial_header,
+                nlohmann::json initial_payload,
+                std::optional<std::chrono::seconds> lifetime)
         : header(std::move(initial_header)),
-          payload(std::move(initial_payload)) {
+          payload(std::move(initial_payload)),
+          lifetime(lifetime) {
         if (!header.contains("typ")) {
             header["typ"] = "JWT";
         }
@@ -61,30 +64,40 @@ class BuilderImpl : public Builder {
         payload[name] = value;
     }
     std::string build() override {
+        if (lifetime.has_value()) {
+            const auto now = std::chrono::system_clock::now();
+            const auto expiration = now + lifetime.value();
+            setIssuedAt(now);
+            setExpiration(expiration);
+        }
         return fmt::format("{}.{}",
                            cb::base64url::encode(header.dump()),
                            cb::base64url::encode(payload.dump()));
     }
 
     nlohmann::json header;
     nlohmann::json payload;
+    std::optional<std::chrono::seconds> lifetime;
 };
 
 class PlainBuilderImpl : public BuilderImpl {
 public:
-    explicit PlainBuilderImpl(nlohmann::json initial)
-        : BuilderImpl({{"alg", "none"}}, std::move(initial)) {
+    explicit PlainBuilderImpl(nlohmann::json initial,
+                              std::optional<std::chrono::seconds> lifetime)
+        : BuilderImpl({{"alg", "none"}}, std::move(initial), lifetime) {
     }
 
     [[nodiscard]] std::unique_ptr<Builder> clone() const override {
-        return std::make_unique<PlainBuilderImpl>(payload);
+        return std::make_unique<PlainBuilderImpl>(payload, lifetime);
     }
 };
 
 class HS256BuilderImpl : public BuilderImpl {
 public:
-    HS256BuilderImpl(std::string passphrase, nlohmann::json initial)
-        : BuilderImpl({{"alg", "HS256"}}, std::move(initial)),
+    HS256BuilderImpl(std::string passphrase,
+                     nlohmann::json initial,
+                     std::optional<std::chrono::seconds> lifetime)
+        : BuilderImpl({{"alg", "HS256"}}, std::move(initial), lifetime),
           passphrase(std::move(passphrase)) {
     }
 
@@ -96,22 +109,25 @@ class HS256BuilderImpl : public BuilderImpl {
     }
 
     [[nodiscard]] std::unique_ptr<Builder> clone() const override {
-        return std::make_unique<HS256BuilderImpl>(passphrase, payload);
+        return std::make_unique<HS256BuilderImpl>(
+                passphrase, payload, lifetime);
     }
 
 protected:
     const std::string passphrase;
 };
 
-std::unique_ptr<Builder> Builder::create(std::string_view alg,
-                                         std::string_view passphrase,
-                                         nlohmann::json payload) {
+std::unique_ptr<Builder> Builder::create(
+        std::string_view alg,
+        std::string_view passphrase,
+        nlohmann::json payload,
+        std::optional<std::chrono::seconds> lifetime) {
     if (alg == "HS256") {
-        return std::make_unique<HS256BuilderImpl>(std::string(passphrase),
-                                                  std::move(payload));
+        return std::make_unique<HS256BuilderImpl>(
+                std::string(passphrase), std::move(payload), lifetime);
     }
     if (alg.empty() || alg == "none") {
-        return std::make_unique<PlainBuilderImpl>(std::move(payload));
+        return std::make_unique<PlainBuilderImpl>(std::move(payload), lifetime);
     }
     throw std::invalid_argument("Invalid Algorithm");
 }

json_web_token/builder.h

@@ -12,6 +12,7 @@
 #include <nlohmann/json.hpp>
 #include <chrono>
 #include <memory>
+#include <optional>
 #include <string_view>
 
 namespace cb::jwt {
@@ -33,7 +34,8 @@ class Builder {
     [[nodiscard]] static std::unique_ptr<Builder> create(
             std::string_view alg = "none",
             std::string_view passphrase = {},
-            nlohmann::json payload = nlohmann::json::object());
+            nlohmann::json payload = nlohmann::json::object(),
+            std::optional<std::chrono::seconds> lifetime = {});
     /// Set the expiration for the token to use
     virtual void setExpiration(std::chrono::system_clock::time_point exp) = 0;
     /// Specify when we can start using the token

protocol/connection/CMakeLists.txt

@@ -24,6 +24,7 @@ target_link_libraries(mc_client_connection PRIVATE
     platform
     ${COUCHBASE_NETWORK_LIBS}
   PRIVATE
+    json_web_token
     folly_io_callbacks
     cbcompress
     json_validator

protocol/connection/client_connection.cc

@@ -17,6 +17,7 @@
 #include <folly/io/IOBuf.h>
 #include <folly/io/async/AsyncSSLSocket.h>
 #include <json/syntax_validator.h>
+#include <json_web_token/builder.h>
 #include <mcbp/codec/dcp_snapshot_marker.h>
 #include <mcbp/codec/frameinfo.h>
 #include <mcbp/mcbp.h>
@@ -38,6 +39,8 @@
 #include <netdb.h>
 #include <netinet/tcp.h> // For TCP_NODELAY etc
 #endif
+#include <json_web_token/token.h>
+
 #include <stdexcept>
 #include <string>
 #include <system_error>
@@ -954,6 +957,9 @@ std::unique_ptr<MemcachedConnection> MemcachedConnection::clone(
     if (!agent_name.empty()) {
         ret->setAgentName(std::move(agent_name));
     }
+    if (tokenBuilder) {
+        ret->tokenBuilder = tokenBuilder->clone();
+    }
 
     if (connect) {
         ret->connect();
@@ -1239,6 +1245,29 @@ void MemcachedConnection::recvResponse(BinprotResponse& response,
     traceData = response.getTracingData();
 }
 
+void MemcachedConnection::setTokenBuilder(
+        std::unique_ptr<cb::jwt::Builder> builder) {
+    if (!ssl) {
+        throw std::logic_error(
+                "MemcachedConnection::setTokenBuilder: "
+                "SSL must be enabled to use JWT");
+    }
+    tokenBuilder = std::move(builder);
+}
+
+void MemcachedConnection::authenticateWithToken() {
+    auto token = tokenBuilder->build();
+
+    auto parsed = cb::jwt::Token::parse(tokenBuilder->build());
+    if (!parsed->payload.contains("sub")) {
+        throw std::logic_error(
+                "MemcachedConnection::authenticateWithToken: "
+                "The token must contain a sub field");
+    }
+
+    doSaslAuthenticate(parsed->payload["sub"], token, "OAUTHBEARER");
+}
+
 void MemcachedConnection::authenticate(
         const std::string& user,
         const std::optional<std::string>& password,
@@ -2257,6 +2286,16 @@ BinprotResponse MemcachedConnection::execute(
                 sendCommand(command);
                 recvResponse(response, command.getOp(), readTimeout);
 
+                if (response.getStatus() == cb::mcbp::Status::AuthStale &&
+                    tokenBuilder) {
+                    if (command.getOp() == cb::mcbp::ClientOpcode::SaslAuth) {
+                        // we don't want to recursively do sasl auth
+                        return true;
+                    }
+                    authenticateWithToken();
+                    return false;
+                }
+
                 bool retry_by_tmpfail =
                         auto_retry_tmpfail &&
                         response.getStatus() == cb::mcbp::Status::Etmpfail;

protocol/connection/client_connection.h

@@ -32,6 +32,9 @@
 #include <utility>
 #include <vector>
 
+namespace cb::jwt {
+class Builder;
+}
 namespace cb::mcbp::request {
 class FrameInfo;
 }
@@ -436,6 +439,16 @@ class MemcachedConnection {
                       const std::optional<std::string>& password = {},
                       const std::string& mech = "PLAIN");
 
+    /**
+     * Set the connection to use JWT (and auto renewal of tokens). Note
+     * that JWT may only be used over TLS
+     *
+     * @param builder The builder used to build JWT tokens
+     */
+    void setTokenBuilder(std::unique_ptr<cb::jwt::Builder> builder);
+
+    void authenticateWithToken();
+
     /**
      * Create a bucket
      *
@@ -1225,6 +1238,7 @@ class MemcachedConnection {
     std::string name;
     std::string serverInterfaceUuid;
     std::optional<std::chrono::microseconds> traceData;
+    std::unique_ptr<cb::jwt::Builder> tokenBuilder;
 
     using Featureset = std::unordered_set<uint16_t>;