1 /*
2 * Copyright (c) Meta Platforms, Inc. and affiliates.
3 * All rights reserved.
4 *
5 * This source code is licensed under the BSD-style license found in the
6 * LICENSE file in the root directory of this source tree.
7 */
8
9 #include <executorch/extension/data_loader/buffer_data_loader.h>
10
11 #include <cstring>
12
13 #include <gtest/gtest.h>
14
15 #include <executorch/runtime/core/result.h>
16 #include <executorch/runtime/platform/runtime.h>
17
18 using namespace ::testing;
19 using executorch::extension::BufferDataLoader;
20 using executorch::runtime::DataLoader;
21 using executorch::runtime::Error;
22 using executorch::runtime::FreeableBuffer;
23 using executorch::runtime::Result;
24
25 class BufferDataLoaderTest : public ::testing::Test {
26 protected:
SetUp()27 void SetUp() override {
28 // Since these tests cause ET_LOG to be called, the PAL must be initialized
29 // first.
30 executorch::runtime::runtime_init();
31 }
32 };
33
TEST_F(BufferDataLoaderTest,InBoundsLoadsSucceed)34 TEST_F(BufferDataLoaderTest, InBoundsLoadsSucceed) {
35 // Create some heterogeneous data.
36 uint8_t data[256];
37 for (int i = 0; i < sizeof(data); ++i) {
38 data[i] = i;
39 }
40
41 // Wrap it in a loader.
42 BufferDataLoader edl(data, sizeof(data));
43
44 // size() should succeed and reflect the total size.
45 Result<size_t> size = edl.size();
46 EXPECT_TRUE(size.ok());
47 EXPECT_EQ(*size, sizeof(data));
48
49 // Load the first bytes of the data.
50 {
51 Result<FreeableBuffer> fb = edl.load(
52 /*offset=*/0,
53 /*size=*/8,
54 /*segment_info=*/
55 DataLoader::SegmentInfo(DataLoader::SegmentInfo::Type::Program));
56 EXPECT_TRUE(fb.ok());
57 EXPECT_EQ(fb->size(), 8);
58 EXPECT_EQ(
59 0,
60 std::memcmp(
61 fb->data(),
62 "\x00\x01\x02\x03"
63 "\x04\x05\x06\x07",
64 fb->size()));
65
66 // Freeing should be a no-op but should still clear out the data/size.
67 fb->Free();
68 EXPECT_EQ(fb->size(), 0);
69 EXPECT_EQ(fb->data(), nullptr);
70
71 // Safe to call multiple times.
72 fb->Free();
73 }
74
75 // Load the last few bytes of the data, a different size than the first time.
76 {
77 Result<FreeableBuffer> fb = edl.load(
78 /*offset=*/sizeof(data) - 3,
79 /*size=*/3,
80 /*segment_info=*/
81 DataLoader::SegmentInfo(DataLoader::SegmentInfo::Type::Program));
82 EXPECT_TRUE(fb.ok());
83 EXPECT_EQ(fb->size(), 3);
84 EXPECT_EQ(0, std::memcmp(fb->data(), "\xfd\xfe\xff", fb->size()));
85 }
86
87 // Loading all of the data succeeds.
88 {
89 Result<FreeableBuffer> fb = edl.load(
90 /*offset=*/0,
91 /*size=*/sizeof(data),
92 /*segment_info*/
93 DataLoader::SegmentInfo(DataLoader::SegmentInfo::Type::Program));
94 EXPECT_TRUE(fb.ok());
95 EXPECT_EQ(fb->size(), sizeof(data));
96 EXPECT_EQ(0, std::memcmp(fb->data(), data, fb->size()));
97 }
98
99 // Loading zero-sized data succeeds, even at the end of the data.
100 {
101 Result<FreeableBuffer> fb = edl.load(
102 /*offset=*/sizeof(data),
103 /*size=*/0,
104 /*segment_info*/
105 DataLoader::SegmentInfo(DataLoader::SegmentInfo::Type::Program));
106 EXPECT_TRUE(fb.ok());
107 EXPECT_EQ(fb->size(), 0);
108 }
109 }
110
TEST_F(BufferDataLoaderTest,OutOfBoundsLoadFails)111 TEST_F(BufferDataLoaderTest, OutOfBoundsLoadFails) {
112 // Wrap some data in a loader.
113 uint8_t data[256] = {};
114 BufferDataLoader edl(data, sizeof(data));
115
116 // Loading beyond the end of the data should fail.
117 {
118 Result<FreeableBuffer> fb = edl.load(
119 /*offset=*/0,
120 /*size=*/sizeof(data) + 1,
121 /*segment_info=*/
122 DataLoader::SegmentInfo(DataLoader::SegmentInfo::Type::Program));
123 EXPECT_NE(fb.error(), Error::Ok);
124 }
125
126 // Loading zero bytes still fails if it's past the end of the data.
127 {
128 Result<FreeableBuffer> fb = edl.load(
129 /*offset=*/sizeof(data) + 1,
130 /*size=*/0,
131 /*segment_info=*/
132 DataLoader::SegmentInfo(DataLoader::SegmentInfo::Type::Program));
133 EXPECT_NE(fb.error(), Error::Ok);
134 }
135 }
136
TEST_F(BufferDataLoaderTest,LoadIntoNullDstFails)137 TEST_F(BufferDataLoaderTest, LoadIntoNullDstFails) {
138 // Wrap some data in a loader.
139 uint8_t data[256] = {};
140 BufferDataLoader edl(data, sizeof(data));
141
142 // Loading beyond the end of the data should fail.
143 {
144 Result<FreeableBuffer> fb = edl.load_into(
145 /*offset=*/0,
146 /*size=*/1,
147 /*segment_info=*/
148 DataLoader::SegmentInfo(DataLoader::SegmentInfo::Type::Program),
149 nullptr);
150 EXPECT_NE(fb.error(), Error::Ok);
151 }
152
153 // Loading zero bytes still fails if dst is null.
154 {
155 Result<FreeableBuffer> fb = edl.load_into(
156 /*offset=*/0,
157 /*size=*/0,
158 /*segment_info=*/
159 DataLoader::SegmentInfo(DataLoader::SegmentInfo::Type::Program),
160 nullptr);
161 EXPECT_NE(fb.error(), Error::Ok);
162 }
163 }
164
TEST_F(BufferDataLoaderTest,InBoundsLoadIntoSucceeds)165 TEST_F(BufferDataLoaderTest, InBoundsLoadIntoSucceeds) {
166 // Wrap some data in a loader.
167 uint8_t data[256] = {};
168 data[0] = 1;
169 uint8_t buffer[256] = {};
170 buffer[0] = 0;
171 BufferDataLoader edl(data, sizeof(data));
172
173 {
174 // Buffer contains 0 before load_into.
175 EXPECT_EQ(buffer[0], 0);
176 Error fb = edl.load_into(
177 /*offset=*/0,
178 /*size=*/1,
179 /*segment_info=*/
180 DataLoader::SegmentInfo(DataLoader::SegmentInfo::Type::Program),
181 buffer);
182 EXPECT_EQ(fb, Error::Ok);
183 // Buffer contains 1 after load_into.
184 EXPECT_EQ(buffer[0], 1);
185 // Data is unaltered.
186 EXPECT_EQ(data[0], 1);
187 }
188 }
189