1 #include <pichi/common/config.hpp>
2 // Include config.hpp first
3 #include <boost/asio/io_context.hpp>
4 #include <boost/asio/ip/tcp.hpp>
5 #include <boost/beast/core/flat_static_buffer.hpp>
6 #include <boost/beast/http/empty_body.hpp>
7 #include <boost/beast/http/message.hpp>
8 #include <boost/beast/http/read.hpp>
9 #include <boost/beast/http/string_body.hpp>
10 #include <boost/beast/http/write.hpp>
11 #include <boost/filesystem/path.hpp>
12 #include <fstream>
13 #include <iostream>
14 #include <pichi/api/server.hpp>
15 #include <pichi/common/asserts.hpp>
16 #include <pichi/common/endpoint.hpp>
17 #include <pichi/net/adapter.hpp>
18 #include <pichi/net/helper.hpp>
19 #include <pichi/net/spawn.hpp>
20 #include <pichi/vo/to_json.hpp>
21 #include <rapidjson/document.h>
22 #include <rapidjson/istreamwrapper.h>
23 #include <vector>
24 
25 #ifdef HAS_SIGNAL_H
26 #include <boost/asio/signal_set.hpp>
27 #include <signal.h>
28 #endif  // HAS_SIGNAL_H
29 
30 using namespace std;
31 using namespace pichi;
32 namespace asio = boost::asio;
33 namespace beast = boost::beast;
34 namespace fs = boost::filesystem;
35 namespace http = boost::beast::http;
36 namespace ip = asio::ip;
37 namespace json = rapidjson;
38 
39 using ip::tcp;
40 
41 static decltype(auto) INGRESSES = "ingresses";
42 static decltype(auto) EGRESSES = "egresses";
43 static decltype(auto) RULES = "rules";
44 static decltype(auto) ROUTE = "route";
45 static decltype(auto) INDENT = "  ";
46 
47 static asio::io_context io{1};
48 
parseJson(char const * str)49 static json::Document parseJson(char const* str)
50 {
51   auto doc = json::Document{};
52   doc.Parse(str);
53   assertFalse(doc.HasParseError());
54   return doc;
55 }
56 
57 class HttpHelper {
58 public:
59   HttpHelper(ResolveResults&& endpoint, asio::yield_context yield);
60 
61   vector<string> get(string const& target);
62   void put(string const& target, string const& body);
63   void del(string const& target);
64 
65 private:
66   ResolveResults endpoint_;
67   asio::yield_context yield_;
68 };
69 
HttpHelper(ResolveResults && endpoint,asio::yield_context yield)70 HttpHelper::HttpHelper(ResolveResults&& endpoint, asio::yield_context yield)
71   : endpoint_{move(endpoint)}, yield_{yield}
72 {
73 }
74 
get(string const & target)75 vector<string> HttpHelper::get(string const& target)
76 {
77   auto s = tcp::socket{io};
78   net::connect(endpoint_, s, yield_);
79 
80   auto req = http::request<http::empty_body>{};
81   req.method(http::verb::get);
82   req.target(target);
83   req.version(11);
84   http::async_write(s, req, yield_);
85 
86   auto buf = beast::flat_static_buffer<1024>{};
87   auto resp = http::response<http::string_body>{};
88   http::async_read(s, buf, resp, yield_);
89   assertTrue(resp.result() == http::status::ok);
90 
91   auto doc = parseJson(resp.body().c_str());
92   auto ret = vector<string>{};
93   transform(doc.MemberBegin(), doc.MemberEnd(), back_inserter(ret),
94             [](auto&& member) { return member.name.GetString(); });
95   return ret;
96 }
97 
put(string const & target,string const & body)98 void HttpHelper::put(string const& target, string const& body)
99 {
100   auto s = tcp::socket{io};
101   net::connect(endpoint_, s, yield_);
102 
103   auto req = http::request<http::string_body>{};
104   req.method(http::verb::put);
105   req.target(target);
106   req.version(11);
107   req.body() = body;
108   req.prepare_payload();
109   http::async_write(s, req, yield_);
110 
111   auto buf = beast::flat_static_buffer<1024>{};
112   auto resp = http::response<http::string_body>{};
113   http::async_read(s, buf, resp, yield_);
114   if (resp.result() != http::status::no_content) cout << INDENT << target << " NOT loaded" << endl;
115 }
116 
del(string const & target)117 void HttpHelper::del(string const& target)
118 {
119   auto s = tcp::socket{io};
120   net::connect(endpoint_, s, yield_);
121 
122   auto req = http::request<http::string_body>{};
123   req.method(http::verb::delete_);
124   req.target(target);
125   req.version(11);
126   http::async_write(s, req, yield_);
127 
128   auto buf = beast::flat_static_buffer<1024>{};
129   auto resp = http::response<http::string_body>{};
130   http::async_read(s, buf, resp, yield_);
131   assertTrue(resp.result() == http::status::no_content);
132 }
133 
readJson(string const & fn)134 static auto readJson(string const& fn)
135 {
136   auto doc = json::Document{};
137   auto ifs = ifstream{fn};
138   auto isw = json::IStreamWrapper{ifs};
139   doc.ParseStream(isw);
140   if (doc.HasParseError() || !doc.IsObject()) {
141     cout << "Invalid JSON configuration" << endl;
142     doc = json::Document{};
143     doc.SetObject();
144   }
145   return doc;
146 }
147 
148 template <typename StringRef>
loadSet(HttpHelper & helper,json::Value const & root,StringRef const & category)149 static void loadSet(HttpHelper& helper, json::Value const& root, StringRef const& category)
150 {
151   auto it = root.FindMember(category);
152   if (it == root.MemberEnd() || !it->value.IsObject()) {
153     cout << INDENT << category << " NOT loaded" << endl;
154     return;
155   }
156 
157   for (auto&& node : it->value.GetObject())
158     helper.put("/"s + category + "/" + node.name.GetString(), vo::toString(node.value));
159 }
160 
load(HttpHelper & helper,string const & fn)161 static void load(HttpHelper& helper, string const& fn)
162 {
163   if (fn.empty()) return;
164   cout << "Loading configuration: " << fn << endl;
165   auto json = readJson(fn);
166   loadSet(helper, json, INGRESSES);
167   loadSet(helper, json, EGRESSES);
168   loadSet(helper, json, RULES);
169   if (json.HasMember(ROUTE))
170     helper.put("/"s + ROUTE, vo::toString(json[ROUTE]));
171   else
172     cout << INDENT << ROUTE << " NOT loaded" << endl;
173   cout << "Configuration " << fn << " loaded" << endl;
174 }
175 
176 #if defined(HAS_SIGNAL_H) && defined(SIGHUP)
flush(HttpHelper & helper)177 static void flush(HttpHelper& helper)
178 {
179   helper.put("/"s + EGRESSES + "/direct", "{\"type\":\"direct\"}"s);
180   helper.put("/"s + ROUTE, "{\"default\":\"direct\",\"rules\":[]}"s);
181 
182   for (auto&& rule : helper.get("/"s + RULES)) helper.del("/"s + RULES + "/" + rule);
183   for (auto&& egress : helper.get("/"s + EGRESSES))
184     if (egress != "direct") helper.del("/"s + EGRESSES + "/" + egress);
185   for (auto&& ingress : helper.get("/"s + INGRESSES)) helper.del("/"s + INGRESSES + "/" + ingress);
186 
187   cout << "Configuration reset" << endl;
188 }
189 #endif  // defined(HAS_SIGNAL_H) && defined(SIGHUP)
190 
run(string const & bind,uint16_t port,string const & fn,string const & mmdb)191 void run(string const& bind, uint16_t port, string const& fn, string const& mmdb)
192 {
193   auto server = api::Server{io, mmdb.c_str()};
194   server.listen(bind, port);
195 
196   // FIXME load & flush aren't designed to be the atomic operations.
197   net::spawn(
198       io,
199       [=](auto yield) {
200         auto resolver = tcp::resolver{io};
201         auto helper = HttpHelper{resolver.async_resolve(bind, to_string(port), yield), yield};
202         load(helper, fn);
203 
204 #if defined(HAS_SIGNAL_H) && defined(SIGHUP)
205         auto ss = asio::signal_set{io};
206         while (true) {
207           ss.add(SIGHUP);
208           ss.async_wait(yield);
209           flush(helper);
210           load(helper, fn);
211         }
212 #endif  // defined(HAS_SIGNAL_H) && defined(SIGHUP)
213       },
214       [](auto, auto) noexcept { io.stop(); });
215 
216   io.run();
217 }
218