Allow passing headers as both list(tuples) and dict()

edenhill · edenhill · commit c55ffc1047a4 · 2018-03-29T19:14:42.000+02:00
This also fixes type checking, the headers were assumed to be a List,
if they were not (such as a dict) it lead to memory corruption and mayhem.

Also added test for binary header values.
diff --git a/confluent_kafka/src/confluent_kafka.c b/confluent_kafka/src/confluent_kafka.c
@@ -939,43 +939,125 @@ rd_kafka_topic_partition_list_t *py_to_c_parts (PyObject *plist) {
 }
 
 #ifdef RD_KAFKA_V_HEADERS
+
+
+/**
+ * @brief Convert Python list of tuples to rd_kafka_headers_t
+ */
+static rd_kafka_headers_t *py_headers_list_to_c (PyObject *hdrs) {
+        int i, len;
+        rd_kafka_headers_t *rd_headers = NULL;
+
+        len = (int)PyList_Size(hdrs);
+        rd_headers = rd_kafka_headers_new(len);
+
+        for (i = 0; i < len; i++) {
+                rd_kafka_resp_err_t err;
+                const char *header_key, *header_value = NULL;
+                int header_key_len = 0, header_value_len = 0;
+
+                if(!PyArg_ParseTuple(PyList_GET_ITEM(hdrs, i), "s#z#",
+                                     &header_key, &header_key_len,
+                                     &header_value, &header_value_len)){
+                        rd_kafka_headers_destroy(rd_headers);
+                        PyErr_SetString(PyExc_TypeError,
+                                        "Headers are expected to be a "
+                                        "tuple of (key, value)");
+                        return NULL;
+                }
+
+                err = rd_kafka_header_add(rd_headers,
+                                          header_key, header_key_len,
+                                          header_value, header_value_len);
+                if (err) {
+                        cfl_PyErr_Format(err,
+                                         "Unable to add message header \"%s\": "
+                                         "%s",
+                                         header_key, rd_kafka_err2str(err));
+                        rd_kafka_headers_destroy(rd_headers);
+                        return NULL;
+                }
+        }
+        return rd_headers;
+}
+
+
+/**
+ * @brief Convert Python dict to rd_kafka_headers_t
+ */
+static rd_kafka_headers_t *py_headers_dict_to_c (PyObject *hdrs) {
+        int len;
+        Py_ssize_t pos = 0;
+        rd_kafka_headers_t *rd_headers = NULL;
+        PyObject *ko, *vo;
+
+        len = (int)PyDict_Size(hdrs);
+        rd_headers = rd_kafka_headers_new(len);
+
+        while (PyDict_Next(hdrs, &pos, &ko, &vo)) {
+                PyObject *ks, *ks8;
+                const char *k;
+                const void *v = NULL;
+                Py_ssize_t vsize = 0;
+                rd_kafka_resp_err_t err;
+
+                if (!(ks = cfl_PyObject_Unistr(ko))) {
+                        PyErr_SetString(PyExc_TypeError,
+                                        "expected header key to be unicode "
+                                        "string");
+                        rd_kafka_headers_destroy(rd_headers);
+                        return NULL;
+                }
+
+                k = cfl_PyUnistr_AsUTF8(ks, &ks8);
+
+                if (vo != Py_None) {
+                        if (PyString_AsStringAndSize(vo, (char **)&v,
+                                                     &vsize) == -1) {
+                                Py_DECREF(ks);
+                                rd_kafka_headers_destroy(rd_headers);
+                                return NULL;
+                        }
+                }
+
+                if ((err = rd_kafka_header_add(rd_headers, k, -1, v, vsize))) {
+                        cfl_PyErr_Format(err,
+                                         "Unable to add message header \"%s\": "
+                                         "%s",
+                                         k, rd_kafka_err2str(err));
+                        Py_DECREF(ks);
+                        rd_kafka_headers_destroy(rd_headers);
+                        return NULL;
+                }
+
+                Py_DECREF(ks);
+        }
+
+        return rd_headers;
+}
+
+
 /**
  * @brief Convert Python list[(header_key, header_value),...]) to C rd_kafka_topic_partition_list_t.
  *
  * @returns The new Python list[(header_key, header_value),...] object.
  */
-rd_kafka_headers_t *py_headers_to_c (PyObject *headers_plist) {
-    int i, len;
-    rd_kafka_headers_t *rd_headers = NULL;
-    rd_kafka_resp_err_t err;
-    const char *header_key, *header_value = NULL;
-    int header_key_len = 0, header_value_len = 0;
-
-    len = PyList_Size(headers_plist);
-    rd_headers = rd_kafka_headers_new(len);
-
-    for (i = 0; i < len; i++) {
-
-        if(!PyArg_ParseTuple(PyList_GET_ITEM(headers_plist, i), "s#z#", &header_key,
-                &header_key_len, &header_value, &header_value_len)){
-            rd_kafka_headers_destroy(rd_headers);
-            PyErr_SetString(PyExc_TypeError,
-                    "Headers are expected to be a tuple of (key, value)");
-            return NULL;
-        }
-
-        err = rd_kafka_header_add(rd_headers, header_key, header_key_len, header_value, header_value_len);
-        if (err) {
-            rd_kafka_headers_destroy(rd_headers);
-            cfl_PyErr_Format(err,
-                     "Unable to create message headers: %s",
-                     rd_kafka_err2str(err));
-            return NULL;
+rd_kafka_headers_t *py_headers_to_c (PyObject *hdrs) {
+
+        if (PyList_Check(hdrs)) {
+                return py_headers_list_to_c(hdrs);
+        } else if (PyDict_Check(hdrs)) {
+                return py_headers_dict_to_c(hdrs);
+        } else {
+                PyErr_Format(PyExc_TypeError,
+                             "expected headers to be "
+                             "dict or list of (key, value) tuples, not %s",
+                             ((PyTypeObject *)PyObject_Type(hdrs))->tp_name);
+                return NULL;
         }
-    }
-    return rd_headers;
 }
 
+
 /**
  * @brief Convert rd_kafka_headers_t to Python list[(header_key, header_value),...])
  *
@@ -989,7 +1071,7 @@ PyObject *c_headers_to_py (rd_kafka_headers_t *headers) {
     size_t header_value_size;
     PyObject *header_list;
 
-    header_size = rd_kafka_header_cnt(headers); 
+    header_size = rd_kafka_header_cnt(headers);
     header_list = PyList_New(header_size);
 
     while (!rd_kafka_header_get_all(headers, idx++,
diff --git a/confluent_kafka/src/confluent_kafka.h b/confluent_kafka/src/confluent_kafka.h
@@ -270,7 +270,7 @@ PyObject *c_parts_to_py (const rd_kafka_topic_partition_list_t *c_parts);
 rd_kafka_topic_partition_list_t *py_to_c_parts (PyObject *plist);
 
 #ifdef RD_KAFKA_V_HEADERS
-rd_kafka_headers_t *py_headers_to_c (PyObject *headers_plist);
+rd_kafka_headers_t *py_headers_to_c (PyObject *hdrs);
 PyObject *c_headers_to_py (rd_kafka_headers_t *headers);
 #endif
 /****************************************************************************
diff --git a/examples/integration_test.py b/examples/integration_test.py
@@ -26,6 +26,7 @@
 import sys
 import json
 import gc
+import struct
 from copy import copy
 
 try:
@@ -117,7 +118,8 @@ def verify_producer():
     p = confluent_kafka.Producer(**conf)
     print('producer at %s' % p)
 
-    headers = [('foo1', 'bar'), ('foo1', 'bar2'), ('foo2', b'1')]
+    headers = [('foo1', 'bar'), ('foo1', 'bar2'), ('foo2', b'1'),
+               ('foobin', struct.pack('hhl', 10, 20, 30))]
 
     # Produce some messages
     p.produce(topic, 'Hello Python!', headers=headers)
@@ -464,6 +466,8 @@ def print_wmark(consumer, parts):
 
     first_msg = None
 
+    example_header = None
+
     while True:
         # Consume until EOF or error
 
@@ -517,12 +521,15 @@ def print_wmark(consumer, parts):
             print('Sync committed offset: %s' % offsets)
 
         msgcnt += 1
-        if msgcnt >= max_msgcnt:
+        if msgcnt >= max_msgcnt and example_header is not None:
             print('max_msgcnt %d reached' % msgcnt)
             break
 
     assert example_header, "We should have received at least one header"
-    assert example_header == [(u'foo1', 'bar'), (u'foo1', 'bar2'), (u'foo2', '1')]
+    assert example_header == [(u'foo1', 'bar'),
+                              (u'foo1', 'bar2'),
+                              (u'foo2', '1'),
+                              ('foobin', struct.pack('hhl', 10, 20, 30))]
 
     # Get current assignment
     assignment = c.assignment()
diff --git a/tests/test_Producer.py b/tests/test_Producer.py
@@ -2,6 +2,7 @@
 import pytest
 
 from confluent_kafka import Producer, KafkaError, KafkaException, libversion
+from struct import pack
 
 
 def error_cb(err):
@@ -65,11 +66,25 @@ def test_produce_headers():
                   'error_cb': error_cb,
                   'default.topic.config': {'message.timeout.ms': 10}})
 
-    p.produce('mytopic', value='somedata', key='a key', headers=[('headerkey', 'headervalue')])
-    p.produce('mytopic', value='somedata', key='a key', headers=[('dupkey', 'dupvalue'), ('dupkey', 'dupvalue')])
-    p.produce('mytopic', value='somedata', key='a key', headers=[('dupkey', 'dupvalue'), ('dupkey', 'diffvalue')])
-    p.produce('mytopic', value='somedata', key='a key', headers=[('key_with_null_value', None)])
-    p.produce('mytopic', value='somedata', key='a key', headers=[])
+    binval = pack('hhl', 1, 2, 3)
+
+    headers_to_test = [
+        [('headerkey', 'headervalue')],
+        [('dupkey', 'dupvalue'), ('empty', ''), ('dupkey', 'dupvalue')],
+        [('dupkey', 'dupvalue'), ('dupkey', 'diffvalue')],
+        [('key_with_null_value', None)],
+        [('binaryval', binval)],
+
+        {'headerkey': 'headervalue'},
+        {'dupkey': 'dupvalue', 'empty': '', 'dupkey': 'dupvalue'},  # noqa: F601
+        {'dupkey': 'dupvalue', 'dupkey': 'diffvalue'},  # noqa: F601
+        {'key_with_null_value': None},
+        {'binaryval': binval}
+        ]
+
+    for headers in headers_to_test:
+        p.produce('mytopic', value='somedata', key='a key', headers=headers)
+        p.produce('mytopic', value='somedata', headers=headers)
 
     with pytest.raises(TypeError) as ex:
         p.produce('mytopic', value='somedata', key='a key', headers=[('malformed_header')])
