From 99fce2dae8371a2cc767f92ba999d3d7fd8fd9a3 Mon Sep 17 00:00:00 2001
From: Seth Ornstein <sornstein@globalcyberalliance.org>
Date: Fri, 2 Jun 2017 15:44:49 -0400
Subject: [PATCH] Added some more Lua commands to dnsdist to make the commands
 more 'balanced'.

---
 pdns/dnsdist-lua.cc  | 14 ++++++++++++--
 pdns/dnsdist-lua2.cc | 28 ++++++++++++++++++++++++----
 2 files changed, 36 insertions(+), 6 deletions(-)

diff --git a/pdns/dnsdist-lua.cc b/pdns/dnsdist-lua.cc
index 99c4c5b9a..e67eeef7f 100644
--- a/pdns/dnsdist-lua.cc
+++ b/pdns/dnsdist-lua.cc
@@ -1588,8 +1588,8 @@ vector<std::function<void(void)>> setupLua(bool client, const std::string& confi
 
 
   // --------------------------------------------------------------------------
-  // GCA - Seth Ornstein added lua callable functions - 5/30/2017
-  // DNSQuestion - setTag, getTagMatch, getTagArray
+  // GCA - Seth Ornstein added lua callable functions - 6/2/2017
+  // DNSQuestion - setTag, setTagArray, getTagMatch, getTagArray
 
     g_lua.registerFunction<void(DNSQuestion::*)(std::string, std::string)>("setTag", [](DNSQuestion& dq, const std::string& strLabel, const std::string& strValue) {
 
@@ -1597,6 +1597,16 @@ vector<std::function<void(void)>> setupLua(bool client, const std::string& confi
 
     });
 
+     g_lua.registerFunction<void(DNSQuestion::*)(vector<pair<string, string>>)>("setTagArray", [](DNSQuestion& dq, const vector<pair<string, string>>&tags) {
+
+      setLuaSideEffect();
+
+      for (const auto& tag : tags)
+        {
+           dq.qTag.add(tag.first, tag.second);
+        }
+     });
+
 
      g_lua.registerFunction<string(DNSQuestion::*)(std::string)>("getTagMatch", [](const DNSQuestion& dq, const std::string& strLabel) {
 
diff --git a/pdns/dnsdist-lua2.cc b/pdns/dnsdist-lua2.cc
index 2a4b47f0d..c00272c95 100644
--- a/pdns/dnsdist-lua2.cc
+++ b/pdns/dnsdist-lua2.cc
@@ -833,8 +833,15 @@ void moreLua(bool client)
 
 
 // --------------------------------------------------------------------------
-// GCA - Seth Ornstein added lua callable functions - 5/30/2017
-// DNSDistProtoBufMessage - setTagArray, setProtobufResponseType
+// GCA - Seth Ornstein added lua callable functions - 6/2/2017
+// DNSDistProtoBufMessage - setTag, setTagArray, setProtobufResponseType, setProtobufResponseTypeQT
+
+     g_lua.registerFunction<void(DNSDistProtoBufMessage::*)(std::string, std::string)>("setTag", [](DNSDistProtoBufMessage& message, const std::string& strLabel, const std::string& strValue) {
+
+      setLuaSideEffect();
+
+      message.addTags(strLabel, strValue);               // add a tag to store text - not used by dnsdist normally
+     });
 
 
      g_lua.registerFunction<void(DNSDistProtoBufMessage::*)(vector<pair<string, string>>)>("setTagArray", [](DNSDistProtoBufMessage& message, const vector<pair<string, string>>&tags) {
@@ -843,11 +850,23 @@ void moreLua(bool client)
 
       for (const auto& tag : tags)
         {
-          message.addTags(tag.first, tag.second);               // add a tag to store text - not used by dnsdist at present?
+          message.addTags(tag.first, tag.second);               // add a tag to store text - not used by dnsdist normally
         }
      });
 
-     g_lua.registerFunction<void(DNSDistProtoBufMessage::*)(const std::string&, time_t sec, uint uSec)>("setProtobufResponseType", [](DNSDistProtoBufMessage& message, const std::string& strQueryName, time_t sec, uint uSec) {
+
+                                                                // setProtobufResponseType - no timestamp
+     g_lua.registerFunction<void(DNSDistProtoBufMessage::*)(const std::string&)>("setProtobufResponseType", [](DNSDistProtoBufMessage& message, const std::string& strQueryName) {
+
+        message.setType(DNSProtoBufMessage::Response);          // set protobuf type to be response - not query
+
+        message.setQueryTime(0, 0);                             // seconds and microseconds
+
+        message.addRRs(strQueryName);                           // add a RR to the response
+     });
+
+                                                                // setProtobufResponseTypeQT - with query time as function parameter
+     g_lua.registerFunction<void(DNSDistProtoBufMessage::*)(const std::string&, time_t sec, uint uSec)>("setProtobufResponseTypeTS", [](DNSDistProtoBufMessage& message, const std::string& strQueryName, time_t sec, uint uSec) {
 
         message.setType(DNSProtoBufMessage::Response);          // set protobuf type to be response - not query
 
@@ -861,6 +880,7 @@ void moreLua(bool client)
         message.addRRs(strQueryName);                           // add a RR to the response
      });
 
+
 // --------------------------------------------------------------------------
 
 
-- 
2.40.0