Merge "Refactor bufferhub token to be more secure"
diff --git a/services/bufferhub/Android.bp b/services/bufferhub/Android.bp
index 72d210c..a4b81fe 100644
--- a/services/bufferhub/Android.bp
+++ b/services/bufferhub/Android.bp
@@ -34,6 +34,7 @@
     ],
     shared_libs: [
         "android.frameworks.bufferhub@1.0",
+        "libcrypto",
         "libcutils",
         "libhidlbase",
         "libhidltransport",
@@ -60,6 +61,7 @@
     shared_libs: [
         "android.frameworks.bufferhub@1.0",
         "libbufferhubservice",
+        "libcrypto",
         "libcutils",
         "libhidltransport",
         "libhwbinder",
diff --git a/services/bufferhub/BufferHubService.cpp b/services/bufferhub/BufferHubService.cpp
index 54796a2..43235f2 100644
--- a/services/bufferhub/BufferHubService.cpp
+++ b/services/bufferhub/BufferHubService.cpp
@@ -14,13 +14,16 @@
  * limitations under the License.
  */
 
+#include <array>
 #include <iomanip>
+#include <random>
 #include <sstream>
 
 #include <android/hardware_buffer.h>
 #include <bufferhub/BufferHubService.h>
 #include <cutils/native_handle.h>
 #include <log/log.h>
+#include <openssl/hmac.h>
 #include <system/graphics-base.h>
 
 using ::android::BufferHubDefs::MetadataHeader;
@@ -32,6 +35,13 @@
 namespace V1_0 {
 namespace implementation {
 
+BufferHubService::BufferHubService() {
+    std::mt19937_64 randomEngine;
+    randomEngine.seed(time(nullptr));
+
+    mKey = randomEngine();
+}
+
 Return<void> BufferHubService::allocateBuffer(const HardwareBufferDescription& description,
                                               const uint32_t userMetadataSize,
                                               allocateBuffer_cb _hidl_cb) {
@@ -66,27 +76,49 @@
 
 Return<void> BufferHubService::importBuffer(const hidl_handle& tokenHandle,
                                             importBuffer_cb _hidl_cb) {
-    if (!tokenHandle.getNativeHandle() || tokenHandle->numFds != 0 || tokenHandle->numInts != 1) {
+    if (!tokenHandle.getNativeHandle() || tokenHandle->numFds != 0 || tokenHandle->numInts <= 1) {
         // nullptr handle or wrong format
         _hidl_cb(/*status=*/BufferHubStatus::INVALID_TOKEN, /*bufferClient=*/nullptr,
                  /*bufferTraits=*/{});
         return Void();
     }
 
-    uint32_t token = tokenHandle->data[0];
+    int tokenId = tokenHandle->data[0];
 
     wp<BufferClient> originClientWp;
     {
-        std::lock_guard<std::mutex> lock(mTokenMapMutex);
-        auto iter = mTokenMap.find(token);
+        std::lock_guard<std::mutex> lock(mTokenMutex);
+        auto iter = mTokenMap.find(tokenId);
         if (iter == mTokenMap.end()) {
-            // Invalid token
+            // Token Id not exist
+            ALOGD("%s: token #%d not found.", __FUNCTION__, tokenId);
             _hidl_cb(/*status=*/BufferHubStatus::INVALID_TOKEN, /*bufferClient=*/nullptr,
                      /*bufferTraits=*/{});
             return Void();
         }
 
-        originClientWp = iter->second;
+        const std::vector<uint8_t>& tokenHMAC = iter->second.first;
+
+        int numIntsForHMAC = (int)ceil(tokenHMAC.size() * sizeof(uint8_t) / (double)sizeof(int));
+        if (tokenHandle->numInts - 1 != numIntsForHMAC) {
+            // HMAC size not match
+            ALOGD("%s: token #%d HMAC size not match. Expected: %d Actual: %d", __FUNCTION__,
+                  tokenId, numIntsForHMAC, tokenHandle->numInts - 1);
+            _hidl_cb(/*status=*/BufferHubStatus::INVALID_TOKEN, /*bufferClient=*/nullptr,
+                     /*bufferTraits=*/{});
+            return Void();
+        }
+
+        size_t hmacSize = tokenHMAC.size() * sizeof(uint8_t);
+        if (memcmp(tokenHMAC.data(), &tokenHandle->data[1], hmacSize) != 0) {
+            // HMAC not match
+            ALOGD("%s: token #%d HMAC not match.", __FUNCTION__, tokenId);
+            _hidl_cb(/*status=*/BufferHubStatus::INVALID_TOKEN, /*bufferClient=*/nullptr,
+                     /*bufferTraits=*/{});
+            return Void();
+        }
+
+        originClientWp = iter->second.second;
         mTokenMap.erase(iter);
     }
 
@@ -225,9 +257,9 @@
     // Map from bufferId to tokenCount
     std::map<int, uint32_t> tokenCount;
     {
-        std::lock_guard<std::mutex> lock(mTokenMapMutex);
+        std::lock_guard<std::mutex> lock(mTokenMutex);
         for (auto iter = mTokenMap.begin(); iter != mTokenMap.end(); ++iter) {
-            sp<BufferClient> client = iter->second.promote();
+            sp<BufferClient> client = iter->second.second.promote();
             if (client != nullptr) {
                 const std::shared_ptr<BufferNode> node = client->getBufferNode();
                 auto mapIter = tokenCount.find(node->id());
@@ -262,21 +294,35 @@
 }
 
 hidl_handle BufferHubService::registerToken(const wp<BufferClient>& client) {
-    uint32_t token;
-    std::lock_guard<std::mutex> lock(mTokenMapMutex);
+    // Find next available token id
+    std::lock_guard<std::mutex> lock(mTokenMutex);
     do {
-        token = mTokenEngine();
-    } while (mTokenMap.find(token) != mTokenMap.end());
+        ++mLastTokenId;
+    } while (mTokenMap.find(mLastTokenId) != mTokenMap.end());
 
-    // native_handle_t use int[], so here need one slots to fit in uint32_t
-    native_handle_t* handle = native_handle_create(/*numFds=*/0, /*numInts=*/1);
-    handle->data[0] = token;
+    std::array<uint8_t, EVP_MAX_MD_SIZE> hmac;
+    uint32_t hmacSize = 0U;
+
+    HMAC(/*evp_md=*/EVP_sha256(), /*key=*/&mKey, /*key_len=*/kKeyLen,
+         /*data=*/(uint8_t*)&mLastTokenId, /*data_len=*/mTokenIdSize,
+         /*out=*/hmac.data(), /*out_len=*/&hmacSize);
+
+    int numIntsForHMAC = (int)ceil(hmacSize / (double)sizeof(int));
+    native_handle_t* handle = native_handle_create(/*numFds=*/0, /*numInts=*/1 + numIntsForHMAC);
+    handle->data[0] = mLastTokenId;
+    // Set all the the bits of last int to 0 since it might not be fully overwritten
+    handle->data[numIntsForHMAC] = 0;
+    memcpy(&handle->data[1], hmac.data(), hmacSize);
 
     // returnToken owns the native_handle_t* thus doing lifecycle management
     hidl_handle returnToken;
     returnToken.setTo(handle, /*shoudOwn=*/true);
 
-    mTokenMap.emplace(token, client);
+    std::vector<uint8_t> hmacVec;
+    hmacVec.resize(hmacSize);
+    memcpy(hmacVec.data(), hmac.data(), hmacSize);
+    mTokenMap.emplace(mLastTokenId, std::pair(hmacVec, client));
+
     return returnToken;
 }
 
@@ -291,10 +337,10 @@
 }
 
 void BufferHubService::removeTokenByClient(const BufferClient* client) {
-    std::lock_guard<std::mutex> lock(mTokenMapMutex);
+    std::lock_guard<std::mutex> lock(mTokenMutex);
     auto iter = mTokenMap.begin();
     while (iter != mTokenMap.end()) {
-        if (iter->second == client) {
+        if (iter->second.second == client) {
             auto oldIter = iter;
             ++iter;
             mTokenMap.erase(oldIter);
diff --git a/services/bufferhub/include/bufferhub/BufferHubService.h b/services/bufferhub/include/bufferhub/BufferHubService.h
index 255f329..0195bcf 100644
--- a/services/bufferhub/include/bufferhub/BufferHubService.h
+++ b/services/bufferhub/include/bufferhub/BufferHubService.h
@@ -17,8 +17,10 @@
 #ifndef ANDROID_FRAMEWORKS_BUFFERHUB_V1_0_BUFFER_HUB_SERVICE_H
 #define ANDROID_FRAMEWORKS_BUFFERHUB_V1_0_BUFFER_HUB_SERVICE_H
 
+#include <map>
 #include <mutex>
-#include <random>
+#include <set>
+#include <vector>
 
 #include <android/frameworks/bufferhub/1.0/IBufferHub.h>
 #include <bufferhub/BufferClient.h>
@@ -39,6 +41,8 @@
 
 class BufferHubService : public IBufferHub {
 public:
+    BufferHubService();
+
     Return<void> allocateBuffer(const HardwareBufferDescription& description,
                                 const uint32_t userMetadataSize,
                                 allocateBuffer_cb _hidl_cb) override;
@@ -60,11 +64,19 @@
     std::mutex mClientSetMutex;
     std::set<wp<BufferClient>> mClientSet GUARDED_BY(mClientSetMutex);
 
-    // TODO(b/118180214): use a more secure implementation
-    std::mt19937 mTokenEngine;
-    // The mapping from token to the client creates it.
-    std::mutex mTokenMapMutex;
-    std::map<uint32_t, const wp<BufferClient>> mTokenMap GUARDED_BY(mTokenMapMutex);
+    // Token generation related
+    // A random number used as private key for HMAC
+    uint64_t mKey;
+    static constexpr size_t kKeyLen = sizeof(uint64_t);
+
+    std::mutex mTokenMutex;
+    // The first TokenId will be 1. TokenId could be negative.
+    int mLastTokenId GUARDED_BY(mTokenMutex) = 0;
+    static constexpr size_t mTokenIdSize = sizeof(int);
+    // A map from token id to the token-buffer_client pair. Using token id as the key to reduce
+    // looking up time
+    std::map<int, std::pair<std::vector<uint8_t>, const wp<BufferClient>>> mTokenMap
+            GUARDED_BY(mTokenMutex);
 };
 
 } // namespace implementation