Skip to content

Commit df80595

Browse files
laramielcopybara-github
authored andcommitted
Add internal/grpc/serverauth
This is the server-side equivalent of internal/grpc/clientauth PiperOrigin-RevId: 716465764 Change-Id: Ib9d0b4bbfa3341e8fbbe4aaab37ac851e1500298
1 parent 5f59990 commit df80595

18 files changed

+286
-63
lines changed

tensorstore/internal/grpc/BUILD

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,8 @@ tensorstore_cc_library(
5555
hdrs = ["server_credentials.h"],
5656
deps = [
5757
"//tensorstore:context",
58+
"//tensorstore/internal/grpc/serverauth:default_strategy",
59+
"//tensorstore/internal/grpc/serverauth:strategy",
5860
"//tensorstore/internal/json_binding",
5961
"//tensorstore/internal/json_binding:bindable",
6062
"//tensorstore/util:result",

tensorstore/internal/grpc/server_credentials.cc

Lines changed: 22 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -20,8 +20,11 @@
2020
#include "absl/base/attributes.h"
2121
#include "absl/base/const_init.h"
2222
#include "absl/synchronization/mutex.h"
23+
#include "grpcpp/security/server_credentials.h" // third_party
2324
#include "tensorstore/context.h"
2425
#include "tensorstore/context_resource_provider.h"
26+
#include "tensorstore/internal/grpc/serverauth/default_strategy.h"
27+
#include "tensorstore/internal/grpc/serverauth/strategy.h"
2528
#include "tensorstore/util/result.h"
2629

2730
namespace tensorstore {
@@ -37,23 +40,32 @@ const internal::ContextResourceRegistration<GrpcServerCredentials>
3740
// of grpc credentials. See grpcpp/security/credentials.h for options, such as:
3841
// ::grpc::experimental::LocalServerCredentials(LOCAL_TCP);
3942

43+
std::shared_ptr<internal_grpc::ServerAuthenticationStrategy>
44+
GrpcServerCredentials::Resource::GetAuthenticationStrategy() {
45+
absl::MutexLock l(&credentials_mu);
46+
if (strategy_) return strategy_;
47+
return internal_grpc::CreateInsecureServerAuthenticationStrategy();
48+
}
49+
4050
/* static */
4151
bool GrpcServerCredentials::Use(
4252
tensorstore::Context context,
4353
std::shared_ptr<::grpc::ServerCredentials> credentials) {
44-
auto resource = context.GetResource<GrpcServerCredentials>().value();
45-
// NOTE: We really want std::atomic<std::shared_ptr<>>.
46-
absl::MutexLock l(&credentials_mu);
47-
bool result = (resource->credentials_ == nullptr);
48-
resource->credentials_ = std::move(credentials);
49-
return result;
54+
return Use(
55+
context,
56+
std::make_shared<internal_grpc::DefaultServerAuthenticationStrategy>(
57+
std::move(credentials)));
5058
}
5159

52-
std::shared_ptr<::grpc::ServerCredentials>
53-
GrpcServerCredentials::Resource::GetCredentials() {
60+
/* static */
61+
bool GrpcServerCredentials::Use(
62+
tensorstore::Context context,
63+
std::shared_ptr<internal_grpc::ServerAuthenticationStrategy> credentials) {
64+
auto resource = context.GetResource<GrpcServerCredentials>().value();
5465
absl::MutexLock l(&credentials_mu);
55-
if (credentials_) return credentials_;
56-
return grpc::InsecureServerCredentials();
66+
bool result = (resource->strategy_ == nullptr);
67+
resource->strategy_ = std::move(credentials);
68+
return result;
5769
}
5870

5971
} // namespace tensorstore

tensorstore/internal/grpc/server_credentials.h

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
#include "grpcpp/security/server_credentials.h" // third_party
2121
#include "tensorstore/context.h"
2222
#include "tensorstore/context_resource_provider.h"
23+
#include "tensorstore/internal/grpc/serverauth/strategy.h"
2324
#include "tensorstore/internal/json_binding/bindable.h"
2425
#include "tensorstore/internal/json_binding/json_binding.h"
2526
#include "tensorstore/util/result.h"
@@ -48,11 +49,12 @@ struct GrpcServerCredentials final
4849

4950
struct Resource {
5051
// Returns either the owned credentials or a new default credential.
51-
std::shared_ptr<::grpc::ServerCredentials> GetCredentials();
52+
std::shared_ptr<internal_grpc::ServerAuthenticationStrategy>
53+
GetAuthenticationStrategy();
5254

5355
private:
5456
friend struct GrpcServerCredentials;
55-
std::shared_ptr<::grpc::ServerCredentials> credentials_;
57+
std::shared_ptr<internal_grpc::ServerAuthenticationStrategy> strategy_;
5658
};
5759

5860
static constexpr Spec Default() { return {}; }
@@ -72,6 +74,9 @@ struct GrpcServerCredentials final
7274
/// Returns true when prior credentials were nullptr.
7375
static bool Use(tensorstore::Context context,
7476
std::shared_ptr<::grpc::ServerCredentials> credentials);
77+
static bool Use(
78+
tensorstore::Context context,
79+
std::shared_ptr<internal_grpc::ServerAuthenticationStrategy> credentials);
7580
};
7681

7782
} // namespace tensorstore

tensorstore/internal/grpc/server_credentials_test.cc

Lines changed: 14 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414

1515
#include "tensorstore/internal/grpc/server_credentials.h"
1616

17+
#include <gmock/gmock.h>
1718
#include <gtest/gtest.h>
1819
#include "grpcpp/security/server_credentials.h" // third_party
1920
#include "tensorstore/context.h"
@@ -22,21 +23,29 @@
2223
namespace {
2324

2425
using ::tensorstore::GrpcServerCredentials;
26+
using ::testing::Eq;
27+
using ::testing::Ne;
2528

2629
TEST(GrpcServerCredentials, Use) {
2730
auto use = grpc::experimental::LocalServerCredentials(LOCAL_TCP);
2831
auto ctx = tensorstore::Context::Default();
2932

3033
EXPECT_TRUE(GrpcServerCredentials::Use(ctx, use));
31-
auto a = ctx.GetResource<GrpcServerCredentials>().value()->GetCredentials();
32-
EXPECT_EQ(a.get(), use.get());
34+
auto a = ctx.GetResource<GrpcServerCredentials>()
35+
.value()
36+
->GetAuthenticationStrategy();
37+
EXPECT_THAT(a->GetServerCredentials().get(), Eq(use.get()));
3338
}
3439

3540
TEST(GrpcServerCredentials, Default) {
3641
auto ctx = tensorstore::Context::Default();
37-
auto a = ctx.GetResource<GrpcServerCredentials>().value()->GetCredentials();
38-
auto b = ctx.GetResource<GrpcServerCredentials>().value()->GetCredentials();
39-
EXPECT_NE(a.get(), b.get());
42+
auto a = ctx.GetResource<GrpcServerCredentials>()
43+
.value()
44+
->GetAuthenticationStrategy();
45+
auto b = ctx.GetResource<GrpcServerCredentials>()
46+
.value()
47+
->GetAuthenticationStrategy();
48+
EXPECT_THAT(a.get(), Ne(b.get()));
4049
}
4150

4251
} // namespace
Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
load("//bazel:tensorstore.bzl", "tensorstore_cc_library")
2+
3+
package(default_visibility = ["//tensorstore:internal_packages"])
4+
5+
tensorstore_cc_library(
6+
name = "strategy",
7+
hdrs = ["strategy.h"],
8+
deps = ["@com_github_grpc_grpc//:grpc++"],
9+
)
10+
11+
tensorstore_cc_library(
12+
name = "default_strategy",
13+
srcs = ["default_strategy.cc"],
14+
hdrs = ["default_strategy.h"],
15+
deps = [
16+
":strategy",
17+
"@com_github_grpc_grpc//:grpc++",
18+
],
19+
)
Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
// Copyright 2025 The TensorStore Authors
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
#include "tensorstore/internal/grpc/serverauth/default_strategy.h"
16+
17+
#include <memory>
18+
19+
#include "grpcpp/security/server_credentials.h" // third_party
20+
#include "tensorstore/internal/grpc/serverauth/strategy.h"
21+
22+
namespace tensorstore {
23+
namespace internal_grpc {
24+
25+
std::shared_ptr<ServerAuthenticationStrategy>
26+
CreateInsecureServerAuthenticationStrategy() {
27+
return std::make_shared<DefaultServerAuthenticationStrategy>(
28+
grpc::InsecureServerCredentials());
29+
}
30+
31+
} // namespace internal_grpc
32+
} // namespace tensorstore
Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
// Copyright 2025 The TensorStore Authors
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
#ifndef TENSORSTORE_INTERNAL_GRPC_SERVERAUTH_DEFAULT_STRATEGY_H_
16+
#define TENSORSTORE_INTERNAL_GRPC_SERVERAUTH_DEFAULT_STRATEGY_H_
17+
18+
#include <memory>
19+
#include <utility>
20+
#include <vector>
21+
22+
#include "grpcpp/security/server_credentials.h" // third_party
23+
#include "grpcpp/server_builder.h" // third_party
24+
#include "tensorstore/internal/grpc/serverauth/strategy.h"
25+
26+
namespace tensorstore {
27+
namespace internal_grpc {
28+
29+
class DefaultServerAuthenticationStrategy
30+
: public ServerAuthenticationStrategy {
31+
public:
32+
DefaultServerAuthenticationStrategy(
33+
std::shared_ptr<grpc::ServerCredentials> credentials)
34+
: credentials_(std::move(credentials)) {}
35+
36+
~DefaultServerAuthenticationStrategy() override = default;
37+
38+
std::shared_ptr<grpc::ServerCredentials> GetServerCredentials()
39+
const override {
40+
return credentials_;
41+
}
42+
43+
void AddBuilderParameters(grpc::ServerBuilder& builder) const override {}
44+
45+
std::shared_ptr<grpc::ServerCredentials> credentials_;
46+
};
47+
48+
/// Creates an "insecure" server authentication strategy.
49+
std::shared_ptr<ServerAuthenticationStrategy>
50+
CreateInsecureServerAuthenticationStrategy();
51+
52+
} // namespace internal_grpc
53+
} // namespace tensorstore
54+
55+
#endif // TENSORSTORE_INTERNAL_GRPC_SERVERAUTH_DEFAULT_STRATEGY_H_
Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
// Copyright 2025 The TensorStore Authors
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
#ifndef TENSORSTORE_INTERNAL_GRPC_SERVERAUTH_STRATEGY_H_
16+
#define TENSORSTORE_INTERNAL_GRPC_SERVERAUTH_STRATEGY_H_
17+
18+
#include <memory>
19+
20+
#include "grpcpp/security/server_credentials.h" // third_party
21+
#include "grpcpp/server_builder.h" // third_party
22+
23+
namespace tensorstore {
24+
namespace internal_grpc {
25+
26+
/// Installs gRPC Server authentication strategies.
27+
///
28+
/// Usage:
29+
/// auto strategy = ...;
30+
/// grpc::ServerBuilder builder;
31+
/// builder.RegisterService(...);
32+
/// strategy->AddBuilderParameters(builder);
33+
/// builder.AddListeningPort(bind_addresses,
34+
/// strategy->GetServerCredentials(),
35+
/// &bound_port);
36+
/// auto server = builder.BuildAndStart();
37+
class ServerAuthenticationStrategy {
38+
public:
39+
virtual ~ServerAuthenticationStrategy() = default;
40+
41+
virtual std::shared_ptr<grpc::ServerCredentials> GetServerCredentials()
42+
const = 0;
43+
44+
virtual void AddBuilderParameters(grpc::ServerBuilder& builder) const = 0;
45+
};
46+
47+
} // namespace internal_grpc
48+
} // namespace tensorstore
49+
50+
#endif // TENSORSTORE_INTERNAL_GRPC_SERVERAUTH_STRATEGY_H_

tensorstore/kvstore/ocdbt/distributed/BUILD

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,8 @@ tensorstore_cc_library(
2525
"//tensorstore/internal/container:heterogeneous_container",
2626
"//tensorstore/internal/container:intrusive_red_black_tree",
2727
"//tensorstore/internal/grpc:peer_address",
28-
"//tensorstore/internal/grpc:utils",
28+
"//tensorstore/internal/grpc/serverauth:default_strategy",
29+
"//tensorstore/internal/grpc/serverauth:strategy",
2930
"//tensorstore/internal/json_binding",
3031
"//tensorstore/internal/json_binding:bindable",
3132
"//tensorstore/internal/log:verbose_flag",
@@ -208,6 +209,7 @@ tensorstore_cc_library(
208209
"//tensorstore/kvstore/ocdbt/non_distributed:create_new_manifest",
209210
"//tensorstore/kvstore/ocdbt/non_distributed:storage_generation",
210211
"//tensorstore/kvstore/ocdbt/non_distributed:write_nodes",
212+
"//tensorstore/util:bit_span",
211213
"//tensorstore/util:bit_vec",
212214
"//tensorstore/util:division",
213215
"//tensorstore/util:executor",
@@ -338,6 +340,8 @@ tensorstore_cc_library(
338340
"//tensorstore/internal/cache_key",
339341
"//tensorstore/internal/grpc/clientauth:authentication_strategy",
340342
"//tensorstore/internal/grpc/clientauth:channel_authentication",
343+
"//tensorstore/internal/grpc/serverauth:default_strategy",
344+
"//tensorstore/internal/grpc/serverauth:strategy",
341345
"//tensorstore/internal/json_binding",
342346
"//tensorstore/internal/json_binding:bindable",
343347
"@com_github_grpc_grpc//:grpc++",

tensorstore/kvstore/ocdbt/distributed/cooperator_get_manifest.cc

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -194,10 +194,7 @@ grpc::ServerUnaryReactor* Cooperator::GetOrCreateManifest(
194194
const grpc_gen::GetOrCreateManifestRequest* request,
195195
grpc_gen::GetOrCreateManifestResponse* response) {
196196
auto* reactor = context->DefaultReactor();
197-
if (auto status = security_->ValidateServerRequest(context); !status.ok()) {
198-
reactor->Finish(internal::AbslStatusToGrpcStatus(status));
199-
return reactor;
200-
}
197+
201198
if (!internal::IncrementReferenceCountIfNonZero(*this)) {
202199
// Shutting down
203200
reactor->Finish(

0 commit comments

Comments
 (0)