Merge branch 'kv-validate'

Author: wallyqs

.github/workflows/check.yml

@@ -2,7 +2,7 @@ name: Check
 on:
   push:
     branches:
-      - main
+      - "*"
   pull_request:
     branches:
       - "*"

.github/workflows/test.yml

@@ -3,9 +3,7 @@ name: test
 on:
     push:
         branches:
-            - main
-            - "release/**"
-            - "dev/**"
+            - "**"
     pull_request:
         branches:
             - "**"
@@ -20,7 +18,7 @@ jobs:
             fail-fast: false
             matrix:
                 python-version: ["3.8", "3.9", "3.10", "3.11", "3.12", "3.13"]
-                nats_version: ["v2.10.26", "v2.11.0", "main"]
+                nats_version: ["v2.10.29", "v2.11.4", "main"]
 
         steps:
             - name: Check out repository

nats/js/errors.py

@@ -254,6 +254,13 @@ def __str__(self) -> str:
         return "nats: history limited to a max of 64"
 
 
+class InvalidKeyError(Error):
+    """
+    Raised when trying to put an object in Key Value with an invalid key.
+    """
+    pass
+
+
 class InvalidBucketNameError(Error):
     """
     Raised when trying to create a KV or OBJ bucket with invalid name.

nats/js/kv.py

@@ -17,6 +17,7 @@
 import asyncio
 import datetime
 import logging
+import re
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, List, Optional
 
@@ -34,6 +35,14 @@
 
 logger = logging.getLogger(__name__)
 
+VALID_KEY_RE = re.compile(r'^[-/_=\.a-zA-Z0-9]+$')
+
+
+def _is_key_valid(key: str) -> bool:
+    if len(key) == 0 or key[0] == '.' or key[-1] == '.':
+        return False
+    return bool(VALID_KEY_RE.match(key))
+
 
 class KeyValue:
     """
@@ -124,10 +133,18 @@ def __init__(
         self._js = js
         self._direct = direct
 
-    async def get(self, key: str, revision: Optional[int] = None) -> Entry:
+    async def get(
+        self,
+        key: str,
+        revision: Optional[int] = None,
+        validate_keys: bool = True
+    ) -> Entry:
         """
         get returns the latest value for the key.
         """
+        if validate_keys and not _is_key_valid(key):
+            raise nats.js.errors.InvalidKeyError
+
         entry = None
         try:
             entry = await self._get(key, revision)
@@ -179,21 +196,33 @@ async def _get(self, key: str, revision: Optional[int] = None) -> Entry:
 
         return entry
 
-    async def put(self, key: str, value: bytes) -> int:
+    async def put(
+        self, key: str, value: bytes, validate_keys: bool = True
+    ) -> int:
         """
         put will place the new value for the key into the store
         and return the revision number.
         """
+        if validate_keys and not _is_key_valid(key):
+            raise nats.js.errors.InvalidKeyError(key)
+
         pa = await self._js.publish(f"{self._pre}{key}", value)
         return pa.seq
 
-    async def create(self, key: str, value: bytes) -> int:
+    async def create(
+        self, key: str, value: bytes, validate_keys: bool = True
+    ) -> int:
         """
         create will add the key/value pair iff it does not exist.
         """
+        if validate_keys and not _is_key_valid(key):
+            raise nats.js.errors.InvalidKeyError(key)
+
         pa = None
         try:
-            pa = await self.update(key, value, last=0)
+            pa = await self.update(
+                key, value, last=0, validate_keys=validate_keys
+            )
         except nats.js.errors.KeyWrongLastSequenceError as err:
             # In case of attempting to recreate an already deleted key,
             # the client would get a KeyWrongLastSequenceError.  When this happens,
@@ -213,16 +242,28 @@ async def create(self, key: str, value: bytes) -> int:
                 # to recreate using the last revision.
                 raise err
             except nats.js.errors.KeyDeletedError as err:
-                pa = await self.update(key, value, last=err.entry.revision)
+                pa = await self.update(
+                    key,
+                    value,
+                    last=err.entry.revision,
+                    validate_keys=validate_keys
+                )
 
         return pa
 
     async def update(
-        self, key: str, value: bytes, last: Optional[int] = None
+        self,
+        key: str,
+        value: bytes,
+        last: Optional[int] = None,
+        validate_keys: bool = True
     ) -> int:
         """
         update will update the value iff the latest revision matches.
         """
+        if validate_keys and not _is_key_valid(key):
+            raise nats.js.errors.InvalidKeyError(key)
+
         hdrs = {}
         if not last:
             last = 0
@@ -243,10 +284,18 @@ async def update(
                 raise err
         return pa.seq
 
-    async def delete(self, key: str, last: Optional[int] = None) -> bool:
+    async def delete(
+        self,
+        key: str,
+        last: Optional[int] = None,
+        validate_keys: bool = True
+    ) -> bool:
         """
         delete will place a delete marker and remove all previous revisions.
         """
+        if validate_keys and not _is_key_valid(key):
+            raise nats.js.errors.InvalidKeyError(key)
+
         hdrs = {}
         hdrs[KV_OP] = KV_DEL
 

tests/test_js.py

@@ -2689,6 +2689,111 @@ async def error_handler(e):
         with pytest.raises(BadBucketError):
             await js.key_value(bucket="TEST3")
 
+    @async_test
+    async def test_bucket_name_validation(self):
+        nc = await nats.connect()
+        js = nc.jetstream()
+
+        invalid_bucket_names = [
+            " x y",
+            "x ",
+            "x!",
+            "xx$",
+            "*",
+            ">",
+            "x.>",
+            "x.*",
+            ".",
+            ".x",
+            ".x.",
+            "x.",
+        ]
+
+        for bucket_name in invalid_bucket_names:
+            with self.subTest(bucket_name):
+                with pytest.raises(InvalidBucketNameError):
+                    await js.create_key_value(
+                        bucket=bucket_name, history=5, ttl=3600
+                    )
+
+                with pytest.raises(InvalidBucketNameError):
+                    await js.key_value(bucket_name)
+
+                with pytest.raises(InvalidBucketNameError):
+                    await js.delete_key_value(bucket_name)
+
+    @async_test
+    async def test_key_validation(self):
+        nc = await nats.connect()
+        js = nc.jetstream()
+
+        kv = await js.create_key_value(bucket="TEST", history=5, ttl=3600)
+        invalid_keys = [
+            " x y",
+            "x ",
+            "x!",
+            "xx$",
+            "*",
+            ">",
+            "x.>",
+            "x.*",
+            ".",
+            ".x",
+            ".x.",
+            "x.",
+        ]
+
+        for key in invalid_keys:
+            with self.subTest(key):
+                # Invalid put (empty)
+                with pytest.raises(InvalidKeyError):
+                    await kv.put(key, b'')
+
+                with pytest.raises(InvalidKeyError):
+                    await kv.get(key)
+
+                with pytest.raises(InvalidKeyError):
+                    await kv.update(key, b'')
+
+    @async_test
+    async def test_key_validation_bypass(self):
+        nc = await nats.connect()
+        js = nc.jetstream()
+
+        kv = await js.create_key_value(bucket="TEST", history=5, ttl=3600)
+        invalid_keys = [
+            "x!",
+            "x.>",
+            "x.*",
+        ]
+
+        for key in invalid_keys:
+            with self.subTest(key):
+                # Should succeed with validate_keys=False
+                seq = await kv.put(key, b'test_value', validate_keys=False)
+                assert seq > 0
+
+                # Should be able to get with validate_keys=False
+                entry = await kv.get(key, validate_keys=False)
+                assert entry.value == b'test_value'
+
+                # Should be able to update with validate_keys=False
+                seq2 = await kv.update(
+                    key, b'updated_value', last=seq, validate_keys=False
+                )
+                assert seq2 > seq
+
+                # Should be able to delete with validate_keys=False
+                result = await kv.delete(key, validate_keys=False)
+                assert result is True
+
+                # Should still fail with default validate_keys=True
+                with pytest.raises(InvalidKeyError):
+                    await kv.put(key, b'fail')
+
+                with pytest.raises(InvalidKeyError):
+                    await kv.get(key)
+
     @async_test
     async def test_kv_basic(self):
         errors = []
@@ -2700,6 +2805,9 @@ async def error_handler(e):
         nc = await nats.connect(error_cb=error_handler)
         js = nc.jetstream()
 
+        with pytest.raises(nats.js.errors.InvalidBucketNameError):
+            await js.create_key_value(bucket="notok!")
+
         bucket = "TEST"
         kv = await js.create_key_value(
             bucket=bucket,