Add OAUTHBEARER token refresh callback

Author: milindl

lib/client.js

@@ -15,7 +15,7 @@ var util = require('util');
 var Kafka = require('../librdkafka.js');
 var assert = require('assert');
 
-const bindingVersion = require('./util').bindingVersion;
+const { bindingVersion, dictToStringList } = require('./util');
 
 var LibrdKafkaError = require('./error');
 
@@ -100,6 +100,40 @@ function Client(globalConf, SubClientType, topicConf) {
     }.bind(this);
   }
 
+  if (Object.hasOwn(this._cb_configs.global, 'oauthbearer_token_refresh_cb')) {
+    const savedCallback = this._cb_configs.global.oauthbearer_token_refresh_cb;
+    this._cb_configs.global.oauthbearer_token_refresh_cb = (oauthbearer_config) => {
+      if (this._isDisconnecting) {
+        // Don't call the callback if we're in the middle of disconnecting.
+        // This is especially important when the credentials are wrong, and
+        // we might want to disconnect without ever completing connection.
+        return;
+      }
+      savedCallback(oauthbearer_config, (err, token) => {
+        try {
+          if (err) {
+            throw err;
+          }
+          let { tokenValue, lifetime, principal, extensions } = token;
+
+          // If the principal isn't there, set an empty principal.
+          if (!principal) {
+            principal = '';
+          }
+
+          // Convert extensions from a Map/object to a list that librdkafka expects.
+          extensions = dictToStringList(extensions);
+
+          this._client.setOAuthBearerToken(tokenValue, lifetime, principal, extensions);
+        } catch (e) {
+          e.message = "oauthbearer_token_refresh_cb: " + e.message;
+          this._client.setOAuthBearerTokenFailure(e.message);
+          this.emit('error', e);
+        }
+      });
+    }
+  }
+
   this.metrics = {};
   this._isConnected = false;
   this.errorCounter = 0;

lib/util.js

@@ -2,6 +2,7 @@
  * confluent-kafka-javascript - Node.js wrapper  for RdKafka C/C++ library
  *
  * Copyright (c) 2016-2023 Blizzard Entertainment
+ *               2024 Confluent, Inc.
  *
  * This software may be modified and distributed under the terms
  * of the MIT license.  See the LICENSE.txt file for details.
@@ -28,4 +29,27 @@ util.isObject = function (obj) {
   return obj && typeof obj === 'object';
 };
 
+// Convert Map or object to a list of [key, value, key, value...].
+util.dictToStringList = function (mapOrObject) {
+  let list = null;
+  if (mapOrObject && (mapOrObject instanceof Map)) {
+    list =
+      Array
+        .from(mapOrObject).reduce((acc, [key, value]) => {
+          acc.push(key, value);
+          return acc;
+        }, [])
+        .map(v => String(v));
+  } else if (util.isObject(mapOrObject)) {
+    list =
+      Object
+        .entries(mapOrObject).reduce((acc, [key, value]) => {
+          acc.push(key, value);
+          return acc;
+        }, [])
+        .map(v => String(v));
+  }
+  return list;
+};
+
 util.bindingVersion = 'v0.1.11-devel';

src/admin.cc

@@ -39,22 +39,41 @@ AdminClient::~AdminClient() {
 }
 
 Baton AdminClient::Connect() {
-  std::string errstr;
+  if (IsConnected()) {
+    return Baton(RdKafka::ERR_NO_ERROR);
+  }
+
+  Baton baton = setupSaslOAuthBearerConfig();
+  if (baton.err() != RdKafka::ERR_NO_ERROR) {
+    return baton;
+  }
 
+  // Activate the dispatchers before the connection, as some callbacks may run
+  // on the background thread.
+  // We will deactivate them if the connection fails.
+  ActivateDispatchers();
+
+  std::string errstr;
   {
     scoped_shared_write_lock lock(m_connection_lock);
     m_client = RdKafka::Producer::create(m_gconfig, errstr);
   }
 
   if (!m_client || !errstr.empty()) {
+    DeactivateDispatchers();
     return Baton(RdKafka::ERR__STATE, errstr);
   }
 
   if (rkqu == NULL) {
     rkqu = rd_kafka_queue_new(m_client->c_ptr());
   }
 
-  return Baton(RdKafka::ERR_NO_ERROR);
+  baton = setupSaslOAuthBearerBackgroundQueue();
+  if (baton.err() != RdKafka::ERR_NO_ERROR) {
+    DeactivateDispatchers();
+  }
+
+  return baton;
 }
 
 Baton AdminClient::Disconnect() {
@@ -66,6 +85,8 @@ Baton AdminClient::Disconnect() {
       rkqu = NULL;
     }
 
+    DeactivateDispatchers();
+
     delete m_client;
     m_client = NULL;
   }
@@ -99,6 +120,9 @@ void AdminClient::Init(v8::Local<v8::Object> exports) {
   Nan::SetPrototypeMethod(tpl, "disconnect", NodeDisconnect);
   Nan::SetPrototypeMethod(tpl, "setSaslCredentials", NodeSetSaslCredentials);
   Nan::SetPrototypeMethod(tpl, "getMetadata", NodeGetMetadata);
+  Nan::SetPrototypeMethod(tpl, "setOAuthBearerToken", NodeSetOAuthBearerToken);
+  Nan::SetPrototypeMethod(tpl, "setOAuthBearerTokenFailure",
+                          NodeSetOAuthBearerTokenFailure);
 
   constructor.Reset(
     (tpl->GetFunction(Nan::GetCurrentContext())).ToLocalChecked());

src/callbacks.cc

@@ -10,9 +10,7 @@
 
 #include <string>
 #include <vector>
-#include <algorithm>
 
-#include "src/callbacks.h"
 #include "src/kafka-consumer.h"
 
 using v8::Local;
@@ -547,6 +545,37 @@ void OffsetCommit::offset_commit_cb(RdKafka::ErrorCode err,
   dispatcher.Execute();
 }
 
+// OAuthBearerTokenRefresh callback
+void OAuthBearerTokenRefreshDispatcher::Add(
+    const std::string &oauthbearer_config) {
+  scoped_mutex_lock lock(async_lock);
+  m_oauthbearer_config = oauthbearer_config;
+}
+
+void OAuthBearerTokenRefreshDispatcher::Flush() {
+  Nan::HandleScope scope;
+
+  const unsigned int argc = 1;
+
+  std::string oauthbearer_config;
+  {
+    scoped_mutex_lock lock(async_lock);
+    oauthbearer_config = m_oauthbearer_config;
+    m_oauthbearer_config.clear();
+  }
+
+  v8::Local<v8::Value> argv[argc] = {};
+  argv[0] = Nan::New<v8::String>(oauthbearer_config.c_str()).ToLocalChecked();
+
+  Dispatch(argc, argv);
+}
+
+void OAuthBearerTokenRefresh::oauthbearer_token_refresh_cb(
+    RdKafka::Handle *handle, const std::string &oauthbearer_config) {
+  dispatcher.Add(oauthbearer_config);
+  dispatcher.Execute();
+}
+
 // Partitioner callback
 
 Partitioner::Partitioner() {}

src/callbacks.h

@@ -248,6 +248,23 @@ class OffsetCommit : public RdKafka::OffsetCommitCb {
   v8::Persistent<v8::Function> m_cb;
 };
 
+class OAuthBearerTokenRefreshDispatcher : public Dispatcher {
+ public:
+  OAuthBearerTokenRefreshDispatcher(){};
+  ~OAuthBearerTokenRefreshDispatcher(){};
+  void Add(const std::string &oauthbearer_config);
+  void Flush();
+
+ private:
+  std::string m_oauthbearer_config;
+};
+
+class OAuthBearerTokenRefresh : public RdKafka::OAuthBearerTokenRefreshCb {
+ public:
+  void oauthbearer_token_refresh_cb(RdKafka::Handle *, const std::string &);
+  OAuthBearerTokenRefreshDispatcher dispatcher;
+};
+
 class Partitioner : public RdKafka::PartitionerCb {
  public:
   Partitioner();

src/common.cc

@@ -143,6 +143,26 @@ std::vector<std::string> v8ArrayToStringVector(v8::Local<v8::Array> parameter) {
   return newItem;
 }
 
+std::list<std::string> v8ArrayToStringList(v8::Local<v8::Array> parameter) {
+  std::list<std::string> newItem;
+  if (parameter->Length() >= 1) {
+    for (unsigned int i = 0; i < parameter->Length(); i++) {
+      v8::Local<v8::Value> v;
+      if (!Nan::Get(parameter, i).ToLocal(&v)) {
+        continue;
+      }
+      Nan::MaybeLocal<v8::String> p = Nan::To<v8::String>(v);
+      if (p.IsEmpty()) {
+        continue;
+      }
+      Nan::Utf8String pVal(p.ToLocalChecked());
+      std::string pString(*pVal);
+      newItem.push_back(pString);
+    }
+  }
+  return newItem;
+}
+
 template<> v8::Local<v8::Array> GetParameter<v8::Local<v8::Array> >(
   v8::Local<v8::Object> object, std::string field_name, v8::Local<v8::Array> def) {
   v8::Local<v8::String> field = Nan::New(field_name.c_str()).ToLocalChecked();

src/common.h

@@ -39,6 +39,7 @@ template<> v8::Local<v8::Array> GetParameter<v8::Local<v8::Array> >(
   v8::Local<v8::Object>, std::string, v8::Local<v8::Array>);
 // template int GetParameter<int>(v8::Local<v8::Object, std::string, int);
 std::vector<std::string> v8ArrayToStringVector(v8::Local<v8::Array>);
+std::list<std::string> v8ArrayToStringList(v8::Local<v8::Array>);
 
 class scoped_mutex_lock {
  public:

src/config.cc

@@ -81,8 +81,8 @@ Conf * Conf::create(RdKafka::Conf::ConfType type, v8::Local<v8::Object> object,
           return NULL;
       }
     } else {
-     // Do nothing - Connection::NodeConfigureCallbacks will handle this for each
-     // of the three client types.
+     // Do nothing - NodeConfigureCallbacks will handle this for each
+     // of the three client types, called from within JavaScript.
     }
   }
 
@@ -118,6 +118,23 @@ void Conf::ConfigureCallback(const std::string &string_key, const v8::Local<v8::
         offset_commit->dispatcher.RemoveCallback(cb);
       }
     }
+  } else if (string_key.compare("oauthbearer_token_refresh_cb") == 0) {
+    NodeKafka::Callbacks::OAuthBearerTokenRefresh *oauthbearer_token_refresh =
+        oauthbearer_token_refresh_cb();
+    if (add) {
+      if (oauthbearer_token_refresh == NULL) {
+        oauthbearer_token_refresh =
+            new NodeKafka::Callbacks::OAuthBearerTokenRefresh();
+        this->set(string_key, oauthbearer_token_refresh, errstr);
+      }
+      oauthbearer_token_refresh->dispatcher.AddCallback(cb);
+    } else {
+      if (oauthbearer_token_refresh != NULL) {
+        oauthbearer_token_refresh->dispatcher.RemoveCallback(cb);
+      }
+    }
+  } else {
+    errstr = "Invalid callback type";
   }
 }
 
@@ -131,6 +148,12 @@ void Conf::listen() {
   if (offset_commit) {
     offset_commit->dispatcher.Activate();
   }
+
+  NodeKafka::Callbacks::OAuthBearerTokenRefresh *oauthbearer_token_refresh =
+      oauthbearer_token_refresh_cb();
+  if (oauthbearer_token_refresh) {
+    oauthbearer_token_refresh->dispatcher.Activate();
+  }
 }
 
 void Conf::stop() {
@@ -143,6 +166,12 @@ void Conf::stop() {
   if (offset_commit) {
     offset_commit->dispatcher.Deactivate();
   }
+
+  NodeKafka::Callbacks::OAuthBearerTokenRefresh *oauthbearer_token_refresh =
+      oauthbearer_token_refresh_cb();
+  if (oauthbearer_token_refresh) {
+    oauthbearer_token_refresh->dispatcher.Deactivate();
+  }
 }
 
 Conf::~Conf() {
@@ -167,4 +196,21 @@ NodeKafka::Callbacks::OffsetCommit* Conf::offset_commit_cb() const {
   return static_cast<NodeKafka::Callbacks::OffsetCommit*>(cb);
 }
 
+NodeKafka::Callbacks::OAuthBearerTokenRefresh *
+Conf::oauthbearer_token_refresh_cb() const {
+  RdKafka::OAuthBearerTokenRefreshCb *cb = NULL;
+  if (this->get(cb) != RdKafka::Conf::CONF_OK) {
+    return NULL;
+  }
+  return static_cast<NodeKafka::Callbacks::OAuthBearerTokenRefresh *>(cb);
+}
+
+bool Conf::is_sasl_oauthbearer() const {
+  std::string sasl_mechanism;
+  if (this->get("sasl.mechanisms", sasl_mechanism) != RdKafka::Conf::CONF_OK) {
+    return false;
+  }
+  return sasl_mechanism.compare("OAUTHBEARER") == 0;
+}
+
 }  // namespace NodeKafka

src/config.h

@@ -34,18 +34,23 @@ class Conf : public RdKafka::Conf {
 
   void ConfigureCallback(const std::string &string_key, const v8::Local<v8::Function> &cb, bool add, std::string &errstr);
 
+  bool is_sasl_oauthbearer() const;
+
  private:
-  NodeKafka::Callbacks::Rebalance* rebalance_cb() const;
+  NodeKafka::Callbacks::Rebalance *rebalance_cb() const;
   NodeKafka::Callbacks::OffsetCommit *offset_commit_cb() const;
+  NodeKafka::Callbacks::OAuthBearerTokenRefresh *oauthbearer_token_refresh_cb()
+      const;
 
   // NOTE: Do NOT add any members to this class.
-  // Internally, to get an instance of this class, we just cast RdKafka::Conf* that we
-  // obtain from RdKafka::Conf::create(). However, that's internally an instance of a sub-class,
-  // ConfImpl.
-  // This means that any members here are aliased to that with the wrong name (for example, the first
-  // member of this class, if it's a pointer, will be aliased to consume_cb_ in the ConfImpl, and
-  // and changing one will change the other!)
-  // TODO: Just don't inherit from RdKafka::Conf, and instead have a member of type RdKafka::Conf*.
+  // Internally, to get an instance of this class, we just cast RdKafka::Conf*
+  // that we obtain from RdKafka::Conf::create(). However, that's internally an
+  // instance of a sub-class, ConfImpl. This means that any members here are
+  // aliased to that with the wrong name (for example, the first member of this
+  // class, if it's a pointer, will be aliased to consume_cb_ in the ConfImpl,
+  // and and changing one will change the other!)
+  // TODO: Just don't inherit from RdKafka::Conf, and instead have a member of
+  // type RdKafka::Conf*.
 };
 
 }  // namespace NodeKafka