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,10 @@ Req.takeError()); return grpc::Status::CANCELLED; } + if (!Req->Limit || *Req->Limit > LimitResults) { + log("[public] Limiting result size for FuzzyFind request."); + Req->Limit = LimitResults; + } unsigned Sent = 0; unsigned FailedToSend = 0; bool HasMore = Index.fuzzyFind(*Req, [&](const clangd::Symbol &Item) { @@ -197,6 +214,10 @@ 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."); + Req->Limit = LimitResults; + } unsigned Sent = 0; unsigned FailedToSend = 0; bool HasMore = Index.refs(*Req, [&](const clangd::Ref &Item) { @@ -236,6 +257,10 @@ Req.takeError()); return grpc::Status::CANCELLED; } + if (!Req->Limit || *Req->Limit > LimitResults) { + log("[public] Limiting result size for Relations request."); + Req->Limit = LimitResults; + } unsigned Sent = 0; unsigned FailedToSend = 0; Index.relations(