ClickHouse/programs/server/fuzzers/tcp_protocol_fuzzer.cpp
2023-09-01 14:20:50 +00:00

120 lines
3.0 KiB
C++

#include <cstdint>
#include <future>
#include <thread>
#include <utility>
#include <vector>
#include <iostream>
#include <chrono>
#include <Poco/Net/PollSet.h>
#include <Poco/Net/SocketAddress.h>
#include <Poco/Net/StreamSocket.h>
#include <Interpreters/Context.h>
int mainEntryClickHouseServer(int argc, char ** argv);
static std::string clickhouse("clickhouse-server");
static std::vector<char *> args{clickhouse.data()};
static std::future<int> main_app;
static std::string s_host("0.0.0.0");
static char * host = s_host.data();
static int64_t port = 9000;
using namespace std::chrono_literals;
extern "C"
int LLVMFuzzerInitialize(int * argc, char ***argv)
{
for (int i = 1; i < *argc; ++i)
{
if ((*argv)[i][0] == '-')
{
if ((*argv)[i][1] == '-')
args.push_back((*argv)[i]);
else
{
if (strncmp((*argv)[i], "-host=", 6) == 0)
{
host = (*argv)[i] + 6;
}
else if (strncmp((*argv)[i], "-port=", 6) == 0)
{
char * p_end = nullptr;
port = strtol((*argv)[i] + 6, &p_end, 10);
}
}
}
}
args.push_back(nullptr);
main_app = std::async(std::launch::async, mainEntryClickHouseServer, args.size() - 1, args.data());
while (!DB::Context::getGlobalContextInstance() || !DB::Context::getGlobalContextInstance()->isServerCompletelyStarted())
{
std::this_thread::sleep_for(100ms);
if (main_app.wait_for(0s) == std::future_status::ready)
exit(-1);
}
return 0;
}
extern "C"
int LLVMFuzzerTestOneInput(const uint8_t * data, size_t size)
{
try
{
if (main_app.wait_for(0s) == std::future_status::ready)
return -1;
if (size == 0)
return -1;
Poco::Net::SocketAddress address(host, port);
Poco::Net::StreamSocket socket;
socket.connectNB(address);
Poco::Net::PollSet ps;
ps.add(socket, Poco::Net::PollSet::POLL_READ | Poco::Net::PollSet::POLL_WRITE);
std::vector<char> buf(1048576);
size_t sent = 0;
while (true)
{
auto m = ps.poll(Poco::Timespan(1000000));
if (m.empty())
continue;
if (m.begin()->second & Poco::Net::PollSet::POLL_READ)
{
if (int n = socket.receiveBytes(buf.data(), static_cast<int>(buf.size())); n == 0)
{
socket.close();
break;
}
continue;
}
if (sent < size && m.begin()->second & Poco::Net::PollSet::POLL_WRITE)
{
sent += socket.sendBytes(data + sent, static_cast<int>(size - sent));
if (sent == size)
{
socket.shutdownSend();
continue;
}
}
}
}
catch (...)
{
}
return 0;
}