From a64c238c10f99acbde2560f1fc2cbb8c27fdbe3f Mon Sep 17 00:00:00 2001
From: Remi Gacogne <remi.gacogne@powerdns.com>
Date: Fri, 5 Nov 2021 09:45:42 +0100
Subject: [PATCH] dnsdist: Add a Lua FFI helper to generate proxy protocol
 payloads

---
 pdns/dnsdistdist/dnsdist-lua-ffi-interface.h |  8 ++++++
 pdns/dnsdistdist/dnsdist-lua-ffi.cc          | 30 ++++++++++++++++++++
 2 files changed, 38 insertions(+)

diff --git a/pdns/dnsdistdist/dnsdist-lua-ffi-interface.h b/pdns/dnsdistdist/dnsdist-lua-ffi-interface.h
index 1dcdfab9e8..7e7dec45b1 100644
--- a/pdns/dnsdistdist/dnsdist-lua-ffi-interface.h
+++ b/pdns/dnsdistdist/dnsdist-lua-ffi-interface.h
@@ -128,3 +128,11 @@ double dnsdist_ffi_server_get_latency(const dnsdist_ffi_server_t* server) __attr
 void dnsdist_ffi_dnsresponse_set_min_ttl(dnsdist_ffi_dnsresponse_t* dr, uint32_t min) __attribute__ ((visibility ("default")));
 void dnsdist_ffi_dnsresponse_set_max_ttl(dnsdist_ffi_dnsresponse_t* dr, uint32_t max) __attribute__ ((visibility ("default")));
 void dnsdist_ffi_dnsresponse_limit_ttl(dnsdist_ffi_dnsresponse_t* dr, uint32_t min, uint32_t max) __attribute__ ((visibility ("default")));
+
+typedef struct dnsdist_ffi_proxy_protocol_value {
+  char* value;
+  uint16_t size;
+  uint8_t type;
+} dnsdist_ffi_proxy_protocol_value_t;
+
+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")));
diff --git a/pdns/dnsdistdist/dnsdist-lua-ffi.cc b/pdns/dnsdistdist/dnsdist-lua-ffi.cc
index ac6676c586..3e455b3baa 100644
--- a/pdns/dnsdistdist/dnsdist-lua-ffi.cc
+++ b/pdns/dnsdistdist/dnsdist-lua-ffi.cc
@@ -598,3 +598,33 @@ void setupLuaFFIPerThreadContext(LuaContext& luaCtx)
   luaCtx.executeCode(getLuaFFIWrappers());
 #endif
 }
+
+size_t dnsdist_ffi_generate_proxy_protocol_payload(const size_t addrSize, const void* srcAddr, const void* dstAddr, const uint16_t srcPort, const uint16_t dstPort, const bool tcp, const size_t valuesCount, const dnsdist_ffi_proxy_protocol_value* values, void* out, const size_t outSize)
+{
+  ComboAddress src, dst;
+  if (addrSize != sizeof(src.sin4.sin_addr) && addrSize != sizeof(src.sin6.sin6_addr.s6_addr)) {
+    return 0;
+  }
+
+  src = makeComboAddressFromRaw(addrSize == sizeof(src.sin4.sin_addr) ? 4 : 6, reinterpret_cast<const char*>(srcAddr), addrSize);
+  src.sin4.sin_port = htons(srcPort);
+  dst = makeComboAddressFromRaw(addrSize == sizeof(dst.sin4.sin_addr) ? 4 : 6, reinterpret_cast<const char*>(dstAddr), addrSize);
+  dst.sin4.sin_port = htons(dstPort);
+
+  std::vector<ProxyProtocolValue> valuesVect;
+  if (valuesCount > 0) {
+    valuesVect.reserve(valuesCount);
+    for (size_t idx = 0; idx < valuesCount; idx++) {
+      valuesVect.push_back({ std::string(values[idx].value, values[idx].size), values[idx].type });
+    }
+  }
+
+  std::string payload = makeProxyHeader(tcp, src, dst, valuesVect);
+  if (payload.size() > outSize) {
+    return 0;
+  }
+
+  memcpy(out, payload.c_str(), payload.size());
+
+  return payload.size();
+}
-- 
2.47.2