diff --git a/clang-tools-extra/clangd/index/remote/server/Server.cpp b/clang-tools-extra/clangd/index/remote/server/Server.cpp --- 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::desc("Maximum time a channel may stay idle until server closes " "the connection, in seconds. Defaults to 480.")); +llvm::cl::opt 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 CurrentRequest; class RemoteIndexServer final : public v1::SymbolIndex::Service { @@ -123,7 +129,12 @@ } 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 @@ 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 @@ 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 @@ 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 @@ 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 --- /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"} +