Add a builder to PrimitiveSet and make the class itself immutable.

This may break users which build a primitive set themselves -- we do not expect anyone to actually do this. If you are broken, please check if you can use the newly added Builder instead. Also, please file an issue if you need support.

#tinkApiChange

PiperOrigin-RevId: 497148275

1 files changed, 175 insertions, 34 deletions

diff --git a/cc/primitive_set.h b/cc/primitive_set.h
index 58926b543..ba3920c6a 100644
--- a/cc/primitive_set.h
+++ b/cc/primitive_set.h
@@ -17,10 +17,14 @@
 #ifndef TINK_PRIMITIVE_SET_H_
 #define TINK_PRIMITIVE_SET_H_
 
+#include <algorithm>
+#include <memory>
 #include <string>
 #include <unordered_map>
+#include <utility>
 #include <vector>
 
+#include "absl/base/thread_annotations.h"
 #include "absl/container/flat_hash_map.h"
 #include "absl/memory/memory.h"
 #include "absl/status/status.h"
@@ -113,34 +117,164 @@ class PrimitiveSet {
   };
 
   typedef std::vector<std::unique_ptr<Entry<P>>> Primitives;
+  typedef std::unordered_map<std::string, Primitives>
+      CiphertextPrefixToPrimitivesMap;
+
+ private:
+  // Helper methods for mutations, used by the Builder and the deprecated
+  // mutation methods on PrimitiveSet.
+
+  static crypto::tink::util::Status SetPrimaryImpl(
+      Entry<P>** output, Entry<P>* primary,
+      const CiphertextPrefixToPrimitivesMap& primitives) {
+    if (!primary) {
+      return util::Status(absl::StatusCode::kInvalidArgument,
+                          "The primary primitive must be non-null.");
+    }
+    if (primary->get_status() != google::crypto::tink::KeyStatusType::ENABLED) {
+      return util::Status(absl::StatusCode::kInvalidArgument,
+                          "Primary has to be enabled.");
+    }
+
+    if (primitives.count(primary->get_identifier()) == 0) {
+      return util::Status(absl::StatusCode::kInvalidArgument,
+                          "Primary cannot be set to an entry which is "
+                          "not held by this primitive set.");
+    }
+
+    *output = primary;
+    return crypto::tink::util::OkStatus();
+  }
+
+  static crypto::tink::util::StatusOr<Entry<P>*> AddPrimitiveImpl(
+      std::unique_ptr<P> primitive,
+      const google::crypto::tink::KeysetInfo::KeyInfo& key_info,
+      CiphertextPrefixToPrimitivesMap& primitives) {
+    auto entry_or = Entry<P>::New(std::move(primitive), key_info);
+    if (!entry_or.ok()) return entry_or.status();
+
+    std::string identifier = entry_or.value()->get_identifier();
+    primitives[identifier].push_back(std::move(entry_or.value()));
+    return primitives[identifier].back().get();
+  }
+
+ public:
+  // Builder is used to construct PrimitiveSet objects. Objects returned by the
+  // builder are immutable. Calling any of the non-const methods on them will
+  // fail.
+  class Builder {
+   public:
+    // Adds 'primitive' to this set for the specified 'key'.
+    Builder& AddPrimitive(
+        std::unique_ptr<P> primitive,
+        const google::crypto::tink::KeysetInfo::KeyInfo& key_info) & {
+      absl::MutexLock lock(&mutex_);
+      if (!status_.ok()) return *this;
+      status_ = AddPrimitiveImpl(std::move(primitive), key_info, primitives_)
+                    .status();
+      return *this;
+    }
+
+    Builder&& AddPrimitive(
+        std::unique_ptr<P> primitive,
+        const google::crypto::tink::KeysetInfo::KeyInfo& key_info) && {
+      return std::move(AddPrimitive(std::move(primitive), key_info));
+    }
+
+    // Adds 'primitive' to this set for the specified 'key' and marks it
+    // primary.
+    Builder& AddPrimaryPrimitive(
+        std::unique_ptr<P> primitive,
+        const google::crypto::tink::KeysetInfo::KeyInfo& key_info) & {
+      absl::MutexLock lock(&mutex_);
+      if (!status_.ok()) return *this;
+      auto entry_result =
+          AddPrimitiveImpl(std::move(primitive), key_info, primitives_);
+      if (!entry_result.ok()) {
+        status_ = entry_result.status();
+        return *this;
+      }
+      status_ = SetPrimaryImpl(&primary_, entry_result.value(), primitives_);
+      return *this;
+    }
+
+    Builder&& AddPrimaryPrimitive(
+        std::unique_ptr<P> primitive,
+        const google::crypto::tink::KeysetInfo::KeyInfo& key_info) && {
+      return std::move(AddPrimaryPrimitive(std::move(primitive), key_info));
+    }
+
+    // Add the given annotations. Existing annotations will not be overwritten.
+    Builder& AddAnnotations(
+        absl::flat_hash_map<std::string, std::string> annotations) & {
+      absl::MutexLock lock(&mutex_);
+      annotations_.merge(std::move(annotations));
+      return *this;
+    }
+
+    Builder&& AddAnnotations(
+        absl::flat_hash_map<std::string, std::string> annotations) && {
+      return std::move(AddAnnotations(std::move(annotations)));
+    }
+
+    crypto::tink::util::StatusOr<PrimitiveSet<P>> Build() && {
+      absl::MutexLock lock(&mutex_);
+      if (!status_.ok()) return status_;
+      return PrimitiveSet<P>(std::move(primitives_), primary_,
+                             std::move(annotations_));
+    }
+
+   private:
+    // The Entry<P> object is owned by primitives_
+    Entry<P>* primary_ ABSL_GUARDED_BY(mutex_) = nullptr;
+    CiphertextPrefixToPrimitivesMap primitives_ ABSL_GUARDED_BY(mutex_);
+    absl::flat_hash_map<std::string, std::string> annotations_
+        ABSL_GUARDED_BY(mutex_);
+    absl::Mutex mutex_;
+    crypto::tink::util::Status status_ ABSL_GUARDED_BY(mutex_);
+  };
+
+  // PrimitiveSet is movable, but not copyable
+  PrimitiveSet(PrimitiveSet&&) = default;
+  PrimitiveSet<P>& operator=(PrimitiveSet&&) = default;
+  PrimitiveSet(const PrimitiveSet&) = delete;
+  PrimitiveSet<P>& operator=(const PrimitiveSet&) = delete;
 
   // Constructs an empty PrimitiveSet.
   // Note: This is equivalent to PrimitiveSet<P>(/*annotations=*/{}).
+  ABSL_DEPRECATED(
+      "Constructing PrimitiveSet using constructors is deprecated. Use "
+      "PrimitiveSet<>::Builder instead.")
   PrimitiveSet<P>() = default;
   // Constructs an empty PrimitiveSet with `annotations`.
+  ABSL_DEPRECATED(
+      "Constructing PrimitiveSet using constructors is deprecated. Use "
+      "PrimitiveSet<>::Builder instead.")
   explicit PrimitiveSet<P>(
       const absl::flat_hash_map<std::string, std::string>& annotations)
       : annotations_(annotations) {}
 
   // Adds 'primitive' to this set for the specified 'key'.
+  ABSL_DEPRECATED(
+      "Mutating PrimitiveSets after construction is deprecated. Use "
+      "PrimitiveSet<>::Builder instead.")
   crypto::tink::util::StatusOr<Entry<P>*> AddPrimitive(
       std::unique_ptr<P> primitive,
       const google::crypto::tink::KeysetInfo::KeyInfo& key_info) {
-    auto entry_or = Entry<P>::New(std::move(primitive), key_info);
-    if (!entry_or.ok()) return entry_or.status();
+    if (!is_mutable()) {
+      return util::Status(absl::StatusCode::kFailedPrecondition,
+                          "PrimitiveSet is not mutable.");
+    }
 
-    absl::MutexLock lock(&primitives_mutex_);
-    std::string identifier = entry_or.value()->get_identifier();
-    primitives_[identifier].push_back(std::move(entry_or.value()));
-    return primitives_[identifier].back().get();
+    absl::MutexLock lock(primitives_mutex_.get());
+    return AddPrimitiveImpl(std::move(primitive), key_info, primitives_);
   }
 
   // Returns the entries with primitives identifed by 'identifier'.
   crypto::tink::util::StatusOr<const Primitives*> get_primitives(
-      absl::string_view identifier) {
-    absl::MutexLock lock(&primitives_mutex_);
-    typename CiphertextPrefixToPrimitivesMap::iterator found =
-        primitives_.find(std::string(identifier));
+      absl::string_view identifier) const {
+    absl::MutexLockMaybe lock(primitives_mutex_.get());
+    auto found = primitives_.find(std::string(identifier));
     if (found == primitives_.end()) {
       return ToStatusF(absl::StatusCode::kNotFound,
                        "No primitives found for identifier '%s'.", identifier);
@@ -149,37 +283,32 @@ class PrimitiveSet {
   }
 
   // Returns all primitives that use RAW prefix.
-  crypto::tink::util::StatusOr<const Primitives*> get_raw_primitives() {
+  crypto::tink::util::StatusOr<const Primitives*> get_raw_primitives() const {
     return get_primitives(CryptoFormat::kRawPrefix);
   }
 
   // Sets the given 'primary' as the primary primitive of this set.
+  ABSL_DEPRECATED(
+      "Mutating PrimitiveSets after construction is deprecated. Use "
+      "PrimitiveSet<>::Builder instead.")
   crypto::tink::util::Status set_primary(Entry<P>* primary) {
-    if (!primary) {
-      return util::Status(absl::StatusCode::kInvalidArgument,
-                          "The primary primitive must be non-null.");
-    }
-    if (primary->get_status() != google::crypto::tink::KeyStatusType::ENABLED) {
-      return util::Status(absl::StatusCode::kInvalidArgument,
-                          "Primary has to be enabled.");
-    }
-    auto entries_result = get_primitives(primary->get_identifier());
-    if (!entries_result.ok()) {
-      return util::Status(absl::StatusCode::kInvalidArgument,
-                          "Primary cannot be set to an entry which is "
-                          "not held by this primitive set.");
+    if (!is_mutable()) {
+      return util::Status(absl::StatusCode::kFailedPrecondition,
+                          "PrimitiveSet is not mutable.");
     }
-
-    primary_ = primary;
-    return crypto::tink::util::OkStatus();
+    absl::MutexLock lock(primitives_mutex_.get());
+    return SetPrimaryImpl(&primary_, primary, primitives_);
   }
 
   // Returns the entry with the primary primitive.
-  const Entry<P>* get_primary() const { return primary_; }
+  const Entry<P>* get_primary() const {
+    absl::MutexLockMaybe lock(primitives_mutex_.get());
+    return primary_;
+  }
 
   // Returns all entries currently in this primitive set.
   std::vector<Entry<P>*> get_all() const {
-    absl::MutexLock lock(&primitives_mutex_);
+    absl::MutexLockMaybe lock(primitives_mutex_.get());
     std::vector<Entry<P>*> result;
     for (const auto& prefix_and_vector : primitives_) {
       for (const auto& primitive : prefix_and_vector.second) {
@@ -193,17 +322,29 @@ class PrimitiveSet {
     return annotations_;
   }
 
+  bool is_mutable() const { return primitives_mutex_ != nullptr; }
+
  private:
-  typedef std::unordered_map<std::string, Primitives>
-      CiphertextPrefixToPrimitivesMap;
+  // Constructs an empty PrimitiveSet.
+  // Note: This is equivalent to PrimitiveSet<P>(/*annotations=*/{}).
+  PrimitiveSet(CiphertextPrefixToPrimitivesMap primitives, Entry<P>* primary,
+               absl::flat_hash_map<std::string, std::string> annotations)
+      : primary_(primary),
+        primitives_mutex_(nullptr),
+        primitives_(std::move(primitives)),
+        annotations_(std::move(annotations)) {}
+
   // The Entry<P> object is owned by primitives_
-  Entry<P>* primary_ = nullptr;
-  mutable absl::Mutex primitives_mutex_;
+  Entry<P>* primary_ ABSL_GUARDED_BY(primitives_mutex_) = nullptr;
+  // If not nullptr, this mutex is used to guard all read and write access to
+  // the primitiveset. If nullptr, the primitiveset is immutable (and lockfree).
+  mutable std::unique_ptr<absl::Mutex> primitives_mutex_ =
+      std::make_unique<absl::Mutex>();
   CiphertextPrefixToPrimitivesMap primitives_
       ABSL_GUARDED_BY(primitives_mutex_);
 
   // Annotations for the set of primitives.
-  const absl::flat_hash_map<std::string, std::string> annotations_;
+  absl::flat_hash_map<std::string, std::string> annotations_;
 };
 
 }  // namespace tink