Add more test cases for EncryptionRunner to improve test coverage.

connections/implementation/encryption_runner.cc

@@ -192,6 +192,7 @@ class ServerRunnable final {
 
   void HandleHandshakeOrIoException(CancelableAlarm* timeout_alarm) {
     timeout_alarm->Cancel();
+    channel_->Close();
     listener_.CallFailureCallback(endpoint_id_, channel_);
   }
 
@@ -325,6 +326,7 @@ class ClientRunnable final {
 
   void HandleHandshakeOrIoException(CancelableAlarm* timeout_alarm) {
     timeout_alarm->Cancel();
+    channel_->Close();
     listener_.CallFailureCallback(endpoint_id_, channel_);
   }
 

connections/implementation/encryption_runner_test.cc

@@ -41,6 +41,8 @@ namespace {
 
 using ::location::nearby::proto::connections::Medium;
 constexpr size_t kChunkSize = 64 * 1024;
+constexpr securegcm::UKey2Handshake::HandshakeCipher kCipher =
+    securegcm::UKey2Handshake::HandshakeCipher::P256_SHA512;
 
 class FakeEndpointChannel : public EndpointChannel {
  public:
@@ -193,6 +195,228 @@ TEST(EncryptionRunnerTest, ReadWrite) {
   EXPECT_EQ(response.client_status, Response::Status::kDone);
 }
 
+TEST(EncryptionRunnerTest, ClientWriteFails) {
+  auto from_a_to_b = CreatePipe();
+  auto from_b_to_a = CreatePipe();
+  User user_a(/*reader=*/from_b_to_a.first.get(),
+              /*writer=*/from_a_to_b.second.get());
+  User user_b(/*reader=*/from_a_to_b.first.get(),
+              /*writer=*/from_b_to_a.second.get());
+  Response response;
+  response.latch = CountDownLatch(1);
+
+  // Close server's input stream, so client can't write to it.
+  from_b_to_a.first->Close();
+
+  user_b.crypto.StartClient(
+      &user_b.client, "endpoint_id", &user_b.channel,
+      {
+          .on_success_cb =
+              [&response](const std::string& endpoint_id,
+                          std::unique_ptr<securegcm::UKey2Handshake> ukey2,
+                          const std::string& auth_token,
+                          const ByteArray& raw_auth_token) {
+                response.client_status = Response::Status::kDone;
+                response.latch.CountDown();
+              },
+          .on_failure_cb =
+              [&response](const std::string& endpoint_id,
+                          EndpointChannel* channel) {
+                response.client_status = Response::Status::kFailed;
+                response.latch.CountDown();
+              },
+      });
+  EXPECT_TRUE(response.latch.Await(absl::Milliseconds(5000)).result());
+  EXPECT_EQ(response.client_status, Response::Status::kFailed);
+}
+
+TEST(EncryptionRunnerTest, ServerWriteFails) {
+  auto from_a_to_b = CreatePipe();
+  auto from_b_to_a = CreatePipe();
+  User user_a(/*reader=*/from_b_to_a.first.get(),
+              /*writer=*/from_a_to_b.second.get());
+  User user_b(/*reader=*/from_a_to_b.first.get(),
+              /*writer=*/from_b_to_a.second.get());
+  Response response;
+  response.latch = CountDownLatch(1);
+
+  // Close client's input stream, so server can't write to it.
+  from_a_to_b.first->Close();
+
+  user_a.crypto.StartServer(
+      &user_a.client, "endpoint_id", &user_a.channel,
+      {
+          .on_success_cb =
+              [&response](const std::string& endpoint_id,
+                          std::unique_ptr<securegcm::UKey2Handshake> ukey2,
+                          const std::string& auth_token,
+                          const ByteArray& raw_auth_token) {
+                response.server_status = Response::Status::kDone;
+                response.latch.CountDown();
+              },
+          .on_failure_cb =
+              [&response](const std::string& endpoint_id,
+                          EndpointChannel* channel) {
+                response.server_status = Response::Status::kFailed;
+                response.latch.CountDown();
+              },
+      });
+  user_b.crypto.StartClient(
+      &user_b.client, "endpoint_id", &user_b.channel,
+      {
+          .on_success_cb =
+              [](const std::string& endpoint_id,
+                 std::unique_ptr<securegcm::UKey2Handshake> ukey2,
+                 const std::string& auth_token,
+                 const ByteArray& raw_auth_token) {},
+          .on_failure_cb =
+              [](const std::string& endpoint_id, EndpointChannel* channel) {},
+      });
+  EXPECT_TRUE(response.latch.Await(absl::Milliseconds(5000)).result());
+  EXPECT_EQ(response.server_status, Response::Status::kFailed);
+}
+
+TEST(EncryptionRunnerTest, ClientSendsGarbageMessage1) {
+  auto from_server_to_client = CreatePipe();
+  auto from_client_to_server = CreatePipe();
+  User user_a(/*reader=*/from_client_to_server.first.get(),
+              /*writer=*/from_server_to_client.second.get());
+  Response response;
+  response.latch = CountDownLatch(1);
+
+  user_a.crypto.StartServer(
+      &user_a.client, "endpoint_id", &user_a.channel,
+      {
+          .on_success_cb =
+              [&response](const std::string& endpoint_id,
+                          std::unique_ptr<securegcm::UKey2Handshake> ukey2,
+                          const std::string& auth_token,
+                          const ByteArray& raw_auth_token) {
+                response.server_status = Response::Status::kDone;
+                response.latch.CountDown();
+              },
+          .on_failure_cb =
+              [&response](const std::string& endpoint_id,
+                          EndpointChannel* channel) {
+                response.server_status = Response::Status::kFailed;
+                response.latch.CountDown();
+              },
+      });
+
+  // Client writes garbage instead of message 1
+  from_client_to_server.second->Write("Garbage");
+
+  EXPECT_TRUE(response.latch.Await(absl::Milliseconds(5000)).result());
+  EXPECT_EQ(response.server_status, Response::Status::kFailed);
+
+  // Check if server sent alert message.
+  // The alert message should be readable from from_server_to_client.first.
+  auto alert = from_server_to_client.first->Read(kChunkSize);
+  EXPECT_TRUE(alert.ok());
+  EXPECT_FALSE(alert.result().Empty());
+}
+
+TEST(EncryptionRunnerTest, ServerSendsGarbageMessage2) {
+  auto from_server_to_client = CreatePipe();
+  auto from_client_to_server = CreatePipe();
+  User user_b(/*reader=*/from_server_to_client.first.get(),
+              /*writer=*/from_client_to_server.second.get());
+  Response response;
+  response.latch = CountDownLatch(1);
+
+  user_b.crypto.StartClient(
+      &user_b.client, "endpoint_id", &user_b.channel,
+      {
+          .on_success_cb =
+              [&response](const std::string& endpoint_id,
+                          std::unique_ptr<securegcm::UKey2Handshake> ukey2,
+                          const std::string& auth_token,
+                          const ByteArray& raw_auth_token) {
+                response.client_status = Response::Status::kDone;
+                response.latch.CountDown();
+              },
+          .on_failure_cb =
+              [&response](const std::string& endpoint_id,
+                          EndpointChannel* channel) {
+                response.client_status = Response::Status::kFailed;
+                response.latch.CountDown();
+              },
+      });
+
+  // Client sends message 1.
+  auto client_init = from_client_to_server.first->Read(kChunkSize);
+  EXPECT_TRUE(client_init.ok());
+
+  // Server writes garbage instead of message 2.
+  from_server_to_client.second->Write("Garbage");
+
+  EXPECT_TRUE(response.latch.Await(absl::Milliseconds(5000)).result());
+  EXPECT_EQ(response.client_status, Response::Status::kFailed);
+
+  // Check if client sent alert message.
+  auto alert = from_client_to_server.first->Read(kChunkSize);
+  EXPECT_TRUE(alert.ok());
+  EXPECT_FALSE(alert.result().Empty());
+}
+
+TEST(EncryptionRunnerTest, ClientSendsGarbageMessage3) {
+  auto from_server_to_client = CreatePipe();
+  auto from_client_to_server = CreatePipe();
+  User user_a(/*reader=*/from_client_to_server.first.get(),
+              /*writer=*/from_server_to_client.second.get());
+  User user_b(/*reader=*/from_server_to_client.first.get(),
+              /*writer=*/from_client_to_server.second.get());
+  Response response;
+  response.latch = CountDownLatch(1);
+
+  user_a.crypto.StartServer(
+      &user_a.client, "endpoint_id", &user_a.channel,
+      {
+          .on_success_cb =
+              [&response](const std::string& endpoint_id,
+                          std::unique_ptr<securegcm::UKey2Handshake> ukey2,
+                          const std::string& auth_token,
+                          const ByteArray& raw_auth_token) {
+                response.server_status = Response::Status::kDone;
+                response.latch.CountDown();
+              },
+          .on_failure_cb =
+              [&response](const std::string& endpoint_id,
+                          EndpointChannel* channel) {
+                response.server_status = Response::Status::kFailed;
+                response.latch.CountDown();
+              },
+      });
+
+  // Client starts, sends message 1
+  std::unique_ptr<securegcm::UKey2Handshake> client_crypto =
+      securegcm::UKey2Handshake::ForInitiator(kCipher);
+  std::unique_ptr<std::string> client_init_str =
+      client_crypto->GetNextHandshakeMessage();
+  from_client_to_server.second->Write(
+      ByteArray(*client_init_str).AsStringView());
+
+  // Server reads message 1, sends message 2.
+  // Read message 2 from server
+  auto server_init = from_server_to_client.first->Read(kChunkSize);
+  EXPECT_TRUE(server_init.ok());
+
+  // Client crypto parses message 2.
+  client_crypto->ParseHandshakeMessage(std::string(server_init.result()));
+
+  // Client sends garbage instead of message 3
+  from_client_to_server.second->Write("Garbage");
+
+  EXPECT_TRUE(response.latch.Await(absl::Milliseconds(5000)).result());
+  EXPECT_EQ(response.server_status, Response::Status::kFailed);
+
+  // Check if server sent alert message.
+  // Message 3 doesn't send alert in current UKEY2 implementation.
+  auto alert = from_server_to_client.first->Read(kChunkSize);
+  EXPECT_TRUE(alert.ok());
+  EXPECT_TRUE(alert.result().Empty());
+}
+
 }  // namespace
 }  // namespace connections
 }  // namespace nearby