127 lines
4.3 KiB
C++
127 lines
4.3 KiB
C++
#include "Connection.hpp"
|
|
|
|
#include <optional>
|
|
#include <algorithm>
|
|
#include <boost/asio.hpp>
|
|
|
|
|
|
|
|
namespace Connection {
|
|
asio::awaitable<Packet> Client::receivePacket() {
|
|
Packet packet;
|
|
// Receive operation
|
|
packet.op = co_await receiveProtobuf<Generic::Operation>();
|
|
// Receive data
|
|
switch (packet.op.code()) {
|
|
case Generic::Operation::Error: {
|
|
packet.data = std::make_unique<Generic::Error>(co_await receiveProtobuf<Generic::Error>());
|
|
} break;
|
|
case Generic::Operation::SimpleAuth: {
|
|
packet.data = std::make_unique<Generic::SimpleAuth>(co_await receiveProtobuf<Generic::SimpleAuth>(30));
|
|
} break;
|
|
case Generic::Operation::Disconnect:
|
|
case Generic::Operation::OK:
|
|
break;
|
|
default: throw FatalError("Invalid opcode: "+std::to_string(packet.op.code()), Generic::Error::InvalidOpcode);
|
|
}
|
|
}
|
|
|
|
asio::awaitable<void> Client::sendProtobuf(const google::protobuf::Message& buffer) {
|
|
// Send size
|
|
uint16_t size = buffer.ByteSizeLong();
|
|
co_await asio::async_write(socket, asio::buffer(&size, sizeof(size)), asio::use_awaitable);
|
|
// Send buffer
|
|
std::vector<char> rawBuffer(size);
|
|
buffer.SerializeToArray(rawBuffer.data(), rawBuffer.size());
|
|
co_await asio::async_write(socket, asio::buffer(rawBuffer.data(), rawBuffer.size()), asio::use_awaitable);
|
|
}
|
|
|
|
asio::awaitable<void> Client::sendPacket(const Packet& packet) {
|
|
// Send operation
|
|
co_await sendProtobuf(packet.op);
|
|
// Send data
|
|
if (packet.data) {
|
|
co_await sendProtobuf(*packet.data);
|
|
}
|
|
}
|
|
|
|
asio::awaitable<void> Client::handlePacket(const Packet& packet) {
|
|
// Check if authentication is required and completed
|
|
if (packet.op.code() >= Generic::Operation::_MinAuth && !isAuthed()) {
|
|
throw FatalError("Authentication required", Generic::Error::AuthenticationRequired);
|
|
}
|
|
// Handle the packet
|
|
switch (packet.op.code()) {
|
|
case Generic::Operation::Disconnect: {
|
|
good = 0;
|
|
} break;
|
|
case Generic::Operation::SimpleAuth: {
|
|
const Generic::SimpleAuth& authData = *static_cast<Generic::SimpleAuth*>(packet.data.get());
|
|
// Check that username is not already taken
|
|
if (std::find_if(clients.begin(), clients.end(), [&, this] (const auto& o) {return o->getUsername() == authData.username();}) != clients.end()) {
|
|
throw Error("Username has already been taken", Generic::Error::UsernameAlreadyTaken);
|
|
}
|
|
// Check that username is "good" (no unprintable characters and stuff like that)
|
|
//TODO...
|
|
// Set username
|
|
username = std::move(authData.username());
|
|
// Send OK
|
|
co_await sendOK();
|
|
} break;
|
|
default: {}
|
|
}
|
|
co_return;
|
|
}
|
|
|
|
asio::awaitable<void> Client::connect() {
|
|
// Send initial "OK"
|
|
try {
|
|
co_await sendOK();
|
|
} catch (...) {
|
|
good = false;
|
|
}
|
|
// Run while connection is good
|
|
while (good) {
|
|
std::optional<Error> error;
|
|
try {
|
|
co_await handlePacket(co_await receivePacket());
|
|
} catch (FatalError& e) {
|
|
error = std::move(e);
|
|
good = false;
|
|
} catch (Error& e) {
|
|
error = std::move(e);
|
|
} catch (std::exception& e) {
|
|
error = Error("Internal server error: "+std::string(e.what()), Generic::Error::Unknown);
|
|
good = false;
|
|
}
|
|
if (error.has_value()) {
|
|
co_await sendPacket(error.value().packet);
|
|
}
|
|
}
|
|
// Try to send disconnect
|
|
try {
|
|
Packet packet;
|
|
packet.op.set_code(Generic::Operation::Disconnect);
|
|
co_await sendPacket(packet);
|
|
} catch (...) {}
|
|
// Close socket
|
|
boost::system::error_code ec;
|
|
socket.close(ec);
|
|
}
|
|
|
|
asio::awaitable<void> Server::listen(int port) {
|
|
try {
|
|
// Create acceptor
|
|
asio::ip::tcp::acceptor acceptor(ioc, asio::ip::tcp::endpoint(asio::ip::tcp::v4(), port));
|
|
// Handle all incoming connections
|
|
while (true) {
|
|
auto client = std::make_shared<Client>(ioc, clients);
|
|
co_await acceptor.async_accept(client->getSocket(), asio::use_awaitable);
|
|
asio::co_spawn(ioc, client->connect(), asio::detached);
|
|
clients.push_back(std::move(client));
|
|
}
|
|
} catch (std::exception& e) {
|
|
std::cerr << "Failed to stay alive: " << e.what() << std::endl;
|
|
}
|
|
}
|
|
}
|