From: Remi Gacogne Date: Mon, 16 Jun 2025 09:01:31 +0000 (+0200) Subject: dnsdist: Add Lua APIs to set Meta tags in protobuf messages X-Git-Tag: dnsdist-2.0.0-beta1~4^2~1 X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=651b98c6755bf164c78586ad1171f246987f4753;p=thirdparty%2Fpdns.git dnsdist: Add Lua APIs to set Meta tags in protobuf messages Signed-off-by: Remi Gacogne --- diff --git a/pdns/dnsdistdist/dnsdist-actions-factory.cc b/pdns/dnsdistdist/dnsdist-actions-factory.cc index 97be48fa5b..07e8ea9b19 100644 --- a/pdns/dnsdistdist/dnsdist-actions-factory.cc +++ b/pdns/dnsdistdist/dnsdist-actions-factory.cc @@ -1637,6 +1637,9 @@ public: static thread_local std::string data; data.clear(); message.serialize(data); + if (!dnsquestion->d_rawProtobufContent.empty()) { + data.insert(data.end(), dnsquestion->d_rawProtobufContent.begin(), dnsquestion->d_rawProtobufContent.end()); + } remoteLoggerQueueData(*d_logger, data); return Action::None; @@ -1798,6 +1801,9 @@ public: static thread_local std::string data; data.clear(); message.serialize(data); + if (!response->d_rawProtobufContent.empty()) { + data.insert(data.end(), response->d_rawProtobufContent.begin(), response->d_rawProtobufContent.end()); + } d_logger->queueData(data); return Action::None; diff --git a/pdns/dnsdistdist/dnsdist-lua-bindings-dnsquestion.cc b/pdns/dnsdistdist/dnsdist-lua-bindings-dnsquestion.cc index 280fd82e40..83f1591ad4 100644 --- a/pdns/dnsdistdist/dnsdist-lua-bindings-dnsquestion.cc +++ b/pdns/dnsdistdist/dnsdist-lua-bindings-dnsquestion.cc @@ -29,6 +29,28 @@ #include "dnsdist-snmp.hh" #include "dnsparser.hh" +#include "protozero.hh" + +static void addMetaKeyAndValuesToProtobufContent(DNSQuestion& dnsQuestion, const std::string& key, const LuaArray>& values) +{ +#if !defined(DISABLE_PROTOBUF) + protozero::pbf_writer pbfWriter{dnsQuestion.d_rawProtobufContent}; + protozero::pbf_writer pbfMetaWriter{pbfWriter, static_cast(pdns::ProtoZero::Message::Field::meta)}; + pbfMetaWriter.add_string(static_cast(pdns::ProtoZero::Message::MetaField::key), key); + protozero::pbf_writer pbfMetaValueWriter{pbfMetaWriter, static_cast(pdns::ProtoZero::Message::MetaField::value)}; + for (const auto& value : values) { + if (value.second.type() == typeid(std::string)) { + pbfMetaValueWriter.add_string(static_cast(pdns::ProtoZero::Message::MetaValueField::stringVal), boost::get(value.second)); + } + else { + pbfMetaValueWriter.add_uint64(static_cast(pdns::ProtoZero::Message::MetaValueField::intVal), boost::get(value.second)); + } + } + pbfMetaValueWriter.commit(); + pbfMetaWriter.commit(); +#endif /* DISABLE_PROTOBUF */ +} + // NOLINTNEXTLINE(readability-function-cognitive-complexity): this function declares Lua bindings, even with a good refactoring it will likely blow up the threshold void setupLuaBindingsDNSQuestion([[maybe_unused]] LuaContext& luaCtx) { @@ -219,6 +241,10 @@ void setupLuaBindingsDNSQuestion([[maybe_unused]] LuaContext& luaCtx) return *dnsQuestion.ids.qTag; }); + luaCtx.registerFunction>)>("setMetaKey", [](DNSQuestion& dnsQuestion, const std::string& key, const LuaArray>& values) { + addMetaKeyAndValuesToProtobufContent(dnsQuestion, key, values); + }); + luaCtx.registerFunction)>("setProxyProtocolValues", [](DNSQuestion& dnsQuestion, const LuaArray& values) { if (!dnsQuestion.proxyProtocolValues) { dnsQuestion.proxyProtocolValues = make_unique>(); @@ -672,5 +698,10 @@ void setupLuaBindingsDNSQuestion([[maybe_unused]] LuaContext& luaCtx) luaCtx.registerFunction("getRestartCount", [](DNSResponse& dnsResponse) { return dnsResponse.ids.restartCount; }); + + luaCtx.registerFunction>)>("setMetaKey", [](DNSResponse& dnsResponse, const std::string& key, const LuaArray>& values) { + addMetaKeyAndValuesToProtobufContent(dnsResponse, key, values); + }); + #endif /* DISABLE_NON_FFI_DQ_BINDINGS */ } diff --git a/pdns/dnsdistdist/dnsdist-lua-ffi-interface.h b/pdns/dnsdistdist/dnsdist-lua-ffi-interface.h index b818054b94..7d5d8bee76 100644 --- a/pdns/dnsdistdist/dnsdist-lua-ffi-interface.h +++ b/pdns/dnsdistdist/dnsdist-lua-ffi-interface.h @@ -306,3 +306,12 @@ void dnsdist_ffi_svc_record_parameters_add_ipv6_hint(dnsdist_ffi_svc_record_para void dnsdist_ffi_svc_record_parameters_free(dnsdist_ffi_svc_record_parameters* parameters) __attribute__ ((visibility ("default"))); bool dnsdist_ffi_dnsquestion_generate_svc_response(dnsdist_ffi_dnsquestion_t* dnsQuestion, const dnsdist_ffi_svc_record_parameters** parametersList, size_t parametersListSize, uint32_t ttl) __attribute__ ((visibility ("default"))); + +/* this function adds a new key to the raw meta buffer. It can only be called with the same key on a given query once, and dnsdist_ffi_dnsquestion_meta_end_key should always be called after values have been added */ +void dnsdist_ffi_dnsquestion_meta_begin_key(dnsdist_ffi_dnsquestion_t* dnsQuestion, const char* key, size_t keyLen) __attribute__ ((visibility ("default"))); +/* this function should never be called if dnsdist_ffi_dnsquestion_meta_begin_key has not been called first */ +void dnsdist_ffi_dnsquestion_meta_add_str_value_to_key(dnsdist_ffi_dnsquestion_t* dnsQuestion, const char* value, size_t valueLen) __attribute__ ((visibility ("default"))); +/* this function should never be called if dnsdist_ffi_dnsquestion_meta_begin_key has not been called first */ +void dnsdist_ffi_dnsquestion_meta_add_int64_value_to_key(dnsdist_ffi_dnsquestion_t* dnsQuestion, int64_t value) __attribute__ ((visibility ("default"))); +/* this function should never be called if dnsdist_ffi_dnsquestion_meta_begin_key has not been called first */ +void dnsdist_ffi_dnsquestion_meta_end_key(dnsdist_ffi_dnsquestion_t* dnsQuestion) __attribute__ ((visibility ("default"))); diff --git a/pdns/dnsdistdist/dnsdist-lua-ffi.cc b/pdns/dnsdistdist/dnsdist-lua-ffi.cc index 33ed5ea692..01838f1343 100644 --- a/pdns/dnsdistdist/dnsdist-lua-ffi.cc +++ b/pdns/dnsdistdist/dnsdist-lua-ffi.cc @@ -2276,3 +2276,77 @@ void dnsdist_ffi_svc_record_parameters_free(dnsdist_ffi_svc_record_parameters* p // NOLINTNEXTLINE(cppcoreguidelines-owning-memory): this is a C API, RAII is not an option delete parameters; } + +void dnsdist_ffi_dnsquestion_meta_begin_key([[maybe_unused]] dnsdist_ffi_dnsquestion_t* dnsQuestion, [[maybe_unused]] const char* key, [[maybe_unused]] size_t keyLen) +{ +#ifndef DISABLE_PROTOBUF + if (dnsQuestion == nullptr || key == nullptr || keyLen == 0) { + return; + } + + if (dnsQuestion->pbfWriter.valid()) { + vinfolog("Error in dnsdist_ffi_dnsquestion_meta_begin_key: the previous key has not been ended"); + return; + } + + dnsQuestion->pbfWriter = protozero::pbf_writer{dnsQuestion->dq->d_rawProtobufContent}; + dnsQuestion->pbfMetaWriter = protozero::pbf_writer{dnsQuestion->pbfWriter, static_cast(pdns::ProtoZero::Message::Field::meta)}; + dnsQuestion->pbfMetaWriter.add_string(static_cast(pdns::ProtoZero::Message::MetaField::key), protozero::data_view(key, keyLen)); + dnsQuestion->pbfMetaValueWriter = protozero::pbf_writer {dnsQuestion->pbfMetaWriter, static_cast(pdns::ProtoZero::Message::MetaField::value)}; +#endif /* DISABLE_PROTOBUF */ +} + +void dnsdist_ffi_dnsquestion_meta_add_str_value_to_key([[maybe_unused]] dnsdist_ffi_dnsquestion_t* dnsQuestion, [[maybe_unused]] const char* value, [[maybe_unused]] size_t valueLen) +{ +#ifndef DISABLE_PROTOBUF + if (dnsQuestion == nullptr || value == nullptr || valueLen == 0) { + return; + } + + if (!dnsQuestion->pbfMetaValueWriter.valid()) { + vinfolog("Error in dnsdist_ffi_dnsquestion_meta_add_str_value_to_key: trying to add a value without starting a key"); + return; + } + + dnsQuestion->pbfMetaValueWriter.add_string(static_cast(pdns::ProtoZero::Message::MetaValueField::stringVal), protozero::data_view(value, valueLen)); +#endif /* DISABLE_PROTOBUF */ +} + +void dnsdist_ffi_dnsquestion_meta_add_int64_value_to_key([[maybe_unused]] dnsdist_ffi_dnsquestion_t* dnsQuestion, [[maybe_unused]] int64_t value) +{ +#ifndef DISABLE_PROTOBUF + if (dnsQuestion == nullptr) { + return; + } + + if (!dnsQuestion->pbfMetaValueWriter.valid()) { + vinfolog("Error in dnsdist_ffi_dnsquestion_meta_add_int64_value_to_key: trying to add a value without starting a key"); + return; + } + + dnsQuestion->pbfMetaValueWriter.add_uint64(static_cast(pdns::ProtoZero::Message::MetaValueField::intVal), value); +#endif /* DISABLE_PROTOBUF */ +} + +void dnsdist_ffi_dnsquestion_meta_end_key([[maybe_unused]] dnsdist_ffi_dnsquestion_t* dnsQuestion) +{ +#ifndef DISABLE_PROTOBUF + if (dnsQuestion == nullptr) { + return; + } + if (!dnsQuestion->pbfWriter.valid()) { + vinfolog("Error in dnsdist_ffi_dnsquestion_meta_end_key: trying to end a key that has not been started"); + return; + } + + try { + /* reset the pbf writer so that the sizes are properly updated */ + dnsQuestion->pbfMetaValueWriter.commit(); + dnsQuestion->pbfMetaWriter.commit(); + dnsQuestion->pbfWriter = protozero::pbf_writer(); + } + catch (const std::exception& exp) { + vinfolog("Error in dnsdist_ffi_dnsquestion_meta_end_key: %s", exp.what()); + } +#endif /* DISABLE_PROTOBUF */ +} diff --git a/pdns/dnsdistdist/dnsdist-lua-ffi.hh b/pdns/dnsdistdist/dnsdist-lua-ffi.hh index 644114d699..b2c7110c15 100644 --- a/pdns/dnsdistdist/dnsdist-lua-ffi.hh +++ b/pdns/dnsdistdist/dnsdist-lua-ffi.hh @@ -22,6 +22,7 @@ #pragma once #include "dnsdist.hh" +#include "protozero.hh" extern "C" { @@ -64,6 +65,11 @@ struct dnsdist_ffi_dnsquestion_t std::unique_ptr> tagsVect; std::unique_ptr> proxyProtocolValuesVect; std::unique_ptr> httpHeaders; +#ifndef DISABLE_PROTOBUF + protozero::pbf_writer pbfWriter; + protozero::pbf_writer pbfMetaWriter; + protozero::pbf_writer pbfMetaValueWriter; +#endif /* DISABLE_PROTOBUF */ }; // dnsdist_ffi_dnsresponse_t is a lightuserdata diff --git a/pdns/dnsdistdist/dnsdist.hh b/pdns/dnsdistdist/dnsdist.hh index 2d10591eab..d2a8095bbc 100644 --- a/pdns/dnsdistdist/dnsdist.hh +++ b/pdns/dnsdistdist/dnsdist.hh @@ -176,6 +176,9 @@ public: InternalQueryState& ids; std::unique_ptr ecs{nullptr}; std::string sni; /* Server Name Indication, if any (DoT or DoH) */ +#if !defined(DISABLE_PROTOBUF) + std::string d_rawProtobufContent; /* protobuf-encoded content to add to protobuf messages */ +#endif /* DISABLE_PROTOBUF */ mutable std::unique_ptr ednsOptions; /* this needs to be mutable because it is parsed just in time, when DNSQuestion is read-only */ std::shared_ptr d_incomingTCPState{nullptr}; std::unique_ptr> proxyProtocolValues{nullptr}; diff --git a/pdns/dnsdistdist/docs/reference/dq.rst b/pdns/dnsdistdist/docs/reference/dq.rst index 6fa21083cd..50b82562fd 100644 --- a/pdns/dnsdistdist/docs/reference/dq.rst +++ b/pdns/dnsdistdist/docs/reference/dq.rst @@ -321,6 +321,15 @@ This state can be modified from the various hooks. :param string body: The body of the HTTP response, or a URL if the status code is a redirect (3xx) :param string contentType: The HTTP Content-Type header to return for a 200 response, ignored otherwise. Default is ``application/dns-message``. + .. method:: DNSQuestion:setMetaKey(key, values) + + .. versionadded:: 2.0.0 + + Set a meta-data entry to be exported in the ``meta`` field of ProtoBuf messages. + + :param string key: The key + :param list values: A list containing strings, integers, or a mix of integers and strings + .. method:: DNSQuestion:setNegativeAndAdditionalSOA(nxd, zone, ttl, mname, rname, serial, refresh, retry, expire, minimum) .. versionadded:: 1.5.0 diff --git a/pdns/dnsdistdist/meson.build b/pdns/dnsdistdist/meson.build index a053232f5b..644e702671 100644 --- a/pdns/dnsdistdist/meson.build +++ b/pdns/dnsdistdist/meson.build @@ -553,6 +553,7 @@ if get_option('unit-tests') dep_boost, dep_boost_test, dep_lua, + dep_protozero, ], ) ) @@ -572,6 +573,7 @@ if get_option('unit-tests') dep_boost_test, dep_ffi_interface, dep_lua, + dep_protozero, ], } } diff --git a/pdns/dnsdistdist/test-dnsdist-lua-ffi.cc b/pdns/dnsdistdist/test-dnsdist-lua-ffi.cc index b4b2333a3c..ca49f6453d 100644 --- a/pdns/dnsdistdist/test-dnsdist-lua-ffi.cc +++ b/pdns/dnsdistdist/test-dnsdist-lua-ffi.cc @@ -28,6 +28,7 @@ #include #include "dnsdist-lua-ffi.hh" +#include "base64.hh" #include "dnsdist-cache.hh" #include "dnsdist-configuration.hh" #include "dnsdist-rings.hh" @@ -1033,4 +1034,60 @@ BOOST_AUTO_TEST_CASE(test_SVC_Generation) dnsdist_ffi_svc_record_parameters_free(parameters); } +#if !defined(DISABLE_PROTOBUF) +BOOST_AUTO_TEST_CASE(test_meta_values) +{ + + InternalQueryState ids; + ids.origRemote = ComboAddress("192.0.2.1:4242"); + ids.origDest = ComboAddress("192.0.2.255:53"); + ids.qtype = QType::A; + ids.qclass = QClass::IN; + ids.protocol = dnsdist::Protocol::DoUDP; + ids.qname = DNSName("www.powerdns.com."); + ids.queryRealTime.start(); + PacketBuffer query; + GenericDNSPacketWriter pwQ(query, ids.qname, QType::A, QClass::IN, 0); + pwQ.getHeader()->rd = 1; + pwQ.getHeader()->id = htons(42); + + DNSQuestion dnsQuestion(ids, query); + dnsdist_ffi_dnsquestion_t lightDQ(&dnsQuestion); + + { + /* check invalid parameters */ + dnsdist_ffi_dnsquestion_meta_begin_key(nullptr, nullptr, 0); + dnsdist_ffi_dnsquestion_meta_begin_key(&lightDQ, nullptr, 0); + dnsdist_ffi_dnsquestion_meta_begin_key(&lightDQ, "some-key", 0); + dnsdist_ffi_dnsquestion_meta_add_str_value_to_key(nullptr, nullptr, 0); + dnsdist_ffi_dnsquestion_meta_add_str_value_to_key(&lightDQ, nullptr, 0); + dnsdist_ffi_dnsquestion_meta_add_str_value_to_key(&lightDQ, "some-str-value", 0); + dnsdist_ffi_dnsquestion_meta_add_int64_value_to_key(nullptr, 0); + dnsdist_ffi_dnsquestion_meta_end_key(nullptr); + } + + { + /* trying to end a key that has not been started */ + dnsdist_ffi_dnsquestion_meta_end_key(&lightDQ); + } + + { + const std::string key{"some-key"}; + const std::string value1{"first value"}; + const std::string value2{"second value"}; + BOOST_CHECK_EQUAL(dnsQuestion.d_rawProtobufContent.size(), 0U); + dnsdist_ffi_dnsquestion_meta_begin_key(&lightDQ, key.data(), key.size()); + /* we should not be able to begin a new key without ending it first */ + dnsdist_ffi_dnsquestion_meta_begin_key(&lightDQ, key.data(), key.size()); + dnsdist_ffi_dnsquestion_meta_add_str_value_to_key(&lightDQ, value1.data(), value1.size()); + dnsdist_ffi_dnsquestion_meta_add_int64_value_to_key(&lightDQ, 42); + dnsdist_ffi_dnsquestion_meta_add_str_value_to_key(&lightDQ, value2.data(), value2.size()); + dnsdist_ffi_dnsquestion_meta_add_int64_value_to_key(&lightDQ, -42); + dnsdist_ffi_dnsquestion_meta_end_key(&lightDQ); + BOOST_CHECK_EQUAL(dnsQuestion.d_rawProtobufContent.size(), 55U); + BOOST_CHECK_EQUAL(Base64Encode(dnsQuestion.d_rawProtobufContent), "sgE0Cghzb21lLWtleRIoCgtmaXJzdCB2YWx1ZRAqCgxzZWNvbmQgdmFsdWUQ1v//////////AQ=="); + } +} +#endif /* DISABLE_PROTOBUF */ + BOOST_AUTO_TEST_SUITE_END(); diff --git a/regression-tests.dnsdist/test_Protobuf.py b/regression-tests.dnsdist/test_Protobuf.py index 77cc97d941..f5717ad3af 100644 --- a/regression-tests.dnsdist/test_Protobuf.py +++ b/regression-tests.dnsdist/test_Protobuf.py @@ -421,9 +421,38 @@ class TestProtobufMetaTags(DNSDistProtobufTest): newServer{address="127.0.0.1:%s"} rl = newRemoteLogger('127.0.0.1:%d') + local ffi = require("ffi") + local C = ffi.C + function add_meta(dq) + local key = "my-meta-key-1" + local key2 = "my-meta-key-2" + C.dnsdist_ffi_dnsquestion_meta_begin_key(dq, key, #key) + C.dnsdist_ffi_dnsquestion_meta_add_str_value_to_key(dq, "test", 4) + C.dnsdist_ffi_dnsquestion_meta_add_int64_value_to_key(dq, -42) + C.dnsdist_ffi_dnsquestion_meta_add_str_value_to_key(dq, "test2", 5) + C.dnsdist_ffi_dnsquestion_meta_end_key(dq) + + local key2 = "my-meta-key-2" + C.dnsdist_ffi_dnsquestion_meta_begin_key(dq, key2, #key2) + C.dnsdist_ffi_dnsquestion_meta_add_str_value_to_key(dq, "foo", 3) + C.dnsdist_ffi_dnsquestion_meta_add_int64_value_to_key(dq, 42) + C.dnsdist_ffi_dnsquestion_meta_add_str_value_to_key(dq, "bar", 3) + C.dnsdist_ffi_dnsquestion_meta_end_key(dq) + + return DNSAction.None + end + + function addMetaToResponse(dr) + dr:setMetaKey('my-meta-key-1', {'test', -42, 'test2'}) + dr:setMetaKey('my-meta-key-2', {'foo', 42, 'bar'}) + return DNSResponseAction.None + end + + addAction(AllRule(), LuaFFIAction(add_meta)) addAction(AllRule(), SetTagAction('my-tag-key', 'my-tag-value')) addAction(AllRule(), SetTagAction('my-empty-key', '')) addAction(AllRule(), RemoteLogAction(rl, nil, {serverID='dnsdist-server-1', exportTags='*'}, {b64='b64-content', ['my-tag-export-name']='tag:my-tag-key'})) + addResponseAction(AllRule(), LuaResponseAction(addMetaToResponse)) addResponseAction(AllRule(), SetTagResponseAction('my-tag-key2', 'my-tag-value2')) addResponseAction(AllRule(), RemoteLogResponseAction(rl, nil, false, {serverID='dnsdist-server-1', exportTags='my-empty-key,my-tag-key2'}, {['my-tag-export-name']='tags'})) """ @@ -462,7 +491,7 @@ class TestProtobufMetaTags(DNSDistProtobufTest): self.assertIn('my-tag-key:my-tag-value', msg.response.tags) self.assertIn('my-empty-key', msg.response.tags) # meta tags - self.assertEqual(len(msg.meta), 2) + self.assertEqual(len(msg.meta), 4) tags = {} for entry in msg.meta: tags[entry.key] = entry.value.stringVal @@ -470,6 +499,18 @@ class TestProtobufMetaTags(DNSDistProtobufTest): self.assertIn('b64', tags) self.assertIn('my-tag-export-name', tags) + self.assertEqual(msg.meta[2].key, 'my-meta-key-1') + self.assertEqual(len(msg.meta[2].value.stringVal), 2) + self.assertIn('test', msg.meta[2].value.stringVal) + self.assertIn('test2', msg.meta[2].value.stringVal) + self.assertIn(-42, msg.meta[2].value.intVal) + + self.assertEqual(msg.meta[3].key, 'my-meta-key-2') + self.assertEqual(len(msg.meta[3].value.stringVal), 2) + self.assertIn('foo', msg.meta[3].value.stringVal) + self.assertIn('bar', msg.meta[3].value.stringVal) + self.assertIn(42, msg.meta[3].value.intVal) + b64EncodedQuery = base64.b64encode(query.to_wire()).decode('ascii') self.assertEqual(tags['b64'], [b64EncodedQuery]) self.assertEqual(tags['my-tag-export-name'], ['my-tag-value']) @@ -482,7 +523,7 @@ class TestProtobufMetaTags(DNSDistProtobufTest): self.assertIn('my-tag-key2:my-tag-value2', msg.response.tags) self.assertIn('my-empty-key', msg.response.tags) # meta tags - self.assertEqual(len(msg.meta), 1) + self.assertEqual(len(msg.meta), 3) self.assertEqual(msg.meta[0].key, 'my-tag-export-name') self.assertEqual(len(msg.meta[0].value.stringVal), 3) self.assertIn('my-tag-key:my-tag-value', msg.meta[0].value.stringVal) @@ -490,6 +531,18 @@ class TestProtobufMetaTags(DNSDistProtobufTest): # no ':' when the value is empty self.assertIn('my-empty-key', msg.meta[0].value.stringVal) + self.assertEqual(msg.meta[1].key, 'my-meta-key-1') + self.assertEqual(len(msg.meta[1].value.stringVal), 2) + self.assertIn('test', msg.meta[1].value.stringVal) + self.assertIn('test2', msg.meta[1].value.stringVal) + self.assertIn(-42, msg.meta[1].value.intVal) + + self.assertEqual(msg.meta[2].key, 'my-meta-key-2') + self.assertEqual(len(msg.meta[2].value.stringVal), 2) + self.assertIn('foo', msg.meta[2].value.stringVal) + self.assertIn('bar', msg.meta[2].value.stringVal) + self.assertIn(42, msg.meta[2].value.intVal) + class TestProtobufExtendedDNSErrorTags(DNSDistProtobufTest): _config_params = ['_testServerPort', '_protobufServerPort'] _config_template = """