Adds size tracking functionality to EnvelopeMaker.
We add to EnvelopeMaker:
- A size() accessor that returns the number of bytes of Observations
currently in the Envelope.
- A status enum for the return value of AddObservation(). Previously
AddObservation() was a void function.
- Two optional parameters to the constructor: max_bytes_each_observation
and max_num_bytes. If the first is set then AddObservation() will return
kObservationTooBig if the Observation being added has size larger
than |max_bytes_each_observation|. If the second is set then
AddObservation() will return kEnvelopeFull if the Observation
being added is not too big but the additional observation would
push the total size of the Envelope above
|max_num_bytes|.
I intend to use this functionality in the implementation of the
ShippingManager. The ShippingManager needs to make sure that
(a) It does not allow its in-memory cache of unsent Observations
to get too big for device memory.
(b) It does not allow any single Envelope get too big to be sent
in a single RPC.
(c) It does not allow any single Observation to be added to an
Envelope if that Observation by itself is too big to be sent
in a single RPC.
Change-Id: If0b78545d0659e9a59ca72638e7531bedf7e7685
diff --git a/encoder/encoder_test.cc b/encoder/encoder_test.cc
index 0770868..6ba0d93 100644
--- a/encoder/encoder_test.cc
+++ b/encoder/encoder_test.cc
@@ -695,9 +695,9 @@
} // namespace cobalt
int main(int argc, char** argv) {
+ ::testing::InitGoogleTest(&argc, argv);
google::ParseCommandLineFlags(&argc, &argv, true);
google::InitGoogleLogging(argv[0]);
google::InstallFailureSignalHandler();
- ::testing::InitGoogleTest(&argc, argv);
return RUN_ALL_TESTS();
}
diff --git a/encoder/envelope_maker.cc b/encoder/envelope_maker.cc
index f1dcafd..e8e2392 100644
--- a/encoder/envelope_maker.cc
+++ b/encoder/envelope_maker.cc
@@ -24,23 +24,46 @@
EnvelopeMaker::EnvelopeMaker(const std::string& analyzer_public_key_pem,
EncryptedMessage::EncryptionScheme analyzer_scheme,
const std::string& shuffler_public_key_pem,
- EncryptedMessage::EncryptionScheme shuffler_scheme)
+ EncryptedMessage::EncryptionScheme shuffler_scheme,
+ size_t max_bytes_each_observation,
+ size_t max_num_bytes)
: encrypt_to_analyzer_(analyzer_public_key_pem, analyzer_scheme),
- encrypt_to_shuffler_(shuffler_public_key_pem, shuffler_scheme) {}
+ encrypt_to_shuffler_(shuffler_public_key_pem, shuffler_scheme),
+ max_bytes_each_observation_(max_bytes_each_observation),
+ max_num_bytes_(max_num_bytes) {}
-void EnvelopeMaker::AddObservation(
+EnvelopeMaker::AddStatus EnvelopeMaker::AddObservation(
const Observation& observation,
std::unique_ptr<ObservationMetadata> metadata) {
EncryptedMessage encrypted_message;
if (encrypt_to_analyzer_.Encrypt(observation, &encrypted_message)) {
+ // "+1" below is for the |scheme| field of EncryptedMessage.
+ size_t obs_size = encrypted_message.ciphertext().size() +
+ encrypted_message.public_key_fingerprint().size() + 1;
+ if (obs_size > max_bytes_each_observation_) {
+ VLOG(1) << "WARNING: An Observation was rejected by "
+ "EnvelopeMaker::AddObservation() because it was too big: "
+ << obs_size;
+ return kObservationTooBig;
+ }
+
+ size_t new_num_bytes = num_bytes_ + obs_size;
+ if (new_num_bytes > max_num_bytes_) {
+ VLOG(4) << "Envelope full.";
+ return kEnvelopeFull;
+ }
+
+ num_bytes_ = new_num_bytes;
// Put the encrypted observation into the appropriate ObservationBatch.
GetBatch(std::move(metadata))
->add_encrypted_observation()
->Swap(&encrypted_message);
+ return kOk;
} else {
VLOG(1)
<< "ERROR: Encryption of Observations failed! Observation not added "
"to batch.";
+ return kEncryptionFailed;
}
}
@@ -96,6 +119,7 @@
batch_map_[other_pair.first] = observation_batch;
}
}
+ num_bytes_ += other->num_bytes_;
other->Clear();
}
diff --git a/encoder/envelope_maker.h b/encoder/envelope_maker.h
index aed454c..321f80c 100644
--- a/encoder/envelope_maker.h
+++ b/encoder/envelope_maker.h
@@ -65,13 +65,44 @@
//
// |shuffler_scheme| The public key encryption scheme to use when encrypting
// Envelopes sent to the Shuffler.
+ //
+ // |max_bytes_each_observation|. If specified then AddObservation() will
+ // return kObservationTooBig if the provided observation's serialized,
+ // encrypted size is greater than this value.
+ //
+ // |max_num_bytes|. If specified then AddObservation() will return
+ // kEnvelopeFull if the provided observation's serialized, encrypted size
+ // is not too large by itself, but adding the additional observation would
+ // cause the sum of the sizes of all added Observations to be greater than
+ // this value.
EnvelopeMaker(const std::string& analyzer_public_key_pem,
EncryptedMessage::EncryptionScheme analyzer_scheme,
const std::string& shuffler_public_key_pem,
- EncryptedMessage::EncryptionScheme shuffler_scheme);
+ EncryptedMessage::EncryptionScheme shuffler_scheme,
+ size_t max_bytes_each_observation = SIZE_MAX,
+ size_t max_num_bytes = SIZE_MAX);
- void AddObservation(const Observation& observation,
- std::unique_ptr<ObservationMetadata> metadata);
+ // The status of an AddObservation() call.
+ enum AddStatus {
+ // AddObservation() succeeded.
+ kOk = 0,
+
+ // The Observation was not added to the Envelope because it is
+ // too big.
+ kObservationTooBig,
+
+ // The Observation was not added to the Envelope because the
+ // envelope is full. The Observation itself is not too big
+ // to be added otherwise.
+ kEnvelopeFull,
+
+ // The Observation was not added to the Envelope because the encryption
+ // failed.
+ kEncryptionFailed
+ };
+
+ AddStatus AddObservation(const Observation& observation,
+ std::unique_ptr<ObservationMetadata> metadata);
// Populates |*encrypted_message| with the encryption of the current
// value of the Envelope. Returns true for success or false for failure.
@@ -80,17 +111,25 @@
// Gives direct read-only access to the internal instance of Envelope.
const Envelope& envelope() const { return envelope_; }
- bool Empty() const {return envelope_.batch_size() == 0;}
+ bool Empty() const { return envelope_.batch_size() == 0; }
void Clear() {
envelope_ = Envelope();
batch_map_.clear();
+ num_bytes_ = 0;
}
// Moves the contents out of |*other| and merges it into |*this|.
// Leaves |*other| empty.
void MergeOutOf(EnvelopeMaker* other);
+ // Returns an approximation to the size of the Envelope in bytes. This value
+ // is the sum of the sizes of the serialized, encrypted Observations contained
+ // in the Envelope. But the size of the EncryptedMessage produced by the
+ // method MakeEncryptedEnvelope() may be somewhat larger than this because
+ // the Envelope itself may be encrypted to the Shuffler.
+ size_t size() { return num_bytes_; }
+
private:
friend class EnvelopeMakerTest;
@@ -106,6 +145,13 @@
// The keys of the map are serialized ObservationMetadata. The values
// are the ObservationBatch containing that Metadata
std::unordered_map<std::string, ObservationBatch*> batch_map_;
+
+ // Keeps a running total of the sum of the sizes of the encrypted Observations
+ // contained in |envelope_|;
+ size_t num_bytes_ = 0;
+
+ const size_t max_bytes_each_observation_;
+ const size_t max_num_bytes_;
};
} // namespace encoder
diff --git a/encoder/envelope_maker_test.cc b/encoder/envelope_maker_test.cc
index da7cd8a..47b71ed 100644
--- a/encoder/envelope_maker_test.cc
+++ b/encoder/envelope_maker_test.cc
@@ -157,22 +157,29 @@
encoder_.set_current_time(kSomeTimestamp);
}
- // Returns the current value of envelope_maker_ and resets envelope_maker_.
- std::unique_ptr<EnvelopeMaker> ResetEnvelopeMaker() {
+ // Returns the current value of envelope_maker_ and resets envelope_maker_
+ // to a new EnvelopeMaker constructed using the given optional arguments.
+ std::unique_ptr<EnvelopeMaker> ResetEnvelopeMaker(
+ size_t max_bytes_each_observation = SIZE_MAX,
+ size_t max_num_bytes = SIZE_MAX) {
std::unique_ptr<EnvelopeMaker> return_val = std::move(envelope_maker_);
- envelope_maker_.reset(
- new EnvelopeMaker(kAnalyzerPublicKey, EncryptedMessage::NONE,
- kShufflerPublicKey, EncryptedMessage::NONE));
+ envelope_maker_.reset(new EnvelopeMaker(
+ kAnalyzerPublicKey, EncryptedMessage::NONE, kShufflerPublicKey,
+ EncryptedMessage::NONE, max_bytes_each_observation, max_num_bytes));
return return_val;
}
// The metric is expected to have a single string part named "Part1" and
// to use the UTC timezone.
+ // expected_size_change: What is the expected change in the size of the
+ // envelope in bytes due to the AddObservation()?
void AddStringObservation(std::string value, uint32_t metric_id,
uint32_t encoding_config_id,
int expected_num_batches,
size_t expected_this_batch_index,
- int expected_this_batch_size) {
+ int expected_this_batch_size,
+ size_t expected_size_change,
+ EnvelopeMaker::AddStatus expected_status) {
// Encode an Observation
Encoder::Result result =
encoder_.EncodeString(metric_id, encoding_config_id, value);
@@ -181,12 +188,20 @@
ASSERT_NE(nullptr, result.metadata);
// Add the Observation to the EnvelopeMaker
- envelope_maker_->AddObservation(*result.observation,
- std::move(result.metadata));
+ size_t size_before_add = envelope_maker_->size();
+ ASSERT_EQ(expected_status,
+ envelope_maker_->AddObservation(*result.observation,
+ std::move(result.metadata)));
+ size_t size_after_add = envelope_maker_->size();
+ EXPECT_EQ(expected_size_change, size_after_add - size_before_add) << value;
// Check the number of batches currently in the envelope.
ASSERT_EQ(expected_num_batches, envelope_maker_->envelope().batch_size());
+ if (expected_status != EnvelopeMaker::kOk) {
+ return;
+ }
+
// Check the ObservationMetadata of the expected batch.
const auto& batch =
envelope_maker_->envelope().batch(expected_this_batch_index);
@@ -235,10 +250,19 @@
for (int i = first; i < limit; i++) {
std::ostringstream stream;
stream << "value " << i;
+ // NOTE(rudominer) The values of expected_observation_num_bytes for
+ // the NoOp encodings in this test are obtained from experimentation
+ // rather than calculation. Notice that there seems to be an overhead
+ // of 20 bytes in addition to the bytes required to store the string
+ // value. Each Observation also needs to store the name of the
+ // MetricPart ("Part 1" in our case) as well as the fact that the NoOp
+ // encoding is being used and the fact that the datatype is string.
+ size_t expected_observation_num_bytes = (i >= 10 ? 28 : 27);
expected_this_batch_size++;
AddStringObservation(stream.str(), metric_id, kEncodingConfigId,
expected_num_batches, expected_this_batch_index,
- expected_this_batch_size);
+ expected_this_batch_size,
+ expected_observation_num_bytes, EnvelopeMaker::kOk);
}
}
@@ -249,39 +273,62 @@
size_t expected_num_batches = 1;
size_t expected_this_batch_index = 0;
size_t expected_this_batch_size = 1;
+ // NOTE(rudominer) The values of expected_observation_num_bytes for
+ // the Forculus and Basic RAPPOR encodings in this test are obtained from
+ // experimentation rather than calculation. We are therefore not testing
+ // that the values are correct but rather testing that there is no
+ // regression in the size() functionality. Also just eybealling the numbers
+ // serves as a sanity test. Notice that the Forculus Observations are
+ // rather large compared to the Basic RAPPOR observations with 3 categoreis.
+ size_t expected_observation_num_bytes = 111;
AddStringObservation("a value", 1, 1, expected_num_batches,
- expected_this_batch_index, expected_this_batch_size);
+ expected_this_batch_index, expected_this_batch_size,
+ expected_observation_num_bytes, EnvelopeMaker::kOk);
expected_this_batch_size = 2;
+ expected_observation_num_bytes = 19;
AddStringObservation("Apple", 1, 2, expected_num_batches,
- expected_this_batch_index, expected_this_batch_size);
+ expected_this_batch_index, expected_this_batch_size,
+ expected_observation_num_bytes, EnvelopeMaker::kOk);
// Add two observations for metric 2
expected_num_batches = 2;
expected_this_batch_index = 1;
expected_this_batch_size = 1;
- AddStringObservation("a value", 2, 1, expected_num_batches,
- expected_this_batch_index, expected_this_batch_size);
+ expected_observation_num_bytes = 112;
+ AddStringObservation("a value2", 2, 1, expected_num_batches,
+ expected_this_batch_index, expected_this_batch_size,
+ expected_observation_num_bytes, EnvelopeMaker::kOk);
expected_this_batch_size = 2;
+ expected_observation_num_bytes = 19;
AddStringObservation("Banana", 2, 2, expected_num_batches,
- expected_this_batch_index, expected_this_batch_size);
+ expected_this_batch_index, expected_this_batch_size,
+ expected_observation_num_bytes, EnvelopeMaker::kOk);
// Add two more observations for metric 1
expected_this_batch_index = 0;
expected_this_batch_size = 3;
- AddStringObservation("a value", 1, 1, expected_num_batches,
- expected_this_batch_index, expected_this_batch_size);
+ expected_observation_num_bytes = 112;
+ AddStringObservation("a value3", 1, 1, expected_num_batches,
+ expected_this_batch_index, expected_this_batch_size,
+ expected_observation_num_bytes, EnvelopeMaker::kOk);
expected_this_batch_size = 4;
+ expected_observation_num_bytes = 19;
AddStringObservation("Banana", 1, 2, expected_num_batches,
- expected_this_batch_index, expected_this_batch_size);
+ expected_this_batch_index, expected_this_batch_size,
+ expected_observation_num_bytes, EnvelopeMaker::kOk);
// Add two more observations for metric 2
expected_this_batch_index = 1;
expected_this_batch_size = 3;
- AddStringObservation("a value", 2, 1, expected_num_batches,
- expected_this_batch_index, expected_this_batch_size);
+ expected_observation_num_bytes = 113;
+ AddStringObservation("a value40", 2, 1, expected_num_batches,
+ expected_this_batch_index, expected_this_batch_size,
+ expected_observation_num_bytes, EnvelopeMaker::kOk);
expected_this_batch_size = 4;
+ expected_observation_num_bytes = 19;
AddStringObservation("Cantaloupe", 2, 2, expected_num_batches,
- expected_this_batch_index, expected_this_batch_size);
+ expected_this_batch_index, expected_this_batch_size,
+ expected_observation_num_bytes, EnvelopeMaker::kOk);
// Make the encrypted Envelope.
EncryptedMessage encrypted_message;
@@ -427,5 +474,98 @@
}
}
+// Tests that EnvelopeMaker returns kObservationTooBig when it is supposed to.
+TEST_F(EnvelopeMakerTest, ObservationTooBig) {
+ static const uint32_t kMetricId = 1;
+ static const uint32_t kEncodingConfigId = 3; // NoOp encoding.
+
+ // Set max_bytes_each_observation = 100.
+ ResetEnvelopeMaker(100);
+
+ // Build an input string of length 75 bytes.
+ std::string value("x", 75);
+
+ // From experimentation we know that the overhead of the NoOp encoding
+ // is 20 bytes so we expect the Observation size to be 95.
+ size_t expected_observation_num_bytes = 95;
+
+ // Invoke AddStringObservation() and expect kOk
+ int expected_num_batches = 1;
+ size_t expected_this_batch_index = 0;
+ int expected_this_batch_size = 1;
+ AddStringObservation(value, kMetricId, kEncodingConfigId,
+ expected_num_batches, expected_this_batch_index,
+ expected_this_batch_size, expected_observation_num_bytes,
+ EnvelopeMaker::kOk);
+
+ // Build an input string of length 101 bytes.
+ value = std::string("x", 101);
+ // We expect the Observation to not be added to the Envelope and so for
+ // the Envelope size to not change.
+ expected_observation_num_bytes = 0;
+
+ // Invoke AddStringObservation() and expect kObservationTooBig
+ AddStringObservation(value, kMetricId, kEncodingConfigId,
+ expected_num_batches, expected_this_batch_index,
+ expected_this_batch_size, expected_observation_num_bytes,
+ EnvelopeMaker::kObservationTooBig);
+
+ // Build an input string of length 75 bytes again.
+ value = std::string("x", 75);
+ expected_observation_num_bytes = 95;
+ expected_this_batch_size = 2;
+ // Invoke AddStringObservation() and expect kOk.
+ AddStringObservation(value, kMetricId, kEncodingConfigId,
+ expected_num_batches, expected_this_batch_index,
+ expected_this_batch_size, expected_observation_num_bytes,
+ EnvelopeMaker::kOk);
+}
+
+// Tests that EnvelopeMaker returns kEnvelopeFull when it is supposed to.
+TEST_F(EnvelopeMakerTest, EnvelopeFull) {
+ static const uint32_t kMetricId = 1;
+ static const uint32_t kEncodingConfigId = 3; // NoOp encoding.
+
+ // Set max_bytes_each_observation = 100, max_num_bytes=1000.
+ ResetEnvelopeMaker(100, 1000);
+
+ int expected_this_batch_size = 1;
+ int expected_num_batches = 1;
+ size_t expected_this_batch_index = 0;
+ for (int i = 0; i < 19; i++) {
+ // Build an input string of length 30 bytes.
+ std::string value("x", 30);
+ // From experimentation we know that the overhead of the NoOp encoding
+ // is 20 bytes so we expect the Observation size to be 50.
+ size_t expected_observation_num_bytes = 50;
+
+ // Invoke AddStringObservation() and expect kOk
+ AddStringObservation(value, kMetricId, kEncodingConfigId,
+ expected_num_batches, expected_this_batch_index,
+ expected_this_batch_size++,
+ expected_observation_num_bytes, EnvelopeMaker::kOk);
+ }
+ EXPECT_EQ(950u, envelope_maker_->size());
+
+ // If we try to add an observation of more than 100 bytes we should
+ // get kObservationTooBig.
+ std::string value("x", 101);
+ // We expect the Observation to not be added to the Envelope and so for
+ // the Envelope size to not change.
+ size_t expected_observation_num_bytes = 0;
+ AddStringObservation(
+ value, kMetricId, kEncodingConfigId, expected_num_batches,
+ expected_this_batch_index, expected_this_batch_size++,
+ expected_observation_num_bytes, EnvelopeMaker::kObservationTooBig);
+
+ // If we try to add an observation of 75 bytes we should
+ // get kEnvelopeFull
+ value = std::string("x", 75);
+ AddStringObservation(
+ value, kMetricId, kEncodingConfigId, expected_num_batches,
+ expected_this_batch_index, expected_this_batch_size++,
+ expected_observation_num_bytes, EnvelopeMaker::kEnvelopeFull);
+}
+
} // namespace encoder
} // namespace cobalt
