diff --git a/src/communication/websocket/auth.hpp b/src/communication/websocket/auth.hpp index 8c1dade44..008d92bd5 100644 --- a/src/communication/websocket/auth.hpp +++ b/src/communication/websocket/auth.hpp @@ -19,15 +19,24 @@ namespace communication::websocket { -class SafeAuth { +class AuthenticationInterface { + public: + virtual bool Authenticate(const std::string &username, const std::string &password) const = 0; + + virtual bool HasUserPermission(const std::string &username, auth::Permission permission) const = 0; + + virtual bool HasAnyUsers() const = 0; +}; + +class SafeAuth : public AuthenticationInterface { public: explicit SafeAuth(utils::Synchronized *auth) : auth_{auth} {} - bool Authenticate(const std::string &username, const std::string &password) const; + bool Authenticate(const std::string &username, const std::string &password) const override; - bool HasUserPermission(const std::string &username, auth::Permission permission) const; + bool HasUserPermission(const std::string &username, auth::Permission permission) const override; - bool HasAnyUsers() const; + bool HasAnyUsers() const override; private: utils::Synchronized *auth_; diff --git a/src/communication/websocket/listener.cpp b/src/communication/websocket/listener.cpp index 167e59a02..ec62b0f44 100644 --- a/src/communication/websocket/listener.cpp +++ b/src/communication/websocket/listener.cpp @@ -27,7 +27,10 @@ void Listener::WriteToAll(std::shared_ptr message) { } } -Listener::Listener(boost::asio::io_context &ioc, ServerContext *context, tcp::endpoint endpoint, SafeAuth auth) +boost::asio::ip::tcp::endpoint Listener::GetEndpoint() const { return acceptor_.local_endpoint(); }; + +Listener::Listener(boost::asio::io_context &ioc, ServerContext *context, tcp::endpoint endpoint, + AuthenticationInterface &auth) : ioc_(ioc), context_(context), acceptor_(ioc), auth_(auth) { boost::beast::error_code ec; diff --git a/src/communication/websocket/listener.hpp b/src/communication/websocket/listener.hpp index c41c38a05..1c91ce8fb 100644 --- a/src/communication/websocket/listener.hpp +++ b/src/communication/websocket/listener.hpp @@ -40,9 +40,10 @@ class Listener : public std::enable_shared_from_this { // Start accepting incoming connections void Run(); void WriteToAll(std::shared_ptr message); + tcp::endpoint GetEndpoint() const; private: - Listener(boost::asio::io_context &ioc, ServerContext *context, tcp::endpoint endpoint, SafeAuth auth); + Listener(boost::asio::io_context &ioc, ServerContext *context, tcp::endpoint endpoint, AuthenticationInterface &auth); void DoAccept(); void OnAccept(boost::beast::error_code ec, tcp::socket socket); @@ -51,6 +52,6 @@ class Listener : public std::enable_shared_from_this { ServerContext *context_; tcp::acceptor acceptor_; utils::Synchronized>, utils::SpinLock> sessions_; - SafeAuth auth_; + AuthenticationInterface &auth_; }; } // namespace communication::websocket diff --git a/src/communication/websocket/server.cpp b/src/communication/websocket/server.cpp index a82a6f863..4b54ff912 100644 --- a/src/communication/websocket/server.cpp +++ b/src/communication/websocket/server.cpp @@ -1,4 +1,4 @@ -// Copyright 2021 Memgraph Ltd. +// Copyright 2022 Memgraph Ltd. // // Use of this software is governed by the Business Source License // included in the file licenses/BSL.txt; by using this file, you agree to be bound by the terms of the Business Source @@ -36,6 +36,8 @@ void Server::AwaitShutdown() { bool Server::IsRunning() const { return background_thread_ && !ioc_.stopped(); } +boost::asio::ip::tcp::endpoint Server::GetEndpoint() const { return listener_->GetEndpoint(); }; + namespace { class QuoteEscapeFormatter : public spdlog::custom_flag_formatter { public: diff --git a/src/communication/websocket/server.hpp b/src/communication/websocket/server.hpp index 5be8bf29e..bda046260 100644 --- a/src/communication/websocket/server.hpp +++ b/src/communication/websocket/server.hpp @@ -28,7 +28,7 @@ class Server final { using tcp = boost::asio::ip::tcp; public: - explicit Server(io::network::Endpoint endpoint, ServerContext *context, SafeAuth auth) + explicit Server(io::network::Endpoint endpoint, ServerContext *context, AuthenticationInterface &auth) : listener_{Listener::Create( ioc_, context, tcp::endpoint{boost::asio::ip::make_address(endpoint.address), endpoint.port}, auth)} {} @@ -43,6 +43,7 @@ class Server final { void Shutdown(); void AwaitShutdown(); bool IsRunning() const; + tcp::endpoint GetEndpoint() const; class LoggingSink : public spdlog::sinks::base_sink { public: diff --git a/src/communication/websocket/session.cpp b/src/communication/websocket/session.cpp index aa86955ae..d7513cfb8 100644 --- a/src/communication/websocket/session.cpp +++ b/src/communication/websocket/session.cpp @@ -23,6 +23,7 @@ #include #include "communication/context.hpp" +#include "communication/websocket/auth.hpp" #include "utils/logging.hpp" namespace communication::websocket { @@ -42,7 +43,7 @@ std::variant Session::CreateWebS return Session::PlainWebSocket{std::move(socket)}; } -Session::Session(tcp::socket &&socket, ServerContext &context, SafeAuth auth) +Session::Session(tcp::socket &&socket, ServerContext &context, AuthenticationInterface &auth) : ws_(CreateWebSocket(std::move(socket), context)), strand_{boost::asio::make_strand(GetExecutor())}, auth_{auth} {} bool Session::Run() { diff --git a/src/communication/websocket/session.hpp b/src/communication/websocket/session.hpp index 899daf366..a0e77992a 100644 --- a/src/communication/websocket/session.hpp +++ b/src/communication/websocket/session.hpp @@ -50,7 +50,7 @@ class Session : public std::enable_shared_from_this { using PlainWebSocket = boost::beast::websocket::stream; using SSLWebSocket = boost::beast::websocket::stream>; - explicit Session(tcp::socket &&socket, ServerContext &context, SafeAuth auth); + explicit Session(tcp::socket &&socket, ServerContext &context, AuthenticationInterface &auth); void DoWrite(); void OnWrite(boost::beast::error_code ec, size_t bytest_transferred); @@ -61,10 +61,10 @@ class Session : public std::enable_shared_from_this { void DoClose(); void OnClose(boost::beast::error_code ec); - utils::BasicResult Authorize(const nlohmann::json &creds); - bool IsAuthenticated() const; + utils::BasicResult Authorize(const nlohmann::json &creds); + void DoShutdown(); auto GetExecutor() { @@ -86,6 +86,6 @@ class Session : public std::enable_shared_from_this { std::atomic connected_{false}; bool authenticated_{false}; bool close_{false}; - SafeAuth auth_; + AuthenticationInterface &auth_; }; } // namespace communication::websocket diff --git a/src/memgraph.cpp b/src/memgraph.cpp index 32a9741ea..dcd714614 100644 --- a/src/memgraph.cpp +++ b/src/memgraph.cpp @@ -1215,8 +1215,8 @@ int main(int argc, char **argv) { []() -> nlohmann::json { return query::plan::CallProcedure::GetAndResetCounters(); }); } - communication::websocket::Server websocket_server{ - {"0.0.0.0", 7444}, &context, communication::websocket::SafeAuth{&auth}}; + communication::websocket::SafeAuth websocket_auth{&auth}; + communication::websocket::Server websocket_server{{"0.0.0.0", 7444}, &context, websocket_auth}; AddLoggerSink(websocket_server.GetLoggingSink()); // Handler for regular termination signals diff --git a/tests/unit/CMakeLists.txt b/tests/unit/CMakeLists.txt index 4994d4711..acfb535b8 100644 --- a/tests/unit/CMakeLists.txt +++ b/tests/unit/CMakeLists.txt @@ -368,3 +368,8 @@ add_custom_target(test_lcp ALL DEPENDS ${CMAKE_CURRENT_BINARY_DIR}/test_lcp) add_test(test_lcp ${CMAKE_CURRENT_BINARY_DIR}/test_lcp) add_dependencies(memgraph__unit test_lcp) +# Test websocket +find_package(Boost REQUIRED) + +add_unit_test(websocket.cpp) +target_link_libraries(${test_prefix}websocket mg-communication Boost::headers) diff --git a/tests/unit/websocket.cpp b/tests/unit/websocket.cpp new file mode 100644 index 000000000..12a627758 --- /dev/null +++ b/tests/unit/websocket.cpp @@ -0,0 +1,305 @@ +// Copyright 2022 Memgraph Ltd. +// +// Use of this software is governed by the Business Source License +// included in the file licenses/BSL.txt; by using this file, you agree to be bound by the terms of the Business Source +// License, and you may not use this file except in compliance with the Business Source License. +// +// As of the Change Date specified in that file, in accordance with +// the Business Source License, use of this software will be governed +// by the Apache License, Version 2.0, included in the file +// licenses/APL.txt.""" + +#define BOOST_ASIO_USE_TS_EXECUTOR_AS_DEFAULT + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +#include "communication/websocket/auth.hpp" +#include "communication/websocket/server.hpp" + +namespace beast = boost::beast; +namespace http = beast::http; +namespace websocket = beast::websocket; +namespace net = boost::asio; +using tcp = boost::asio::ip::tcp; + +constexpr auto kResponseSuccess{"success"}; +constexpr auto kResponseMessage{"message"}; + +struct MockAuth : public communication::websocket::AuthenticationInterface { + MockAuth() = default; + + bool Authenticate(const std::string & /*username*/, const std::string & /*password*/) const override { + return authentication; + } + + bool HasUserPermission(const std::string & /*username*/, auth::Permission /*permission*/) const override { + return authorization; + } + + bool HasAnyUsers() const override { return has_any_users; } + + bool authentication{true}; + bool authorization{true}; + bool has_any_users{true}; +}; + +class WebSocketServerTest : public ::testing::Test { + public: + protected: + WebSocketServerTest() : websocket_server{{"0.0.0.0", 0}, &context, auth} { + EXPECT_NO_THROW(websocket_server.Start()); + } + + void TearDown() override { + EXPECT_NO_THROW(websocket_server.Shutdown()); + EXPECT_NO_THROW(websocket_server.AwaitShutdown()); + } + + std::string ServerPort() const { return std::to_string(websocket_server.GetEndpoint().port()); } + + std::string ServerAddress() const { return websocket_server.GetEndpoint().address().to_string(); } + + MockAuth auth; + communication::ServerContext context{}; + communication::websocket::Server websocket_server; +}; + +class Client { + public: + ~Client() { ws_.close(websocket::close_code::normal); } + + void Connect(const std::string &host, const std::string &port) { + tcp::resolver resolver{ioc_}; + auto endpoint_ = resolver.resolve(host, port); + auto ep = net::connect(ws_.next_layer(), endpoint_); + const auto server = fmt::format("{}:{}", host, ep.port()); + ws_.set_option(websocket::stream_base::decorator([](websocket::request_type &req) { + req.set(http::field::user_agent, std::string(BOOST_BEAST_VERSION_STRING) + " websocket-client-coro"); + })); + + // Perform the websocket handshake + ws_.handshake(host, "/"); + } + + void Write(const std::string &msg) { ws_.write(net::buffer(msg)); } + + std::string Read() { + ws_.read(buffer_); + const std::string response = beast::buffers_to_string(buffer_.data()); + buffer_.consume(buffer_.size()); + return response; + } + + private: + net::io_context ioc_{}; + websocket::stream ws_{ioc_}; + beast::flat_buffer buffer_; +}; + +TEST(WebSocketServer, WebsocketWorkflow) { + /** + * Notice how there is no port management for the clients + * and the servers, that is because when using "0.0.0.0" as address and + * and 0 as port number we delegate port assignment to the OS + * and it is the keeper of all available port numbers and + * assigns them automatically. + */ + MockAuth auth{}; + communication::ServerContext context{}; + communication::websocket::Server websocket_server({"0.0.0.0", 0}, &context, auth); + const auto port = websocket_server.GetEndpoint().port(); + + SCOPED_TRACE(fmt::format("Checking port number different then 0: {}", port)); + EXPECT_NE(port, 0); + EXPECT_NO_THROW(websocket_server.Start()); + EXPECT_TRUE(websocket_server.IsRunning()); + + EXPECT_NO_THROW(websocket_server.Shutdown()); + EXPECT_FALSE(websocket_server.IsRunning()); + + EXPECT_NO_THROW(websocket_server.AwaitShutdown()); + EXPECT_FALSE(websocket_server.IsRunning()); +} + +TEST_F(WebSocketServerTest, WebsocketConnection) { + { + auto client = Client{}; + EXPECT_NO_THROW(client.Connect("0.0.0.0", ServerPort())); + } + + websocket_server.Shutdown(); + websocket_server.AwaitShutdown(); +} + +TEST_F(WebSocketServerTest, WebsocketLogging) { + auth.has_any_users = false; + // Set up the websocket logger as one of the defaults for spdlog + { + auto default_logger = spdlog::default_logger(); + auto sinks = default_logger->sinks(); + sinks.push_back(websocket_server.GetLoggingSink()); + + auto logger = std::make_shared("memgraph_log", sinks.begin(), sinks.end()); + logger->set_level(default_logger->level()); + logger->flush_on(spdlog::level::trace); + spdlog::set_default_logger(std::move(logger)); + } + { + auto client = Client(); + client.Connect(ServerAddress(), ServerPort()); + + auto log_message = [](spdlog::level::level_enum level, std::string_view message) { + spdlog::log(level, message); + spdlog::default_logger()->flush(); + }; + auto log_and_check = [log_message, &client](spdlog::level::level_enum level, std::string_view message, + std::string_view log_level_received) { + std::thread(log_message, level, message).detach(); + const auto received_message = client.Read(); + EXPECT_EQ(received_message, fmt::format("{{\"event\": \"log\", \"level\": \"{}\", \"message\": \"{}\"}}\n", + log_level_received, message)); + }; + + log_and_check(spdlog::level::err, "Sending error message!", "error"); + log_and_check(spdlog::level::warn, "Sending warn message!", "warning"); + log_and_check(spdlog::level::info, "Sending info message!", "info"); + log_and_check(spdlog::level::trace, "Sending trace message!", "trace"); + } +} + +TEST_F(WebSocketServerTest, WebsocketAuthenticationParsingError) { + constexpr auto auth_fail = "Cannot parse JSON for WebSocket authentication"; + + { + SCOPED_TRACE("Checking handling of first request parsing error."); + auto client = Client(); + EXPECT_NO_THROW(client.Connect(ServerAddress(), ServerPort())); + EXPECT_NO_THROW(client.Write("Test")); + const auto response = nlohmann::json::parse(client.Read()); + const auto message_header = response[kResponseMessage].get(); + const auto message_first_part = message_header.substr(0, message_header.find(": ")); + + EXPECT_FALSE(response[kResponseSuccess]); + EXPECT_EQ(message_first_part, auth_fail); + } + { + SCOPED_TRACE("Checking handling of JSON parsing error."); + auto client = Client(); + EXPECT_NO_THROW(client.Connect(ServerAddress(), ServerPort())); + const std::string json_without_comma = R"({"username": "user" "password": "123"})"; + EXPECT_NO_THROW(client.Write(json_without_comma)); + const auto response = nlohmann::json::parse(client.Read()); + const auto message_header = response[kResponseMessage].get(); + const auto message_first_part = message_header.substr(0, message_header.find(": ")); + + EXPECT_FALSE(response[kResponseSuccess]); + EXPECT_EQ(message_first_part, auth_fail); + } +} + +TEST_F(WebSocketServerTest, WebsocketAuthenticationWhenAuthPasses) { + constexpr auto auth_success = R"({"message":"User has been successfully authenticated!","success":true})"; + + { + SCOPED_TRACE("Checking successful authentication response."); + auto client = Client(); + EXPECT_NO_THROW(client.Connect(ServerAddress(), ServerPort())); + EXPECT_NO_THROW(client.Write(R"({"username": "user", "password": "123"})")); + const auto response = client.Read(); + + EXPECT_EQ(response, auth_success); + } +} + +TEST_F(WebSocketServerTest, WebsocketAuthenticationWithMultipleAttempts) { + constexpr auto auth_success = R"({"message":"User has been successfully authenticated!","success":true})"; + constexpr auto auth_fail = "Cannot parse JSON for WebSocket authentication"; + + { + SCOPED_TRACE("Checking multiple authentication tries from same client"); + auto client = Client(); + EXPECT_NO_THROW(client.Connect(ServerAddress(), ServerPort())); + EXPECT_NO_THROW(client.Write(R"({"username": "user" "password": "123"})")); + + { + const auto response = nlohmann::json::parse(client.Read()); + const auto message_header = response[kResponseMessage].get(); + const auto message_first_part = message_header.substr(0, message_header.find(": ")); + + EXPECT_FALSE(response[kResponseSuccess]); + EXPECT_EQ(message_first_part, auth_fail); + } + { + EXPECT_NO_THROW(client.Connect(ServerAddress(), ServerPort())); + EXPECT_NO_THROW(client.Write(R"({"username": "user", "password": "123"})")); + const auto response = client.Read(); + EXPECT_EQ(response, auth_success); + } + } + { + SCOPED_TRACE("Checking multiple authentication tries from different clients"); + auto client1 = Client(); + auto client2 = Client(); + + EXPECT_NO_THROW(client1.Connect(ServerAddress(), ServerPort())); + EXPECT_NO_THROW(client2.Connect(ServerAddress(), ServerPort())); + + EXPECT_NO_THROW(client1.Write(R"({"username": "user" "password": "123"})")); + EXPECT_NO_THROW(client2.Write(R"({"username": "user", "password": "123"})")); + + { + const auto response = nlohmann::json::parse(client1.Read()); + const auto message_header = response[kResponseMessage].get(); + const auto message_first_part = message_header.substr(0, message_header.find(": ")); + + EXPECT_FALSE(response[kResponseSuccess]); + EXPECT_EQ(message_first_part, auth_fail); + } + { + const auto response = client2.Read(); + EXPECT_EQ(response, auth_success); + } + } +} + +TEST_F(WebSocketServerTest, WebsocketAuthenticationFails) { + auth.authentication = false; + + constexpr auto auth_fail = R"({"message":"Authentication failed!","success":false})"; + { + auto client = Client(); + EXPECT_NO_THROW(client.Connect(ServerAddress(), ServerPort())); + EXPECT_NO_THROW(client.Write(R"({"username": "user", "password": "123"})")); + + const auto response = client.Read(); + EXPECT_EQ(response, auth_fail); + } +} + +#ifdef MG_ENTERPRISE +TEST_F(WebSocketServerTest, WebsocketAuthorizationFails) { + auth.authorization = false; + constexpr auto auth_fail = R"({"message":"Authorization failed!","success":false})"; + + { + auto client = Client(); + EXPECT_NO_THROW(client.Connect(ServerAddress(), ServerPort())); + EXPECT_NO_THROW(client.Write(R"({"username": "user", "password": "123"})")); + + const auto response = client.Read(); + EXPECT_EQ(response, auth_fail); + } +} +#endif