Refactor test contact to use new session interface and transform base…

… class

Change-Id: Id74e21c1f54a5a01251bce718686e93c4b53efe6

containers/confidential_transform_test_concat/BUILD

@@ -18,10 +18,10 @@ load("@rules_pkg//pkg:tar.bzl", "pkg_tar")
 
 cc_library(
     name = "confidential_transform_server",
-    srcs = ["confidential_transform_server.cc"],
     hdrs = ["confidential_transform_server.h"],
     deps = [
         "//containers:blob_metadata",
+        "//containers:confidential_transform_server_base",
         "//containers:crypto",
         "//containers:session",
         "@com_github_grpc_grpc//:grpc++",
@@ -52,6 +52,7 @@ cc_test(
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
+        "@federated-compute//fcp/confidentialcompute:crypto",
         "@federated-compute//fcp/protos/confidentialcompute:confidential_transform_cc_grpc",
         "@federated-compute//fcp/protos/confidentialcompute:confidential_transform_cc_proto",
         "@googletest//:gtest_main",

containers/confidential_transform_test_concat/confidential_transform_server.h

@@ -21,7 +21,9 @@
 #include "absl/log/die_if_null.h"
 #include "absl/status/status.h"
 #include "absl/synchronization/mutex.h"
+#include "containers/confidential_transform_server_base.h"
 #include "containers/crypto.h"
+#include "containers/session.h"
 #include "fcp/protos/confidentialcompute/confidential_transform.grpc.pb.h"
 #include "fcp/protos/confidentialcompute/confidential_transform.pb.h"
 #include "grpcpp/server_context.h"
@@ -30,44 +32,70 @@
 
 namespace confidential_federated_compute::confidential_transform_test_concat {
 
-// Test ConfidentialTransform service that concatenates inputs. This test
-// service doesn't manage the number of sessions.
-class TestConcatConfidentialTransform final
-    : public fcp::confidentialcompute::ConfidentialTransform::Service {
+// TestConcat implementation of Session interface. Not threadsafe.
+class TestConcatSession final : public confidential_federated_compute::Session {
  public:
-  // The OrchestratorCrypto stub must not be NULL and must outlive this object.
-  explicit TestConcatConfidentialTransform(
-      oak::containers::v1::OrchestratorCrypto::StubInterface* crypto_stub)
-      : crypto_stub_(*ABSL_DIE_IF_NULL(crypto_stub)) {}
-
-  grpc::Status Initialize(
-      grpc::ServerContext* context,
-      const fcp::confidentialcompute::InitializeRequest* request,
-      fcp::confidentialcompute::InitializeResponse* response) override;
+  TestConcatSession() {};
+  // Currently no per-session configuration.
+  absl::Status ConfigureSession(
+      fcp::confidentialcompute::SessionRequest configure_request) override {
+    return absl::OkStatus();
+  }
+  // Concatenates the unencrypted data to the result string.
+  absl::StatusOr<fcp::confidentialcompute::SessionResponse> SessionWrite(
+      const fcp::confidentialcompute::WriteRequest& write_request,
+      std::string unencrypted_data) override {
+    absl::StrAppend(&state_, unencrypted_data);
+    return confidential_federated_compute::ToSessionWriteFinishedResponse(
+        absl::OkStatus(),
+        write_request.first_request_metadata().total_size_bytes());
+  }
+  // Run any session finalization logic and complete the session.
+  // After finalization, the session state is no longer mutable.
+  absl::StatusOr<fcp::confidentialcompute::SessionResponse> FinalizeSession(
+      const fcp::confidentialcompute::FinalizeRequest& request,
+      const fcp::confidentialcompute::BlobMetadata& input_metadata) override {
+    fcp::confidentialcompute::SessionResponse response;
+    fcp::confidentialcompute::ReadResponse* read_response =
+        response.mutable_read();
+    read_response->set_finish_read(true);
+    *(read_response->mutable_data()) = state_;
 
-  grpc::Status Session(
-      grpc::ServerContext* context,
-      grpc::ServerReaderWriter<fcp::confidentialcompute::SessionResponse,
-                               fcp::confidentialcompute::SessionRequest>*
-          stream) override;
+    fcp::confidentialcompute::BlobMetadata result_metadata;
+    result_metadata.mutable_unencrypted();
+    result_metadata.set_total_size_bytes(state_.length());
+    result_metadata.set_compression_type(
+        fcp::confidentialcompute::BlobMetadata::COMPRESSION_TYPE_NONE);
+    *(read_response->mutable_first_response_metadata()) = result_metadata;
+    return response;
+  }
 
  private:
-  absl::Status Initialize(
-      const fcp::confidentialcompute::InitializeRequest* request,
-      fcp::confidentialcompute::InitializeResponse* response);
+  std::string state_ = "";
+};
 
-  absl::Status Session(
-      grpc::ServerReaderWriter<fcp::confidentialcompute::SessionResponse,
-                               fcp::confidentialcompute::SessionRequest>*
-          stream);
+// Test ConfidentialTransform service that concatenates inputs.
+class TestConcatConfidentialTransform final
+    : public confidential_federated_compute::ConfidentialTransformBase {
+ public:
+  TestConcatConfidentialTransform(
+      oak::containers::v1::OrchestratorCrypto::StubInterface* crypto_stub,
+      int max_num_sessions = 1)
+      : ConfidentialTransformBase(crypto_stub, max_num_sessions) {};
 
-  oak::containers::v1::OrchestratorCrypto::StubInterface& crypto_stub_;
-  absl::Mutex mutex_;
-  // The mutex is used to protect the optional wrapping blob_decryptor_ to
-  // ensure the BlobDecryptor is initialized, but the BlobDecryptor is itself
-  // threadsafe.
-  std::optional<confidential_federated_compute::BlobDecryptor> blob_decryptor_
-      ABSL_GUARDED_BY(mutex_);
+ protected:
+  virtual absl::StatusOr<google::protobuf::Struct> InitializeTransform(
+      const fcp::confidentialcompute::InitializeRequest* request) override {
+    google::protobuf::Struct config_properties;
+    return config_properties;
+  }
+  virtual absl::StatusOr<
+      std::unique_ptr<confidential_federated_compute::Session>>
+  CreateSession() override {
+    return std::make_unique<
+        confidential_federated_compute::confidential_transform_test_concat::
+            TestConcatSession>();
+  };
 };
 
 }  // namespace

containers/confidential_transform_test_concat/confidential_transform_server_test.cc

@@ -21,6 +21,7 @@
 #include "containers/blob_metadata.h"
 #include "containers/crypto.h"
 #include "containers/crypto_test_utils.h"
+#include "fcp/confidentialcompute/crypto.h"
 #include "fcp/protos/confidentialcompute/confidential_transform.grpc.pb.h"
 #include "fcp/protos/confidentialcompute/confidential_transform.pb.h"
 #include "gmock/gmock.h"
@@ -39,6 +40,7 @@ namespace confidential_federated_compute::confidential_transform_test_concat {
 
 namespace {
 
+using ::fcp::confidential_compute::MessageDecryptor;
 using ::fcp::confidential_compute::NonceAndCounter;
 using ::fcp::confidential_compute::NonceGenerator;
 using ::fcp::confidentialcompute::BlobHeader;
@@ -211,7 +213,10 @@ TEST_F(TestConcatServerSessionTest, SessionWritesAndFinalizesUnencryptedBlobs) {
 }
 
 TEST_F(TestConcatServerSessionTest, SessionDecryptsMultipleBlobsAndFinalizes) {
-  std::string reencryption_public_key = "reencryption key";
+  MessageDecryptor decryptor;
+  absl::StatusOr<std::string> reencryption_public_key =
+      decryptor.GetPublicKey([](absl::string_view) { return ""; }, 0);
+  ASSERT_TRUE(reencryption_public_key.ok());
   std::string ciphertext_associated_data =
       BlobHeader::default_instance().SerializeAsString();
 
@@ -222,7 +227,7 @@ TEST_F(TestConcatServerSessionTest, SessionDecryptsMultipleBlobsAndFinalizes) {
   absl::StatusOr<Record> rewrapped_record_0 =
       crypto_test_utils::CreateRewrappedRecord(
           message_0, ciphertext_associated_data, public_key_,
-          nonce_0->blob_nonce, reencryption_public_key);
+          nonce_0->blob_nonce, *reencryption_public_key);
   ASSERT_TRUE(rewrapped_record_0.ok()) << rewrapped_record_0.status();
 
   SessionRequest request_0;
@@ -249,7 +254,7 @@ TEST_F(TestConcatServerSessionTest, SessionDecryptsMultipleBlobsAndFinalizes) {
   absl::StatusOr<Record> rewrapped_record_1 =
       crypto_test_utils::CreateRewrappedRecord(
           message_1, ciphertext_associated_data, public_key_,
-          nonce_1->blob_nonce, reencryption_public_key);
+          nonce_1->blob_nonce, *reencryption_public_key);
   ASSERT_TRUE(rewrapped_record_1.ok()) << rewrapped_record_1.status();
 
   SessionRequest request_1;
@@ -285,7 +290,10 @@ TEST_F(TestConcatServerSessionTest, SessionDecryptsMultipleBlobsAndFinalizes) {
 }
 
 TEST_F(TestConcatServerSessionTest, SessionIgnoresUndecryptableInputs) {
-  std::string reencryption_public_key = "reencryption key";
+  MessageDecryptor decryptor;
+  absl::StatusOr<std::string> reencryption_public_key =
+      decryptor.GetPublicKey([](absl::string_view) { return ""; }, 0);
+  ASSERT_TRUE(reencryption_public_key.ok());
   std::string ciphertext_associated_data =
       BlobHeader::default_instance().SerializeAsString();
 
@@ -296,7 +304,7 @@ TEST_F(TestConcatServerSessionTest, SessionIgnoresUndecryptableInputs) {
   absl::StatusOr<Record> rewrapped_record_0 =
       crypto_test_utils::CreateRewrappedRecord(
           message_0, ciphertext_associated_data, public_key_,
-          nonce_0->blob_nonce, reencryption_public_key);
+          nonce_0->blob_nonce, *reencryption_public_key);
   ASSERT_TRUE(rewrapped_record_0.ok()) << rewrapped_record_0.status();
 
   SessionRequest request_0;
@@ -321,7 +329,7 @@ TEST_F(TestConcatServerSessionTest, SessionIgnoresUndecryptableInputs) {
   absl::StatusOr<Record> rewrapped_record_1 =
       crypto_test_utils::CreateRewrappedRecord(
           "unused message", ciphertext_associated_data, public_key_,
-          nonce_1->blob_nonce, reencryption_public_key);
+          nonce_1->blob_nonce, *reencryption_public_key);
   ASSERT_TRUE(rewrapped_record_1.ok()) << rewrapped_record_1.status();
 
   SessionRequest invalid_request;