Fix AIVO security findings

Author: AmeliaYe

.project/spec.yaml

@@ -101,7 +101,7 @@ execution:
         - name: Chat Frontend
           type: custom
           class: webapp
-          start_command: export PROXY_PREFIX && cd /project/code && uvicorn --log-level info --reload frontend.server:app --reload-dir frontend --port 7070 --host 0.0.0.0 --timeout-graceful-shutdown 10
+          start_command: export PROXY_PREFIX NIM_ANYWHERE_API_KEY && cd /project/code && uvicorn --log-level info --reload frontend.server:app --reload-dir frontend --port 7070 --host 0.0.0.0 --timeout-graceful-shutdown 10
           health_check_command: curl localhost:7070/healthz > /dev/null
           stop_command: ps aux | grep frontend.server:app | grep -v grep | awk '{print $2}' | xargs kill
           user_msg: ""
@@ -117,7 +117,7 @@ execution:
         - name: Chain Server
           type: custom
           class: webapp
-          start_command: export PROXY_PREFIX && export NGC_API_KEY && cd /project/code && APP_LOG_LEVEL=DEBUG uvicorn --log-level info chain_server.server:app --port 3030 --host 0.0.0.0 --timeout-graceful-shutdown 10 --reload-exclude 'frontend/*' --reload-exclude 'evaluation/*' --reload-include '*.html' --reload-include '*.css' --reload-include '*.yaml' --reload-include '*.js' --reload --reload-include reload
+          start_command: export PROXY_PREFIX NGC_API_KEY NIM_ANYWHERE_API_KEY && cd /project/code && APP_LOG_LEVEL=DEBUG uvicorn --log-level info chain_server.server:app --port 3030 --host 0.0.0.0 --timeout-graceful-shutdown 10 --reload-exclude 'frontend/*' --reload-exclude 'evaluation/*' --reload-include '*.html' --reload-include '*.css' --reload-include '*.yaml' --reload-include '*.js' --reload --reload-include reload
           health_check_command: ps aux |grep -v grep | grep chain_server.server:app
           stop_command: ps aux | grep chain_server.server:app | grep -v grep | awk '{print $2}' | xargs kill
           user_msg: ""
@@ -135,7 +135,7 @@ execution:
           class: webapp
           start_command: |-
             cd /project/code/tutorial_app && \
-            export PROXY_PREFIX && \
+            export PROXY_PREFIX NIM_ANYWHERE_API_KEY && \
             streamlit run app.py --server.baseUrlPath=$PROXY_PREFIX
           health_check_command: curl -I http://localhost:8501
           stop_command: |
@@ -157,6 +157,8 @@ execution:
     secrets:
         - variable: NGC_API_KEY
           description: NGC Personal Key from https://org.ngc.nvidia.com/setup/personal-keys
+        - variable: NIM_ANYWHERE_API_KEY
+          description: A random 32+ character secret used to authenticate the apps and Redis (for example, openssl rand -hex 32)
     mounts:
         - type: project
           target: /project/

README.md

@@ -22,7 +22,7 @@ to help our partners and customers build effective RAG pipeline with
 ease.
 
 NIM Anywhere contains all the tooling required to start integrating NIMs
-for RAG. It natively scales out to full-sized labs and up to production 
+for RAG. It natively scales out to full-sized labs and up to production
 environments. This is great news for building a RAG architecture and
 easily adding NIMs as needed. If you're unfamiliar with RAG, it
 dynamically retrieves relevant external information during inference
@@ -399,14 +399,15 @@ section.
 
 ## Configure this project
 
-The project must be configured to use your NGC personal key.
+The project must be configured with an application secret and your NGC
+personal key.
 
 <details>
 <summary>
 <b>Expand this section for a details on configuring this project.</b>
 </summary>
 
-1.  Before running for the first time, your NGC personal key must be
+1.  Before running for the first time, the project secrets must be
     configured in Workbench. This is done using the *Environment* tab
     from the left-hand panel.
 
@@ -416,6 +417,15 @@ The project must be configured to use your NGC personal key.
     entry. Press *Configure* and provide the personal key for NGC that
     was generated earlier.
 
+3.  Configure *NIM_ANYWHERE_API_KEY* with a random secret of at least 32
+    characters. For example, generate one with `openssl rand -hex 32`.
+    This key protects the frontend, chain server, tutorial, sessions,
+    and Redis. Do not commit it to the repository.
+
+    The frontend uses HTTP Basic authentication with `nim-anywhere` as
+    the username and this secret as the password. The tutorial asks for
+    the secret directly.
+
 </details>
 
 ## Start This Project
@@ -504,8 +514,12 @@ without GPUs.
 
 5.  Once the *Chain Server* is up, the *Chat Frontend* can be started.
     Starting the interface will automatically open it in a browser
-    window. If you are running any local NIMs, you can edit the config
-    to connect to them via the *Chat Frontend*
+    window. Sign in with `nim-anywhere` as the username and your
+    *NIM_ANYWHERE_API_KEY* as the password. If you are running any local
+    NIMs, you can select their preconfigured model names through the
+    *Chat Frontend*. Service URLs, credentials, and data-store addresses
+    can only be changed in the project configuration, not through the
+    web UI.
 
 ![NIM Anywhere Frontend](.static/_static/na_frontend.png)
 
@@ -654,6 +668,11 @@ milvus:
     # Type: string
     collection_name: collection_1
 
+    # Milvus authentication token in username:password format.
+    # ENV Variables: APP_MILVUS__TOKEN
+    # Type: string, null
+    token: ~
+
 
 log_level: 
 

code/chain_server/chain.py

@@ -26,6 +26,7 @@
 from langchain_core.runnables.history import RunnableWithMessageHistory
 from langchain_milvus.vectorstores.milvus import Milvus
 from langchain_nvidia_ai_endpoints import ChatNVIDIA, NVIDIAEmbeddings, NVIDIARerank
+from nim_anywhere_security import validate_session_id
 from pydantic import BaseModel
 
 from . import prompts
@@ -40,7 +41,7 @@
 )
 vector_store = Milvus(
     embedding_function=embedding_model,
-    connection_args={"uri": app_config.milvus.url},
+    connection_args=app_config.milvus.connection_args(),
     collection_name=app_config.milvus.collection_name,
     auto_id=True,
 )
@@ -125,9 +126,15 @@ class ChainInputs(BaseModel):
 ChainOutputs = str
 
 
+def _get_message_history(session_id: str) -> RedisChatMessageHistory:
+    """Return history only for an unmodified server-signed session id."""
+    validated_session_id = validate_session_id(session_id)
+    return RedisChatMessageHistory(validated_session_id, url=str(app_config.redis_dsn))
+
+
 my_chain = RunnableWithMessageHistory(
     my_chain,
-    lambda session_id: RedisChatMessageHistory(session_id, url=str(app_config.redis_dsn)),
+    _get_message_history,
     input_messages_key="question",
     output_messages_key="output",
     history_messages_key="history",

code/chain_server/configuration.py

@@ -48,6 +48,7 @@
 Values specified as environment variables will take precedence over all values from files.
 
 """
+
 import logging
 import os
 from enum import Enum
@@ -86,6 +87,17 @@ class MilvusConfig(BaseModel):
         description="The host machine running Milvus vector DB.",
     )
     collection_name: str = Field("collection_1", description="The name of the Milvus collection.")
+    token: Optional[str] = Field(
+        None,
+        description="Milvus authentication token in username:password format.",
+    )
+
+    def connection_args(self) -> dict[str, str]:
+        """Build authenticated Milvus connection arguments."""
+        arguments = {"uri": self.url}
+        if self.token:
+            arguments["token"] = self.token
+        return arguments
 
 
 class LLMModelConfig(BaseModel):
@@ -145,7 +157,7 @@ class EmbeddingModelConfig(BaseModel):
     ]
 
 
-class Configuration(BaseConfig):
+class Configuration(BaseConfig):  # type: ignore[metaclass]
     """Configuration for this microservice."""
 
     # configuration fields

code/chain_server/prompts.py

@@ -21,16 +21,14 @@
     PromptTemplate,
 )
 
-CONDENSE_QUESTION_TEMPLATE = PromptTemplate.from_template(
-    """Given a chat history and the latest user question
+CONDENSE_QUESTION_TEMPLATE = PromptTemplate.from_template("""Given a chat history and the latest user question
     which might reference context in the chat history, formulate a standalone question
     which can be understood without the chat history. Do NOT answer the question,
     just reformulate it if needed and otherwise return it as is. Don't frame your response.
     Chat History:
     {history}
     Follow Up question: {question}
-    Standalone question:"""
-)
+    Standalone question:""")
 
 
 # Primary Chat Prompt template

code/chain_server/server.py

@@ -23,6 +23,8 @@
 from langserve import add_routes
 from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
 
+from nim_anywhere_security import AuthenticationMiddleware
+
 from . import errors
 from .chain import my_chain  # type: ignore
 
@@ -34,6 +36,7 @@
     root_path=PROXY_PREFIX or "",
     middleware=[Middleware(errors.ErrorHandlerMiddleware)],
 )
+app.add_middleware(AuthenticationMiddleware)
 
 
 add_routes(

code/frontend/_assets/index.html

@@ -19,8 +19,6 @@
 <head>
     <title>NVIDIA NIM Anywhere</title>
     <link rel="stylesheet" href="static/css/style.css">
-    <link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/5.12.1/css/all.min.css">
-    <script src="https://cdn.jsdelivr.net/npm/mermaid@10.9.0/dist/mermaid.min.js"></script>
     <script src="static/js/page_router.js"></script>
     <script src="static/js/events.js"></script>
 </head>

code/frontend/configuration.py

@@ -37,7 +37,7 @@ class LogLevels(Enum):
     CRITICAL = "CRITICAL"
 
 
-class Configuration(BaseConfig):
+class Configuration(BaseConfig):  # type: ignore[metaclass]
     """Configuration for this microservice."""
 
     # configuration fields

code/frontend/mermaid.py

@@ -18,8 +18,12 @@
 import gradio as gr
 
 HEAD = """
-<script src="https://cdn.jsdelivr.net/npm/mermaid@10.9.0/dist/mermaid.min.js"></script>
-<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/5.12.1/css/all.min.css">
+<script
+  src="https://cdn.jsdelivr.net/npm/mermaid@10.9.0/dist/mermaid.min.js"
+  integrity="sha384-6F4Ibv/ylL12O35KFWTeGTHuBKDz5L6yjKsgv3QHQ8s4NTqlDXq7kMlYXGs7MHFc"
+  crossorigin="anonymous"
+  referrerpolicy="no-referrer"
+></script>
 """
 
 _ONLOAD = """

code/frontend/pages/chat/__init__.py

@@ -15,13 +15,13 @@
 
 """The gradio chat app in the frontend."""
 
-import uuid
 from pathlib import Path
 from typing import AsyncGenerator, Any
 
 import gradio as gr
 from httpx import ConnectError, HTTPStatusError
 from langserve import RemoteRunnable
+from nim_anywhere_security import bearer_headers, create_session_id
 
 from ...common import IMG_DIR, THEME, USE_KB_INITIAL, USE_RERANKER_INITIAL
 from ...configuration import config
@@ -38,7 +38,7 @@
 with open(_JS_FILE, "r", encoding="UTF-8") as js_file:
     _HEAD = "<script>" + js_file.read() + "</script>"
 _ONLOAD = "async()=>{ registerHandlers(); chatResize(); }"
-_CHAIN: RemoteRunnable[dict[str, Any], str] = RemoteRunnable(str(config.chain_url))
+_CHAIN: RemoteRunnable[dict[str, Any], str] = RemoteRunnable(str(config.chain_url), headers=bearer_headers())
 
 
 # web ui definition
@@ -65,8 +65,8 @@
 
     @page.load(outputs=session_id)
     def sid_generator() -> str:
-        """Generate a new session id."""
-        return str(uuid.uuid4())
+        """Generate a new signed session id."""
+        return create_session_id()
 
     @gr.on(
         triggers=[msg.submit, submit.click],