Skip to content

Commit

Permalink
Fix C++ server/client wrapper undefined behavior
Browse files Browse the repository at this point in the history 
Switches to using a SessionConfig-provider model instead of trying to
keep re-using the same SessionConfig pointer (which results in invalid
memory accesses)

In this change, we keep the old constructor around so that copybara
imports will pass. Once clients are migrated, we can remove this.

Fixed: b/389917088

Change-Id: Idfb83faaaca72b5028716626875351d5a3151e96
  • Loading branch information
@jblebrun
jblebrun committed Jan 17, 2025
1 parent f26388b commit 42fcfac
Show file tree
Hide file tree
Showing 6 changed files with 88 additions and 29 deletions.
2 changes: 1 addition & 1 deletion cc/client/session_client.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@ namespace oak::client {
absl::StatusOr<std::unique_ptr<OakSessionClient::Channel>>
OakSessionClient::NewChannel(std::unique_ptr<Channel::Transport> transport) {
absl::StatusOr<std::unique_ptr<session::ClientSession>> session =
session::ClientSession::Create(config_);
session::ClientSession::Create(config_provider_());

while (!(*session)->IsOpen()) {
absl::StatusOr<std::optional<session::v1::SessionRequest>> init_request =
Expand Down
31 changes: 22 additions & 9 deletions cc/client/session_client.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -37,16 +37,29 @@ class OakSessionClient {
session::ClientSession>;

// A valid `SessionConfig` can be obtained using
// oak::session::SessionConfigBuilder.
OakSessionClient(session::SessionConfig* config) : config_(config) {}
// oak::session::SessionConfigBuilder. Each session needs its own unique
// SessionConfig instance, so a function to create a new SessionConfig should
// be provided here.
OakSessionClient(
absl::AnyInvocable<session::SessionConfig*()> config_provider)
: config_provider_(std::move(config_provider)) {}

// Use a default configuration, Unattested + NoiseNN
ABSL_DEPRECATED("Use the config-providing variant.")
// Use a default configuration provider, Unattested + NoiseNN
ABSL_DEPRECATED("Use the config-provider-providing variant.")
OakSessionClient()
: OakSessionClient(
session::SessionConfigBuilder(session::AttestationType::kUnattested,
session::HandshakeType::kNoiseNN)
.Build()) {}
: OakSessionClient([] {
return session::SessionConfigBuilder(
session::AttestationType::kUnattested,
session::HandshakeType::kNoiseNN)
.Build();
}) {}

// Keeping this around briefly until we transition existing clients.
ABSL_DEPRECATED(
"This constructor will lead to UB. Use the config-provider-providing "
"variant.")
OakSessionClient(session::SessionConfig* config)
: OakSessionClient([config] { return config; }) {}

// Create a new OakClientChannel instance with the provided session and
// transport.
Expand All @@ -61,7 +74,7 @@ class OakSessionClient {
std::unique_ptr<OakSessionClient::Channel::Transport> transport);

private:
session::SessionConfig* config_;
absl::AnyInvocable<session::SessionConfig*()> config_provider_;
};

} // namespace oak::client
Expand Down
26 changes: 21 additions & 5 deletions cc/client/session_client_test.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -65,17 +65,33 @@ session::SessionConfig* TestSessionConfig() {
TEST(OakSessionClientTest, CreateSuccessFullyHandshakes) {
auto server_session = session::ServerSession::Create(TestSessionConfig());
ASSERT_THAT(server_session, IsOk());
auto _ = OakSessionClient(TestSessionConfig())
.NewChannel(
std::make_unique<TestTransport>(std::move(*server_session)));
auto channel = OakSessionClient(TestSessionConfig)
.NewChannel(std::make_unique<TestTransport>(
std::move(*server_session)));
EXPECT_THAT(channel, IsOk());
}

TEST(OakSessionClientTest, CreateSecondClientSuccessFullyHandshakes) {
auto server_session = session::ServerSession::Create(TestSessionConfig());
ASSERT_THAT(server_session, IsOk());
auto client = OakSessionClient(TestSessionConfig);

auto channel = client.NewChannel(
std::make_unique<TestTransport>(std::move(*server_session)));
EXPECT_THAT(channel, IsOk());

auto server_session2 = session::ServerSession::Create(TestSessionConfig());
auto channel2 = client.NewChannel(
std::make_unique<TestTransport>(std::move(*server_session2)));
EXPECT_THAT(channel2, IsOk());
}

TEST(OakSessionClientTest, CreatedSessionCanSend) {
auto server_session = session::ServerSession::Create(TestSessionConfig());
// Hold a pointer for testing behavior below.
session::ServerSession* server_session_ptr = server_session->get();
ASSERT_THAT(server_session, IsOk());
auto channel = OakSessionClient(TestSessionConfig())
auto channel = OakSessionClient(TestSessionConfig)
.NewChannel(std::make_unique<TestTransport>(
std::move(*server_session)));

Expand All @@ -93,7 +109,7 @@ TEST(OakSessionClientTest, CreatedSessionCanReceive) {
// Hold a pointer for testing behavior below.
session::ServerSession* server_session_ptr = server_session->get();
ASSERT_THAT(server_session, IsOk());
auto channel = OakSessionClient(TestSessionConfig())
auto channel = OakSessionClient(TestSessionConfig)
.NewChannel(std::make_unique<TestTransport>(
std::move(*server_session)));

Expand Down
2 changes: 1 addition & 1 deletion cc/server/session_server.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@ namespace oak::server {

absl::StatusOr<std::unique_ptr<OakSessionServer::Channel>>
OakSessionServer::NewChannel(std::unique_ptr<Channel::Transport> transport) {
auto session = session::ServerSession::Create(config_);
auto session = session::ServerSession::Create(config_provider_());
if (!session.ok()) {
return util::status::Annotate(session.status(),
"Failed to create server session");
Expand Down
29 changes: 21 additions & 8 deletions cc/server/session_server.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -37,16 +37,29 @@ class OakSessionServer {
session::ServerSession>;

// A valid `SessionConfig` can be obtained using
// oak::session::SessionConfigBuilder.
OakSessionServer(session::SessionConfig* config) : config_(config) {}
// oak::session::SessionConfigBuilder. Each session needs its own unique
// SessionConfig instance, so a function to create a new SessionConfig should
// be provided.
OakSessionServer(
absl::AnyInvocable<session::SessionConfig*()> config_provider)
: config_provider_(std::move(config_provider)) {}

// Use a default configuration, Unattested + NoiseNN
// Use a default configuration provider, Unattested + NoiseNN
ABSL_DEPRECATED("Use the config-providing variant.")
OakSessionServer()
: OakSessionServer(
session::SessionConfigBuilder(session::AttestationType::kUnattested,
session::HandshakeType::kNoiseNN)
.Build()) {}
: OakSessionServer([] {
return session::SessionConfigBuilder(
session::AttestationType::kUnattested,
session::HandshakeType::kNoiseNN)
.Build();
}) {}

// Keeping this around briefly until we transition existing clients.
ABSL_DEPRECATED(
"This constructor will lead to UB. Use the config-provider-providing "
"variant.")
OakSessionServer(session::SessionConfig* config)
: OakSessionServer([config] { return config; }) {}

// Create a new OakServerChannel instance with the provided session and
// transport.
Expand All @@ -61,7 +74,7 @@ class OakSessionServer {
std::unique_ptr<OakSessionServer::Channel::Transport> transport);

private:
session::SessionConfig* config_;
absl::AnyInvocable<session::SessionConfig*()> config_provider_;
};

} // namespace oak::server
Expand Down
27 changes: 22 additions & 5 deletions cc/server/session_server_test.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -68,17 +68,34 @@ session::SessionConfig* TestSessionConfig() {
TEST(OakSessionServerTest, CreateSuccessFullyHandshakes) {
auto client_session = session::ClientSession::Create(TestSessionConfig());
ASSERT_THAT(client_session, IsOk());
auto _ = OakSessionServer(TestSessionConfig())
.NewChannel(
std::make_unique<TestTransport>(std::move(*client_session)));
auto channel = OakSessionServer(TestSessionConfig)
.NewChannel(std::make_unique<TestTransport>(
std::move(*client_session)));
ASSERT_THAT(channel, IsOk());
}

TEST(OakSessionServerTest, SecondCreateSuccessFullyHandshakes) {
auto server = OakSessionServer(TestSessionConfig);

auto client_session = session::ClientSession::Create(TestSessionConfig());
ASSERT_THAT(client_session, IsOk());
auto channel = server.NewChannel(
std::make_unique<TestTransport>(std::move(*client_session)));
ASSERT_THAT(channel, IsOk());

auto client_session2 = session::ClientSession::Create(TestSessionConfig());
ASSERT_THAT(client_session2, IsOk());
auto channel2 = server.NewChannel(
std::make_unique<TestTransport>(std::move(*client_session2)));
ASSERT_THAT(channel2, IsOk());
}

TEST(OakSessionServerTest, CreatedSessionCanSend) {
auto client_session = session::ClientSession::Create(TestSessionConfig());
// Hold a pointer for testing behavior below.
session::ClientSession* client_session_ptr = client_session->get();
ASSERT_THAT(client_session, IsOk());
auto channel = OakSessionServer(TestSessionConfig())
auto channel = OakSessionServer(TestSessionConfig)
.NewChannel(std::make_unique<TestTransport>(
std::move(*client_session)));

Expand All @@ -96,7 +113,7 @@ TEST(OakSessionServerTest, CreatedSessionCanReceive) {
// Hold a pointer for testing behavior below.
session::ClientSession* client_session_ptr = client_session->get();
ASSERT_THAT(client_session, IsOk());
auto channel = OakSessionServer(TestSessionConfig())
auto channel = OakSessionServer(TestSessionConfig)
.NewChannel(std::make_unique<TestTransport>(
std::move(*client_session)));

Expand Down

0 comments on commit 42fcfac

Please sign in to comment.