Fix for "Core dump when providing extensions to oauthbearer_set_token" #719 (#733)

* Attempting to fix issue #719

* Tests

* Optimization

* Comments

* Fix code review issue (empty but non-nil extensions)

* Fix unsafe memory management

* Fix for code review issue, close producer in tests

---------

Co-authored-by: Daniel Sjöblom <[email protected]>

Author: dssjoblom

lib/rdkafka/helpers/oauth.rb

@@ -12,12 +12,18 @@ module OAuth
       # @return [Integer] 0 on success
       def oauthbearer_set_token(token:, lifetime_ms:, principal_name:, extensions: nil)
         error_buffer = FFI::MemoryPointer.from_string(" " * 256)
+        extensions_ptr, extensions_str_ptrs = map_extensions(extensions)
 
-        response = @native_kafka.with_inner do |inner|
-          Rdkafka::Bindings.rd_kafka_oauthbearer_set_token(
-            inner, token, lifetime_ms, principal_name,
-            flatten_extensions(extensions), extension_size(extensions), error_buffer, 256
-          )
+        begin
+          response = @native_kafka.with_inner do |inner|
+            Rdkafka::Bindings.rd_kafka_oauthbearer_set_token(
+              inner, token, lifetime_ms, principal_name,
+              extensions_ptr, extension_size(extensions), error_buffer, 256
+            )
+          end
+        ensure
+          extensions_str_ptrs&.each { |ptr| ptr.free }
+          extensions_ptr&.free
         end
 
         return response if response.zero?
@@ -41,10 +47,31 @@ def oauthbearer_set_token_failure(reason)
 
       private
 
-      # Flatten the extensions hash into a string according to the spec, https://datatracker.ietf.org/doc/html/rfc7628#section-3.1
-      def flatten_extensions(extensions)
-        return nil unless extensions
-        "\x01#{extensions.map { |e| e.join("=") }.join("\x01")}"
+      # Convert extensions hash to FFI::MemoryPointer (const char **).
+      # Note: the returned pointers must be freed manually (autorelease = false).
+      def map_extensions(extensions)
+        return [nil, nil] if extensions.nil? || extensions.empty?
+
+        # https://github.com/confluentinc/librdkafka/blob/master/src/rdkafka_sasl_oauthbearer.c#L327-L347
+
+        # The method argument is const char **
+        array_ptr = FFI::MemoryPointer.new(:pointer, extension_size(extensions))
+        array_ptr.autorelease = false
+        str_ptrs = []
+
+        # Element i is the key, i + 1 is the value.
+        extensions.each_with_index do |(k, v), i|
+          k_ptr = FFI::MemoryPointer.from_string(k.to_s)
+          k_ptr.autorelease = false
+          str_ptrs << k_ptr
+          v_ptr = FFI::MemoryPointer.from_string(v.to_s)
+          v_ptr.autorelease = false
+          str_ptrs << v_ptr
+          array_ptr[i * 2].put_pointer(0, k_ptr)
+          array_ptr[i * 2 + 1].put_pointer(0, v_ptr)
+        end
+
+        [array_ptr, str_ptrs]
       end
 
       # extension_size is the number of keys + values which should be a non-negative even number

spec/lib/rdkafka/admin_spec.rb

@@ -953,14 +953,29 @@
         $admin_sasl.close
       end
 
-      it 'should succeed' do
+      context 'without extensions' do
+        it 'should succeed' do
+          response = $admin_sasl.oauthbearer_set_token(
+            token: "foo",
+            lifetime_ms: Time.now.to_i*1000 + 900 * 1000,
+            principal_name: "kafka-cluster"
+          )
+          expect(response).to eq(0)
+        end
+      end
 
-        response = $admin_sasl.oauthbearer_set_token(
-          token: "foo",
-          lifetime_ms: Time.now.to_i*1000 + 900 * 1000,
-          principal_name: "kafka-cluster"
-        )
-        expect(response).to eq(0)
+      context 'with extensions' do
+        it 'should succeed' do
+          response = $admin_sasl.oauthbearer_set_token(
+            token: "foo",
+            lifetime_ms: Time.now.to_i*1000 + 900 * 1000,
+            principal_name: "kafka-cluster",
+            extensions: {
+              "foo" => "bar"
+            }
+          )
+          expect(response).to eq(0)
+        end
       end
     end
   end

spec/lib/rdkafka/consumer_spec.rb

@@ -1230,14 +1230,29 @@ def collect(name, list)
         $consumer_sasl.close
       end
 
-      it 'should succeed' do
+      context 'without extensions' do
+        it 'should succeed' do
+          response = $consumer_sasl.oauthbearer_set_token(
+            token: "foo",
+            lifetime_ms: Time.now.to_i*1000 + 900 * 1000,
+            principal_name: "kafka-cluster"
+          )
+          expect(response).to eq(0)
+        end
+      end
 
-        response = $consumer_sasl.oauthbearer_set_token(
-          token: "foo",
-          lifetime_ms: Time.now.to_i*1000 + 900 * 1000,
-          principal_name: "kafka-cluster"
-        )
-        expect(response).to eq(0)
+      context 'with extensions' do
+        it 'should succeed' do
+          response = $consumer_sasl.oauthbearer_set_token(
+            token: "foo",
+            lifetime_ms: Time.now.to_i*1000 + 900 * 1000,
+            principal_name: "kafka-cluster",
+            extensions: {
+              "foo" => "bar"
+            }
+          )
+          expect(response).to eq(0)
+        end
       end
     end
   end

spec/lib/rdkafka/producer_spec.rb

@@ -880,26 +880,47 @@ def call(_, handle)
         response = producer.oauthbearer_set_token(
               token: "foo",
               lifetime_ms: Time.now.to_i*1000 + 900 * 1000,
-              principal_name: "kafka-cluster"
+              principal_name: "kafka-cluster",
             )
         expect(response).to eq(Rdkafka::Bindings::RD_KAFKA_RESP_ERR__STATE)
       end
     end
 
     context 'when sasl configured' do
-      it 'should succeed' do
-        producer_sasl = rdkafka_producer_config(
-          {
-            "security.protocol": "sasl_ssl",
-            "sasl.mechanisms": 'OAUTHBEARER'
-          }
+      before do
+        $producer_sasl = rdkafka_producer_config(
+          "security.protocol": "sasl_ssl",
+          "sasl.mechanisms": 'OAUTHBEARER'
         ).producer
-        response = producer_sasl.oauthbearer_set_token(
-          token: "foo",
-          lifetime_ms: Time.now.to_i*1000 + 900 * 1000,
-          principal_name: "kafka-cluster"
-        )
-        expect(response).to eq(0)
+      end
+
+      after do
+        $producer_sasl.close
+      end
+
+      context 'without extensions' do
+        it 'should succeed' do
+          response = $producer_sasl.oauthbearer_set_token(
+            token: "foo",
+            lifetime_ms: Time.now.to_i*1000 + 900 * 1000,
+            principal_name: "kafka-cluster"
+          )
+          expect(response).to eq(0)
+        end
+      end
+
+      context 'with extensions' do
+        it 'should succeed' do
+          response = $producer_sasl.oauthbearer_set_token(
+            token: "foo",
+            lifetime_ms: Time.now.to_i*1000 + 900 * 1000,
+            principal_name: "kafka-cluster",
+            extensions: {
+              "foo" => "bar"
+            }
+          )
+          expect(response).to eq(0)
+        end
       end
     end
   end