Skip to content

Commit 4a5a7ac

Browse files
[FIX] Prevent keys from starting with slashes (#172)
* [FIX] Prevent keys from starting with slashes If a key starts with a slash, then it becomes undeletable and prevents database purges from working properly as well. This prevents that from occuring by stripping slashes from the left of the key name. * Double newlines for flake8 * flake8 wanted another newline here * Force `set_bulk_raw` to handle keys with slashes as well * Add tests for keys starting with a slash * Fix a typo I made twice * flake8 * `del self.db[k]` not `self.db.delete(k)` in non-Async * One space for flake8 * These were also wrong * These shouldn't be using `get` * Match format of some of the other tests in TestDatabase * Perhaps the key is corrupted? * Have to `get_raw` for `_raw` calls. * Reassociate _dumps with def dumps * Only call keyStrip at the root of the .set function hierarchy * Clarify that keyStrip is an internal method --------- Co-authored-by: Devon Stewart <[email protected]>
1 parent 5c0ab7a commit 4a5a7ac

File tree

2 files changed

+69
-0
lines changed

2 files changed

+69
-0
lines changed

src/replit/database/database.py

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,18 @@ def dumps(val: Any) -> str:
6161
_dumps = dumps
6262

6363

64+
def _sanitize_key(key: str) -> str:
65+
"""Strip slashes from the beginning of keys.
66+
67+
Args:
68+
key (str): The key to strip
69+
70+
Returns:
71+
str: The stripped key
72+
"""
73+
return key.lstrip("/")
74+
75+
6476
class AsyncDatabase:
6577
"""Async interface for Replit Database.
6678
@@ -195,6 +207,7 @@ async def set_bulk_raw(self, values: Dict[str, str]) -> None:
195207
Args:
196208
values (Dict[str, str]): The key-value pairs to set.
197209
"""
210+
values = {_sanitize_key(k): v for k, v in values.items()}
198211
async with self.client.post(self.db_url, data=values) as response:
199212
response.raise_for_status()
200213

@@ -629,6 +642,7 @@ def set_bulk_raw(self, values: Dict[str, str]) -> None:
629642
Args:
630643
values (Dict[str, str]): The key-value pairs to set.
631644
"""
645+
values = {_sanitize_key(k): v for k, v in values.items()}
632646
r = self.sess.post(self.db_url, data=values)
633647
r.raise_for_status()
634648

tests/test_database.py

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -125,6 +125,33 @@ async def test_bulk_raw(self) -> None:
125125
self.assertEqual(await self.db.get_raw("bulk1"), "val1")
126126
self.assertEqual(await self.db.get_raw("bulk2"), "val2")
127127

128+
async def test_slash_keys(self) -> None:
129+
"""Test that slash keys work."""
130+
k = "/key"
131+
# set
132+
await self.db.set(k,"val1")
133+
self.assertEqual(await self.db.get(k), "val1")
134+
await self.db.delete(k)
135+
with self.assertRaises(KeyError):
136+
await self.db.get(k)
137+
# set_raw
138+
await self.db.set_raw(k,"val1")
139+
self.assertEqual(await self.db.get_raw(k), "val1")
140+
await self.db.delete(k)
141+
with self.assertRaises(KeyError):
142+
await self.db.get(k)
143+
# set_bulk
144+
await self.db.set_bulk({k: "val1"})
145+
self.assertEqual(await self.db.get(k), "val1")
146+
await self.db.delete(k)
147+
with self.assertRaises(KeyError):
148+
await self.db.get(k)
149+
# set_bulk_raw
150+
await self.db.set_bulk_raw({k: "val1"})
151+
self.assertEqual(await self.db.get_raw(k), "val1")
152+
await self.db.delete(k)
153+
with self.assertRaises(KeyError):
154+
await self.db.get(k)
128155

129156
class TestDatabase(unittest.TestCase):
130157
"""Tests for replit.database.Database."""
@@ -259,3 +286,31 @@ def test_bulk_raw(self) -> None:
259286
self.db.set_bulk_raw({"bulk1": "val1", "bulk2": "val2"})
260287
self.assertEqual(self.db.get_raw("bulk1"), "val1")
261288
self.assertEqual(self.db.get_raw("bulk2"), "val2")
289+
290+
def test_slash_keys(self) -> None:
291+
"""Test that slash keys work."""
292+
k = "/key"
293+
# set
294+
self.db.set(k,"val1")
295+
self.assertEqual(self.db[k], "val1")
296+
del self.db[k]
297+
with self.assertRaises(KeyError):
298+
self.db[k]
299+
# set_raw
300+
self.db.set_raw(k,"val1")
301+
self.assertEqual(self.db.get_raw(k), "val1")
302+
del self.db[k]
303+
with self.assertRaises(KeyError):
304+
self.db[k]
305+
# set_bulk
306+
self.db.set_bulk({k: "val1"})
307+
self.assertEqual(self.db.get(k), "val1")
308+
del self.db[k]
309+
with self.assertRaises(KeyError):
310+
self.db[k]
311+
# set_bulk_raw
312+
self.db.set_bulk_raw({k: "val1"})
313+
self.assertEqual(self.db.get_raw(k), "val1")
314+
del self.db[k]
315+
with self.assertRaises(KeyError):
316+
self.db[k]

0 commit comments

Comments
 (0)