add protocol check

apache · Jan 9, 2024 · 2dfd003 · 2dfd003
1 parent de5ce5c
commit 2dfd003
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 3 deletions.
diff --git a/src/brpc/server.cpp b/src/brpc/server.cpp
@@ -565,13 +565,22 @@ int Server::AddBuiltinServices() {
         LOG(ERROR) << "Fail to add GetJsService";
         return -1;
     }
-    if (AddBuiltinService(new (std::nothrow) GrpcHealthCheckService)) {
+    if (ProtocolEnabled("h2") &&
+        AddBuiltinService(new (std::nothrow) GrpcHealthCheckService)) {
         LOG(ERROR) << "Fail to add GrpcHealthCheckService";
         return -1;
     }
     return 0;
 }
 
+bool Server::ProtocolEnabled(const char * name) {
+    if (!_options.enabled_protocols.empty() &&
+        _options.enabled_protocols.find(name) == std::string::npos) {
+        return false;
+    }
+    return true;
+}
+
 bool is_http_protocol(const char* name) {
     if (name[0] != 'h') {
         return false;

diff --git a/src/brpc/server.h b/src/brpc/server.h
@@ -673,6 +673,7 @@ friend class Controller;
 
     AdaptiveMaxConcurrency& MaxConcurrencyOf(MethodProperty*);
     int MaxConcurrencyOf(const MethodProperty*) const;
+    bool ProtocolEnabled(const char *);
 
     DISALLOW_COPY_AND_ASSIGN(Server);
 

diff --git a/test/brpc_builtin_service_unittest.cpp b/test/brpc_builtin_service_unittest.cpp
@@ -598,6 +598,25 @@ TEST_F(BuiltinServiceTest, normal_grpc_health) {
     EXPECT_EQ(response.status(), grpc::health::v1::HealthCheckResponse_ServingStatus_SERVING);
 }
 
+TEST_F(BuiltinServiceTest, normal_grpc_health_with_protocol_check) {
+    brpc::Server server;
+    brpc::ServerOptions opt;
+    opt.enabled_protocols = "baidu_std";
+    ASSERT_EQ(0, server.Start(9798, &opt));
+
+    grpc::health::v1::HealthCheckResponse response;
+    grpc::health::v1::HealthCheckRequest request;
+    brpc::Controller cntl;
+    brpc::ChannelOptions copt;
+    copt.protocol = "h2:grpc";
+    brpc::Channel chan;
+    ASSERT_EQ(0, chan.Init("127.0.0.1:9798", &copt));
+    grpc::health::v1::Health_Stub stub(&chan);
+    stub.Check(&cntl, &request, &response, NULL);
+    EXPECT_TRUE(cntl.Failed());
+    EXPECT_EQ(1002, cntl.ErrorCode()); // ENOMETHOD Fail to find method /grpc.health.v1.Health/Check
+}
+
 TEST_F(BuiltinServiceTest, customized_grpc_health) {
     brpc::ServerOptions opt;
     MyGrpcHealthReporter hr;
@@ -606,7 +625,6 @@ TEST_F(BuiltinServiceTest, customized_grpc_health) {
 
     grpc::health::v1::HealthCheckResponse response;
     grpc::health::v1::HealthCheckRequest request;
-    request.set_service("grpc_req_from_brpc");
     brpc::Controller cntl;
 
     brpc::ChannelOptions copt;
@@ -616,7 +634,6 @@ TEST_F(BuiltinServiceTest, customized_grpc_health) {
 
     grpc::health::v1::Health_Stub stub(&chan);
     stub.Check(&cntl, &request, &response, NULL);
-
     EXPECT_FALSE(cntl.Failed()) << cntl.ErrorText();
     EXPECT_EQ(response.status(), grpc::health::v1::HealthCheckResponse_ServingStatus_UNKNOWN);
 }