FIX: decode variation selector payloads next to emoji with presentation selectors

pyrit/converter/token_smuggling/variation_selector_smuggler_converter.py

@@ -10,6 +10,11 @@
 logger = logging.getLogger(__name__)
 
 
+def _is_variation_selector(char: str) -> bool:
+    code_point = ord(char)
+    return 0xFE00 <= code_point <= 0xFE0F or 0xE0100 <= code_point <= 0xE01EF
+
+
 class VariationSelectorSmugglerConverter(SmugglerConverter):
     """
     Encodes and decodes text using Unicode Variation Selectors.
@@ -54,8 +59,7 @@ def __init__(
         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:
+        if _is_variation_selector(base_char):
             raise ValueError("base_char_utf8 must not be a variation selector.")
         self.utf8_base_char = base_char
         self.embed_in_base = embed_in_base
@@ -111,14 +115,28 @@ def encode_message(self, message: str) -> tuple[str, str]:
     def decode_message(self, message: str) -> str:
         """
         Decode a message encoded using Unicode variation selectors.
-        The decoder scans the string for variation selectors, ignoring any visible separator.
+
+        The decoder reads the variation selectors that follow the base character, with or without the
+        space separator. Selectors elsewhere in the text, such as the U+FE0F in an emoji like ❤️, are not part
+        of the payload. If no base character carries a payload, the first run of variation selectors
+        in the string is decoded.
 
         Args:
             message (str): The encoded message.
 
         Returns:
             str: The decoded message.
         """
+        start = message.find(self.utf8_base_char)
+        while start != -1:
+            payload_start = start + 1
+            if message[payload_start : payload_start + 1] == " ":
+                payload_start += 1
+            if payload_start < len(message) and _is_variation_selector(message[payload_start]):
+                message = message[payload_start:]
+                break
+            start = message.find(self.utf8_base_char, start + 1)
+
         bytes_out = bytearray()
         started = False
         for char in message:
@@ -169,17 +187,20 @@ def decode_visible_hidden(self, combined: str) -> tuple[str, str]:
         """
         Extract the visible text and decodes the hidden text from a combined string.
 
-        It searches for the first occurrence of the base character (``self.utf8_base_char``) and treats everything
-        from that point on as the hidden payload.
+        ``encode_visible_hidden`` appends the base character and the hidden payload after the visible text, so the
+        last occurrence of the base character starts the hidden payload. The visible text may contain the base
+        character itself.
 
         Args:
             combined (str): The combined text containing visible and hidden parts.
 
         Returns:
             tuple[str, str]: A tuple with the visible text and the decoded hidden text.
         """
-        base_char = self.utf8_base_char
-        index = combined.find(base_char)
+        index = combined.rfind(self.utf8_base_char)
+        if not self.embed_in_base and index > 0 and combined[index - 1] == self.utf8_base_char == " ":
+            # With a space as the base character, the last space found is the separator after it.
+            index -= 1
         if index == -1:
             return combined, ""
         visible = combined[:index]

tests/unit/converter/test_token_smuggler_converter.py

@@ -2,6 +2,8 @@
 # Licensed under the MIT license.
 
 
+import pytest
+
 from pyrit.converter import ConverterResult
 from pyrit.converter.token_smuggling import (
     AsciiSmugglerConverter,
@@ -192,6 +194,64 @@ async def test_encode_decode_visible_hidden():
     assert decoded_hidden == hidden_text
 
 
+@pytest.mark.parametrize("embed_in_base", [True, False])
+@pytest.mark.parametrize("carrier", ["I ❤️ this ", "Done ✔️ ", "Hi \U0001f60a there "])
+async def test_decode_variation_selector_smuggler_ignores_emoji_selectors_before_payload(carrier, embed_in_base):
+    # Emoji such as U+2764 U+FE0F carry a presentation selector that is not part of the payload.
+    encode_converter = VariationSelectorSmugglerConverter(action="encode", embed_in_base=embed_in_base)
+    encoded_result = await encode_converter.convert_async(prompt="secret", input_type="text")
+
+    decode_converter = VariationSelectorSmugglerConverter(action="decode", embed_in_base=embed_in_base)
+    decoded_result = await decode_converter.convert_async(
+        prompt=carrier + encoded_result.output_text + " bye", input_type="text"
+    )
+
+    assert decoded_result.output_text == "secret"
+
+
+async def test_decode_variation_selector_smuggler_reads_payload_on_other_base_char():
+    encode_converter = VariationSelectorSmugglerConverter(action="encode", base_char_utf8="A")
+    encoded_result = await encode_converter.convert_async(prompt="secret", input_type="text")
+
+    decode_converter = VariationSelectorSmugglerConverter(action="decode")
+    decoded_result = await decode_converter.convert_async(prompt="Hi " + encoded_result.output_text, input_type="text")
+
+    assert decoded_result.output_text == "secret"
+
+
+@pytest.mark.parametrize("encode_embed, decode_embed", [(True, False), (False, True)])
+async def test_decode_variation_selector_smuggler_across_separator_settings(encode_embed, decode_embed):
+    encoder = VariationSelectorSmugglerConverter(embed_in_base=encode_embed)
+    decoder = VariationSelectorSmugglerConverter(embed_in_base=decode_embed)
+
+    _, encoded = encoder.encode_message("secret")
+    _, combined = encoder.encode_visible_hidden(visible="I \u2764\ufe0f this ", hidden="secret")
+
+    assert decoder.decode_message("I \u2764\ufe0f this " + encoded) == "secret"
+    assert decoder.decode_visible_hidden(combined) == ("I \u2764\ufe0f this ", "secret")
+
+
+@pytest.mark.parametrize("embed_in_base", [True, False])
+@pytest.mark.parametrize("visible", ["Hi \U0001f60a there ", "Hi \U0001f60a❤️ there ", "\U0001f60a "])
+@pytest.mark.parametrize("hidden", ["secret", ""])
+async def test_encode_decode_visible_hidden_with_base_char_in_visible_text(visible, hidden, embed_in_base):
+    converter = VariationSelectorSmugglerConverter(embed_in_base=embed_in_base)
+
+    _, combined = converter.encode_visible_hidden(visible=visible, hidden=hidden)
+
+    assert converter.decode_visible_hidden(combined) == (visible, hidden)
+
+
+@pytest.mark.parametrize("visible", ["Hello", "Hello "])
+@pytest.mark.parametrize("hidden", ["secret", ""])
+async def test_encode_decode_visible_hidden_with_space_base_char_and_separator(visible, hidden):
+    converter = VariationSelectorSmugglerConverter(base_char_utf8=" ", embed_in_base=False)
+
+    _, combined = converter.encode_visible_hidden(visible=visible, hidden=hidden)
+
+    assert converter.decode_visible_hidden(combined) == (visible, hidden)
+
+
 async def test_embed_in_base_false_inserts_separator():
     # Test that when embed_in_base is False, space is inserted after base char.
     converter = VariationSelectorSmugglerConverter(action="encode", embed_in_base=False)