chaps: Add OnLoad to ObjectImpl

BUG=chromium:1032100
TEST=manual && FEATURES=test emerge-$BOARD chaps

Change-Id: I8a1b1ade62bbb34c42771580c67f1ce19f5d76a7
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform2/+/1967780
Tested-by: John L Chen <zuan@chromium.org>
Commit-Queue: John L Chen <zuan@chromium.org>
Reviewed-by: Andrey Pronin <apronin@chromium.org>
(cherry picked from commit f4dec3d7ed3085f24c7d2e204fff9fa98865dcc7)
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform2/+/1992671
Commit-Queue: Andrey Pronin <apronin@chromium.org>
Tested-by: Andrey Pronin <apronin@chromium.org>
diff --git a/chaps/object.h b/chaps/object.h
index 10d2c51..04e88af 100644
--- a/chaps/object.h
+++ b/chaps/object.h
@@ -94,6 +94,9 @@
   virtual void RemoveAttribute(CK_ATTRIBUTE_TYPE type) = 0;
   // Provides a read-only map of all existing attributes.
   virtual const AttributeMap* GetAttributeMap() const = 0;
+  // This should be called after an object is loaded from disk. If this returns
+  // false, then object loading should be considered as failed.
+  virtual bool OnLoad() = 0;
   // Get / set handle as seen by PKCS #11 clients.
   virtual int handle() const = 0;
   virtual void set_handle(int handle) = 0;
diff --git a/chaps/object_impl.cc b/chaps/object_impl.cc
index a1f6915..3664f36 100644
--- a/chaps/object_impl.cc
+++ b/chaps/object_impl.cc
@@ -197,6 +197,15 @@
   return &attributes_;
 }
 
+bool ObjectImpl::OnLoad() {
+  if (!SetPolicyByClass()) {
+    LOG(ERROR) << "Failed to set attribute access policy.";
+    return false;
+  }
+  stage_ = kModify;
+  return true;
+}
+
 bool ObjectImpl::SetPolicyByClass() {
   if (!IsAttributePresent(CKA_CLASS)) {
     LOG(ERROR) << "Missing object class attribute.";
diff --git a/chaps/object_impl.h b/chaps/object_impl.h
index 0281d97..e9cdc0a 100644
--- a/chaps/object_impl.h
+++ b/chaps/object_impl.h
@@ -48,6 +48,7 @@
                           const std::string& value) override;
   void RemoveAttribute(CK_ATTRIBUTE_TYPE type) override;
   const AttributeMap* GetAttributeMap() const override;
+  bool OnLoad() override;
   int handle() const override { return handle_; }
   void set_handle(int handle) override { handle_ = handle; }
   int store_id() const override { return store_id_; }
diff --git a/chaps/object_mock.h b/chaps/object_mock.h
index 5ffe704..933ab3a 100644
--- a/chaps/object_mock.h
+++ b/chaps/object_mock.h
@@ -42,6 +42,7 @@
   MOCK_METHOD2(SetAttributeString, void(CK_ATTRIBUTE_TYPE, const std::string&));
   MOCK_METHOD1(RemoveAttribute, void(CK_ATTRIBUTE_TYPE));
   MOCK_CONST_METHOD0(GetAttributeMap, const AttributeMap*());
+  MOCK_METHOD0(OnLoad, bool());
   MOCK_CONST_METHOD0(handle, int());
   MOCK_METHOD1(set_handle, void(int));
   MOCK_CONST_METHOD0(store_id, int());
@@ -81,6 +82,7 @@
         .WillByDefault(testing::Invoke(this, &ObjectMock::FakeRemoveAttribute));
     ON_CALL(*this, GetAttributeMap())
         .WillByDefault(testing::Return(&attributes_));
+    ON_CALL(*this, OnLoad()).WillByDefault(testing::Return(true));
     ON_CALL(*this, set_handle(testing::_))
         .WillByDefault(testing::Invoke(this, &ObjectMock::FakeSetHandle));
     ON_CALL(*this, set_store_id(testing::_))
diff --git a/chaps/object_pool_impl.cc b/chaps/object_pool_impl.cc
index 40b7fa7..6fe9038 100644
--- a/chaps/object_pool_impl.cc
+++ b/chaps/object_pool_impl.cc
@@ -259,6 +259,12 @@
       return false;
     }
   }
+
+  if (!object->OnLoad()) {
+    LOG(ERROR) << "Object's OnLoad failed.";
+    return false;
+  }
+
   return true;
 }