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

[kadircet]

Summary

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.

[kbobyrev]

Good point, thanks!

I wonder if this can somehow make the experience worse in _some_ specific cases but I can't think of anything on top of my mind, so this should be all good.

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

152

nit (here and elsewhere): maybe add something like (requested X, sending Y) to the message for clarity. I don't expect this to be common anyway so probably verbosity isn't an issue here.

[kadircet]

• Change logs to mention original limits.
• Add the test i forgot first time.

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

152

doesn't really apply here, as lookup request doesn't have a limit. practically it triggers when user makes a request with more than LimitResults request ids in a single batch, but in theory could kick in earlier as an index can return multiple results for the same symbol id (but they never do).
changing the rest though.

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

llvm::cl::opt<std::string> ServerAddress(
     "server-address", llvm::cl::init("0.0.0.0:50051"),
     llvm::cl::desc("Address of the invoked server. Defaults to 0.0.0.0:50051"));
 
 llvm::cl::opt<size_t> IdleTimeoutSeconds(
     "idle-timeout", llvm::cl::init(8 * 60),
     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 {
 public:
   RemoteIndexServer(clangd::SymbolIndex &Index, llvm::StringRef IndexRoot)
       : Index(Index) {
     llvm::SmallString<256> NativePath = IndexRoot;
     llvm::sys::path::native(NativePath);
grpc::Status Lookup(grpc::ServerContext *Context,
     trace::Span Tracer("LookupRequest");
     auto Req = ProtobufMarshaller->fromProtobuf(Request);
     if (!Req) {
       elog("Can not parse LookupRequest from protobuf: {0}", Req.takeError());
       return grpc::Status::CANCELLED;
     }
     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}",
              SerializedItem.takeError());
         ++FailedToSend;
         return;
       }
       LookupReply NextMessage;
       *NextMessage.mutable_stream_result() = *SerializedItem;
       logResponse(NextMessage);
       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);
     SPAN_ATTACH(Tracer, "Failed to send", FailedToSend);
     logRequestSummary("v1/Lookup", Sent, StartTime);
     return grpc::Status::OK;
   }
 
   grpc::Status FuzzyFind(grpc::ServerContext *Context,
                          const FuzzyFindRequest *Request,
                          grpc::ServerWriter<FuzzyFindReply> *Reply) override {
     auto StartTime = stopwatch::now();
     WithContextValue WithRequestContext(CurrentRequest, Context);
     logRequest(*Request);
     trace::Span Tracer("FuzzyFindRequest");
     auto Req = ProtobufMarshaller->fromProtobuf(Request);
     if (!Req) {
       elog("Can not parse FuzzyFindRequest from protobuf: {0}",
            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) {
       auto SerializedItem = ProtobufMarshaller->toProtobuf(Item);
       if (!SerializedItem) {
         elog("Unable to convert Symbol to protobuf: {0}",
              SerializedItem.takeError());
         ++FailedToSend;
grpc::Status Refs(grpc::ServerContext *Context, const RefsRequest *Request,
     WithContextValue WithRequestContext(CurrentRequest, Context);
     logRequest(*Request);
     trace::Span Tracer("RefsRequest");
     auto Req = ProtobufMarshaller->fromProtobuf(Request);
     if (!Req) {
       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) {
       auto SerializedItem = ProtobufMarshaller->toProtobuf(Item);
       if (!SerializedItem) {
         elog("Unable to convert Ref to protobuf: {0}",
              SerializedItem.takeError());
         ++FailedToSend;
grpc::Status Relations(grpc::ServerContext *Context,
     logRequest(*Request);
     trace::Span Tracer("RelationsRequest");
     auto Req = ProtobufMarshaller->fromProtobuf(Request);
     if (!Req) {
       elog("Can not parse RelationsRequest from protobuf: {0}",
            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(
         *Req, [&](const SymbolID &Subject, const clangd::Symbol &Object) {
           auto SerializedItem = ProtobufMarshaller->toProtobuf(Subject, Object);
           if (!SerializedItem) {
             elog("Unable to convert Relation to protobuf: {0}",
                  SerializedItem.takeError());

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

+# 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"}
+