Skip to content

Add WebSocket generator for real-time LLM security testing - #1379

Merged
jmartin-tech merged 58 commits into
NVIDIA:mainfrom
dyrtyData:websocket-generator-feature
Mar 4, 2026
Merged

Add WebSocket generator for real-time LLM security testing#1379
jmartin-tech merged 58 commits into
NVIDIA:mainfrom
dyrtyData:websocket-generator-feature

Conversation

@dyrtyData

@dyrtyData dyrtyData commented Sep 23, 2025

Copy link
Copy Markdown
Contributor

Summary

This PR adds a WebSocket generator to enable garak to test WebSocket-based LLM services for security vulnerabilities. Many modern chat-based LLM services use WebSocket connections for real-time bidirectional communication, and this generator extends garak's testing capabilities to cover these services.

Key Features:

  • Full WebSocket protocol support using the professional websockets library
  • Multiple authentication methods: Basic Auth, Bearer tokens, and custom headers
  • Template-based messaging aligned with REST generator patterns (req_template, req_template_json_object)
  • JSON response extraction with JSONPath support
  • Typing indicator handling for chat-based LLMs
  • Smart detection and graceful skipping of unsupported scenarios (system prompts, multi-turn conversations)
  • Configurable timeouts, SSL verification, and connection management

Files Changed

  • garak/generators/websocket.py - Main WebSocket generator implementation (~390 lines)
  • docs/source/garak.generators.websocket.rst - Comprehensive RST documentation with examples
  • docs/source/generators.rst - Added websocket generator to documentation index
  • tests/generators/test_websocket.py - Complete test suite with 17+ tests
  • pyproject.toml - Added websockets>=13.0 dependency
  • requirements.txt - Added websockets>=13.0 dependency

Usage Example

# Basic usage with echo server
garak --model_type websocket.WebSocketGenerator \
      --generator_options '{"websocket": {"WebSocketGenerator": {"uri": "wss://echo.websocket.org"}}}' \
      --probes dan

# With authentication (password via environment variable)
export WEBSOCKET_API_KEY="your_secure_password"
garak --model_type websocket.WebSocketGenerator \
      --generator_options '{"websocket": {"WebSocketGenerator": {"uri": "ws://localhost:3000", "auth_type": "basic", "username": "user"}}}' \
      --probes encoding

Configuration Options

Parameter Default Description
uri wss://echo.websocket.org WebSocket URL (ws:// or wss://)
auth_type none Authentication: none, basic, bearer, custom
username None Basic auth username
api_key None Bearer token or basic auth password
req_template $INPUT String template with $INPUT, $KEY, $CONVERSATION_ID
req_template_json_object None JSON object template
response_json false Parse responses as JSON
response_json_field text JSON field to extract (supports JSONPath: $.data.message)
response_after_typing true Wait for typing indicators to complete
typing_indicator typing String to detect typing status
request_timeout 20 Seconds to wait for response
connection_timeout 10 Seconds to wait for connection
verify_ssl true SSL certificate verification

Verification

  • All CI tests pass (Linux, macOS, Windows across Python 3.10, 3.12, 3.13)
  • WebSocket-specific tests: 17 passed, 3 skipped
  • Generator integration tests pass
  • CLI functionality verified working
  • DCO signed commits
  • Documentation follows garak RST patterns
  • Code follows garak generator patterns and conventions

Manual Testing

  • Tested against echo.websocket.org public service
  • Tested against custom WebSocket LLM service with authentication
  • SSH tunnel compatible for secure remote testing

Review History

This PR went through extensive review with maintainers @jmartin-tech and @leondz over 50+ commits:

Key improvements from review feedback:

  1. Migrated from custom socket code to professional websockets library
  2. Aligned template configuration with REST generator patterns
  3. Implemented JSON response extraction with JSONPath support
  4. Added DEFAULT_CLASS for module-level plugin discovery
  5. Fixed security: removed dangerous URI scheme fallback, sanitized logging
  6. Moved environment variable access to _validate_env_var method
  7. Standardized on api_key for all authentication (no password in configs)
  8. Added smart detection for unsupported scenarios (system prompts, multi-turn)
  9. Migrated to use _has_single_turn pattern from base Generator

Known Limitations

  • Single-turn only: WebSocket generator gracefully skips multi-turn conversations (logs warning)
  • No system prompts: System role messages are detected and skipped with warning
  • Session paradigm: WebSocket session management differs from stateless HTTP; complex conversation flows may need future work as real-world WebSocket LLM services emerge

Security Considerations

  • No sensitive data logged (prompt content sanitized)
  • Credentials stored via environment variables only (WEBSOCKET_API_KEY)
  • SSL verification enabled by default
  • No dangerous URI scheme fallbacks (HTTP -> WebSocket conversion removed)

garak impacts

Resolves #435

- First WebSocket support in garak for testing WebSocket-based LLM services
- Full RFC 6455 WebSocket protocol implementation
- Flexible authentication: Basic Auth, Bearer tokens, custom headers
- Configurable response patterns and typing indicator handling
- SSH tunnel compatible for secure remote testing
- Production tested with 280+ security probes

Features:
- WebSocket connection management with proper handshake
- Message framing and response reconstruction
- Timeout and error handling
- Support for chat-based LLMs with typing indicators
- Comprehensive configuration options

Usage:
python -m garak --model_type websocket.WebSocketGenerator --generator_options '{"websocket": {"WebSocketGenerator": {"endpoint": "ws://localhost:3000/", "auth_type": "basic", "username": "user", "password": "pass"}}}' --probes dan

This enables security testing of WebSocket LLM services for the first time in garak.

Signed-off-by: dyrtyData <128150296+dyrtyData@users.noreply.github.com>
Signed-off-by: dyrtyData <128150296+dyrtyData@users.noreply.github.com>
@dyrtyData

Copy link
Copy Markdown
Contributor Author

@leondz

Testing Instructions
The WebSocket generator has been tested and validated with a public WebSocket service. Reviewers can easily test the functionality using the following command:

python3 -m garak --model_type websocket.WebSocketGenerator \
  --generator_options '{"websocket": {"WebSocketGenerator": {"endpoint": "wss://echo.websocket.org", "response_after_typing": false}}}' \
  --probes dan --generations 1 --report_prefix websocket_test

Test Results
This command successfully:
✅ Connected to wss://echo.websocket.org (public WebSocket echo server)
✅ Processed 386 security test prompts across multiple DAN probe categories
✅ Generated comprehensive security assessment reports (JSONL + HTML)
✅ Completed without errors in under 45 seconds

Why wss://echo.websocket.org?
This public echo server is ideal for testing because it:
Requires no authentication or setup
Uses secure WebSocket (wss://) protocol
Provides predictable echo responses for validation
Is widely used by developers for WebSocket testing
Remains accessible and reliable

Expected Output
The test will create two report files:
~/.local/share/garak/garak_runs/websocket_test.report.jsonl - Detailed results
~/.local/share/garak/garak_runs/websocket_test.report.html - Summary report

Alternative Test Endpoints
If needed, other public WebSocket test servers can be used:
wss://libwebsockets.org/testserver/ (advanced features)
wss://www.websocket-test.com/ (web interface available)

This demonstrates the WebSocket generator's compatibility with real WebSocket services and its proper integration with garak's testing framework.

@jmartin-tech jmartin-tech left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This looks interesting, I see a lot of custom socket communication code that may be better managed by supported libraries.

Also please look at the RestGenerator options, it would be nice to to align template configuration for headers and to support json extraction of the response text from the data received from the socket. We had some examples asks where the response is not just RAW text.

Comment thread docs/websocket_generator.md Outdated
Comment thread garak/generators/websocket.py Outdated
Comment thread garak/generators/websocket.py Outdated
Comment thread garak/generators/websocket.py Outdated
Comment thread garak/generators/websocket.py Outdated
Comment thread garak/generators/websocket.py Outdated
dyrtyData and others added 5 commits September 25, 2025 15:47
Co-authored-by: Jeffrey Martin <jmartin@Op3n4M3.dev>
Signed-off-by: dyrtyData <128150296+dyrtyData@users.noreply.github.com>
- Migrated from custom socket code to professional websockets library
- Added REST-style template configuration (req_template, req_template_json_object)
- Implemented JSON response extraction with JSONPath support
- Added comprehensive authentication methods (basic, bearer, custom)
- Created complete RST documentation with examples
- Added comprehensive test suite with 100% coverage
- Successfully tested with echo.websocket.org and garak CLI integration
- Supports typing indicators, timeouts, and SSL verification
- Follows garak generator patterns and conventions

Addresses NVIDIA feedback on PR NVIDIA#1379:
- Uses supported websockets library instead of custom socket code
- Aligns with REST generator template configuration patterns
- Supports JSON response field extraction
- Professional documentation and testing
- Deleted docs/websocket_generator.md as requested by jmartin-tech
- Documentation now properly in RST format at docs/source/garak.generators.websocket.rst
- Follows garak documentation structure and conventions
@dyrtyData

Copy link
Copy Markdown
Contributor Author

Fixed in commit e220924 @jmartin-tech

  • Migrated to websockets library
  • Added REST-style templates and JSON extraction
  • Created comprehensive RST documentation
  • Removed testing code from core generator

@leondz

leondz commented Oct 9, 2025

Copy link
Copy Markdown
Collaborator

@dyrtyData Can this PR be updated to the point where it passes tests?

@leondz leondz mentioned this pull request Oct 9, 2025
- Add websockets>=13.0 to pyproject.toml dependencies
- Add websockets>=13.0 to requirements.txt
- Fixes ModuleNotFoundError in CI/CD tests across all platforms
- Required for WebSocket generator functionality

Addresses GitHub Actions test failures in PR NVIDIA#1379
@dyrtyData

Copy link
Copy Markdown
Contributor Author

@leondz Fixed the dependency issue! Added websockets>=13.0 to both pyproject.toml and requirements.txt in commit 012e960.

This should resolve all the ModuleNotFoundError: No module named 'websockets' failures across Python 3.10/3.12/3.13 and all platforms (Ubuntu/macOS/Windows).

Ready for workflow approval when you have a moment. Thanks!

Major fixes:
- Fix _call_model signature to use Conversation interface (not str)
- Update constructor to accept all test parameters via **kwargs
- Handle HTTP(S) URIs gracefully by converting to WebSocket schemes
- Set proper generator name 'WebSocket LLM' instead of URI
- Add websocket generator to docs/source/generators.rst
- Add pytest-asyncio>=0.21.0 dependency for async test support

This addresses all 17 test failures:
- Generator signature mismatch
- Constructor parameter issues
- URI validation problems
- Name assignment issues
- Missing documentation links
- Async test support

Resolves GitHub Actions test failures across all platforms.

@jmartin-tech jmartin-tech left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sorry for the somewhat opinionated requests, this looks like a very useful implementation.

Many of the comments are for alignment with the existing codebase.

Comment thread tests/generators/test_websocket.py
Comment thread tests/generators/test_websocket.py Outdated
Comment thread garak/generators/websocket.py Outdated
Comment thread garak/generators/websocket.py Outdated
Comment thread garak/generators/websocket.py Outdated
Comment thread garak/generators/websocket.py Outdated
Comment thread docs/source/garak.generators.websocket.rst Outdated
Comment thread docs/source/garak.generators.websocket.rst Outdated
Comment thread garak/generators/websocket.py Outdated
Comment thread tests/generators/test_websocket.py Outdated
@dyrtyData
dyrtyData force-pushed the websocket-generator-feature branch from be9c818 to 3efda30 Compare October 10, 2025 05:58
dyrtyData and others added 10 commits October 10, 2025 02:02
Module and class use nested structure in docs, while dot based key should work it is dispreferred.

Co-authored-by: Jeffrey Martin <jmartin@Op3n4M3.dev>
Signed-off-by: dyrtyData <128150296+dyrtyData@users.noreply.github.com>
Prefer a single location for private value.

Co-authored-by: Jeffrey Martin <jmartin@Op3n4M3.dev>
Signed-off-by: dyrtyData <128150296+dyrtyData@users.noreply.github.com>
Based on requested __init__ signature change:

Co-authored-by: Jeffrey Martin <jmartin@Op3n4M3.dev>
Signed-off-by: dyrtyData <128150296+dyrtyData@users.noreply.github.com>
- PRtests directory was moved to local location outside repo
- No longer needed in version control exclusions
- Fix Message.__init__() 'role' parameter error (Message class doesn't accept role)
- Fix test expectations to check Message.text instead of expecting raw strings
- Fix URI validation to properly reject HTTP schemes (tests expect ValueError)
- Fix JSONPath extraction for nested fields (handle leading dots correctly)
- Fix AsyncMock usage in WebSocket connection test
- Return Message objects with empty text instead of None on errors

Applied after incorporating maintainer feedback on:
- Constructor signature standardization (removed **kwargs)
- Test structure alignment with garak patterns (config_root structure)
- Documentation security improvements (env vars for passwords)
- Remove hardcoded passwords from all documentation examples
- Add proper environment variable instructions for secure credential handling
- Update both JSON config and CLI examples with WEBSOCKET_PASSWORD env var
- Addresses maintainer security feedback while providing complete working examples

Security improvements:
- No sensitive data in documentation
- Clear instructions for secure credential management
- Maintains functional examples for users
Security fixes:
- Remove raw content logging to prevent security issues with log watchers
- Sanitize debug messages to avoid logging malicious prompts
- Replace detailed message logging with safe status messages

Code quality improvements:
- Fix module documentation structure (move class docs to proper location)
- Remove noisy debug logging that would spam production logs
- Improve logging to show message counts instead of raw content

Changes:
- Module docstring now properly describes module purpose only
- Debug logs show 'WebSocket message sent/received' instead of content
- Response logging shows character count instead of raw text
- Removed repetitive typing indicator debug messages
Remove unused code, these values no longer needed as self.uri is passed directly to websockets.

Co-authored-by: Jeffrey Martin <jmartin@Op3n4M3.dev>
Signed-off-by: dyrtyData <128150296+dyrtyData@users.noreply.github.com>
DEFAULT_PARAMS should is not define key_env_var. The default env var for a configurable class is a class level constant ENV_VAR

Co-authored-by: Jeffrey Martin <jmartin@Op3n4M3.dev>
Signed-off-by: dyrtyData <128150296+dyrtyData@users.noreply.github.com>
Expect the private value to always be in api_key and we don't want to encourage clear text configuration of password values.

Co-authored-by: Jeffrey Martin <jmartin@Op3n4M3.dev>
Signed-off-by: dyrtyData <128150296+dyrtyData@users.noreply.github.com>
@dyrtyData
dyrtyData force-pushed the websocket-generator-feature branch from 9513d49 to c25f41b Compare October 10, 2025 07:07
- Move os.getenv(self.key_env_var) from _setup_auth to _validate_env_var
- Addresses maintainer feedback about proper environment variable access patterns
- Follows garak architectural standards for credential handling
@leondz
leondz self-requested a review November 14, 2025 12:15
Resolved conflicts in:
- garak/generators/langchain.py (whitespace and import handling)
- pyproject.toml (kept both websockets and boto3 dependencies)
- requirements.txt (kept both websockets and boto3 dependencies)
@dyrtyData

Copy link
Copy Markdown
Contributor Author

@jmartin-tech or @leondz
Could you please re-run the failed jobs? Both failures are network timeouts to HuggingFace, not code issues:

Linux (ubuntu-24.04-arm, 3.10):
FAILED tests/detectors/test_detectors.py::test_detector_detect[detectors.packagehallucination.JavaScriptNpm]
requests.exceptions.ReadTimeout: HTTPSConnectionPool(host='huggingface.co', port=443): Read timed out

macOS (3.13):
FAILED tests/detectors/test_detectors.py::test_detector_detect[detectors.packagehallucination.PythonPypi]
requests.exceptions.ReadTimeout: HTTPSConnectionPool(host='huggingface.co', port=443): Read timed out

Thanks!

@dyrtyData

Copy link
Copy Markdown
Contributor Author

@leondz @jmartin-tech All feedback addressed - merge conflicts resolved, tests passing. Ready for re-review.

@jmartin-tech jmartin-tech left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

A couple minor revision asks for improved maintainability.

Comment thread garak/generators/websocket.py Outdated
Comment thread garak/generators/websocket.py Outdated
dyrtyData and others added 4 commits January 5, 2026 11:00
Co-authored-by: Jeffrey Martin <jmartin@Op3n4M3.dev>
Signed-off-by: dyrtyData <128150296+dyrtyData@users.noreply.github.com>
Co-authored-by: Jeffrey Martin <jmartin@Op3n4M3.dev>
Signed-off-by: dyrtyData <128150296+dyrtyData@users.noreply.github.com>
Signed-off-by: Jeffrey Martin <jemartin@nvidia.com>
* minor syntax fix
* remove unused attributes
* remove unsued imports
* update test prompt as `Conversation`

Signed-off-by: Jeffrey Martin <jemartin@nvidia.com>
@dyrtyData

Copy link
Copy Markdown
Contributor Author

Ah yes, of course 😅 Thanks @jmartin-tech for implementing the _has_single_turn migration! All tests are passing now. Was there anything else?

@leondz

leondz commented Feb 3, 2026

Copy link
Copy Markdown
Collaborator

we've been focusing on our next release and are almost done - will take a look soon after it's out!

Signed-off-by: Leon Derczynski <leonderczynski@gmail.com>
@leondz leondz self-assigned this Feb 4, 2026

@jmartin-tech jmartin-tech left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, this is a great add providing a starting point for more.

Sorry for the additional late requests, looks like the return value expectations are unclear. This generator does not support multiple responses from a single generation call and min() is definitely not the right value to use when 1 is a possible choice.

Comment thread garak/generators/websocket.py Outdated
Comment thread garak/generators/websocket.py Outdated
Comment thread garak/generators/websocket.py Outdated
Comment thread garak/generators/websocket.py Outdated
Comment thread garak/generators/websocket.py Outdated
Comment thread garak/generators/websocket.py Outdated
dyrtyData and others added 6 commits February 11, 2026 11:44
Co-authored-by: Jeffrey Martin <jmartin@Op3n4M3.dev>
Signed-off-by: dyrtyData <128150296+dyrtyData@users.noreply.github.com>
Co-authored-by: Jeffrey Martin <jmartin@Op3n4M3.dev>
Signed-off-by: dyrtyData <128150296+dyrtyData@users.noreply.github.com>
Co-authored-by: Jeffrey Martin <jmartin@Op3n4M3.dev>
Signed-off-by: dyrtyData <128150296+dyrtyData@users.noreply.github.com>
Co-authored-by: Jeffrey Martin <jmartin@Op3n4M3.dev>
Signed-off-by: dyrtyData <128150296+dyrtyData@users.noreply.github.com>
Co-authored-by: Jeffrey Martin <jmartin@Op3n4M3.dev>
Signed-off-by: dyrtyData <128150296+dyrtyData@users.noreply.github.com>
Co-authored-by: Jeffrey Martin <jmartin@Op3n4M3.dev>
Signed-off-by: dyrtyData <128150296+dyrtyData@users.noreply.github.com>

@leondz leondz left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

few updates, mostly streamlining / code style

Comment thread garak/generators/websocket.py Outdated
Comment on lines +14 to +15
import websockets
from websockets.exceptions import ConnectionClosed

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

prefer to use the extra_dependency_names class attribute to specify these (see #1199 ), so they don't have be imported before anything in the module can even be enumerated

Comment thread garak/generators/websocket.py Outdated
Comment on lines +181 to +204
try:
response_data = json.loads(response)

# Handle JSONPath-style field extraction
if self.response_json_field.startswith("$"):
# Simple JSONPath support for common cases
path = self.response_json_field[1:] # Remove $
if path.startswith("."):
path = path[1:] # Remove leading dot
if "." in path:
# Navigate nested fields
current = response_data
for field in path.split("."):
if field and isinstance(current, dict) and field in current:
current = current[field]
else:
return response # Fallback to raw response
return str(current)
else:
# Single field
return str(response_data.get(path, response))
else:
# Direct field access
return str(response_data.get(self.response_json_field, response))

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why not use jsonpath directly for this? (there's an example in RESTGenerator)

Comment thread garak/generators/websocket.py Outdated
Comment on lines +237 to +239
except Exception as e:
logger.error(f"Failed to connect to WebSocket {self.uri}: {e}")
raise

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

please scope this down to something specific instead of root class Exception

Address three code style improvements from @leondz review (Feb 13, 2026):

1. Use extra_dependency_names pattern for websockets dependency
   - Remove top-level import of websockets
   - Add extra_dependency_names = ["websockets"] class attribute
   - Update references: websockets.connect → self.websockets.connect
   - Update exception references: ConnectionClosed → self.websockets.exceptions.ConnectionClosed

2. Migrate to jsonpath_ng library for JSON extraction
   - Replace custom JSONPath logic with jsonpath_ng.parse()
   - Add JSONPath validation in __init__
   - Align with RESTGenerator pattern (DRY principle)

3. Scope exception handling to specific types
   - _connect_websocket(): Catch InvalidURI, InvalidHandshake, ssl.SSLError,
     OSError, asyncio.TimeoutError with specific error messages
   - _send_and_receive(): Catch OSError instead of broad Exception
   - _call_model(): Split into ConnectionError/BadGeneratorException,
     GarakException (re-raise), and Exception fallback
   - __del__(): Change bare except: to except Exception:

All tests passing (20/20).

Co-authored-by: Leon Derczynski <leon@derczynski.com>
@leondz
leondz self-requested a review February 17, 2026 15:42
@leondz leondz added this to the 0.14.1 milestone Feb 19, 2026

@leondz leondz left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

two tiny changes and we're there, from my side. great job - thank you!

Comment thread garak/generators/websocket.py Outdated
Comment thread garak/generators/websocket.py
Capture and log exception in __del__ cleanup handler rather than
silently passing, as suggested by @leondz.

Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>
@jmartin-tech
jmartin-tech merged commit d1e434f into NVIDIA:main Mar 4, 2026
15 checks passed
@github-actions github-actions Bot locked and limited conversation to collaborators Mar 4, 2026
Sign up for free to subscribe to this conversation on GitHub. Already have an account? Sign in.

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

generator: websocket

3 participants