From: Remi Gacogne <remi.gacogne@powerdns.com>
Date: Tue, 28 May 2024 07:49:58 +0000 (+0200)
Subject: dnsdist: Add a Lua FFI function to set proxy protocol values
X-Git-Tag: rec-5.2.0-alpha0~38^2~1
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=060a00a1ecfd5d064294c26b152ee6c1c1e7c295;p=thirdparty%2Fpdns.git

dnsdist: Add a Lua FFI function to set proxy protocol values
---

diff --git a/pdns/dnsdistdist/dnsdist-lua-ffi-interface.h b/pdns/dnsdistdist/dnsdist-lua-ffi-interface.h
index c254b6ad73..03ea1b3315 100644
--- a/pdns/dnsdistdist/dnsdist-lua-ffi-interface.h
+++ b/pdns/dnsdistdist/dnsdist-lua-ffi-interface.h
@@ -177,6 +177,7 @@ typedef struct dnsdist_ffi_proxy_protocol_value {
 
 size_t dnsdist_ffi_generate_proxy_protocol_payload(size_t addrSize, const void* srcAddr, const void* dstAddr, uint16_t srcPort, uint16_t dstPort, bool tcp, size_t valuesCount, const dnsdist_ffi_proxy_protocol_value_t* values, void* out, size_t outSize) __attribute__ ((visibility ("default")));
 size_t dnsdist_ffi_dnsquestion_generate_proxy_protocol_payload(const dnsdist_ffi_dnsquestion_t* dq, const size_t valuesCount, const dnsdist_ffi_proxy_protocol_value_t* values, void* out, const size_t outSize) __attribute__ ((visibility ("default")));
+bool dnsdist_ffi_dnsquestion_add_proxy_protocol_values(dnsdist_ffi_dnsquestion_t* dnsQuestion, const size_t valuesCount, const dnsdist_ffi_proxy_protocol_value_t* values) __attribute__ ((visibility ("default")));
 
 typedef struct dnsdist_ffi_domain_list_t dnsdist_ffi_domain_list_t;
 typedef struct dnsdist_ffi_address_list_t dnsdist_ffi_address_list_t;
diff --git a/pdns/dnsdistdist/dnsdist-lua-ffi.cc b/pdns/dnsdistdist/dnsdist-lua-ffi.cc
index 2c77b35d7f..48d6bb5a5c 100644
--- a/pdns/dnsdistdist/dnsdist-lua-ffi.cc
+++ b/pdns/dnsdistdist/dnsdist-lua-ffi.cc
@@ -1083,6 +1083,25 @@ size_t dnsdist_ffi_dnsquestion_generate_proxy_protocol_payload(const dnsdist_ffi
   return payload.size();
 }
 
+bool dnsdist_ffi_dnsquestion_add_proxy_protocol_values(dnsdist_ffi_dnsquestion_t* dnsQuestion, const size_t valuesCount, const dnsdist_ffi_proxy_protocol_value_t* values)
+{
+  if (dnsQuestion == nullptr || dnsQuestion->dq == nullptr || values == nullptr || valuesCount == 0) {
+    return false;
+  }
+
+  if (!dnsQuestion->dq->proxyProtocolValues) {
+    dnsQuestion->dq->proxyProtocolValues = make_unique<std::vector<ProxyProtocolValue>>();
+  }
+  if (valuesCount > 0) {
+    dnsQuestion->dq->proxyProtocolValues->reserve(dnsQuestion->dq->proxyProtocolValues->size() + valuesCount);
+    for (size_t idx = 0; idx < valuesCount; idx++) {
+      dnsQuestion->dq->proxyProtocolValues->push_back({ std::string(values[idx].value, values[idx].size), values[idx].type });
+    }
+  }
+
+  return true;
+}
+
 struct dnsdist_ffi_domain_list_t
 {
   std::vector<std::string> d_domains;
diff --git a/pdns/dnsdistdist/test-dnsdist-lua-ffi.cc b/pdns/dnsdistdist/test-dnsdist-lua-ffi.cc
index df8f4d87fa..c58a205bfb 100644
--- a/pdns/dnsdistdist/test-dnsdist-lua-ffi.cc
+++ b/pdns/dnsdistdist/test-dnsdist-lua-ffi.cc
@@ -575,6 +575,53 @@ BOOST_AUTO_TEST_CASE(test_ProxyProtocol)
   }
 }
 
+BOOST_AUTO_TEST_CASE(test_ProxyProtocolQuery)
+{
+  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<PacketBuffer> pwQ(query, ids.qname, QType::A, QClass::IN, 0);
+  pwQ.getHeader()->rd = 1;
+  pwQ.getHeader()->id = htons(42);
+
+  DNSQuestion dq(ids, query);
+  dnsdist_ffi_dnsquestion_t lightDQ(&dq);
+
+  std::vector<dnsdist_ffi_proxy_protocol_value> values;
+  values.push_back({"test-value", 10U, 1U});
+
+  {
+    auto added = dnsdist_ffi_dnsquestion_add_proxy_protocol_values(nullptr, values.size(), values.data());
+    BOOST_CHECK_EQUAL(added, false);
+  }
+
+  {
+    auto added = dnsdist_ffi_dnsquestion_add_proxy_protocol_values(&lightDQ, 0, values.data());
+    BOOST_CHECK_EQUAL(added, false);
+  }
+
+  {
+    auto added = dnsdist_ffi_dnsquestion_add_proxy_protocol_values(&lightDQ, values.size(), nullptr);
+    BOOST_CHECK_EQUAL(added, false);
+  }
+
+  {
+    auto added = dnsdist_ffi_dnsquestion_add_proxy_protocol_values(&lightDQ, values.size(), values.data());
+    BOOST_CHECK_EQUAL(added, true);
+    BOOST_REQUIRE(dq.proxyProtocolValues != nullptr);
+    BOOST_REQUIRE_EQUAL(dq.proxyProtocolValues->size(), values.size());
+    BOOST_CHECK_EQUAL(dq.proxyProtocolValues->at(0).type, values.at(0).type);
+    BOOST_REQUIRE_EQUAL(dq.proxyProtocolValues->at(0).content.size(), values.at(0).size);
+    BOOST_CHECK_EQUAL(memcmp(dq.proxyProtocolValues->at(0).content.data(), values.at(0).value, values.at(0).size), 0);
+  }
+}
+
 BOOST_AUTO_TEST_CASE(test_PacketOverlay)
 {
   const DNSName target("powerdns.com.");