From bc766d7794923ae58b5bf3a91f265c7b500ffaa2 Mon Sep 17 00:00:00 2001 From: Jason Weill Date: Wed, 25 Oct 2023 14:25:58 -0700 Subject: [PATCH] Validates slash ID --- packages/jupyter-ai/jupyter_ai/extension.py | 22 ++++++++++++++++----- 1 file changed, 17 insertions(+), 5 deletions(-) diff --git a/packages/jupyter-ai/jupyter_ai/extension.py b/packages/jupyter-ai/jupyter_ai/extension.py index 24af16524..6b7077ffe 100644 --- a/packages/jupyter-ai/jupyter_ai/extension.py +++ b/packages/jupyter-ai/jupyter_ai/extension.py @@ -1,4 +1,5 @@ import logging +import re import time from dask.distributed import Client as DaskClient @@ -142,6 +143,7 @@ def initialize_settings(self): "/help": help_chat_handler, } + slash_command_pattern = r'^[a-zA-Z0-9_]+$' for chat_handler_ep in chat_handler_eps: try: chat_handler = chat_handler_ep.load() @@ -156,11 +158,21 @@ def initialize_settings(self): # Slash IDs may contain only alphanumerics and underscores. slash_id = chat_handler.slash_id - # TODO: Validate slash ID (/^[A-Za-z0-9_]+$/) - # The "default" handler above takes precedence over any "default" command here - command_name = ( - f"/{chat_handler.slash_id}" if chat_handler.slash_id else "default" - ) + if chat_handler.routing_method == 'slash_command': + if chat_handler.slash_id: + # Validate slash ID (/^[A-Za-z0-9_]+$/) + if re.match(slash_command_pattern, chat_handler.slash_id): + command_name = f"/{chat_handler.slash_id}" + else: + self.log.error( + f"Handler `{chat_handler_ep.name}` has an invalid slash command " + + f"`{chat_handler.slash_id}`; must contain only letters, numbers, " + + "and underscores" + ) + continue + else: # No slash ID provided; assume default + # The "default" handler above takes precedence over any "default" command here + command_name = "default" if command_name in jai_chat_handlers: self.log.error(