extensions/splunk/tests/MockSplunkHEC.h - nifi-minifi-cpp - Git at Google

 /**
  * Licensed to the Apache Software Foundation (ASF) under one or more
  * contributor license agreements.  See the NOTICE file distributed with
  * this work for additional information regarding copyright ownership.
  * The ASF licenses this file to You under the Apache License, Version 2.0
  * (the "License"); you may not use this file except in compliance with
  * the License.  You may obtain a copy of the License at
  *
  *     http://www.apache.org/licenses/LICENSE-2.0
  *
  * Unless required by applicable law or agreed to in writing, software
  * distributed under the License is distributed on an "AS IS" BASIS,
  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */

 #pragma once

 #include <memory>
 #include <string>
 #include <utility>
 #include <vector>
 #include <CivetServer.h>
 #include "integration/CivetLibrary.h"
 #include "core/logging/Logger.h"
 #include "core/logging/LoggerFactory.h"
 #include "rapidjson/document.h"
 #include "rapidjson/writer.h"
 #include "rapidjson/stringbuffer.h"


 class MockSplunkHandler : public CivetHandler {
  public:
   explicit MockSplunkHandler(std::string token, std::function<void(const struct mg_request_info *request_info)>& assertions) : token_(std::move(token)), assertions_(assertions) {
   }

   enum HeaderResult {
     MissingAuth,
     InvalidAuth,
     MissingReqChannel,
     HeadersOk
   };

   bool handlePost(CivetServer*, struct mg_connection *conn) override {
     switch (checkHeaders(conn)) {
       case MissingAuth:
         return send401(conn);
       case InvalidAuth:
         return send403(conn);
       case MissingReqChannel:
         return send400(conn);
       case HeadersOk:
         return handlePostImpl(conn);
     }
     return false;
   }

   HeaderResult checkHeaders(struct mg_connection *conn) const {
     const struct mg_request_info* req_info = mg_get_request_info(conn);
     assertions_(req_info);
     auto auth_header = std::find_if(std::begin(req_info->http_headers),
                                     std::end(req_info->http_headers),
                                     [](auto header) -> bool {return strcmp(header.name, "Authorization") == 0;});
     if (auth_header == std::end(req_info->http_headers))
       return MissingAuth;
     if (strcmp(auth_header->value, token_.c_str()) != 0)
       return InvalidAuth;

     auto request_channel_header = std::find_if(std::begin(req_info->http_headers),
                                                std::end(req_info->http_headers),
                                                [](auto header) -> bool {return strcmp(header.name, "X-Splunk-Request-Channel") == 0;});

     if (request_channel_header == std::end(req_info->http_headers))
       return MissingReqChannel;
     return HeadersOk;
   }

   bool send400(struct mg_connection *conn) const {
     constexpr const char * body = "{\"text\":\"Data channel is missing\",\"code\":10}";
     mg_printf(conn, "HTTP/1.1 400 Bad Request\r\n");
     mg_printf(conn, "Content-length: %lu", strlen(body));
     mg_printf(conn, "\r\n\r\n");
     mg_printf(conn, body);
     return true;
   }

   bool send401(struct mg_connection *conn) const {
     constexpr const char * body = "{\"text\":\"Token is required\",\"code\":2}";
     mg_printf(conn, "HTTP/1.1 401 Unauthorized\r\n");
     mg_printf(conn, "Content-length: %lu", strlen(body));
     mg_printf(conn, "\r\n\r\n");
     mg_printf(conn, body);
     return true;
   }

   bool send403(struct mg_connection *conn) const {
     constexpr const char * body = "{\"text\":\"Invalid token\",\"code\":4}";
     mg_printf(conn, "HTTP/1.1 403 Forbidden\r\n");
     mg_printf(conn, "Content-length: %lu", strlen(body));
     mg_printf(conn, "\r\n\r\n");
     mg_printf(conn, body);
     return true;
   }

  protected:
   virtual bool handlePostImpl(struct mg_connection *conn) = 0;
   std::string token_;
   std::function<void(const struct mg_request_info *request_info)>& assertions_;
 };

 class RawCollectorHandler : public MockSplunkHandler {
  public:
   explicit RawCollectorHandler(std::string token, std::function<void(const struct mg_request_info *request_info)>& assertions) : MockSplunkHandler(std::move(token), assertions) {}
  protected:
   bool handlePostImpl(struct mg_connection* conn) override {
     constexpr const char * body = "{\"text\":\"Success\",\"code\":0,\"ackId\":808}";
     mg_printf(conn, "HTTP/1.1 200 OK\r\n");
     mg_printf(conn, "Content-length: %lu", strlen(body));
     mg_printf(conn, "\r\n\r\n");
     mg_printf(conn, body);
     return true;
   }
 };

 class AckIndexerHandler : public MockSplunkHandler {
  public:
   explicit AckIndexerHandler(std::string token, std::vector<uint64_t> indexed_events, std::function<void(const struct mg_request_info *request_info)>& assertions)
       : MockSplunkHandler(std::move(token), assertions), indexed_events_(indexed_events) {}

  protected:
   bool handlePostImpl(struct mg_connection* conn) override {
     std::vector<char> data;
     data.reserve(2048);
     mg_read(conn, data.data(), 2048);
     rapidjson::Document post_data;

     rapidjson::ParseResult parse_result = post_data.Parse<rapidjson::kParseStopWhenDoneFlag>(data.data());
     if (parse_result.IsError())
       return sendInvalidFormat(conn);
     if (!post_data.HasMember("acks") || !post_data["acks"].IsArray())
       return sendInvalidFormat(conn);
     std::vector<uint64_t> ids;
     for (auto& id : post_data["acks"].GetArray()) {
       ids.push_back(id.GetUint64());
     }
     rapidjson::Document reply = rapidjson::Document(rapidjson::kObjectType);
     reply.AddMember("acks", rapidjson::kObjectType, reply.GetAllocator());
     for (auto& id : ids) {
       rapidjson::Value key(std::to_string(id).c_str(), reply.GetAllocator());
       reply["acks"].AddMember(key, std::find(indexed_events_.begin(), indexed_events_.end(), id) != indexed_events_.end() ? true : false, reply.GetAllocator());
     }
     rapidjson::StringBuffer buffer;
     rapidjson::Writer<rapidjson::StringBuffer> writer(buffer);
     reply.Accept(writer);

     mg_printf(conn, "HTTP/1.1 200 OK\r\n");
     mg_printf(conn, "Content-length: %lu", buffer.GetSize());
     mg_printf(conn, "\r\n\r\n");
     mg_printf(conn, "%s" , buffer.GetString());
     return true;
   }

   bool sendInvalidFormat(struct mg_connection* conn) {
     constexpr const char * body = "{\"text\":\"Invalid data format\",\"code\":6}";
     mg_printf(conn, "HTTP/1.1 400 Bad Request\r\n");
     mg_printf(conn, "Content-length: %lu", strlen(body));
     mg_printf(conn, "\r\n\r\n");
     mg_printf(conn, body);
     return true;
   }

   std::vector<uint64_t> indexed_events_;
 };

 class MockSplunkHEC {
  public:
   static constexpr const char* TOKEN = "Splunk 822f7d13-2b70-4f8c-848b-86edfc251222";

   static inline std::vector<uint64_t> indexed_events = {0, 1};

   explicit MockSplunkHEC(std::string port) : port_(std::move(port)) {
     std::vector<std::string> options;
     options.emplace_back("listening_ports");
     options.emplace_back(port_);
     server_.reset(new CivetServer(options, &callbacks_, &logger_));
     {
       MockSplunkHandler* raw_collector_handler = new RawCollectorHandler(TOKEN, assertions_);
       server_->addHandler("/services/collector/raw", raw_collector_handler);
       handlers_.emplace_back(std::move(raw_collector_handler));
     }
     {
       MockSplunkHandler* ack_indexer_handler = new AckIndexerHandler(TOKEN, indexed_events, assertions_);
       server_->addHandler("/services/collector/ack", ack_indexer_handler);
       handlers_.emplace_back(std::move(ack_indexer_handler));
     }
   }

   const std::string& getPort() const {
     return port_;
   }

   void setAssertions(std::function<void(const struct mg_request_info *request_info)> assertions) {
     assertions_ = assertions;
   }


  private:
   CivetLibrary lib_;
   std::string port_;
   std::unique_ptr<CivetServer> server_;
   std::vector<std::unique_ptr<MockSplunkHandler>> handlers_;
   CivetCallbacks callbacks_;
   std::function<void(const struct mg_request_info *request_info)> assertions_ = [](const struct mg_request_info*) {};
   std::shared_ptr<org::apache::nifi::minifi::core::logging::Logger> logger_ = org::apache::nifi::minifi::core::logging::LoggerFactory<MockSplunkHEC>::getLogger();
 };
	/**
	* Licensed to the Apache Software Foundation (ASF) under one or more
	* contributor license agreements. See the NOTICE file distributed with
	* this work for additional information regarding copyright ownership.
	* The ASF licenses this file to You under the Apache License, Version 2.0
	* (the "License"); you may not use this file except in compliance with
	* the License. You may obtain a copy of the License at
	*
	* http://www.apache.org/licenses/LICENSE-2.0
	*
	* Unless required by applicable law or agreed to in writing, software
	* distributed under the License is distributed on an "AS IS" BASIS,
	* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	* See the License for the specific language governing permissions and
	* limitations under the License.
	*/

	#pragma once

	#include <memory>
	#include <string>
	#include <utility>
	#include <vector>
	#include <CivetServer.h>
	#include "integration/CivetLibrary.h"
	#include "core/logging/Logger.h"
	#include "core/logging/LoggerFactory.h"
	#include "rapidjson/document.h"
	#include "rapidjson/writer.h"
	#include "rapidjson/stringbuffer.h"


	class MockSplunkHandler : public CivetHandler {
	public:
	explicit MockSplunkHandler(std::string token, std::function<void(const struct mg_request_info *request_info)>& assertions) : token_(std::move(token)), assertions_(assertions) {
	}

	enum HeaderResult {
	MissingAuth,
	InvalidAuth,
	MissingReqChannel,
	HeadersOk
	};

	bool handlePost(CivetServer, struct mg_connection conn) override {
	switch (checkHeaders(conn)) {
	case MissingAuth:
	return send401(conn);
	case InvalidAuth:
	return send403(conn);
	case MissingReqChannel:
	return send400(conn);
	case HeadersOk:
	return handlePostImpl(conn);
	}
	return false;
	}

	HeaderResult checkHeaders(struct mg_connection *conn) const {
	const struct mg_request_info* req_info = mg_get_request_info(conn);
	assertions_(req_info);
	auto auth_header = std::find_if(std::begin(req_info->http_headers),
	std::end(req_info->http_headers),
	[](auto header) -> bool {return strcmp(header.name, "Authorization") == 0;});
	if (auth_header == std::end(req_info->http_headers))
	return MissingAuth;
	if (strcmp(auth_header->value, token_.c_str()) != 0)
	return InvalidAuth;

	auto request_channel_header = std::find_if(std::begin(req_info->http_headers),
	std::end(req_info->http_headers),
	[](auto header) -> bool {return strcmp(header.name, "X-Splunk-Request-Channel") == 0;});

	if (request_channel_header == std::end(req_info->http_headers))
	return MissingReqChannel;
	return HeadersOk;
	}

	bool send400(struct mg_connection *conn) const {
	constexpr const char * body = "{\"text\":\"Data channel is missing\",\"code\":10}";
	mg_printf(conn, "HTTP/1.1 400 Bad Request\r\n");
	mg_printf(conn, "Content-length: %lu", strlen(body));
	mg_printf(conn, "\r\n\r\n");
	mg_printf(conn, body);
	return true;
	}

	bool send401(struct mg_connection *conn) const {
	constexpr const char * body = "{\"text\":\"Token is required\",\"code\":2}";
	mg_printf(conn, "HTTP/1.1 401 Unauthorized\r\n");
	mg_printf(conn, "Content-length: %lu", strlen(body));
	mg_printf(conn, "\r\n\r\n");
	mg_printf(conn, body);
	return true;
	}

	bool send403(struct mg_connection *conn) const {
	constexpr const char * body = "{\"text\":\"Invalid token\",\"code\":4}";
	mg_printf(conn, "HTTP/1.1 403 Forbidden\r\n");
	mg_printf(conn, "Content-length: %lu", strlen(body));
	mg_printf(conn, "\r\n\r\n");
	mg_printf(conn, body);
	return true;
	}

	protected:
	virtual bool handlePostImpl(struct mg_connection *conn) = 0;
	std::string token_;
	std::function<void(const struct mg_request_info *request_info)>& assertions_;
	};

	class RawCollectorHandler : public MockSplunkHandler {
	public:
	explicit RawCollectorHandler(std::string token, std::function<void(const struct mg_request_info *request_info)>& assertions) : MockSplunkHandler(std::move(token), assertions) {}
	protected:
	bool handlePostImpl(struct mg_connection* conn) override {
	constexpr const char * body = "{\"text\":\"Success\",\"code\":0,\"ackId\":808}";
	mg_printf(conn, "HTTP/1.1 200 OK\r\n");
	mg_printf(conn, "Content-length: %lu", strlen(body));
	mg_printf(conn, "\r\n\r\n");
	mg_printf(conn, body);
	return true;
	}
	};

	class AckIndexerHandler : public MockSplunkHandler {
	public:
	explicit AckIndexerHandler(std::string token, std::vector<uint64_t> indexed_events, std::function<void(const struct mg_request_info *request_info)>& assertions)
	: MockSplunkHandler(std::move(token), assertions), indexed_events_(indexed_events) {}

	protected:
	bool handlePostImpl(struct mg_connection* conn) override {
	std::vector<char> data;
	data.reserve(2048);
	mg_read(conn, data.data(), 2048);
	rapidjson::Document post_data;

	rapidjson::ParseResult parse_result = post_data.Parse<rapidjson::kParseStopWhenDoneFlag>(data.data());
	if (parse_result.IsError())
	return sendInvalidFormat(conn);
	if (!post_data.HasMember("acks") \|\| !post_data["acks"].IsArray())
	return sendInvalidFormat(conn);
	std::vector<uint64_t> ids;
	for (auto& id : post_data["acks"].GetArray()) {
	ids.push_back(id.GetUint64());
	}
	rapidjson::Document reply = rapidjson::Document(rapidjson::kObjectType);
	reply.AddMember("acks", rapidjson::kObjectType, reply.GetAllocator());
	for (auto& id : ids) {
	rapidjson::Value key(std::to_string(id).c_str(), reply.GetAllocator());
	reply["acks"].AddMember(key, std::find(indexed_events_.begin(), indexed_events_.end(), id) != indexed_events_.end() ? true : false, reply.GetAllocator());
	}
	rapidjson::StringBuffer buffer;
	rapidjson::Writer<rapidjson::StringBuffer> writer(buffer);
	reply.Accept(writer);

	mg_printf(conn, "HTTP/1.1 200 OK\r\n");
	mg_printf(conn, "Content-length: %lu", buffer.GetSize());
	mg_printf(conn, "\r\n\r\n");
	mg_printf(conn, "%s" , buffer.GetString());
	return true;
	}

	bool sendInvalidFormat(struct mg_connection* conn) {
	constexpr const char * body = "{\"text\":\"Invalid data format\",\"code\":6}";
	mg_printf(conn, "HTTP/1.1 400 Bad Request\r\n");
	mg_printf(conn, "Content-length: %lu", strlen(body));
	mg_printf(conn, "\r\n\r\n");
	mg_printf(conn, body);
	return true;
	}

	std::vector<uint64_t> indexed_events_;
	};

	class MockSplunkHEC {
	public:
	static constexpr const char* TOKEN = "Splunk 822f7d13-2b70-4f8c-848b-86edfc251222";

	static inline std::vector<uint64_t> indexed_events = {0, 1};

	explicit MockSplunkHEC(std::string port) : port_(std::move(port)) {
	std::vector<std::string> options;
	options.emplace_back("listening_ports");
	options.emplace_back(port_);
	server_.reset(new CivetServer(options, &callbacks_, &logger_));
	{
	MockSplunkHandler* raw_collector_handler = new RawCollectorHandler(TOKEN, assertions_);
	server_->addHandler("/services/collector/raw", raw_collector_handler);
	handlers_.emplace_back(std::move(raw_collector_handler));
	}
	{
	MockSplunkHandler* ack_indexer_handler = new AckIndexerHandler(TOKEN, indexed_events, assertions_);
	server_->addHandler("/services/collector/ack", ack_indexer_handler);
	handlers_.emplace_back(std::move(ack_indexer_handler));
	}
	}

	const std::string& getPort() const {
	return port_;
	}

	void setAssertions(std::function<void(const struct mg_request_info *request_info)> assertions) {
	assertions_ = assertions;
	}


	private:
	CivetLibrary lib_;
	std::string port_;
	std::unique_ptr<CivetServer> server_;
	std::vector<std::unique_ptr<MockSplunkHandler>> handlers_;
	CivetCallbacks callbacks_;
	std::function<void(const struct mg_request_info request_info)> assertions_ = [](const struct mg_request_info) {};
	std::shared_ptr<org::apache::nifi::minifi::core::logging::Logger> logger_ = org::apache::nifi::minifi::core::logging::LoggerFactory<MockSplunkHEC>::getLogger();
	};