diff --git a/python/python/lancedb/functions.py b/python/python/lancedb/functions.py index 0e0e6bccf..c26a681d3 100644 --- a/python/python/lancedb/functions.py +++ b/python/python/lancedb/functions.py @@ -488,6 +488,7 @@ _FUNCTION_NAME = re.compile(r"^[A-Za-z_][A-Za-z0-9_.-]*$") _SECRET_NAME = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") # Keep this byte limit aligned with Sophon's MAX_FUNCTION_SECRET_VALUE_BYTES. _MAX_FUNCTION_SECRET_VALUE_BYTES = 64 * 1024 +_MAX_FUNCTION_SECRET_VALUES_BYTES = 512 * 1024 def _validate_secret_value(name: str, value: Any) -> str: @@ -1038,8 +1039,16 @@ class UdfDefinition: ) canonical_values = {} + total_bytes = 0 for name in sorted(secret_values): - canonical_values[name] = _validate_secret_value(name, secret_values[name]) + value = _validate_secret_value(name, secret_values[name]) + total_bytes += len(value.encode("utf-8")) + if total_bytes > _MAX_FUNCTION_SECRET_VALUES_BYTES: + raise ValueError( + "Function secret values exceed the " + f"{_MAX_FUNCTION_SECRET_VALUES_BYTES}-byte request limit" + ) + canonical_values[name] = value submission = self._request._known_dict() if canonical_values: diff --git a/python/python/tests/test_first_class_function_slice2.py b/python/python/tests/test_first_class_function_slice2.py index 983cbf780..d826aba7a 100644 --- a/python/python/tests/test_first_class_function_slice2.py +++ b/python/python/tests/test_first_class_function_slice2.py @@ -21,6 +21,7 @@ import pytest import lancedb from lancedb.functions import ( _MAX_FUNCTION_SECRET_VALUE_BYTES, + _MAX_FUNCTION_SECRET_VALUES_BYTES, FunctionRegistrationRequest, UdfDefinition, udf, @@ -782,6 +783,43 @@ def test_secret_value_rejects_over_utf8_byte_limit_before_json_construction( normalize_score._submission_json({"API_TOKEN": value}) +def test_secret_values_accept_exact_aggregate_utf8_byte_limit(monkeypatch): + names = tuple(f"SECRET_{index}" for index in range(8)) + value = "é" * (_MAX_FUNCTION_SECRET_VALUE_BYTES // len("é".encode("utf-8"))) + values = {name: value for name in names} + monkeypatch.setattr( + normalize_score, + "_request", + normalize_score._request._copy(update={"required_secrets": names}), + ) + + submission = json.loads(normalize_score._submission_json(values)) + + assert submission["secret_values"] == values + assert sum(len(item.encode("utf-8")) for item in values.values()) == ( + _MAX_FUNCTION_SECRET_VALUES_BYTES + ) + + +def test_secret_values_reject_aggregate_over_limit_before_construction(monkeypatch): + names = tuple(f"SECRET_{index}" for index in range(9)) + values = {name: "x" * _MAX_FUNCTION_SECRET_VALUE_BYTES for name in names} + monkeypatch.setattr( + normalize_score, + "_request", + normalize_score._request._copy(update={"required_secrets": names}), + ) + + def fail_if_json_construction_starts(self): + pytest.fail("oversized aggregate reached JSON construction") + + monkeypatch.setattr( + FunctionRegistrationRequest, "_known_dict", fail_if_json_construction_starts + ) + with pytest.raises(ValueError, match=r"exceed.*524288-byte request limit"): + normalize_score._submission_json(values) + + @pytest.mark.asyncio async def test_async_remote_registration_submits_secret_values_only_once(): with _mock_remote_function_catalog() as (host, state): diff --git a/rust/lancedb/src/function.rs b/rust/lancedb/src/function.rs index 2e0be3a88..5c06a4c22 100644 --- a/rust/lancedb/src/function.rs +++ b/rust/lancedb/src/function.rs @@ -19,6 +19,12 @@ use crate::{Error, Result}; pub(crate) const MAX_FUNCTION_SECRET_VALUE_BYTES: usize = 64 * 1024; const MAX_FUNCTION_SECRET_VALUES_BYTES: usize = 512 * 1024; +fn is_portable_environment_name(name: &str) -> bool { + let mut bytes = name.bytes(); + matches!(bytes.next(), Some(b'A'..=b'Z' | b'a'..=b'z' | b'_')) + && bytes.all(|byte| matches!(byte, b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'_')) +} + fn invalid_json(error: impl std::fmt::Display) -> Error { Error::InvalidInput { message: format!("invalid remote Function JSON: {error}"), @@ -433,7 +439,32 @@ pub struct FunctionRegistrationRequest { impl FunctionRegistrationRequest { pub(crate) fn validate_secret_values(&self) -> Result<()> { - let required = self.required_secrets.iter().collect::>(); + let mut required = BTreeSet::new(); + for name in &self.required_secrets { + if !is_portable_environment_name(name) { + return Err(Error::InvalidInput { + message: format!( + "Function secret name {name:?} must be a portable environment variable name" + ), + }); + } + if !required.insert(name) { + return Err(Error::InvalidInput { + message: format!("Function required_secrets contains duplicate name {name:?}"), + }); + } + } + + if let PythonRuntimeSpec::Python { env, .. } = &self.runtime + && let Some(name) = required.iter().find(|name| env.contains_key(**name)) + { + return Err(Error::InvalidInput { + message: format!( + "Function runtime env and secret names must be disjoint: {name:?}" + ), + }); + } + let provided = self.secret_values.keys().collect::>(); if required != provided { return Err(Error::InvalidInput { @@ -686,7 +717,7 @@ impl_json!(RefreshColumnResult); mod secret_value_tests { use super::{ FunctionRegistrationRequest, MAX_FUNCTION_SECRET_VALUE_BYTES, - MAX_FUNCTION_SECRET_VALUES_BYTES, + MAX_FUNCTION_SECRET_VALUES_BYTES, PythonRuntimeSpec, }; use crate::Error; @@ -732,6 +763,67 @@ mod secret_value_tests { )); } + #[test] + fn rejects_invalid_duplicate_and_overlapping_secret_declarations() { + let mut invalid_name = request(); + invalid_name.required_secrets = vec!["BAD=NAME".to_string()]; + invalid_name + .secret_values + .insert("BAD=NAME".to_string(), "secret".to_string()); + + let mut duplicate = request(); + duplicate.required_secrets = vec!["API_TOKEN".to_string(), "API_TOKEN".to_string()]; + duplicate + .secret_values + .insert("API_TOKEN".to_string(), "secret".to_string()); + + let mut overlap = request(); + overlap + .secret_values + .insert("API_TOKEN".to_string(), "secret".to_string()); + if let PythonRuntimeSpec::Python { env, .. } = &mut overlap.runtime { + env.insert("API_TOKEN".to_string(), "public".to_string()); + } + + assert!(matches!( + invalid_name.validate_secret_values(), + Err(Error::InvalidInput { message }) if message.contains("portable environment variable") + )); + assert!(matches!( + duplicate.validate_secret_values(), + Err(Error::InvalidInput { message }) if message.contains("duplicate") + )); + assert!(matches!( + overlap.validate_secret_values(), + Err(Error::InvalidInput { message }) if message.contains("must be disjoint") + )); + } + + #[test] + fn enforces_portable_secret_name_boundaries() { + for name in ["A", "_", "A0_"] { + let mut request = request(); + request.required_secrets = vec![name.to_string()]; + request + .secret_values + .insert(name.to_string(), "secret".to_string()); + request.validate_secret_values().unwrap(); + } + + for name in ["", "0TOKEN", "BAD-NAME", "TÖKEN"] { + let mut request = request(); + request.required_secrets = vec![name.to_string()]; + request + .secret_values + .insert(name.to_string(), "secret".to_string()); + assert!(matches!( + request.validate_secret_values(), + Err(Error::InvalidInput { message }) + if message.contains("portable environment variable") + )); + } + } + #[test] fn accepts_exact_secret_value_utf8_byte_limit() { for value in [ @@ -777,6 +869,27 @@ mod secret_value_tests { if message.contains(&format!("{MAX_FUNCTION_SECRET_VALUES_BYTES}-byte request limit")) )); } + + #[test] + fn accepts_exact_aggregate_secret_value_byte_limit() { + let mut request = request(); + request.required_secrets = (0..8).map(|index| format!("SECRET_{index}")).collect(); + request.secret_values = request + .required_secrets + .iter() + .map(|name| (name.clone(), "x".repeat(MAX_FUNCTION_SECRET_VALUE_BYTES))) + .collect(); + + assert_eq!( + request + .secret_values + .values() + .map(String::len) + .sum::(), + MAX_FUNCTION_SECRET_VALUES_BYTES + ); + request.validate_secret_values().unwrap(); + } } #[cfg(test)] diff --git a/rust/lancedb/src/remote/client.rs b/rust/lancedb/src/remote/client.rs index 2e0b410f6..da4e139b1 100644 --- a/rust/lancedb/src/remote/client.rs +++ b/rust/lancedb/src/remote/client.rs @@ -894,7 +894,7 @@ impl RestfulLanceDbClient { } } - pub(crate) fn log_request(&self, request: &Request, request_id: &String) { + pub(crate) fn log_request(&self, request: &Request, request_id: &str) { if log::log_enabled!(log::Level::Debug) { debug!("{}", request_log_message(request, request_id)); }