diff --git a/pyrit/converter/token_smuggling/variation_selector_smuggler_converter.py b/pyrit/converter/token_smuggling/variation_selector_smuggler_converter.py index ad2ecd1a57..5b874a98ce 100644 --- a/pyrit/converter/token_smuggling/variation_selector_smuggler_converter.py +++ b/pyrit/converter/token_smuggling/variation_selector_smuggler_converter.py @@ -47,10 +47,17 @@ def __init__( Default is True. Raises: - ValueError: If an unsupported action or ``encoding_mode`` is provided. + ValueError: If an unsupported action is provided, or ``base_char_utf8`` is not exactly one character + or is a variation selector used to encode payload bytes. """ super().__init__(action=action) - self.utf8_base_char = base_char_utf8 if base_char_utf8 is not None else "😊" + base_char = base_char_utf8 if base_char_utf8 is not None else "😊" + if len(base_char) != 1: + raise ValueError("base_char_utf8 must be exactly one character.") + code_point = ord(base_char) + if 0xFE00 <= code_point <= 0xFE0F or 0xE0100 <= code_point <= 0xE01EF: + raise ValueError("base_char_utf8 must not be a variation selector.") + self.utf8_base_char = base_char self.embed_in_base = embed_in_base def _build_identifier(self) -> ComponentIdentifier: diff --git a/tests/unit/converter/test_variation_selector_smuggler_converter.py b/tests/unit/converter/test_variation_selector_smuggler_converter.py index 59e242093f..527fda532b 100644 --- a/tests/unit/converter/test_variation_selector_smuggler_converter.py +++ b/tests/unit/converter/test_variation_selector_smuggler_converter.py @@ -43,6 +43,31 @@ def test_variation_selector_invalid_action(): VariationSelectorSmugglerConverter(action="invalid") +@pytest.mark.parametrize("base_char", ["", "ab", "😊x"]) +def test_variation_selector_invalid_base_char(base_char: str) -> None: + with pytest.raises(ValueError, match="base_char_utf8 must be exactly one character"): + VariationSelectorSmugglerConverter(base_char_utf8=base_char) + + +@pytest.mark.parametrize("base_char", ["\ufe00", "\ufe0f", "\U000e0100", "\U000e01ef"]) +def test_variation_selector_rejects_selector_base_char(base_char: str) -> None: + with pytest.raises(ValueError, match="base_char_utf8 must not be a variation selector"): + VariationSelectorSmugglerConverter(base_char_utf8=base_char) + + +@pytest.mark.parametrize("base_char", ["A", "\ufdff", "\ufe10", "\U000e00ff", "\U000e01f0"]) +@pytest.mark.parametrize("embed_in_base", [True, False]) +async def test_variation_selector_custom_base_char_roundtrip(base_char: str, embed_in_base: bool) -> None: + encoder = VariationSelectorSmugglerConverter(action="encode", base_char_utf8=base_char, embed_in_base=embed_in_base) + encoded = await encoder.convert_async(prompt="test", input_type="text") + + decoder = VariationSelectorSmugglerConverter(action="decode", base_char_utf8=base_char, embed_in_base=embed_in_base) + decoded = await decoder.convert_async(prompt=encoded.output_text, input_type="text") + + assert encoded.output_text.startswith(base_char) + assert decoded.output_text == "test" + + async def test_variation_selector_input_not_supported(): converter = VariationSelectorSmugglerConverter(action="encode") with pytest.raises(ValueError, match="Input type not supported"):