[clangd][index-sever] Limit results in repsonse

This is to prevent server from being DOS'd by possible malicious
parties issuing requests that can yield huge responses.

One possible drawback is on rename workflow. As it really requests all
occurences, but it has an internal limit on 50 files currently.
We are putting the limit on 10000 elements per response So for rename to regress
one should have 10k refs to a symbol in less than 50 files. This seems unlikely
and we fix it if there are complaints by giving up on the response based on the
number of files covered instead.

Differential Revision: https://reviews.llvm.org/D101914

Author: kadircet

clang-tools-extra/clangd/index/remote/server/Server.cpp

@@ -93,6 +93,12 @@ llvm::cl::opt<size_t> IdleTimeoutSeconds(
     llvm::cl::desc("Maximum time a channel may stay idle until server closes "
                    "the connection, in seconds. Defaults to 480."));
 
+llvm::cl::opt<size_t> LimitResults(
+    "limit-results", llvm::cl::init(10000),
+    llvm::cl::desc("Maximum number of results to stream as a response to "
+                   "single request. Limit is to keep the server from being "
+                   "DOS'd. Defaults to 10000."));
+
 static Key<grpc::ServerContext *> CurrentRequest;
 
 class RemoteIndexServer final : public v1::SymbolIndex::Service {
@@ -123,7 +129,12 @@ class RemoteIndexServer final : public v1::SymbolIndex::Service {
     }
     unsigned Sent = 0;
     unsigned FailedToSend = 0;
+    bool HasMore = false;
     Index.lookup(*Req, [&](const clangd::Symbol &Item) {
+      if (Sent >= LimitResults) {
+        HasMore = true;
+        return;
+      }
       auto SerializedItem = ProtobufMarshaller->toProtobuf(Item);
       if (!SerializedItem) {
         elog("Unable to convert Symbol to protobuf: {0}",
@@ -137,8 +148,10 @@ class RemoteIndexServer final : public v1::SymbolIndex::Service {
       Reply->Write(NextMessage);
       ++Sent;
     });
+    if (HasMore)
+      log("[public] Limiting result size for Lookup request.");
     LookupReply LastMessage;
-    LastMessage.mutable_final_result()->set_has_more(true);
+    LastMessage.mutable_final_result()->set_has_more(HasMore);
     logResponse(LastMessage);
     Reply->Write(LastMessage);
     SPAN_ATTACH(Tracer, "Sent", Sent);
@@ -160,6 +173,11 @@ class RemoteIndexServer final : public v1::SymbolIndex::Service {
            Req.takeError());
       return grpc::Status::CANCELLED;
     }
+    if (!Req->Limit || *Req->Limit > LimitResults) {
+      log("[public] Limiting result size for FuzzyFind request from {0} to {1}",
+          Req->Limit, LimitResults);
+      Req->Limit = LimitResults;
+    }
     unsigned Sent = 0;
     unsigned FailedToSend = 0;
     bool HasMore = Index.fuzzyFind(*Req, [&](const clangd::Symbol &Item) {
@@ -197,6 +215,11 @@ class RemoteIndexServer final : public v1::SymbolIndex::Service {
       elog("Can not parse RefsRequest from protobuf: {0}", Req.takeError());
       return grpc::Status::CANCELLED;
     }
+    if (!Req->Limit || *Req->Limit > LimitResults) {
+      log("[public] Limiting result size for Refs request from {0} to {1}.",
+          Req->Limit, LimitResults);
+      Req->Limit = LimitResults;
+    }
     unsigned Sent = 0;
     unsigned FailedToSend = 0;
     bool HasMore = Index.refs(*Req, [&](const clangd::Ref &Item) {
@@ -236,6 +259,12 @@ class RemoteIndexServer final : public v1::SymbolIndex::Service {
            Req.takeError());
       return grpc::Status::CANCELLED;
     }
+    if (!Req->Limit || *Req->Limit > LimitResults) {
+      log("[public] Limiting result size for Relations request from {0} to "
+          "{1}.",
+          Req->Limit, LimitResults);
+      Req->Limit = LimitResults;
+    }
     unsigned Sent = 0;
     unsigned FailedToSend = 0;
     Index.relations(

clang-tools-extra/clangd/test/remote-index/result-limiting.test

@@ -0,0 +1,39 @@
+# REQUIRES: clangd-remote-index
+# RUN: rm -rf %t
+# RUN: clangd-indexer %S/Inputs/Source.cpp > %t.idx
+# RUN: %python %S/pipeline_helper.py --input-file-name=%s --server-arg=--limit-results=1 --project-root=%S --index-file=%t.idx | FileCheck %s
+
+{"jsonrpc":"2.0","id":0,"method":"initialize","params":{"processId":123,"rootPath":"clangd","capabilities":{},"trace":"off"}}
+---
+# Ensure there's a single result.
+{"jsonrpc":"2.0","id":1,"method":"workspace/symbol","params":{"query":"c"}}
+# CHECK: {
+# CHECK:   "id": 1,
+# CHECK:   "jsonrpc": "2.0",
+# CHECK:   "result": [
+# CHECK:     {
+# CHECK:       "containerName": "{{.*}}",
+# CHECK:       "kind": {{.*}},
+# CHECK:       "location": {
+# CHECK:         "range": {
+# CHECK:           "end": {
+# CHECK:             "character": {{.*}},
+# CHECK:             "line": {{.*}}
+# CHECK:           },
+# CHECK:           "start": {
+# CHECK:             "character": {{.*}},
+# CHECK:             "line": {{.*}}
+# CHECK:           }
+# CHECK:         },
+# CHECK:         "uri": "{{.*}}"
+# CHECK:       },
+# CHECK:       "name": "{{.*}}",
+# CHECK:       "score": {{.*}}
+# CHECK:     }
+# CHECK:   ]
+# CHECK: }
+---
+{"jsonrpc":"2.0","id":4,"method":"shutdown"}
+---
+{"jsonrpc":"2.0","method":"exit"}
+