diff --git a/src/grpc_server.cpp b/src/grpc_server.cpp index 50ac445..07a39df 100644 --- a/src/grpc_server.cpp +++ b/src/grpc_server.cpp @@ -6,14 +6,11 @@ namespace kv { namespace { constexpr std::size_t max_key_size = 64 * 1024; // 64 KiB constexpr std::size_t max_value_size = 1024 * 1024; // 1 MiB -} // namespace - -KvStoreServiceImpl::KvStoreServiceImpl(Engine *engine) : engine_(engine) {} -grpc::Status KvStoreServiceImpl::Put(grpc::ServerContext *, - const kv::v1::PutRequest *request, - kv::v1::PutResponse *) { - const std::string &key = request->key(); +// Check a key passed to Put / Get / Delete: must be non-empty and within +// the size cap. Returns OK if valid, otherwise INVALID_ARGUMENT with a +// descriptive message. +grpc::Status validate_key(const std::string &key) { if (key.empty()) { return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, "Key cannot be empty"); @@ -23,16 +20,38 @@ grpc::Status KvStoreServiceImpl::Put(grpc::ServerContext *, "Key size exceeds maximum allowed size of " + std::to_string(max_key_size) + " bytes"); } + return grpc::Status::OK; +} - const std::string &value = request->value(); - if (value.size() > max_value_size) { +// Generic size check with a customisable field name for the error message. +// Used for Scan's start_key / end_key (empty is allowed) and Put's value. +grpc::Status validate_size(const std::string &field_name, std::size_t actual, + std::size_t max) { + if (actual > max) { return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, - "Value size exceeds maximum allowed size of " + - std::to_string(max_value_size) + " bytes"); + field_name + " size exceeds maximum allowed size of " + + std::to_string(max) + " bytes"); + } + return grpc::Status::OK; +} +} // namespace + +KvStoreServiceImpl::KvStoreServiceImpl(Engine *engine) : engine_(engine) {} + +grpc::Status KvStoreServiceImpl::Put(grpc::ServerContext *, + const kv::v1::PutRequest *request, + kv::v1::PutResponse *) { + if (auto status = validate_key(request->key()); !status.ok()) { + return status; + } + if (auto status = + validate_size("Value", request->value().size(), max_value_size); + !status.ok()) { + return status; } try { - engine_->put(key, value); + engine_->put(request->key(), request->value()); } catch (const std::exception &e) { return grpc::Status(grpc::StatusCode::INTERNAL, e.what()); } @@ -42,20 +61,13 @@ grpc::Status KvStoreServiceImpl::Put(grpc::ServerContext *, grpc::Status KvStoreServiceImpl::Get(grpc::ServerContext *, const kv::v1::GetRequest *request, kv::v1::GetResponse *response) { - const std::string &key = request->key(); - if (key.empty()) { - return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, - "Key cannot be empty"); - } - if (key.size() > max_key_size) { - return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, - "Key size exceeds maximum allowed size of " + - std::to_string(max_key_size) + " bytes"); + if (auto status = validate_key(request->key()); !status.ok()) { + return status; } std::optional value; try { - value = engine_->get(key); + value = engine_->get(request->key()); } catch (const std::exception &e) { return grpc::Status(grpc::StatusCode::INTERNAL, e.what()); } @@ -70,19 +82,12 @@ grpc::Status KvStoreServiceImpl::Get(grpc::ServerContext *, grpc::Status KvStoreServiceImpl::Delete(grpc::ServerContext *, const kv::v1::DeleteRequest *request, kv::v1::DeleteResponse *) { - const std::string &key = request->key(); - if (key.empty()) { - return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, - "Key cannot be empty"); - } - if (key.size() > max_key_size) { - return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, - "Key size exceeds maximum allowed size of " + - std::to_string(max_key_size) + " bytes"); + if (auto status = validate_key(request->key()); !status.ok()) { + return status; } try { - engine_->remove(key); + engine_->remove(request->key()); } catch (const std::exception &e) { return grpc::Status(grpc::StatusCode::INTERNAL, e.what()); } @@ -93,16 +98,17 @@ grpc::Status KvStoreServiceImpl::Scan(grpc::ServerContext *, const kv::v1::ScanRequest *request, grpc::ServerWriter *writer) { - // Empty start_key / end_key are allowed (mean unbounded). - if (request->start_key().size() > max_key_size) { - return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, - "start_key size exceeds maximum allowed size of " + - std::to_string(max_key_size) + " bytes"); + // Empty start_key / end_key are allowed (mean unbounded), so size-only + // validation rather than validate_key. + if (auto status = + validate_size("start_key", request->start_key().size(), max_key_size); + !status.ok()) { + return status; } - if (request->end_key().size() > max_key_size) { - return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, - "end_key size exceeds maximum allowed size of " + - std::to_string(max_key_size) + " bytes"); + if (auto status = + validate_size("end_key", request->end_key().size(), max_key_size); + !status.ok()) { + return status; } try {