[clang-tools-extra] daf3cb3 - [clangd][index-sever] Limit results in repsonse

[Kadir Cetinkaya via cfe-commits]

Author: Kadir Cetinkaya
Date: 2021-05-11T08:22:23+02:00
New Revision: daf3cb3b8a5868d9089a69025c556b564615b844

URL: https://github.com/llvm/llvm-project/commit/daf3cb3b8a5868d9089a69025c556b564615b844
DIFF: https://github.com/llvm/llvm-project/commit/daf3cb3b8a5868d9089a69025c556b564615b844.diff

LOG: [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

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

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

Removed: 
    


################################################################################
diff  --git a/clang-tools-extra/clangd/index/remote/server/Server.cpp b/clang-tools-extra/clangd/index/remote/server/Server.cpp
index f3cf131bb8a58..7073cc0dc5670 100644
--- a/clang-tools-extra/clangd/index/remote/server/Server.cpp
+++ b/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(

diff  --git a/clang-tools-extra/clangd/test/remote-index/result-limiting.test b/clang-tools-extra/clangd/test/remote-index/result-limiting.test
new file mode 100644
index 0000000000000..d0f6f6f0a5287
--- /dev/null
+++ b/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"}
+