fix: Extract validate_key() and validate_size() in gRPC handlers

Author: IsaacCheng9

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<std::string> 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<kv::v1::ScanResponse> *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 {