Add mutex for stream in ProtoReflectionDescriptorDatabase, fix headers check

2 files changed, 14 insertions, 20 deletions

diff --git a/test/cpp/util/proto_reflection_descriptor_database.cc b/test/cpp/util/proto_reflection_descriptor_database.cc
index 3963b0c093..6513a14992 100644
--- a/test/cpp/util/proto_reflection_descriptor_database.cc
+++ b/test/cpp/util/proto_reflection_descriptor_database.cc
@@ -31,7 +31,7 @@
  *
  */
 
-#include "proto_reflection_descriptor_database.h"
+#include "test/cpp/util/proto_reflection_descriptor_database.h"
 
 #include <vector>
 
@@ -69,16 +69,14 @@ bool ProtoReflectionDescriptorDatabase::FindFileByName(
   request.set_file_by_filename(filename);
   ServerReflectionResponse response;
 
+  stream_mutex_.lock();
   GetStream()->Write(request);
   GetStream()->Read(&response);
+  stream_mutex_.unlock();
 
   if (response.message_response_case() ==
       ServerReflectionResponse::MessageResponseCase::kFileDescriptorResponse) {
     AddFileFromResponse(response.file_descriptor_response());
-    // const google::protobuf::FileDescriptorProto file_proto =
-    //     ParseFileDescriptorProtoResponse(response.file_descriptor_response());
-    // known_files_.insert(file_proto.name());
-    // cached_db_.Add(file_proto);
   } else if (response.message_response_case() ==
              ServerReflectionResponse::MessageResponseCase::kErrorResponse) {
     const ErrorResponse error = response.error_response();
@@ -119,19 +117,14 @@ bool ProtoReflectionDescriptorDatabase::FindFileContainingSymbol(
   request.set_file_containing_symbol(symbol_name);
   ServerReflectionResponse response;
 
+  stream_mutex_.lock();
   GetStream()->Write(request);
   GetStream()->Read(&response);
+  stream_mutex_.unlock();
 
-  // Status status = stub_->GetFileContainingSymbol(&ctx, request, &response);
   if (response.message_response_case() ==
       ServerReflectionResponse::MessageResponseCase::kFileDescriptorResponse) {
     AddFileFromResponse(response.file_descriptor_response());
-    // const google::protobuf::FileDescriptorProto file_proto =
-    //     ParseFileDescriptorProtoResponse(response.file_descriptor_response());
-    // if (known_files_.find(file_proto.name()) == known_files_.end()) {
-    //   known_files_.insert(file_proto.name());
-    //   cached_db_.Add(file_proto);
-    // }
   } else if (response.message_response_case() ==
              ServerReflectionResponse::MessageResponseCase::kErrorResponse) {
     const ErrorResponse error = response.error_response();
@@ -181,20 +174,14 @@ bool ProtoReflectionDescriptorDatabase::FindFileContainingExtension(
       field_number);
   ServerReflectionResponse response;
 
+  stream_mutex_.lock();
   GetStream()->Write(request);
   GetStream()->Read(&response);
+  stream_mutex_.unlock();
 
-  // Status status = stub_->GetFileContainingExtension(&ctx, request,
-  // &response);
   if (response.message_response_case() ==
       ServerReflectionResponse::MessageResponseCase::kFileDescriptorResponse) {
     AddFileFromResponse(response.file_descriptor_response());
-    // const google::protobuf::FileDescriptorProto file_proto =
-    //     ParseFileDescriptorProtoResponse(response.file_descriptor_response());
-    // if (known_files_.find(file_proto.name()) == known_files_.end()) {
-    //   known_files_.insert(file_proto.name());
-    //   cached_db_.Add(file_proto);
-    // }
   } else if (response.message_response_case() ==
              ServerReflectionResponse::MessageResponseCase::kErrorResponse) {
     const ErrorResponse error = response.error_response();
@@ -240,8 +227,10 @@ bool ProtoReflectionDescriptorDatabase::FindAllExtensionNumbers(
   request.set_all_extension_numbers_of_type(extendee_type);
   ServerReflectionResponse response;
 
+  stream_mutex_.lock();
   GetStream()->Write(request);
   GetStream()->Read(&response);
+  stream_mutex_.unlock();
 
   if (response.message_response_case() ==
       ServerReflectionResponse::MessageResponseCase::
@@ -272,8 +261,11 @@ bool ProtoReflectionDescriptorDatabase::GetServices(
   ServerReflectionRequest request;
   request.set_list_services("");
   ServerReflectionResponse response;
+
+  stream_mutex_.lock();
   GetStream()->Write(request);
   GetStream()->Read(&response);
+  stream_mutex_.unlock();
 
   if (response.message_response_case() ==
       ServerReflectionResponse::MessageResponseCase::kListServicesResponse) {
diff --git a/test/cpp/util/proto_reflection_descriptor_database.h b/test/cpp/util/proto_reflection_descriptor_database.h
index 4bb9c21a92..bc25fb0f5c 100644
--- a/test/cpp/util/proto_reflection_descriptor_database.h
+++ b/test/cpp/util/proto_reflection_descriptor_database.h
@@ -32,6 +32,7 @@
  */
 
 #include <memory>
+#include <mutex>
 #include <unordered_map>
 #include <unordered_set>
 #include <vector>
@@ -98,6 +99,7 @@ class ProtoReflectionDescriptorDatabase
   std::unordered_set<string> missing_symbols_;
   std::unordered_map<string, std::unordered_set<int>> missing_extensions_;
   std::unordered_map<string, std::vector<int>> cached_extension_numbers_;
+  std::mutex stream_mutex_;
 
   google::protobuf::SimpleDescriptorDatabase cached_db_;
 };