Home | History | Annotate | Download | only in util 
      1 /*
      2  * Copyright (C) 2017 The Android Open Source Project
      3  *
      4  * Licensed under the Apache License, Version 2.0 (the "License");
      5  * you may not use this file except in compliance with the License.
      6  * You may obtain a copy of the License at
      7  *
      8  *      http://www.apache.org/licenses/LICENSE-2.0
      9  *
     10  * Unless required by applicable law or agreed to in writing, software
     11  * distributed under the License is distributed on an "AS IS" BASIS,
     12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
     13  * See the License for the specific language governing permissions and
     14  * limitations under the License.
     15  */
     16 
     17 // Utility functions for working with FlatBuffers.
     18 
     19 #ifndef LIBTEXTCLASSIFIER_UTIL_FLATBUFFERS_H_
     20 #define LIBTEXTCLASSIFIER_UTIL_FLATBUFFERS_H_
     21 
     22 #include <memory>
     23 #include <string>
     24 
     25 #include "model_generated.h"
     26 #include "flatbuffers/flatbuffers.h"
     27 
     28 namespace libtextclassifier2 {
     29 
     30 // Loads and interprets the buffer as 'FlatbufferMessage' and verifies its
     31 // integrity.
     32 template <typename FlatbufferMessage>
     33 const FlatbufferMessage* LoadAndVerifyFlatbuffer(const void* buffer, int size) {
     34   const FlatbufferMessage* message =
     35       flatbuffers::GetRoot<FlatbufferMessage>(buffer);
     36   if (message == nullptr) {
     37     return nullptr;
     38   }
     39   flatbuffers::Verifier verifier(reinterpret_cast<const uint8_t*>(buffer),
     40                                  size);
     41   if (message->Verify(verifier)) {
     42     return message;
     43   } else {
     44     return nullptr;
     45   }
     46 }
     47 
     48 // Same as above but takes string.
     49 template <typename FlatbufferMessage>
     50 const FlatbufferMessage* LoadAndVerifyFlatbuffer(const std::string& buffer) {
     51   return LoadAndVerifyFlatbuffer<FlatbufferMessage>(buffer.c_str(),
     52                                                     buffer.size());
     53 }
     54 
     55 // Loads and interprets the buffer as 'FlatbufferMessage', verifies its
     56 // integrity and returns its mutable version.
     57 template <typename FlatbufferMessage>
     58 std::unique_ptr<typename FlatbufferMessage::NativeTableType>
     59 LoadAndVerifyMutableFlatbuffer(const void* buffer, int size) {
     60   const FlatbufferMessage* message =
     61       LoadAndVerifyFlatbuffer<FlatbufferMessage>(buffer, size);
     62   if (message == nullptr) {
     63     return nullptr;
     64   }
     65   return std::unique_ptr<typename FlatbufferMessage::NativeTableType>(
     66       message->UnPack());
     67 }
     68 
     69 // Same as above but takes string.
     70 template <typename FlatbufferMessage>
     71 std::unique_ptr<typename FlatbufferMessage::NativeTableType>
     72 LoadAndVerifyMutableFlatbuffer(const std::string& buffer) {
     73   return LoadAndVerifyMutableFlatbuffer<FlatbufferMessage>(buffer.c_str(),
     74                                                            buffer.size());
     75 }
     76 
     77 template <typename FlatbufferMessage>
     78 const char* FlatbufferFileIdentifier() {
     79   return nullptr;
     80 }
     81 
     82 template <>
     83 const char* FlatbufferFileIdentifier<Model>();
     84 
     85 // Packs the mutable flatbuffer message to string.
     86 template <typename FlatbufferMessage>
     87 std::string PackFlatbuffer(
     88     const typename FlatbufferMessage::NativeTableType* mutable_message) {
     89   flatbuffers::FlatBufferBuilder builder;
     90   builder.Finish(FlatbufferMessage::Pack(builder, mutable_message),
     91                  FlatbufferFileIdentifier<FlatbufferMessage>());
     92   return std::string(reinterpret_cast<const char*>(builder.GetBufferPointer()),
     93                      builder.GetSize());
     94 }
     95 
     96 }  // namespace libtextclassifier2
     97 
     98 #endif  // LIBTEXTCLASSIFIER_UTIL_FLATBUFFERS_H_
     99