diff options
Diffstat (limited to 'server_main.cpp')
-rw-r--r-- | server_main.cpp | 230 |
1 files changed, 230 insertions, 0 deletions
diff --git a/server_main.cpp b/server_main.cpp new file mode 100644 index 0000000..8da6477 --- /dev/null +++ b/server_main.cpp | |||
@@ -0,0 +1,230 @@ | |||
1 | #define ASIO_STANDALONE | ||
2 | |||
3 | #include <base64.hpp> | ||
4 | #include <chrono> | ||
5 | #include <fstream> | ||
6 | #include <functional> | ||
7 | #include <json.hpp> | ||
8 | #include <mutex> | ||
9 | #include <random> | ||
10 | #include <thread> | ||
11 | #include <websocketpp/config/asio_no_tls.hpp> | ||
12 | #include <websocketpp/server.hpp> | ||
13 | |||
14 | #include "cardset.h" | ||
15 | #include "database.h" | ||
16 | #include "imagestore.h" | ||
17 | #include "wizard.h" | ||
18 | |||
19 | namespace { | ||
20 | |||
21 | using socket_type = websocketpp::server<websocketpp::config::asio>; | ||
22 | |||
23 | class server { | ||
24 | public: | ||
25 | server(const cardset& cards, const imagestore& images, std::mt19937& rng) | ||
26 | : cards_(cards), images_(images), rng_(rng) { | ||
27 | socket_.init_asio(); | ||
28 | |||
29 | socket_.set_message_handler([this](websocketpp::connection_hdl connection, | ||
30 | socket_type::message_ptr message) { | ||
31 | on_message(connection, message); | ||
32 | }); | ||
33 | } | ||
34 | |||
35 | void run() { | ||
36 | socket_.listen(9002); | ||
37 | socket_.start_accept(); | ||
38 | |||
39 | asio::post(std::bind(&server::cleanup_thread, this)); | ||
40 | |||
41 | socket_.run(); | ||
42 | } | ||
43 | |||
44 | private: | ||
45 | void on_message(websocketpp::connection_hdl connection, | ||
46 | socket_type::message_ptr message) { | ||
47 | nlohmann::json msgJson; | ||
48 | |||
49 | try { | ||
50 | msgJson = nlohmann::json::parse(message->get_payload()); | ||
51 | } catch (const std::exception& ex) { | ||
52 | std::string response = R"( | ||
53 | { | ||
54 | "type": "error", | ||
55 | "msg": "Could not parse message." | ||
56 | } | ||
57 | )"; | ||
58 | |||
59 | socket_.send(connection, response, | ||
60 | websocketpp::frame::opcode::value::TEXT); | ||
61 | return; | ||
62 | } | ||
63 | |||
64 | std::string cmd = msgJson["cmd"]; | ||
65 | if (cmd == "generate") { | ||
66 | cmd_generate(connection, msgJson["text"]); | ||
67 | } else if (cmd == "check") { | ||
68 | cmd_check(connection, msgJson["token"]); | ||
69 | } else { | ||
70 | std::string response = R"( | ||
71 | { | ||
72 | "type": "error", | ||
73 | "msg": "No command in message." | ||
74 | } | ||
75 | )"; | ||
76 | |||
77 | socket_.send(connection, response, | ||
78 | websocketpp::frame::opcode::value::TEXT); | ||
79 | return; | ||
80 | } | ||
81 | } | ||
82 | |||
83 | void cmd_generate(websocketpp::connection_hdl connection, std::string text) { | ||
84 | std::string token; | ||
85 | |||
86 | { | ||
87 | std::lock_guard rng_guard(rng_mutex_); | ||
88 | token = database_.create(rng_); | ||
89 | } | ||
90 | |||
91 | nlohmann::json tokenMsg; | ||
92 | tokenMsg["type"] = "token"; | ||
93 | tokenMsg["token"] = token; | ||
94 | socket_.send(connection, tokenMsg.dump(), | ||
95 | websocketpp::frame::opcode::value::TEXT); | ||
96 | |||
97 | database_.subscribe(token, [this, connection](const std::string& msg) { | ||
98 | socket_.send(connection, msg, websocketpp::frame::opcode::value::TEXT); | ||
99 | }); | ||
100 | |||
101 | asio::post(std::bind(&server::generate_thread, this, token, text)); | ||
102 | } | ||
103 | |||
104 | void cmd_check(websocketpp::connection_hdl connection, std::string token) { | ||
105 | bool failed = false; | ||
106 | |||
107 | try { | ||
108 | database_.subscribe(token, [this, connection](const std::string& msg) { | ||
109 | socket_.send(connection, msg, websocketpp::frame::opcode::value::TEXT); | ||
110 | }); | ||
111 | |||
112 | if (!database_.is_done(token)) { | ||
113 | return; | ||
114 | } | ||
115 | |||
116 | std::string result = database_.getResult(token); | ||
117 | nlohmann::json resultMsg; | ||
118 | if (result.empty()) { | ||
119 | resultMsg["type"] = "error"; | ||
120 | resultMsg["msg"] = "Unknown error occurred."; | ||
121 | } else { | ||
122 | resultMsg["type"] = "result"; | ||
123 | resultMsg["image"] = database_.getResult(token); | ||
124 | resultMsg["msg"] = "Success!"; | ||
125 | } | ||
126 | |||
127 | socket_.send(connection, resultMsg.dump(), | ||
128 | websocketpp::frame::opcode::value::TEXT); | ||
129 | } catch (const std::exception& ex) { | ||
130 | failed = true; | ||
131 | } | ||
132 | |||
133 | if (failed) { | ||
134 | try { | ||
135 | socket_.send(connection, R"( | ||
136 | { | ||
137 | "type": "error", | ||
138 | "msg": "Error retrieving request." | ||
139 | })", | ||
140 | websocketpp::frame::opcode::value::TEXT); | ||
141 | } catch (const std::exception& ex) { | ||
142 | // Well, okay | ||
143 | std::cout << ex.what() << std::endl; | ||
144 | } | ||
145 | } | ||
146 | } | ||
147 | |||
148 | void generate_thread(std::string token, std::string text) { | ||
149 | std::unique_ptr<wizard> generator; | ||
150 | |||
151 | { | ||
152 | std::lock_guard rng_guard(rng_mutex_); | ||
153 | generator = std::make_unique<wizard>(cards_, images_, text, rng_); | ||
154 | } | ||
155 | |||
156 | try { | ||
157 | generator->set_status_callback([this, token](const std::string& status) { | ||
158 | nlohmann::json msg; | ||
159 | msg["type"] = "status"; | ||
160 | msg["msg"] = status; | ||
161 | |||
162 | database_.post(token, msg.dump()); | ||
163 | }); | ||
164 | |||
165 | Magick::Image resultImage = generator->run(); | ||
166 | Magick::Blob resultBlob; | ||
167 | resultImage.write(&resultBlob); | ||
168 | |||
169 | std::string resultBytes((const char*)resultBlob.data(), | ||
170 | resultBlob.length()); | ||
171 | std::string resultEncoded = base64::to_base64(resultBytes); | ||
172 | |||
173 | database_.setResult(token, resultEncoded); | ||
174 | |||
175 | nlohmann::json resultMsg; | ||
176 | resultMsg["type"] = "result"; | ||
177 | resultMsg["image"] = resultEncoded; | ||
178 | resultMsg["msg"] = "Success!"; | ||
179 | |||
180 | database_.post(token, resultMsg.dump()); | ||
181 | } catch (const std::exception& ex) { | ||
182 | nlohmann::json response; | ||
183 | response["type"] = "error"; | ||
184 | response["msg"] = | ||
185 | std::string("Error generating card (") + ex.what() + ")"; | ||
186 | |||
187 | database_.post(token, response.dump()); | ||
188 | } | ||
189 | |||
190 | database_.mark_done(token); | ||
191 | } | ||
192 | |||
193 | void cleanup_thread() { | ||
194 | for (;;) { | ||
195 | // sleep | ||
196 | } | ||
197 | } | ||
198 | |||
199 | const cardset& cards_; | ||
200 | const imagestore& images_; | ||
201 | |||
202 | std::mutex rng_mutex_; | ||
203 | std::mt19937& rng_; | ||
204 | |||
205 | socket_type socket_; | ||
206 | database database_; | ||
207 | }; | ||
208 | |||
209 | } // namespace | ||
210 | |||
211 | int main(int argc, char** argv) { | ||
212 | Magick::InitializeMagick(nullptr); | ||
213 | |||
214 | std::random_device randomDevice; | ||
215 | std::mt19937 rng(5); // randomDevice()); | ||
216 | |||
217 | if (argc != 2) { | ||
218 | std::cout << "usage: wizard_server [configfile]" << std::endl; | ||
219 | return -1; | ||
220 | } | ||
221 | |||
222 | std::ifstream config_file(argv[1]); | ||
223 | nlohmann::json config_data = nlohmann::json::parse(config_file); | ||
224 | |||
225 | cardset cards(config_data["cards_path"]); | ||
226 | imagestore images(config_data["cache_path"]); | ||
227 | |||
228 | server app(cards, images, rng); | ||
229 | app.run(); | ||
230 | } | ||