| 194 | constexpr const char MemmappedFileSystem::kMemmappedPackageDefaultGraphDef[]; |
| 195 | |
| 196 | Status MemmappedFileSystem::InitializeFromFile(Env* env, |
| 197 | const string& filename) { |
| 198 | TF_RETURN_IF_ERROR( |
| 199 | env->NewReadOnlyMemoryRegionFromFile(filename, &mapped_memory_)); |
| 200 | directory_.clear(); |
| 201 | if (mapped_memory_->length() <= sizeof(uint64)) { |
| 202 | return errors::DataLoss("Corrupted memmapped model file: ", filename, |
| 203 | " Invalid package size"); |
| 204 | } |
| 205 | const auto memory_start = |
| 206 | reinterpret_cast<const uint8*>(mapped_memory_->data()); |
| 207 | const uint64 directory_offset = DecodeUint64LittleEndian( |
| 208 | memory_start + mapped_memory_->length() - sizeof(uint64)); |
| 209 | if (directory_offset > mapped_memory_->length() - sizeof(uint64)) { |
| 210 | return errors::DataLoss("Corrupted memmapped model file: ", filename, |
| 211 | " Invalid directory offset"); |
| 212 | } |
| 213 | MemmappedFileSystemDirectory proto_directory; |
| 214 | if (!ParseProtoUnlimited( |
| 215 | &proto_directory, memory_start + directory_offset, |
| 216 | mapped_memory_->length() - directory_offset - sizeof(uint64))) { |
| 217 | return errors::DataLoss("Corrupted memmapped model file: ", filename, |
| 218 | " Can't parse its internal directory"); |
| 219 | } |
| 220 | |
| 221 | // Iterating in reverse order to get lengths of elements; |
| 222 | uint64 prev_element_offset = directory_offset; |
| 223 | for (auto element_iter = proto_directory.element().rbegin(); |
| 224 | element_iter != proto_directory.element().rend(); ++element_iter) { |
| 225 | // Check that the element offset is in the right range. |
| 226 | if (element_iter->offset() >= prev_element_offset) { |
| 227 | return errors::DataLoss("Corrupted memmapped model file: ", filename, |
| 228 | " Invalid offset of internal component"); |
| 229 | } |
| 230 | if (!directory_ |
| 231 | .insert(std::make_pair( |
| 232 | element_iter->name(), |
| 233 | FileRegion(element_iter->offset(), element_iter->length()))) |
| 234 | .second) { |
| 235 | return errors::DataLoss("Corrupted memmapped model file: ", filename, |
| 236 | " Duplicate name of internal component ", |
| 237 | element_iter->name()); |
| 238 | } |
| 239 | prev_element_offset = element_iter->offset(); |
| 240 | } |
| 241 | return Status::OK(); |
| 242 | } |
| 243 | |
| 244 | bool MemmappedFileSystem::IsMemmappedPackageFilename(const string& filename) { |
| 245 | return absl::StartsWith(filename, kMemmappedPackagePrefix); |