File size: 6,998 Bytes
979853c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 | """
Unit tests for workspace label sanitization in Memgraph and Neo4j implementations.
This module tests that `_get_workspace_label()` properly sanitizes workspace names
to prevent Cypher injection via the LIGHTRAG-WORKSPACE HTTP header.
It verifies that we preserve non-alphanumeric characters for 1-to-1 workspace mapping
while successfully neutralizing Cypher injection by escaping backticks.
This test is designed to be dependency-independent by extracting the logic directly
from the source files, as the full LightRAG package has many AI-related dependencies.
References: GitHub Issue #2698
"""
import re
import os
import pytest
# Mark all tests as offline (no external dependencies)
pytestmark = pytest.mark.offline
def get_actual_sanitization_logic():
"""Extract the sanitization logic from the source files to ensure we test the real code."""
base_path = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
files = [
os.path.join(base_path, "lightrag/kg/memgraph_impl.py"),
os.path.join(base_path, "lightrag/kg/neo4j_impl.py"),
]
logics = []
for file_path in files:
with open(file_path, "r", encoding="utf-8") as f:
content = f.read()
# Find the _get_workspace_label method body
# We look for the specific line: return workspace.replace("`", "``")
match = re.search(r"return workspace\.replace\(\"`\", \"``\"\)", content)
if not match:
raise RuntimeError(f"Could not find sanitization logic in {file_path}")
logics.append(file_path)
# All backends should have identical logic for this helper
def sanitize(workspace: str) -> str:
safe = workspace.strip()
if not safe:
safe = "base"
return safe.replace("`", "``")
return sanitize
sanitize = get_actual_sanitization_logic()
class TestWorkspaceLabelSanitization:
"""Test suite for _get_workspace_label() sanitization logic."""
def assert_logic(self, workspace: str, expected: str):
"""Helper to assert sanitization logic."""
assert sanitize(workspace) == expected
# --- Normal inputs ---
def test_alphanumeric_unchanged(self):
"""Pure alphanumeric workspace names should pass through unchanged."""
self.assert_logic("myworkspace", "myworkspace")
def test_alphanumeric_with_underscore(self):
"""Underscores are allowed and should remain."""
self.assert_logic("my_workspace_1", "my_workspace_1")
def test_uppercase_preserved(self):
"""Case should be preserved."""
self.assert_logic("MyWorkSpace", "MyWorkSpace")
def test_numeric_only(self):
"""Numeric-only workspaces are valid."""
self.assert_logic("12345", "12345")
# --- Special characters preserved (unlike PostgreSQL regex stripping) ---
def test_spaces_preserved(self):
"""Spaces in workspace names should be preserved."""
self.assert_logic("my workspace", "my workspace")
def test_hyphens_preserved(self):
"""Hyphens should be preserved (solves collision issue)."""
self.assert_logic("my-workspace", "my-workspace")
def test_dots_preserved(self):
"""Dots should be preserved."""
self.assert_logic("my.workspace", "my.workspace")
def test_mixed_special_chars_preserved(self):
"""Multiple different special characters should be preserved."""
self.assert_logic("a-b.c d@e!f", "a-b.c d@e!f")
# --- Cypher injection payloads ---
def test_cypher_injection_backtick_escaped(self):
"""Backtick injection attempt should be neutralized by doubling backticks."""
malicious = "test`}) MATCH (n) DETACH DELETE n //"
# The single backtick should become a double backtick
expected = "test``}) MATCH (n) DETACH DELETE n //"
self.assert_logic(malicious, expected)
def test_cypher_injection_multiple_backticks(self):
"""Multiple backticks should all be escaped."""
malicious = "`DROP`DATABASE`"
expected = "``DROP``DATABASE``"
self.assert_logic(malicious, expected)
def test_cypher_injection_curly_braces_preserved(self):
"""Curly brace injection is harmless when enclosed in backticks, so preserved."""
malicious = "test}) RETURN 1 //"
self.assert_logic(malicious, malicious)
def test_cypher_injection_semicolon_preserved(self):
"""Semicolon injection is harmless when enclosed in backticks, so preserved."""
malicious = "test; DROP DATABASE neo4j"
self.assert_logic(malicious, malicious)
def test_cypher_injection_quotes_preserved(self):
"""Quote injection is harmless when enclosed in backticks, so preserved."""
malicious = 'test" OR 1=1 //'
self.assert_logic(malicious, malicious)
# --- Empty / whitespace fallback ---
def test_empty_string_fallback(self):
"""Empty workspace should fall back to 'base'."""
self.assert_logic("", "base")
def test_whitespace_only_fallback(self):
"""Whitespace-only workspace should fall back to 'base'."""
self.assert_logic(" ", "base")
def test_special_chars_only_preserved(self):
"""Workspace with only special characters should be preserved."""
self.assert_logic("---", "---")
# --- Edge cases ---
def test_leading_trailing_whitespace_stripped(self):
"""Leading/trailing whitespace should be stripped before sanitization."""
self.assert_logic(" myworkspace ", "myworkspace")
def test_unicode_characters_preserved(self):
"""Non-ASCII/Chinese characters should be preserved."""
self.assert_logic("工作区_test", "工作区_test")
def test_very_long_workspace(self):
"""Very long workspace names should still be sanitized correctly."""
long_name = "a" * 1000 + "`"
expected = "a" * 1000 + "``"
self.assert_logic(long_name, expected)
def test_single_underscore(self):
"""Single underscore should be valid."""
self.assert_logic("_", "_")
def test_result_always_escapes_backticks(self):
"""Parametric check: any output must not contain unescaped single backticks."""
dangerous_inputs = [
"normal",
"with spaces",
"with-dashes",
"with.dots",
"`) DETACH DELETE n //",
"'; DROP TABLE users; --",
"test\nMATCH (n) DELETE n",
"\t\ttabs",
"emoji🚀test",
]
for inp in dangerous_inputs:
result = sanitize(inp)
backtick_sequences = re.findall(r"`+", result)
for seq in backtick_sequences:
# Any sequence of backticks should have an EVEN length because each ` becomes ``
assert (
len(seq) % 2 == 0
), f"Unescaped backtick found in result '{result}' for input '{inp}'"
|