lua: add dynamic metadata API

This patch add access to request info's dynamic metadata from lua. Signed-off-by: Dhi Aurrahman <[email protected]>
dio · Jun 13, 2018 · 293fcac · 293fcac
1 parent 6d3d5d7
commit 293fcac
Show file tree

Hide file tree

Showing 7 changed files with 207 additions and 9 deletions.
diff --git a/source/extensions/filters/common/lua/BUILD b/source/extensions/filters/common/lua/BUILD
@@ -30,6 +30,9 @@ envoy_cc_library(
     deps = [
         ":lua_lib",
         "//include/envoy/buffer:buffer_interface",
+        "//include/envoy/request_info:request_info_interface",
         "//source/common/protobuf",
+        "//source/common/protobuf:utility_lib",
+        "//source/common/request_info:request_info_lib",
     ],
 )
diff --git a/source/extensions/filters/common/lua/wrappers.cc b/source/extensions/filters/common/lua/wrappers.cc
@@ -1,5 +1,7 @@
 #include "extensions/filters/common/lua/wrappers.h"
 
+#include "common/protobuf/utility.h"
+
 namespace Envoy {
 namespace Extensions {
 namespace Filters {
@@ -26,7 +28,7 @@ int BufferWrapper::luaGetBytes(lua_State* state) {
   return 1;
 }
 
-void MetadataMapWrapper::setValue(lua_State* state, const ProtobufWkt::Value& value) {
+void MetadataMapHelper::setValue(lua_State* state, const ProtobufWkt::Value& value) {
   ProtobufWkt::Value::KindCase kind = value.kind_case();
 
   switch (kind) {
@@ -76,7 +78,7 @@ void MetadataMapWrapper::setValue(lua_State* state, const ProtobufWkt::Value& va
   }
 }
 
-void MetadataMapWrapper::createTable(
+void MetadataMapHelper::createTable(
     lua_State* state,
     const Protobuf::Map<Envoy::ProtobufTypes::String, ProtobufWkt::Value>& fields) {
   lua_createtable(state, 0, fields.size());
@@ -98,7 +100,7 @@ int MetadataMapIterator::luaPairsIterator(lua_State* state) {
   }
 
   lua_pushstring(state, current_->first.c_str());
-  parent_.setValue(state, current_->second);
+  MetadataMapHelper::setValue(state, current_->second);
 
   current_++;
   return 2;
@@ -111,7 +113,7 @@ int MetadataMapWrapper::luaGet(lua_State* state) {
     return 0;
   }
 
-  setValue(state, filter_it->second);
+  MetadataMapHelper::setValue(state, filter_it->second);
   return 1;
 }
 
@@ -125,6 +127,63 @@ int MetadataMapWrapper::luaPairs(lua_State* state) {
   return 1;
 }
 
+int RequestInfoWrapper::luaDynamicMetadata(lua_State* state) {
+  if (metadata_wrapper_.get() != nullptr) {
+    metadata_wrapper_.pushStack();
+  } else {
+    metadata_wrapper_.reset(DynamicMetadataMapWrapper::create(state, request_info_), true);
+  }
+  return 1;
+}
+
+DynamicMetadataMapIterator::DynamicMetadataMapIterator(DynamicMetadataMapWrapper& parent)
+    : parent_{parent}, current_{parent.request_info_.dynamicMetadata().filter_metadata().begin()} {}
+
+int DynamicMetadataMapIterator::luaPairsIterator(lua_State* state) {
+  if (current_ == parent_.request_info_.dynamicMetadata().filter_metadata().end()) {
+    parent_.iterator_.reset();
+    return 0;
+  }
+
+  lua_pushstring(state, current_->first.c_str());
+  MetadataMapHelper::createTable(state, current_->second.fields());
+
+  current_++;
+  return 2;
+}
+
+int DynamicMetadataMapWrapper::luaGet(lua_State* state) {
+  const char* filter_name = luaL_checkstring(state, 2);
+  const auto& metadata = request_info_.dynamicMetadata().filter_metadata();
+  const auto filter_it = metadata.find(filter_name);
+  if (filter_it == metadata.end()) {
+    return 0;
+  }
+
+  MetadataMapHelper::createTable(state, filter_it->second.fields());
+  return 1;
+}
+
+int DynamicMetadataMapWrapper::luaSet(lua_State* state) {
+  const char* filter_name = luaL_checkstring(state, 2);
+  const char* key = luaL_checkstring(state, 3);
+
+  // TODO(dio): Allow to set other than string.
+  const char* value = luaL_checkstring(state, 4);
+  request_info_.setDynamicMetadata(filter_name, MessageUtil::keyValueStruct(key, value));
+  return 0;
+}
+
+int DynamicMetadataMapWrapper::luaPairs(lua_State* state) {
+  if (iterator_.get() != nullptr) {
+    luaL_error(state, "cannot create a second iterator before completing the first");
+  }
+
+  iterator_.reset(DynamicMetadataMapIterator::create(state, *this), true);
+  lua_pushcclosure(state, DynamicMetadataMapIterator::static_luaPairsIterator, 1);
+  return 1;
+}
+
 } // namespace Lua
 } // namespace Common
 } // namespace Filters

diff --git a/source/extensions/filters/common/lua/wrappers.h b/source/extensions/filters/common/lua/wrappers.h
@@ -1,6 +1,7 @@
 #pragma once
 
 #include "envoy/buffer/buffer.h"
+#include "envoy/request_info/request_info.h"
 
 #include "common/protobuf/protobuf.h"
 
@@ -42,6 +43,13 @@ class BufferWrapper : public BaseLuaObject<BufferWrapper> {
 
 class MetadataMapWrapper;
 
+struct MetadataMapHelper {
+  static void setValue(lua_State* state, const ProtobufWkt::Value& value);
+  static void
+  createTable(lua_State* state,
+              const Protobuf::Map<Envoy::ProtobufTypes::String, ProtobufWkt::Value>& fields);
+};
+
 /**
  * Iterator over a metadata map.
  */
@@ -89,16 +97,97 @@ class MetadataMapWrapper : public BaseLuaObject<MetadataMapWrapper> {
     iterator_.reset();
   }
 
-  void setValue(lua_State* state, const ProtobufWkt::Value& value);
-  void createTable(lua_State* state,
-                   const Protobuf::Map<Envoy::ProtobufTypes::String, ProtobufWkt::Value>& fields);
-
   const ProtobufWkt::Struct metadata_;
   LuaDeathRef<MetadataMapIterator> iterator_;
 
   friend class MetadataMapIterator;
 };
 
+class DynamicMetadataMapWrapper;
+
+/**
+ * Iterator over a dynamic metadata map.
+ */
+class DynamicMetadataMapIterator : public BaseLuaObject<DynamicMetadataMapIterator> {
+public:
+  DynamicMetadataMapIterator(DynamicMetadataMapWrapper& parent);
+
+  static ExportedFunctions exportedFunctions() { return {}; }
+
+  DECLARE_LUA_CLOSURE(DynamicMetadataMapIterator, luaPairsIterator);
+
+private:
+  DynamicMetadataMapWrapper& parent_;
+  Protobuf::Map<Envoy::ProtobufTypes::String, ProtobufWkt::Struct>::const_iterator current_;
+};
+
+class DynamicMetadataMapWrapper : public BaseLuaObject<DynamicMetadataMapWrapper> {
+public:
+  DynamicMetadataMapWrapper(RequestInfo::RequestInfo& request_info) : request_info_{request_info} {}
+
+  static ExportedFunctions exportedFunctions() {
+    return {{"get", static_luaGet}, {"set", static_luaSet}, {"__pairs", static_luaPairs}};
+  }
+
+private:
+  /**
+   * Get a metadata value from the map.
+   * @param 1 (string): filter name.
+   * @return value if found or nil.
+   */
+  DECLARE_LUA_FUNCTION(DynamicMetadataMapWrapper, luaGet);
+
+  /**
+   * Get a metadata value from the map.
+   * @param 1 (string): filter name.
+   * @param 2 (string): key.
+   * @param 1 (string): value.
+   * @return nil.
+   */
+  DECLARE_LUA_FUNCTION(DynamicMetadataMapWrapper, luaSet);
+
+  /**
+   * Implementation of the __pairs metamethod so a dynamic metadata wrapper can be iterated over
+   * using pairs().
+   */
+  DECLARE_LUA_FUNCTION(DynamicMetadataMapWrapper, luaPairs);
+
+  // Envoy::Lua::BaseLuaObject
+  void onMarkDead() override {
+    // Iterators do not survive yields.
+    iterator_.reset();
+  }
+
+  RequestInfo::RequestInfo& request_info_;
+  LuaDeathRef<DynamicMetadataMapIterator> iterator_;
+
+  friend class DynamicMetadataMapIterator;
+};
+
+class RequestInfoWrapper : public BaseLuaObject<RequestInfoWrapper> {
+public:
+  RequestInfoWrapper(RequestInfo::RequestInfo& request_info) : request_info_{request_info} {}
+  static ExportedFunctions exportedFunctions() {
+    return {{"dynamicMetadata", static_luaDynamicMetadata}};
+  }
+
+private:
+  /**
+   * Get a dynamic metadata value from the map.
+   * @return value if found or nil.
+   */
+  DECLARE_LUA_FUNCTION(RequestInfoWrapper, luaDynamicMetadata);
+
+  // Envoy::Lua::BaseLuaObject
+  void onMarkDead() override {
+    // TODO(dio): Check if it is required to always reset in here.
+    metadata_wrapper_.reset();
+  }
+
+  LuaDeathRef<DynamicMetadataMapWrapper> metadata_wrapper_;
+  RequestInfo::RequestInfo& request_info_;
+};
+
 } // namespace Lua
 } // namespace Common
 } // namespace Filters

diff --git a/source/extensions/filters/http/lua/lua_filter.cc b/source/extensions/filters/http/lua/lua_filter.cc
@@ -367,6 +367,17 @@ int StreamHandleWrapper::luaMetadata(lua_State* state) {
   return 1;
 }
 
+int StreamHandleWrapper::luaRequestInfo(lua_State* state) {
+  ASSERT(state_ == State::Running);
+  if (metadata_wrapper_.get() != nullptr) {
+    request_info_wrapper_.pushStack();
+  } else {
+    request_info_wrapper_.reset(
+        Filters::Common::Lua::RequestInfoWrapper::create(state, callbacks_.requestInfo()), true);
+  }
+  return 1;
+}
+
 int StreamHandleWrapper::luaLogTrace(lua_State* state) {
   const char* message = luaL_checkstring(state, 2);
   filter_.scriptLog(spdlog::level::trace, message);
@@ -409,6 +420,9 @@ FilterConfig::FilterConfig(const std::string& lua_code, ThreadLocal::SlotAllocat
   lua_state_.registerType<Filters::Common::Lua::BufferWrapper>();
   lua_state_.registerType<Filters::Common::Lua::MetadataMapWrapper>();
   lua_state_.registerType<Filters::Common::Lua::MetadataMapIterator>();
+  lua_state_.registerType<Filters::Common::Lua::DynamicMetadataMapWrapper>();
+  lua_state_.registerType<Filters::Common::Lua::DynamicMetadataMapIterator>();
+  lua_state_.registerType<Filters::Common::Lua::RequestInfoWrapper>();
   lua_state_.registerType<HeaderMapWrapper>();
   lua_state_.registerType<HeaderMapIterator>();
   lua_state_.registerType<StreamHandleWrapper>();

diff --git a/source/extensions/filters/http/lua/lua_filter.h b/source/extensions/filters/http/lua/lua_filter.h
@@ -68,6 +68,8 @@ class FilterCallbacks {
    * route entry.
    */
   virtual const ProtobufWkt::Struct& metadata() const PURE;
+
+  virtual RequestInfo::RequestInfo& requestInfo() PURE;
 };
 
 class Filter;
@@ -122,7 +124,7 @@ class StreamHandleWrapper : public Filters::Common::Lua::BaseLuaObject<StreamHan
             {"logDebug", static_luaLogDebug},       {"logInfo", static_luaLogInfo},
             {"logWarn", static_luaLogWarn},         {"logErr", static_luaLogErr},
             {"logCritical", static_luaLogCritical}, {"httpCall", static_luaHttpCall},
-            {"respond", static_luaRespond}};
+            {"respond", static_luaRespond},         {"requestInfo", static_luaRequestInfo}};
   }
 
 private:
@@ -176,6 +178,8 @@ class StreamHandleWrapper : public Filters::Common::Lua::BaseLuaObject<StreamHan
    */
   DECLARE_LUA_FUNCTION(StreamHandleWrapper, luaMetadata);
 
+  DECLARE_LUA_FUNCTION(StreamHandleWrapper, luaRequestInfo);
+
   /**
    * Log a message to the Envoy log.
    * @param 1 (string): The log message.
@@ -202,6 +206,7 @@ class StreamHandleWrapper : public Filters::Common::Lua::BaseLuaObject<StreamHan
     body_wrapper_.reset();
     trailers_wrapper_.reset();
     metadata_wrapper_.reset();
+    request_info_wrapper_.reset();
   }
 
   // Http::AsyncClient::Callbacks
@@ -221,6 +226,7 @@ class StreamHandleWrapper : public Filters::Common::Lua::BaseLuaObject<StreamHan
   Filters::Common::Lua::LuaDeathRef<Filters::Common::Lua::BufferWrapper> body_wrapper_;
   Filters::Common::Lua::LuaDeathRef<HeaderMapWrapper> trailers_wrapper_;
   Filters::Common::Lua::LuaDeathRef<Filters::Common::Lua::MetadataMapWrapper> metadata_wrapper_;
+  Filters::Common::Lua::LuaDeathRef<Filters::Common::Lua::RequestInfoWrapper> request_info_wrapper_;
   State state_{State::Running};
   std::function<void()> yield_callback_;
   Http::AsyncClient::Request* http_request_{};
@@ -310,6 +316,7 @@ class Filter : public Http::StreamFilter, Logger::Loggable<Logger::Id::lua> {
     void respond(Http::HeaderMapPtr&& headers, Buffer::Instance* body, lua_State* state) override;
 
     const ProtobufWkt::Struct& metadata() const override { return getMetadata(callbacks_); }
+    RequestInfo::RequestInfo& requestInfo() override { return callbacks_->requestInfo(); }
 
     Filter& parent_;
     Http::StreamDecoderFilterCallbacks* callbacks_{};
@@ -328,6 +335,7 @@ class Filter : public Http::StreamFilter, Logger::Loggable<Logger::Id::lua> {
     void respond(Http::HeaderMapPtr&& headers, Buffer::Instance* body, lua_State* state) override;
 
     const ProtobufWkt::Struct& metadata() const override { return getMetadata(callbacks_); }
+    RequestInfo::RequestInfo& requestInfo() override { return callbacks_->requestInfo(); }
 
     Filter& parent_;
     Http::StreamEncoderFilterCallbacks* callbacks_{};

diff --git a/test/extensions/filters/http/lua/BUILD b/test/extensions/filters/http/lua/BUILD
@@ -16,6 +16,7 @@ envoy_extension_cc_test(
     srcs = ["lua_filter_test.cc"],
     extension_name = "envoy.filters.http.lua",
     deps = [
+        "//source/common/request_info:request_info_lib",
         "//source/extensions/filters/http/lua:lua_filter_lib",
         "//test/mocks/http:http_mocks",
         "//test/mocks/thread_local:thread_local_mocks",

diff --git a/test/extensions/filters/http/lua/lua_filter_test.cc b/test/extensions/filters/http/lua/lua_filter_test.cc
@@ -1,5 +1,6 @@
 #include "common/buffer/buffer_impl.h"
 #include "common/http/message_impl.h"
+#include "common/request_info/request_info_impl.h"
 
 #include "extensions/filters/http/lua/lua_filter.h"
 
@@ -1460,6 +1461,29 @@ TEST_F(LuaHttpFilterTest, GetMetadataFromHandleNoLuaMetadata) {
   EXPECT_EQ(Http::FilterHeadersStatus::Continue, filter_->decodeHeaders(request_headers, true));
 }
 
+// Set request's request-info dynamic metadata.
+TEST_F(LuaHttpFilterTest, RequestInfoSetDynamicMetadata) {
+  const std::string SCRIPT{R"EOF(
+    function envoy_on_request(request_handle)
+      request_handle:requestInfo():dynamicMetadata():set("envoy.lb", "foo", "bar")
+      request_handle:logTrace(request_handle:requestInfo():dynamicMetadata():get("envoy.lb")["foo"])
+    end
+  )EOF"};
+
+  InSequence s;
+  setup(SCRIPT);
+
+  RequestInfo::RequestInfoImpl request_info(Http::Protocol::Http2);
+  EXPECT_EQ(0, request_info.dynamicMetadata().filter_metadata_size());
+  EXPECT_CALL(decoder_callbacks_, requestInfo()).WillRepeatedly(testing::ReturnRef(request_info));
+
+  Http::TestHeaderMapImpl request_headers{{":path", "/"}};
+  EXPECT_CALL(*filter_, scriptLog(spdlog::level::trace, StrEq("bar")));
+  EXPECT_EQ(Http::FilterHeadersStatus::Continue, filter_->decodeHeaders(request_headers, true));
+
+  EXPECT_EQ(1, request_info.dynamicMetadata().filter_metadata_size());
+}
+
 } // namespace Lua
 } // namespace HttpFilters
 } // namespace Extensions