From e79f2196497a2088d609dd35ce199b0dd973ce0e Mon Sep 17 00:00:00 2001 From: Remi Gacogne Date: Fri, 3 Mar 2023 16:36:06 +0100 Subject: [PATCH] dnsdist: Add Lua bindings for PB requestorID, deviceName and deviceID --- pdns/dnsdist-idstate.hh | 14 +++++++-- pdns/dnsdist-lua-actions.cc | 14 ++++++--- pdns/dnsdist-lua-bindings-dnsquestion.cc | 33 ++++++++++++++++++++ pdns/dnsdist-protobuf.cc | 15 ++++++++- pdns/dnsdistdist/dnsdist-lua-ffi-interface.h | 4 +++ pdns/dnsdistdist/dnsdist-lua-ffi.cc | 33 ++++++++++++++++++++ pdns/dnsdistdist/docs/reference/dq.rst | 24 ++++++++++++-- pdns/dnsdistdist/test-dnsdist-lua-ffi.cc | 17 ++++++++++ 8 files changed, 143 insertions(+), 11 deletions(-) diff --git a/pdns/dnsdist-idstate.hh b/pdns/dnsdist-idstate.hh index bf3d76c332..b61984249e 100644 --- a/pdns/dnsdist-idstate.hh +++ b/pdns/dnsdist-idstate.hh @@ -91,6 +91,14 @@ private: struct InternalQueryState { + struct ProtoBufData + { + std::optional uniqueId{std::nullopt}; // 17 + std::string d_deviceName; + std::string d_deviceID; + std::string d_requestorID; + }; + static void DeleterPlaceHolder(DOHUnit*) { } @@ -114,11 +122,12 @@ struct InternalQueryState ComboAddress hopLocal; DNSName qname; // 24 std::string poolName; // 24 - StopWatch queryRealTime{true}; // 16 + StopWatch queryRealTime{true}; // 24 std::shared_ptr packetCache{nullptr}; // 16 std::unique_ptr dnsCryptQuery{nullptr}; // 8 std::unique_ptr qTag{nullptr}; // 8 - std::unique_ptr d_packet; // Initial packet, so we can restart the query from the response path if needed // 8 + std::unique_ptr d_packet{nullptr}; // Initial packet, so we can restart the query from the response path if needed // 8 + std::unique_ptr d_protoBufData{nullptr}; boost::optional tempFailureTTL{boost::none}; // 8 ClientState* cs{nullptr}; // 8 std::unique_ptr du; // 8 @@ -137,7 +146,6 @@ struct InternalQueryState uint16_t cacheFlags{0}; // DNS flags as sent to the backend // 2 uint16_t udpPayloadSize{0}; // Max UDP payload size from the query // 2 dnsdist::Protocol protocol; // 1 - boost::optional uniqueId{boost::none}; // 17 (placed here to reduce the space lost to padding) bool ednsAdded{false}; bool ecsAdded{false}; bool skipCache{false}; diff --git a/pdns/dnsdist-lua-actions.cc b/pdns/dnsdist-lua-actions.cc index b6b65fb29a..acbada4e14 100644 --- a/pdns/dnsdist-lua-actions.cc +++ b/pdns/dnsdist-lua-actions.cc @@ -1519,8 +1519,11 @@ public: DNSAction::Action operator()(DNSQuestion* dq, std::string* ruleresult) const override { - if (!dq->ids.uniqueId) { - dq->ids.uniqueId = getUniqueID(); + if (!dq->ids.d_protoBufData) { + dq->ids.d_protoBufData = std::make_unique(); + } + if (!dq->ids.d_protoBufData->uniqueId) { + dq->ids.d_protoBufData->uniqueId = getUniqueID(); } DNSDistProtoBufMessage message(*dq); @@ -1660,8 +1663,11 @@ public: } DNSResponseAction::Action operator()(DNSResponse* dr, std::string* ruleresult) const override { - if (!dr->ids.uniqueId) { - dr->ids.uniqueId = getUniqueID(); + if (!dr->ids.d_protoBufData) { + dr->ids.d_protoBufData = std::make_unique(); + } + if (!dr->ids.d_protoBufData->uniqueId) { + dr->ids.d_protoBufData->uniqueId = getUniqueID(); } DNSDistProtoBufMessage message(*dr, d_includeCNAME); diff --git a/pdns/dnsdist-lua-bindings-dnsquestion.cc b/pdns/dnsdist-lua-bindings-dnsquestion.cc index ec21b70d30..9d5da2c736 100644 --- a/pdns/dnsdist-lua-bindings-dnsquestion.cc +++ b/pdns/dnsdist-lua-bindings-dnsquestion.cc @@ -56,6 +56,39 @@ void setupLuaBindingsDNSQuestion(LuaContext& luaCtx) dq.ids.tempFailureTTL = newValue; } ); + luaCtx.registerMember("deviceID", [](const DNSQuestion& dq) -> std::string { + if (dq.ids.d_protoBufData) { + return dq.ids.d_protoBufData->d_deviceID; + } + return std::string(); + }, [](DNSQuestion& dq, const std::string& newValue) { + if (!dq.ids.d_protoBufData) { + dq.ids.d_protoBufData = std::make_unique(); + } + dq.ids.d_protoBufData->d_deviceID = newValue; + }); + luaCtx.registerMember("deviceName", [](const DNSQuestion& dq) -> std::string { + if (dq.ids.d_protoBufData) { + return dq.ids.d_protoBufData->d_deviceName; + } + return std::string(); + }, [](DNSQuestion& dq, const std::string& newValue) { + if (!dq.ids.d_protoBufData) { + dq.ids.d_protoBufData = std::make_unique(); + } + dq.ids.d_protoBufData->d_deviceName = newValue; + }); + luaCtx.registerMember("requestorID", [](const DNSQuestion& dq) -> std::string { + if (dq.ids.d_protoBufData) { + return dq.ids.d_protoBufData->d_requestorID; + } + return std::string(); + }, [](DNSQuestion& dq, const std::string& newValue) { + if (!dq.ids.d_protoBufData) { + dq.ids.d_protoBufData = std::make_unique(); + } + dq.ids.d_protoBufData->d_requestorID = newValue; + }); luaCtx.registerFunction("getDO", [](const DNSQuestion& dq) { return getEDNSZ(dq) & EDNS_HEADER_FLAG_DO; }); diff --git a/pdns/dnsdist-protobuf.cc b/pdns/dnsdist-protobuf.cc index a2d24622fc..c38529cb67 100644 --- a/pdns/dnsdist-protobuf.cc +++ b/pdns/dnsdist-protobuf.cc @@ -154,7 +154,7 @@ void DNSDistProtoBufMessage::serialize(std::string& data) const protocol = pdns::ProtoZero::Message::TransportProtocol::DNSCryptTCP; } - m.setRequest(d_dq.ids.uniqueId ? *d_dq.ids.uniqueId : getUniqueID(), d_requestor ? *d_requestor : d_dq.ids.origRemote, d_responder ? *d_responder : d_dq.ids.origDest, d_question ? d_question->d_name : d_dq.ids.qname, d_question ? d_question->d_type : d_dq.ids.qtype, d_question ? d_question->d_class : d_dq.ids.qclass, d_dq.getHeader()->id, protocol, d_bytes ? *d_bytes : d_dq.getData().size()); + m.setRequest(d_dq.ids.d_protoBufData && d_dq.ids.d_protoBufData->uniqueId ? *d_dq.ids.d_protoBufData->uniqueId : getUniqueID(), d_requestor ? *d_requestor : d_dq.ids.origRemote, d_responder ? *d_responder : d_dq.ids.origDest, d_question ? d_question->d_name : d_dq.ids.qname, d_question ? d_question->d_type : d_dq.ids.qtype, d_question ? d_question->d_class : d_dq.ids.qclass, d_dq.getHeader()->id, protocol, d_bytes ? *d_bytes : d_dq.getData().size()); if (d_serverIdentity) { m.setServerIdentity(*d_serverIdentity); @@ -196,6 +196,19 @@ void DNSDistProtoBufMessage::serialize(std::string& data) const m.commitResponse(); + if (d_dq.ids.d_protoBufData) { + const auto& pbData = d_dq.ids.d_protoBufData; + if (!pbData->d_deviceName.empty()) { + m.setDeviceName(pbData->d_deviceName); + } + if (!pbData->d_deviceID.empty()) { + m.setDeviceId(pbData->d_deviceID); + } + if (!pbData->d_requestorID.empty()) { + m.setRequestorId(pbData->d_requestorID); + } + } + for (const auto& [key, values] : d_metaTags) { if (!values.empty()) { m.setMeta(key, values, {}); diff --git a/pdns/dnsdistdist/dnsdist-lua-ffi-interface.h b/pdns/dnsdistdist/dnsdist-lua-ffi-interface.h index 69dab22602..c8d9397815 100644 --- a/pdns/dnsdistdist/dnsdist-lua-ffi-interface.h +++ b/pdns/dnsdistdist/dnsdist-lua-ffi-interface.h @@ -111,6 +111,10 @@ void dnsdist_ffi_dnsquestion_unset_temp_failure_ttl(dnsdist_ffi_dnsquestion_t* d void dnsdist_ffi_dnsquestion_set_tag(dnsdist_ffi_dnsquestion_t* dq, const char* label, const char* value) __attribute__ ((visibility ("default"))); void dnsdist_ffi_dnsquestion_set_tag_raw(dnsdist_ffi_dnsquestion_t* dq, const char* label, const char* value, size_t valueSize) __attribute__ ((visibility ("default"))); +void dnsdist_ffi_dnsquestion_set_requestor_id(dnsdist_ffi_dnsquestion_t* dq, const char* value, size_t valueSize) __attribute__ ((visibility ("default"))); +void dnsdist_ffi_dnsquestion_set_device_id(dnsdist_ffi_dnsquestion_t* dq, const char* value, size_t valueSize) __attribute__ ((visibility ("default"))); +void dnsdist_ffi_dnsquestion_set_device_name(dnsdist_ffi_dnsquestion_t* dq, const char* value, size_t valueSize) __attribute__ ((visibility ("default"))); + void dnsdist_ffi_dnsquestion_set_http_response(dnsdist_ffi_dnsquestion_t* dq, uint16_t statusCode, const char* body, size_t bodyLen, const char* contentType) __attribute__ ((visibility ("default"))); size_t dnsdist_ffi_dnsquestion_get_trailing_data(dnsdist_ffi_dnsquestion_t* dq, const char** out) __attribute__ ((visibility ("default"))); diff --git a/pdns/dnsdistdist/dnsdist-lua-ffi.cc b/pdns/dnsdistdist/dnsdist-lua-ffi.cc index 2283078fd4..b690dd368b 100644 --- a/pdns/dnsdistdist/dnsdist-lua-ffi.cc +++ b/pdns/dnsdistdist/dnsdist-lua-ffi.cc @@ -512,6 +512,39 @@ void dnsdist_ffi_dnsquestion_set_tag_raw(dnsdist_ffi_dnsquestion_t* dq, const ch dq->dq->setTag(label, std::string(value, valueSize)); } +void dnsdist_ffi_dnsquestion_set_requestor_id(dnsdist_ffi_dnsquestion_t* dq, const char* value, size_t valueSize) +{ + if (!dq || !dq->dq || !value) { + return; + } + if (!dq->dq->ids.d_protoBufData) { + dq->dq->ids.d_protoBufData = std::make_unique(); + } + dq->dq->ids.d_protoBufData->d_requestorID = std::string(value, valueSize); +} + +void dnsdist_ffi_dnsquestion_set_device_id(dnsdist_ffi_dnsquestion_t* dq, const char* value, size_t valueSize) +{ + if (!dq || !dq->dq || !value) { + return; + } + if (!dq->dq->ids.d_protoBufData) { + dq->dq->ids.d_protoBufData = std::make_unique(); + } + dq->dq->ids.d_protoBufData->d_deviceID = std::string(value, valueSize); +} + +void dnsdist_ffi_dnsquestion_set_device_name(dnsdist_ffi_dnsquestion_t* dq, const char* value, size_t valueSize) +{ + if (!dq || !dq->dq || !value) { + return; + } + if (!dq->dq->ids.d_protoBufData) { + dq->dq->ids.d_protoBufData = std::make_unique(); + } + dq->dq->ids.d_protoBufData->d_deviceName = std::string(value, valueSize); +} + size_t dnsdist_ffi_dnsquestion_get_trailing_data(dnsdist_ffi_dnsquestion_t* dq, const char** out) { dq->trailingData = dq->dq->getTrailingData(); diff --git a/pdns/dnsdistdist/docs/reference/dq.rst b/pdns/dnsdistdist/docs/reference/dq.rst index 7a87c6aeb5..6025c64e16 100644 --- a/pdns/dnsdistdist/docs/reference/dq.rst +++ b/pdns/dnsdistdist/docs/reference/dq.rst @@ -11,6 +11,18 @@ This state can be modified from the various hooks. The DNSQuestion object has several attributes, many of them read-only: + .. attribute:: DNSQuestion.deviceID + + .. versionadded:: 1.8.0 + + The identifier of the remote device, which will be exported via ProtoBuf if set. + + .. attribute:: DNSQuestion.deviceName + + .. versionadded:: 1.8.0 + + The name of the remote device, which will be exported via ProtoBuf if set. + .. attribute:: DNSQuestion.dh The :ref:`DNSHeader` of this query. @@ -59,6 +71,12 @@ This state can be modified from the various hooks. :ref:`ComboAddress` of the remote client. + .. attribute:: DNSQuestion.requestorID + + .. versionadded:: 1.8.0 + + The identifier of the requestor, which will be exported via ProtoBuf if set. + .. attribute:: DNSQuestion.rcode RCode (as an unsigned integer) of this question. @@ -73,7 +91,7 @@ This state can be modified from the various hooks. Whether to skip cache lookup / storing the answer for this question, settable. .. attribute:: DNSQuestion.tempFailureTTL - + On a SERVFAIL or REFUSED from the backend, cache for this amount of seconds, settable. .. attribute:: DNSQuestion.tcp @@ -313,7 +331,7 @@ This state can be modified from the various hooks. Prior to 1.7.0 calling :func:`DNSQuestion:setTag` would not overwrite an existing tag value if already set. Set a tag into the DNSQuestion object. Overwrites the value if any already exists. - + :param string key: The tag's key :param string value: The tag's value @@ -323,7 +341,7 @@ This state can be modified from the various hooks. Prior to 1.7.0 calling :func:`DNSQuestion:setTagArray` would not overwrite existing tag values if already set. Set an array of tags into the DNSQuestion object. Overwrites the values if any already exist. - + :param table tags: A table of tags, using strings as keys and values .. method:: DNSQuestion:setTrailingData(tail) -> bool diff --git a/pdns/dnsdistdist/test-dnsdist-lua-ffi.cc b/pdns/dnsdistdist/test-dnsdist-lua-ffi.cc index fcd53608da..d4367ea50e 100644 --- a/pdns/dnsdistdist/test-dnsdist-lua-ffi.cc +++ b/pdns/dnsdistdist/test-dnsdist-lua-ffi.cc @@ -345,6 +345,23 @@ BOOST_AUTO_TEST_CASE(test_Query) dnsdist_ffi_dnsquestion_set_http_response(&lightDQ, 0U, nullptr, 0U, nullptr); } } + + const std::string deviceID{"my-device-id"}; + const std::string deviceName{"my-device-name"}; + const std::string requestorID{"my-requestor-ID"}; + dnsdist_ffi_dnsquestion_set_device_id(nullptr, nullptr, 0); + dnsdist_ffi_dnsquestion_set_device_id(&lightDQ, nullptr, 0); + dnsdist_ffi_dnsquestion_set_device_id(&lightDQ, deviceID.c_str(), deviceID.size()); + dnsdist_ffi_dnsquestion_set_device_name(nullptr, nullptr, 0); + dnsdist_ffi_dnsquestion_set_device_name(&lightDQ, nullptr, 0); + dnsdist_ffi_dnsquestion_set_device_name(&lightDQ, deviceName.c_str(), deviceName.size()); + dnsdist_ffi_dnsquestion_set_requestor_id(nullptr, nullptr, 0); + dnsdist_ffi_dnsquestion_set_requestor_id(&lightDQ, nullptr, 0); + dnsdist_ffi_dnsquestion_set_requestor_id(&lightDQ, requestorID.c_str(), requestorID.size()); + BOOST_REQUIRE(ids.d_protoBufData != nullptr); + BOOST_CHECK_EQUAL(ids.d_protoBufData->d_deviceID, deviceID); + BOOST_CHECK_EQUAL(ids.d_protoBufData->d_deviceName, deviceName); + BOOST_CHECK_EQUAL(ids.d_protoBufData->d_requestorID, requestorID); } BOOST_AUTO_TEST_CASE(test_Response) -- 2.47.2