1
// Copyright (c) 2026 Cloudflare, Inc.
2
// Licensed under the Apache 2.0 license found in the LICENSE file or at:
3
//     https://opensource.org/licenses/Apache-2.0
4
 
5
#include "sqlite-metering.h"
6
 
7
#include <kj/test.h>
8
#include <kj/thread.h>
9
 
10
#include <atomic>
11
 
12
namespace workerd {
13
namespace {
14
 
15
void test_sqlite3_mem_methods(bool expectEnforcedLimits) {
16
  // sqliteMemMalloc with a negative size returns a nullptr.
17
  KJ_EXPECT(sqliteMemMalloc(-1) == nullptr);
18
 
19
  // sqliteMemMalloc with a zero size returns a nullptr.
20
  KJ_EXPECT(sqliteMemMalloc(0) == nullptr);
21
 
22
  // sqliteMemMalloc with a positive size returns a non-nullptr.
23
  void* ptr = sqliteMemMalloc(1024);
24
  KJ_EXPECT(ptr != nullptr);
25
 
26
  // sqliteMemFree on a non-nullptr does not throw.
27
  sqliteMemFree(ptr);
28
 
29
  // sqliteMemMalloc with a larger size. Note that the size allocated here needs to be less than 1
30
  // MiB below, or the pointer address will be reused when requesting 1 MiB.
31
  ptr = sqliteMemMalloc(1024 * 512 + 1);
32
  if (expectEnforcedLimits) {
33
    KJ_EXPECT(ptr == nullptr);
34
  } else {
35
    KJ_EXPECT(ptr != nullptr);
36
  }
37
  sqliteMemFree(ptr);
38
 
39
  // sqliteMemFree on a nullptr does not throw.
40
  sqliteMemFree(nullptr);
41
 
42
  // sqliteMemRealloc(nullptr, <n>) should behave as sqliteMemMalloc(<n>).
43
  ptr = sqliteMemRealloc(nullptr, -1);
44
  KJ_EXPECT(ptr == nullptr);
45
  ptr = sqliteMemRealloc(nullptr, 0);
46
  KJ_EXPECT(ptr == nullptr);
47
  ptr = sqliteMemRealloc(nullptr, 1024);
48
  KJ_EXPECT(ptr != nullptr);
49
  sqliteMemFree(ptr);
50
 
51
  // sqliteMemRealloc should return a different pointer if a larger size is requested.
52
  ptr = sqliteMemMalloc(1024);
53
  void* new_ptr = sqliteMemRealloc(ptr, 1024 * 1024);
54
  KJ_EXPECT(ptr != new_ptr);
55
  KJ_EXPECT(ptr != nullptr);
56
  if (expectEnforcedLimits) {
57
    // Realloc was rejected: new_ptr is nullptr and ptr still owns the original allocation.
58
    KJ_EXPECT(new_ptr == nullptr);
59
    sqliteMemFree(ptr);
60
  } else {
61
    KJ_EXPECT(new_ptr != nullptr);
62
    sqliteMemFree(new_ptr);
63
  }
64
 
65
  // sqliteMemRealloc with a valid pointer and negative size returns a nullptr. Note that ptr is
66
  // implicitly freed by the preceding call to sqliteMemRealloc.
67
  ptr = sqliteMemMalloc(1024);
68
  KJ_EXPECT(ptr != nullptr);
69
  new_ptr = sqliteMemRealloc(ptr, -1);
70
  KJ_EXPECT(new_ptr == nullptr);
71
  // sqliteMemRealloc with a valid pointer and zero size returns a nullptr. Note again that ptr is
72
  // implicitly freed by the preceding call to sqliteMemRealloc.
73
  ptr = sqliteMemMalloc(1024);
74
  KJ_EXPECT(ptr != nullptr);
75
  new_ptr = sqliteMemRealloc(ptr, 0);
76
  KJ_EXPECT(new_ptr == nullptr);
77
}
78
 
79
KJ_TEST("sqlite3_mem_methods work without SqliteMemoryScope") {
80
  // We expect a bunch of error logs, but we do not crash.
81
  KJ_EXPECT_LOG(ERROR, "sqliteMemMalloc did not find a valid scope.");
82
  KJ_EXPECT_LOG(ERROR, "sqliteMemFree did not find a valid scope.");
83
  KJ_EXPECT_LOG(ERROR, "sqliteMemMalloc did not find a valid scope.");
84
  KJ_EXPECT_LOG(ERROR, "sqliteMemFree did not find a valid scope.");
85
  KJ_EXPECT_LOG(ERROR, "sqliteMemMalloc did not find a valid scope.");
86
  KJ_EXPECT_LOG(ERROR, "sqliteMemFree did not find a valid scope.");
87
  KJ_EXPECT_LOG(ERROR, "sqliteMemMalloc did not find a valid scope.");
88
  KJ_EXPECT_LOG(ERROR, "sqliteMemRealloc did not find a valid scope.");
89
  KJ_EXPECT_LOG(ERROR, "sqliteMemFree did not find a valid scope.");
90
  KJ_EXPECT_LOG(ERROR, "sqliteMemMalloc did not find a valid scope.");
91
  KJ_EXPECT_LOG(ERROR, "sqliteMemFree did not find a valid scope.");
92
  KJ_EXPECT_LOG(ERROR, "sqliteMemMalloc did not find a valid scope.");
93
  KJ_EXPECT_LOG(ERROR, "sqliteMemFree did not find a valid scope.");
94
  test_sqlite3_mem_methods(/*expectEnforcedLimits*/ false);
95
}
96
 
97
KJ_TEST("sqlite3_mem_methods work with SqliteMemoryScope and high limits") {
98
  size_t memoryBytes = 0;
99
  SqliteMemoryScope scope(memoryBytes, 2 * 1024 * 1024);
100
  test_sqlite3_mem_methods(/*expectEnforcedLimits*/ false);
101
}
102
 
103
KJ_TEST("sqlite3_mem_methods work with SqliteMemoryScope and low limits") {
104
  size_t memoryBytes = 0;
105
  SqliteMemoryScope scope(memoryBytes, 512 * 1024);
106
  test_sqlite3_mem_methods(/*expectEnforcedLimits*/ true);
107
}
108
 
109
KJ_TEST("sqlite3_mem_methods correctly modify the SqliteMemoryScope memoryBytes") {
110
  auto runThread = [&]() {
111
    size_t memoryBytes = 0;
112
    SqliteMemoryScope scope(memoryBytes, 1024 * 1024);
113
 
114
    // sqliteMemMalloc followed by sqliteMemFree should result in a zero memoryBytes.
115
    void* p1 = sqliteMemMalloc(1024);
116
    KJ_EXPECT(p1 != nullptr);
117
    KJ_EXPECT(memoryBytes >= 1024);
118
    sqliteMemFree(p1);
119
    KJ_EXPECT(memoryBytes == 0);
120
 
121
    // Multiple sqliteMemMalloc's should increase memoryBytes respectively.
122
    void* p2 = sqliteMemMalloc(4096);
123
    KJ_EXPECT(p2 != nullptr);
124
    KJ_EXPECT(memoryBytes >= 4096);
125
    void* p3 = sqliteMemMalloc(4096);
126
    KJ_EXPECT(p3 != nullptr);
127
    KJ_EXPECT(memoryBytes >= 8192);
128
    size_t memoryBytesSnapshot = memoryBytes;
129
 
130
    // sqliteMemMalloc leaves memoryBytes unchanged when hardLimitBytes is exceeded.
131
    KJ_EXPECT(sqliteMemMalloc(1024 * 1024) == nullptr);
132
    KJ_EXPECT(memoryBytes == memoryBytesSnapshot);
133
 
134
    // sqliteMemRealloc with an increase in size should increase memoryBytes respectively, but less than
135
    // a sqliteMemMalloc of equivalent size.
136
    void* p4 = sqliteMemRealloc(p2, 8192);
137
    KJ_EXPECT(p4 != nullptr);
138
    KJ_EXPECT(memoryBytes >= 12288);
139
    KJ_EXPECT(memoryBytes <= 16384);
140
 
141
    // sqliteMemRealloc with an decrease in size should decrease memoryBytes respectively.
142
    void* p5 = sqliteMemRealloc(p3, 1024);
143
    KJ_EXPECT(p5 != nullptr);
144
    KJ_EXPECT(memoryBytes >= 9216);
145
    KJ_EXPECT(memoryBytes <= 12288);
146
    memoryBytesSnapshot = memoryBytes;
147
 
148
    // sqliteMemRealloc leaves memoryBytes unchanged when hardLimitBytes is exceeded.
149
    KJ_EXPECT(sqliteMemRealloc(p5, 1024 * 1024) == nullptr);
150
    KJ_EXPECT(memoryBytes == memoryBytesSnapshot);
151
 
152
    // sqliteMemFree on all active pointers should result in a zero memoryBytes.
153
    sqliteMemFree(p4);
154
    sqliteMemFree(p5);
155
    KJ_EXPECT(memoryBytes == 0);
156
  };
157
 
158
  kj::Thread thread1(runThread);
159
  kj::Thread thread2(runThread);
160
}
161
 
162
KJ_TEST("SqliteMemoryScope is idempotent and handles nested scopes") {
163
  size_t memoryBytes = 0;
164
  SqliteMemoryScope outerScope(memoryBytes, 1024 * 1024);
165
 
166
  void* p1 = sqliteMemMalloc(1024);
167
  KJ_EXPECT(p1 != nullptr);
168
  size_t memoryBytesSnapshot = memoryBytes;
169
  KJ_EXPECT(memoryBytesSnapshot >= 1024);
170
 
171
  {
172
    size_t innerMemoryBytes = 0;
173
    SqliteMemoryScope innerScope(innerMemoryBytes, 512);
174
 
175
    void* p2 = sqliteMemMalloc(2048);
176
    KJ_EXPECT(p2 != nullptr);
177
    KJ_EXPECT(memoryBytes > memoryBytesSnapshot);
178
    KJ_EXPECT(innerMemoryBytes == 0);
179
 
180
    sqliteMemFree(p2);
181
  }
182
 
183
  void* p3 = sqliteMemMalloc(512);
184
  KJ_EXPECT(p3 != nullptr);
185
  KJ_EXPECT(memoryBytes >= memoryBytesSnapshot);
186
 
187
  sqliteMemFree(p1);
188
  sqliteMemFree(p3);
189
  KJ_EXPECT(memoryBytes == 0);
190
}
191
 
192
}  // namespace
193
}  // namespace workerd