turbod.cpp

#include <iostream>
#include <fstream>
#include <sstream>
#include <memory>
#include <string>

#include <grpcpp/grpcpp.h>
#include <turbo-rpc/turbo.grpc.pb.h>

#include <pqxx/pqxx>

using namespace grpc;
using namespace Turbo;
using namespace std;

#define TURBODB "turbodb"
// Logic and data behind the server's behavior.
class TurboServiceImpl final : public TurboService::Service 
{
	public:

		TurboServiceImpl()
		{
			pqxx::connection conn{db_connect_options.c_str()};
			pqxx::work txn(conn);


			cout << "Creating turbod tables" << endl;
			txn.exec(turbod_schema.c_str());
			txn.commit();

		}

		Status Ping(ServerContext *context, const Empty_t *request, Empty_t *reply) override final  
		{
			return Status::OK;
		}

		Status BoostAdd(ServerContext *context, const BoostAddRequest_t *request, BoostAddReply_t *reply) override final  
		{
			try
			{
				pqxx::connection conn{db_connect_options.c_str()};
				pqxx::work txn(conn);

				cout << "Requested creation of new boost named " << request->name() << endl;

				conn.prepare("boostcheck", "select boost_id from boosts where boost_name = $1");
				const auto check_res =  txn.prepared("boostcheck")(request->name()).exec();

				if(check_res.size() != 0)
				{
					return Status(ABORTED, "Boost already exists");
				}

				conn.prepare("boostadd", "insert into boosts (boost_name, create_time) values ($1, NOW()) returning boost_id ");
				const auto add_res =  txn.prepared("boostadd")(request->name()).exec();

				if(add_res.size() != 1)
				{
					return Status(ABORTED, "Unexpected sql reply.");
				}

				reply->set_boost_id(add_res[0]["boost_id"].as<int>());
				txn.commit();
				return Status::OK;
			}
			catch(const std::exception& e)
			{
				return Status(ABORTED, e.what());
			}

			
		}

		Status GetZafld  (ServerContext *context, const GetZafldRequest_t *request, ZafldBinary_t *reply) override final  
		{
			try
			{
				pqxx::connection conn{db_connect_options.c_str()};
				pqxx::work txn(conn);

				cout << "Requested zafld binary for " << request->ver_id() << endl;

				conn.prepare("getzafld", 
						R"""(select 
							case when zafl_binary is null then '' else zafl_binary end as zafl_binary ,
							case when err is null then false else err end as err 
						from boost_versions where ver_id = $1)"""
						);
				const auto res_table =  txn.prepared("getzafld")(request->ver_id()).exec();

				if(res_table.size() != 1)
				{
					return Status(ABORTED, "Invalid version id");
				}
				const auto zafl_binary_contents = pqxx::binarystring(res_table[0]["zafl_binary"]).str();
				cout << "Zafl version is " << zafl_binary_contents.size() << " bytes" << endl;

				reply->set_zafld_bin(zafl_binary_contents);
				reply->set_err(res_table[0]["err"].as<bool>() );

				return Status::OK;
			}
			catch(const std::exception& e)
			{
				return Status(ABORTED, e.what());
			}
		}

		Status BoostList(ServerContext *context, const Empty_t *request, ServerWriter<BoostInfo_t> *writer) override final  
		{
			try
			{
				pqxx::connection conn{db_connect_options.c_str()};
				pqxx::work txn(conn);

				cout << "Requested boost list " << endl;

				conn.prepare("boostlist", "select boost_id,encode(boost_name,'escape'), to_char(create_time, 'MM-DD-YYYY HH24:MI:SS') as create_time from boosts");
				const auto res_table =  txn.prepared("boostlist").exec();

				if(res_table.size() == 0)
				{
					return Status(ABORTED, "No Boosts found");
				}

				for(const auto &res : res_table)
				{
					auto reply = BoostInfo_t();
					reply.set_id          (res["boost_id"   ].as<int>   ());
					reply.set_name        (res["encode"     ].as<string>());
					reply.set_create_time (res["create_time"].as<string>());

					writer->Write(reply);
				}

				return Status::OK;
			}
			catch(const std::exception& e)
			{
				return Status(ABORTED, e.what());
			}
		}

		Status VersionList(ServerContext *context, const VersionListRequest_t *request, ServerWriter<VersionInfo_t> *writer) override final  
		{
			try
			{
				pqxx::connection conn{db_connect_options.c_str()};
				pqxx::work txn(conn);

				cout << "Requested version list for boost id " << request->boost_id() << endl;

				conn.prepare("boostlist", 
						R"""(
							select 
								ver_id,
								encode(orig_bin_sha,'text') as orig_sha,
								encode(zafl_bin_sha,'text') as zafl_sha, 
								COALESCE(to_char(submit_time, 'MM-DD-YYYY HH24:MI:SS'), '') as submit_time,
								COALESCE(to_char(start_time , 'MM-DD-YYYY HH24:MI:SS'), '') as start_time,
								COALESCE(to_char(end_time   , 'MM-DD-YYYY HH24:MI:SS'), '') as end_time,
								err 
								from boost_versions where boost_id = $1)""" 
						);
				const auto res_table =  txn.prepared("boostlist")(request->boost_id()).exec();

				if(res_table.size() == 0)
				{
					return Status(ABORTED, "No versions found");
				}

				const auto to_string = [](const pqxx::field &res) -> string
					{
						return res.is_null() ? "" : res.as<string>();
					};

				for(const auto &res : res_table)
				{
					auto reply = VersionInfo_t();
					reply.set_ver_id       (res["ver_id"     ].as<int>   ());
					reply.set_orig_bin_sha (to_string(res["orig_sha"   ]));
					reply.set_zafl_bin_sha (to_string(res["zafl_sha"   ]));
					reply.set_submit_time  (res["submit_time"].as<string>());
					reply.set_start_time   (res["start_time" ].as<string>());
					reply.set_end_time     (res["end_time"   ].as<string>());
					reply.set_err          (res["err"].is_null() ? false : res["err"].as<bool>());
					writer->Write(reply);
				}

				return Status::OK;
			}
			catch(const std::exception& e)
			{
				return Status(ABORTED, e.what());
			}
		}

		Status VersionAdd(ServerContext *context, const VersionAddRequest_t *request, VersionAddReply_t *reply) override final  
		{
			try
			{
				pqxx::connection conn{db_connect_options.c_str()};
				pqxx::work txn(conn);

				cout << "Requested creation of new version for boost id " << request->boost_id() << endl;

				conn.prepare("boostcheck", "select boost_id,boost_name from boosts where boost_id = $1");
				const auto check_res =  txn.prepared("boostcheck")(request->boost_id()).exec();

				if(check_res.size() != 1)
				{
					return Status(ABORTED, "Boost id not found");
				}

				conn.prepare("versionadd", "insert into boost_versions ( boost_id, orig_binary, submit_time, err )"
						" values ($1, $2, NOW(), false) returning ver_id ");

				const auto boost_id    = check_res[0]["boost_id"].as<int>();
				const auto binary_data = pqxx::binarystring(request->version_contents().c_str(), request->version_contents().size());
				const auto add_res     = txn.prepared("versionadd")(boost_id)(binary_data).exec();

				if(add_res.size() != 1)
				{
					return Status(ABORTED, "Unexpected sql reply.");
				}

				reply->set_version_id(add_res[0]["ver_id"].as<int>());
				txn.commit();
				return Status::OK;
			}
			catch(const std::exception& e)
			{
				return Status(ABORTED, e.what());
			}

		}

		Status GetZaflLog(ServerContext *context, const GetZaflLogRequest_t *request, LogFile_t *reply) override final  
		{
			try
			{
				pqxx::connection conn{db_connect_options.c_str()};
				pqxx::work txn(conn);

				cout << "Requested zafl log  for " << request->ver_id() << endl;

				conn.prepare("getzafllog", 
						R"""(select 
							case when log is null then '' else encode(log,'escape') end as log 
						from zafl_logs where ver_id = $1)"""
						);
				const auto res_table =  txn.prepared("getzafllog")(request->ver_id()).exec();

				if(res_table.size() != 1)
				{
					return Status(ABORTED, "Invalid version id");
				}

				const auto zafl_log = res_table[0]["log"].as<string>();
				reply->set_log(zafl_log);

				return Status::OK;
			}
			catch(const std::exception& e)
			{
				return Status(ABORTED, e.what());
			}
		}

		Status GetFuzzLog(ServerContext *context, const GetFuzzLogRequest_t *request, LogFile_t *reply) override final  
		{
			try
			{
				pqxx::connection conn{db_connect_options.c_str()};
				pqxx::work txn(conn);

				cout << "Requested Fuzz log  for " << request->ver_id() << endl;

				conn.prepare("getFuzzlog", 
						R"""(select 
							case when log is null then '' else encode(log,'escape') end as log 
						from fuzz_logs where ver_id = $1)"""
						);
				const auto res_table =  txn.prepared("getFuzzlog")(request->ver_id()).exec();

				if(res_table.size() != 1)
				{
					return Status(ABORTED, "Invalid version id");
				}

				const auto zafl_log = res_table[0]["log"].as<string>();
				reply->set_log(zafl_log);

				return Status::OK;
			}
			catch(const std::exception& e)
			{
				return Status(ABORTED, e.what());
			}
		}


		Status SeedAdd(ServerContext *context, const SeedAddRequest_t *request, SeedAddReply_t *reply) override final  
		{
			try
			{
				pqxx::connection conn{db_connect_options.c_str()};
				pqxx::work txn(conn);

				cout << "Requested seed add " << request->boost_id() << endl;

				// prepare and execute a command to insert an input
				conn.prepare("input_add", 
						R"""(insert into inputs ( boost_id, create_time, prov, input_data) values ($1, NOW(), 'User',  $2) returning input_id )""");
				const auto res_input =  txn.prepared("input_add")(request->boost_id())(request->seed_file()).exec();
				if(res_input.size() != 1)
				{
					return Status(ABORTED, "Invalid boost id");
				}

				const auto input_id = res_input[0]["input_id"].as<int32_t>();

				// prepare and execute a command to mark the input as a seed
				conn.prepare("seed_add", R"""(insert into seeds ( input_id ) values ($1) returning seed_id )""");
				const auto res_seed =  txn.prepared("seed_add")(input_id).exec();
				if(res_seed.size() != 1)
				{
					return Status(ABORTED, "Unexpected error");
				}

				const auto seed_id = res_seed[0]["seed_id"].as<int32_t>();

				// finalize the transaction
				txn.commit();

				// prepare a reply
				reply->set_input_id(input_id);
				reply->set_seed_id (seed_id );
				return Status::OK;
			}
			catch(const std::exception& e)
			{
				return Status(ABORTED, e.what());
			}
		}

		Status InputAdd(ServerContext *context, const InputAddRequest_t *request, InputAddReply_t *reply) override final  
		{
			try
			{
				pqxx::connection conn{db_connect_options.c_str()};
				pqxx::work txn(conn);

				cout << "Requested input add for ver-id=" << request->ver_id() << endl;


				conn.prepare("findboost", "select boost_id from boost_versions where ver_id = $1");
				const auto boost_id_res =  txn.prepared("findboost")(request->ver_id()).exec();
				if(boost_id_res.size() != 1)
				{
					return Status(ABORTED, "Invalid version id");
				}

				const auto boost_id = boost_id_res[0]["boost_id"].as<int>();

				// prepare and execute a command to insert an input
				conn.prepare("aflinputadd", R"""(insert into inputs ( boost_id, create_time, prov, input_data) values ($1, NOW(), 'AFL',  $2) returning input_id )""");
				const auto res_input =  txn.prepared("aflinputadd")(boost_id)(request->input_file()).exec();
				if(res_input.size() != 1)
				{
					return Status(ABORTED, "Invalid boost id");
				}

				const auto input_id = res_input[0]["input_id"].as<int32_t>();

				// finalize the transaction
				txn.commit();

				// prepare a reply
				reply->set_input_id(input_id);
				return Status::OK;
			}
			catch(const std::exception& e)
			{
				return Status(ABORTED, e.what());
			}
		}

		Status DownloadInputs(ServerContext *context, const DownloadInputsRequest_t *request, ServerWriter<DownloadInputsReply_t> *writer) override final  
		{
			try
			{
				pqxx::connection conn{db_connect_options.c_str()};
				pqxx::work txn(conn);

				cout << "Requested boost list " << endl;

				conn.prepare("input_download", "select input_id,input_data from inputs where boost_id in (select boost_id from boost_versions where ver_id = $1) ");
				const auto res_table =  txn.prepared("input_download")(request->ver_id()).exec();

				if(res_table.size() == 0)
				{
					return Status(ABORTED, "No inputs found");
				}

				for(const auto &res : res_table)
				{
					// create a response
					auto reply = DownloadInputsReply_t();

					// get input id
					const auto input_id = res["input_id"].as<uint32_t>();
					reply.set_input_id(input_id);

					// get yaml string
					const auto input_contents = pqxx::binarystring(res["input_data"]).str();
					reply.set_yaml_string(input_contents);

					writer->Write(reply);
				}

				return Status::OK;
			}
			catch(const std::exception& e)
			{
				return Status(ABORTED, e.what());
			}
		}

		Status DownloadCrashes(ServerContext *context, const DownloadInputsRequest_t *request, ServerWriter<DownloadInputsReply_t> *writer) override final  
		{
			try
			{
				pqxx::connection conn{db_connect_options.c_str()};
				pqxx::work txn(conn);

				cout << "Requested boost list " << endl;

				conn.prepare("input_download", "select input_id,input_data from inputs where input_id in (select input_id from forensics where ver_id = $1) ");
				const auto res_table =  txn.prepared("input_download")(request->ver_id()).exec();

				// OK to have 0 results, just means we haven't found a crash yet. 
				for(const auto &res : res_table)
				{
					// create a response
					auto reply = DownloadInputsReply_t();

					// get input id
					const auto input_id = res["input_id"].as<uint32_t>();
					reply.set_input_id(input_id);

					// get yaml string
					const auto input_contents = pqxx::binarystring(res["input_data"]).str();
					reply.set_yaml_string(input_contents);

					writer->Write(reply);
				}

				return Status::OK;
			}
			catch(const std::exception& e)
			{
				return Status(ABORTED, e.what());
			}
		}

		Status AddFuzzLog(ServerContext *context, const AddFuzzLogRequest_t *request, Empty_t *reply) override final  
		{
			try
			{
				pqxx::connection conn{db_connect_options.c_str()};
				pqxx::work txn(conn);

				cout << "Add to Fuzz request" << endl;

				conn.prepare("update_fuzz_log", 
						R"""(
							with upsert as (update fuzz_logs set log = log || $2 where ver_id = $1 returning *) 
							insert into fuzz_logs (ver_id,log) select $1,$2 where not exists (select * from upsert)
						    )""");

				txn.prepared("update_fuzz_log")(request->ver_id())(request->log()).exec();
				txn.commit();
				return Status::OK;
			}
			catch(const std::exception& e)
			{
				return Status(ABORTED, e.what());
			}
			
		}
		Status AddForensics(ServerContext *context, const AddForensicsRequest_t *request, AddForensicsReply_t *reply) override final  
		{
			try
			{
				pqxx::connection conn{db_connect_options.c_str()};
				pqxx::work txn(conn);

				cout << "Add to crash request" << endl;

				conn.prepare("addfor", 
						R"""(
							insert into forensics (ver_id,input_id) values ($1,$2) returning forensic_id
						    )""");

				const auto res = txn.prepared("addfor")(request->ver_id())(request->input_id()).exec();
				txn.commit();
				if(res.size() != 1)
				{
					return Status(ABORTED, "Internal error");
				}
				reply->set_forensic_id(res[0]["forensic_id"].as<uint32_t>());
				return Status::OK;
			}
			catch(const std::exception& e)
			{
				return Status(ABORTED, e.what());
			}
		}

	private:
		string db_connect_options = string("dbname=" TURBODB " host=localhost");
		const string turbod_schema = 
			R"""(
				-- avoid NOTICES if tables already exist.
				SET client_min_messages = error;

				-- avoid error if type already exists
				DO $$ BEGIN
					create type InputProvenance as ENUM ('User', 'AFL', 'QSym', 'Unknown');
				EXCEPTION
					when duplicate_object then null;
				END $$;

				create table if not exists boosts
				(
					boost_id            serial primary key,
					boost_name          bytea,
					create_time         timestamp
				);

				create table if not exists boost_versions
				(
					ver_id              serial primary key,
					boost_id            integer references boosts(boost_id),
					orig_binary         bytea,
					orig_bin_sha        bytea,
					zafl_binary         bytea,
					zafl_bin_sha        bytea,
					submit_time         timestamp,
					start_time          timestamp,
					end_time            timestamp,
					err                 bool
				);

				create table if not exists inputs
				(
					input_id            serial primary key,
					boost_id            integer references boosts(boost_id),
					create_time         timestamp,
					prov                InputProvenance,
					input_data          bytea,
					sha256              bytea
				);

				create table if not exists forensics
				(
					forensic_id         serial primary key,
					input_id            integer references inputs(input_id),
					ver_id              integer references boost_versions(ver_id),
					exit_code           integer,
					backtrace           bytea,
					prov                InputProvenance
				);

				create table if not exists seeds
				(
					seed_id             serial primary key,
					input_id            integer references inputs(input_id)
				);


				create table if not exists zafl_logs
				(
					log_id              serial primary key,
					log                 bytea,
					ver_id              integer references boost_versions(ver_id)
				);

				create table if not exists fuzz_logs
				(
					log_id              serial primary key,
					log                 bytea,
					ver_id              integer references boost_versions(ver_id),
					prov                InputProvenance
				);
			)""";



};

void RunServer() 
{
	// listening port
	const auto server_address=string("0.0.0.0:55155");

	// create the service, build it, bind to listening port, and register.
	TurboServiceImpl service;
	ServerBuilder builder;
	builder.AddListeningPort(server_address, grpc::InsecureServerCredentials());
	builder.RegisterService(&service);

	// Finally start the server.
	auto server = builder.BuildAndStart();

	if(server == nullptr)
	{
		cout << "Could not start server! (Port already in use?)" << endl;
		exit(1);
	}

	// log
	cout << "Server listening on " << server_address << endl;

	// Wait for the server to shutdown. Note that some other thread must be
	// responsible for shutting down the server for this call to ever return.
	server->Wait();
}

int main(int argc, char **argv) 
{
	const auto create_db_res = system("createdb " TURBODB);
	if(create_db_res == -1)
	{
		cout << "Could not create process for creating database" << endl;
		perror("main"); 
	}
	RunServer();
	return 0;
}