[storage] Support reading unencrypted observations

Bug: 3842
Bug: 3752
Change-Id: Ib6fe69f75b84214bd4d36399e94c8c1bd6b37d50
diff --git a/src/observation_store/BUILD.gn b/src/observation_store/BUILD.gn
index 0e6b193..dd13b7d 100644
--- a/src/observation_store/BUILD.gn
+++ b/src/observation_store/BUILD.gn
@@ -48,6 +48,7 @@
     "$cobalt_root/src/lib/util:proto_serialization",
     "$cobalt_root/src/logger:internal_metrics",
     "$cobalt_root/src/logger:logger_interface",
+    "//third_party/googletest:gtest_prod",
   ]
   public_deps = [
     ":observation_store_internal_proto",
@@ -69,6 +70,7 @@
   deps = [
     ":observation_store",
     "$cobalt_root/src/lib/util:posix_file_system",
+    "$cobalt_root/src/lib/util:proto_serialization",
     "$cobalt_root/src/system_data:client_secret",
     "//third_party/gflags",
     "//third_party/googletest:gmock",
diff --git a/src/observation_store/envelope_maker.cc b/src/observation_store/envelope_maker.cc
index bc8c38d..e834794 100644
--- a/src/observation_store/envelope_maker.cc
+++ b/src/observation_store/envelope_maker.cc
@@ -54,11 +54,39 @@
   // "+1" below is for the |scheme| field of EncryptedMessage.
   num_bytes_ += message->ciphertext().size() + message->public_key_fingerprint().size() + 1;
   // Put the encrypted observation into the appropriate ObservationBatch.
-  GetBatch(std::move(metadata))->add_encrypted_observation()->Swap(message.get());
+  GetBatch(std::move(metadata))->add_observation()->mutable_encrypted()->Swap(message.get());
   return ObservationStore::kOk;
 }
 
-ObservationBatch* EnvelopeMaker::GetBatch(std::unique_ptr<ObservationMetadata> metadata) {
+const Envelope& EnvelopeMaker::GetEnvelope(util::EncryptedMessageMaker* encrypter) {
+  envelope_.Clear();
+
+  auto num_batches = batch_map_.size();
+  for (auto i = 0; i < num_batches; i++) {
+    envelope_.add_batch();
+  }
+
+  auto i = num_batches;
+  for (const auto& batch : batch_map_) {
+    i -= 1;
+    ObservationBatch* observation_batch = envelope_.mutable_batch(i);
+    *observation_batch->mutable_meta_data() = batch.second->meta_data();
+    for (const auto& obs : batch.second->observation()) {
+      auto obs_out = observation_batch->add_encrypted_observation();
+      if (obs.has_encrypted()) {
+        *obs_out = obs.encrypted();
+      } else if (obs.has_unencrypted()) {
+        if (!encrypter->Encrypt(obs.unencrypted(), obs_out)) {
+          LOG_FIRST_N(ERROR, 10) << "ERROR: Unable to encrypt observation on read.";
+        }
+      }
+    }
+  }
+
+  return envelope_;
+}
+
+StoredObservationBatch* EnvelopeMaker::GetBatch(std::unique_ptr<ObservationMetadata> metadata) {
   // Serialize metadata.
   std::string serialized_metadata;
   (*metadata).SerializeToString(&serialized_metadata);
@@ -66,14 +94,12 @@
   // See if metadata is already in batch_map_. If so return it.
   auto iter = batch_map_.find(serialized_metadata);
   if (iter != batch_map_.end()) {
-    return iter->second;
+    return iter->second.get();
   }
 
-  // Create a new ObservationBatch and add it to both envelope_ and batch_map_.
-  ObservationBatch* observation_batch = envelope_.add_batch();
-  observation_batch->set_allocated_meta_data(metadata.release());
-  batch_map_[serialized_metadata] = observation_batch;
-  return observation_batch;
+  batch_map_[serialized_metadata] = std::make_unique<StoredObservationBatch>();
+  batch_map_[serialized_metadata]->set_allocated_meta_data(metadata.release());
+  return batch_map_[serialized_metadata].get();
 }
 
 void EnvelopeMaker::MergeWith(std::unique_ptr<ObservationStore::EnvelopeHolder> other_ref) {
@@ -88,17 +114,17 @@
       // from the other's batch into our batch. Note that this process
       // reverses the order of the messages in other but the order of
       // the messages in a batch has no meaning so this doesn't matter.
-      auto* other_messages = other_pair.second->mutable_encrypted_observation();
-      auto* this_messages = iter->second->mutable_encrypted_observation();
+      auto* other_messages = other_pair.second->mutable_observation();
+      auto* this_messages = iter->second->mutable_observation();
       while (!other_messages->empty()) {
         this_messages->AddAllocated(other_messages->ReleaseLast());
       }
     } else {
       // We do not have a pair with the same key. Make one and swap the
       // contents of the others batch into it.
-      ObservationBatch* observation_batch = envelope_.add_batch();
-      observation_batch->Swap(other_pair.second);
-      batch_map_[other_pair.first] = observation_batch;
+      auto observation_batch = std::make_unique<StoredObservationBatch>();
+      observation_batch->Swap(other_pair.second.get());
+      batch_map_[other_pair.first] = std::move(observation_batch);
     }
   }
   num_bytes_ += other->num_bytes_;
diff --git a/src/observation_store/envelope_maker.h b/src/observation_store/envelope_maker.h
index cb317bd..1dcc4fb 100644
--- a/src/observation_store/envelope_maker.h
+++ b/src/observation_store/envelope_maker.h
@@ -14,9 +14,9 @@
 #include "src/observation_store/observation_store.h"
 #include "src/pb/encrypted_message.pb.h"
 #include "src/pb/observation.pb.h"
+#include "third_party/googletest/googletest/include/gtest/gtest_prod.h"
 
-namespace cobalt {
-namespace observation_store {
+namespace cobalt::observation_store {
 
 // EnvelopeMaker is an implementation of ObservationStore::EnvelopeHolder that
 // holds its Envelope in memory. This implementation is used by
@@ -51,11 +51,9 @@
   ObservationStore::StoreStatus AddEncryptedObservation(
       std::unique_ptr<EncryptedMessage> message, std::unique_ptr<ObservationMetadata> metadata);
 
-  const Envelope& GetEnvelope(util::EncryptedMessageMaker* /*encrypter*/) override {
-    return envelope_;
-  }
+  const Envelope& GetEnvelope(util::EncryptedMessageMaker* encrypter) override;
 
-  bool Empty() const { return envelope_.batch_size() == 0; }
+  bool Empty() const { return batch_map_.empty(); }
 
   void Clear() {
     envelope_.Clear();
@@ -74,17 +72,18 @@
 
  private:
   friend class EnvelopeMakerTest;
+  FRIEND_TEST(EnvelopeMakerTest, CanReadUnencrypted);
 
   // Returns the ObservationBatch containing the given |metadata|. If
   // this is the first time we have seen the given |metadata| then a
   // new ObservationBatch is created.
-  ObservationBatch* GetBatch(std::unique_ptr<ObservationMetadata> metadata);
+  StoredObservationBatch* GetBatch(std::unique_ptr<ObservationMetadata> metadata);
 
   Envelope envelope_;
 
   // The keys of the map are serialized ObservationMetadata. The values
   // are the ObservationBatch containing that Metadata
-  std::unordered_map<std::string, ObservationBatch*> batch_map_;
+  std::unordered_map<std::string, std::unique_ptr<StoredObservationBatch>> batch_map_;
 
   // Keeps a running total of the sum of the sizes of the encrypted Observations
   // contained in |envelope_|;
@@ -94,7 +93,6 @@
   const size_t max_num_bytes_;
 };
 
-}  // namespace observation_store
-}  // namespace cobalt
+}  // namespace cobalt::observation_store
 
 #endif  // COBALT_SRC_OBSERVATION_STORE_ENVELOPE_MAKER_H_
diff --git a/src/observation_store/envelope_maker_test.cc b/src/observation_store/envelope_maker_test.cc
index 9cc96e0..f80320a 100644
--- a/src/observation_store/envelope_maker_test.cc
+++ b/src/observation_store/envelope_maker_test.cc
@@ -331,4 +331,33 @@
                  expected_this_batch_size, ObservationStore::kOk);
 }
 
+TEST_F(EnvelopeMakerTest, CanReadUnencrypted) {
+  auto observation = std::make_unique<Observation2>();
+  observation->set_random_id("test123");
+  observation->mutable_basic_rappor()->set_data("test");
+
+  auto encrypted_obs = std::make_unique<EncryptedMessage>();
+  encrypt_->Encrypt(*observation, encrypted_obs.get());
+
+  // Verify that our encrypted observation is non-trivial.
+  ASSERT_GT(encrypted_obs->ciphertext().size(), 0);
+
+  auto metadata = std::make_unique<ObservationMetadata>();
+  metadata->set_customer_id(kCustomerId);
+  metadata->set_project_id(kProjectId);
+  metadata->set_metric_id(10);
+
+  envelope_maker_->GetBatch(std::move(metadata))
+      ->add_observation()
+      ->mutable_unencrypted()
+      ->Swap(observation.get());
+
+  auto read_env = envelope_maker_->GetEnvelope(encrypt_.get());
+  ASSERT_EQ(read_env.batch_size(), 1);
+  ASSERT_EQ(read_env.batch(0).encrypted_observation_size(), 1);
+
+  // Verify that we got the observation we expected.
+  ASSERT_EQ(read_env.batch(0).encrypted_observation(0).ciphertext(), encrypted_obs->ciphertext());
+}
+
 }  // namespace cobalt::observation_store
diff --git a/src/observation_store/file_observation_store.cc b/src/observation_store/file_observation_store.cc
index af656e4..5df0ffa 100644
--- a/src/observation_store/file_observation_store.cc
+++ b/src/observation_store/file_observation_store.cc
@@ -349,7 +349,7 @@
 }
 
 const Envelope &FileObservationStore::FileEnvelopeHolder::GetEnvelope(
-    util::EncryptedMessageMaker * /*encrypter*/) {
+    util::EncryptedMessageMaker *encrypter) {
   if (envelope_read_) {
     return envelope_;
   }
@@ -386,6 +386,11 @@
         }
       } else if (stored.has_encrypted_observation()) {
         current_batch->add_encrypted_observation()->Swap(stored.mutable_encrypted_observation());
+      } else if (stored.has_unencrypted_observation()) {
+        if (!encrypter->Encrypt(stored.unencrypted_observation(),
+                                current_batch->add_encrypted_observation())) {
+          LOG_FIRST_N(ERROR, 10) << "ERROR: Unable to encrypt observation on read.";
+        }
       } else {
         clean_eof = false;
         break;
diff --git a/src/observation_store/file_observation_store_test.cc b/src/observation_store/file_observation_store_test.cc
index 613600c..7bfb7b9 100644
--- a/src/observation_store/file_observation_store_test.cc
+++ b/src/observation_store/file_observation_store_test.cc
@@ -8,9 +8,11 @@
 #include <random>
 #include <utility>
 
+#include "google/protobuf/util/delimited_message_util.h"
 #include "src/lib/util/posix_file_system.h"
 #include "src/logging.h"
 #include "src/observation_store/observation_store.h"
+#include "src/observation_store/observation_store_internal.pb.h"
 #include "src/system_data/client_secret.h"
 #include "third_party/googletest/googlemock/include/gmock/gmock.h"
 #include "third_party/googletest/googletest/include/gtest/gtest.h"
@@ -350,6 +352,45 @@
   }
 }
 
+TEST_F(FileObservationStoreTest, CanReadUnencrypted) {
+  auto observation = std::make_unique<Observation2>();
+  observation->set_random_id("test123");
+  observation->mutable_basic_rappor()->set_data("test");
+
+  auto encrypted_obs = std::make_unique<EncryptedMessage>();
+  encrypt_->Encrypt(*observation, encrypted_obs.get());
+
+  // Verify that our encrypted observation is non-trivial.
+  ASSERT_GT(encrypted_obs->ciphertext().size(), 0);
+
+  auto metadata = std::make_unique<ObservationMetadata>();
+  metadata->set_customer_id(kCustomerId);
+  metadata->set_project_id(kProjectId);
+  metadata->set_metric_id(10);
+
+  {
+    std::ofstream manual_data(test_dir_name_ + "/1234567890123-1234567890.data");
+    google::protobuf::io::OstreamOutputStream outstream(&manual_data);
+    FileObservationStoreRecord stored_metadata;
+    stored_metadata.mutable_meta_data()->Swap(metadata.get());
+    google::protobuf::util::SerializeDelimitedToZeroCopyStream(stored_metadata, &outstream);
+
+    FileObservationStoreRecord stored_observation;
+    stored_observation.mutable_unencrypted_observation()->Swap(observation.get());
+    google::protobuf::util::SerializeDelimitedToZeroCopyStream(stored_observation, &outstream);
+  }
+
+  ASSERT_EQ(store_->ListFinalizedFiles().size(), 1u);
+  auto envelope = store_->TakeNextEnvelopeHolder();
+  ASSERT_NE(envelope, nullptr);
+  auto read_env = envelope->GetEnvelope(encrypt_.get());
+  ASSERT_EQ(read_env.batch_size(), 1);
+  ASSERT_EQ(read_env.batch(0).encrypted_observation_size(), 1);
+
+  // Verify that we got the observation we expected.
+  ASSERT_EQ(read_env.batch(0).encrypted_observation(0).ciphertext(), encrypted_obs->ciphertext());
+}
+
 TEST(FilenameGenerator, PadsTimestamp) {
   EXPECT_THAT(FileObservationStore::FilenameGenerator([] { return 1234; }).GenerateFilename(),
               MatchesRegex(R"(0000000001234-[0-9]{10}.data)"));