feat: implement WebSocket authentication system with UUID tokens (#56)

Author: ThomasK33

CLAUDE.md

@@ -22,10 +22,16 @@ claudecode.nvim - A Neovim plugin that implements the same WebSocket-based MCP p
 
 ### Build Commands
 
+- `make` - **RECOMMENDED**: Run formatting, linting, and testing (complete validation)
 - `make all` - Run check and format (default target)
+- `make test` - Run all tests using busted with coverage
+- `make check` - Check Lua syntax and run luacheck
+- `make format` - Format code with stylua (or nix fmt if available)
 - `make clean` - Remove generated test files
 - `make help` - Show available commands
 
+**Best Practice**: Always use `make` at the end of editing sessions for complete validation.
+
 ### Development with Nix
 
 - `nix develop` - Enter development shell with all dependencies
@@ -45,11 +51,20 @@ claudecode.nvim - A Neovim plugin that implements the same WebSocket-based MCP p
 ### WebSocket Server Implementation
 
 - **TCP Server**: `server/tcp.lua` handles port binding and connections
-- **Handshake**: `server/handshake.lua` processes HTTP upgrade requests
+- **Handshake**: `server/handshake.lua` processes HTTP upgrade requests with authentication
 - **Frame Processing**: `server/frame.lua` implements RFC 6455 WebSocket frames
 - **Client Management**: `server/client.lua` manages individual connections
 - **Utils**: `server/utils.lua` provides base64, SHA-1, XOR operations in pure Lua
 
+#### Authentication System
+
+The WebSocket server implements secure authentication using:
+
+- **UUID v4 Tokens**: Generated per session with enhanced entropy
+- **Header-based Auth**: Uses `x-claude-code-ide-authorization` header
+- **Lock File Discovery**: Tokens stored in `~/.claude/ide/[port].lock` for Claude CLI
+- **MCP Compliance**: Follows official Claude Code IDE authentication protocol
+
 ### MCP Tools Architecture
 
 Tools are registered with JSON schemas and handlers. MCP-exposed tools include:
@@ -76,14 +91,125 @@ Tests are organized in three layers:
 
 Test files follow the pattern `*_spec.lua` or `*_test.lua` and use the busted framework.
 
+### Test Organization Principles
+
+- **Isolation**: Each test should be independent and not rely on external state
+- **Mocking**: Use comprehensive mocking for vim APIs and external dependencies
+- **Coverage**: Aim for both positive and negative test cases, edge cases included
+- **Performance**: Tests should run quickly to encourage frequent execution
+- **Clarity**: Test names should clearly describe what behavior is being verified
+
+## Authentication Testing
+
+The plugin implements authentication using UUID v4 tokens that are generated for each server session and stored in lock files. This ensures secure connections between Claude CLI and the Neovim WebSocket server.
+
+### Testing Authentication Features
+
+**Lock File Authentication Tests** (`tests/lockfile_test.lua`):
+
+- Auth token generation and uniqueness validation
+- Lock file creation with authentication tokens
+- Reading auth tokens from existing lock files
+- Error handling for missing or invalid tokens
+
+**WebSocket Handshake Authentication Tests** (`tests/unit/server/handshake_spec.lua`):
+
+- Valid authentication token acceptance
+- Invalid/missing token rejection
+- Edge cases (empty tokens, malformed headers, length limits)
+- Case-insensitive header handling
+
+**Server Integration Tests** (`tests/unit/server_spec.lua`):
+
+- Server startup with authentication tokens
+- Auth token state management during server lifecycle
+- Token validation throughout server operations
+
+**End-to-End Authentication Tests** (`tests/integration/mcp_tools_spec.lua`):
+
+- Complete authentication flow from server start to tool execution
+- Authentication state persistence across operations
+- Concurrent operations with authentication enabled
+
+### Manual Authentication Testing
+
+**Test Script Authentication Support**:
+
+```bash
+# Test scripts automatically detect and use authentication tokens
+cd scripts/
+./claude_interactive.sh  # Automatically reads auth token from lock file
+```
+
+**Authentication Flow Testing**:
+
+1. Start the plugin: `:ClaudeCodeStart`
+2. Check lock file contains `authToken`: `cat ~/.claude/ide/*.lock | jq .authToken`
+3. Test WebSocket connection with auth: Use test scripts in `scripts/` directory
+4. Verify authentication in logs: Set `log_level = "debug"` in config
+
+**Testing Authentication Failures**:
+
+```bash
+# Test invalid auth token (should fail)
+websocat ws://localhost:PORT --header "x-claude-code-ide-authorization: invalid-token"
+
+# Test missing auth header (should fail)
+websocat ws://localhost:PORT
+
+# Test valid auth token (should succeed)
+websocat ws://localhost:PORT --header "x-claude-code-ide-authorization: $(cat ~/.claude/ide/*.lock | jq -r .authToken)"
+```
+
+### Authentication Logging
+
+Enable detailed authentication logging by setting:
+
+```lua
+require("claudecode").setup({
+  log_level = "debug"  -- Shows auth token generation, validation, and failures
+})
+```
+
+Log levels for authentication events:
+
+- **DEBUG**: Server startup authentication state, client connections, handshake processing, auth token details
+- **WARN**: Authentication failures during handshake
+- **ERROR**: Auth token generation failures, handshake response errors
+
+### Logging Best Practices
+
+- **Connection Events**: Use DEBUG level for routine connection establishment/teardown
+- **Authentication Flow**: Use DEBUG for successful auth, WARN for failures
+- **User-Facing Events**: Use INFO sparingly for events users need to know about
+- **System Errors**: Use ERROR for failures that require user attention
+
 ## Development Notes
 
+### Technical Requirements
+
 - Plugin requires Neovim >= 0.8.0
 - Uses only Neovim built-ins for WebSocket implementation (vim.loop, vim.json, vim.schedule)
-- Lock files are created at `~/.claude/ide/[port].lock` for Claude CLI discovery
-- WebSocket server only accepts local connections for security
+- Zero external dependencies for core functionality
+
+### Security Considerations
+
+- WebSocket server only accepts local connections (127.0.0.1) for security
+- Authentication tokens are UUID v4 with enhanced entropy
+- Lock files created at `~/.claude/ide/[port].lock` for Claude CLI discovery
+- All authentication events are logged for security auditing
+
+### Performance Optimizations
+
 - Selection tracking is debounced to reduce overhead
+- WebSocket frame processing optimized for JSON-RPC payload sizes
+- Connection pooling and cleanup to prevent resource leaks
+
+### Integration Support
+
 - Terminal integration supports both snacks.nvim and native Neovim terminal
+- Compatible with popular file explorers (nvim-tree, oil.nvim)
+- Visual selection tracking across different selection modes
 
 ## Release Process
 
@@ -134,6 +260,23 @@ make
 rg "0\.1\.0" .  # Should only show CHANGELOG.md historical entries
 ```
 
-## CRITICAL: Pre-commit Requirements
+## Development Workflow
+
+### Pre-commit Requirements
 
 **ALWAYS run `make` before committing any changes.** This runs code quality checks and formatting that must pass for CI to succeed. Never skip this step - many PRs fail CI because contributors don't run the build commands before committing.
+
+### Recommended Development Flow
+
+1. **Start Development**: Use existing tests and documentation to understand the system
+2. **Make Changes**: Follow existing patterns and conventions in the codebase
+3. **Validate Work**: Run `make` to ensure formatting, linting, and tests pass
+4. **Document Changes**: Update relevant documentation (this file, PROTOCOL.md, etc.)
+5. **Commit**: Only commit after successful `make` execution
+
+### Code Quality Standards
+
+- **Test Coverage**: Maintain comprehensive test coverage (currently 314+ tests)
+- **Zero Warnings**: All code must pass luacheck with 0 warnings/errors
+- **Consistent Formatting**: Use `nix fmt` or `stylua` for consistent code style
+- **Documentation**: Update CLAUDE.md for architectural changes, PROTOCOL.md for protocol changes

PROTOCOL.md

@@ -23,7 +23,8 @@ The IDE writes a discovery file to `~/.claude/ide/[port].lock`:
   "pid": 12345, // IDE process ID
   "workspaceFolders": ["/path/to/project"], // Open folders
   "ideName": "VS Code", // or "Neovim", "IntelliJ", etc.
-  "transport": "ws" // WebSocket transport
+  "transport": "ws", // WebSocket transport
+  "authToken": "550e8400-e29b-41d4-a716-446655440000" // Random UUID for authentication
 }
 ```
 
@@ -38,6 +39,16 @@ When launching Claude, the IDE sets:
 
 Claude reads the lock files, finds the matching port from the environment, and connects to the WebSocket server.
 
+## Authentication
+
+When Claude connects to the IDE's WebSocket server, it must authenticate using the token from the lock file. The authentication happens via a custom WebSocket header:
+
+```
+x-claude-code-ide-authorization: 550e8400-e29b-41d4-a716-446655440000
+```
+
+The IDE validates this header against the `authToken` value from the lock file. If the token doesn't match, the connection is rejected.
+
 ## The Protocol
 
 Communication uses WebSocket with JSON-RPC 2.0 messages:
@@ -503,11 +514,13 @@ local server = create_websocket_server("127.0.0.1", random_port)
 
 ```lua
 -- ~/.claude/ide/[port].lock
+local auth_token = generate_uuid() -- Generate random UUID
 local lock_data = {
   pid = vim.fn.getpid(),
   workspaceFolders = { vim.fn.getcwd() },
   ideName = "YourEditor",
-  transport = "ws"
+  transport = "ws",
+  authToken = auth_token
 }
 write_json(lock_path, lock_data)
 ```
@@ -523,6 +536,12 @@ claude  # Claude will now connect!
 ### 4. Handle Messages
 
 ```lua
+-- Validate authentication on WebSocket handshake
+function validate_auth(headers)
+  local auth_header = headers["x-claude-code-ide-authorization"]
+  return auth_header == auth_token
+end
+
 -- Send selection updates
 send_message({
   jsonrpc = "2.0",

lua/claudecode/init.lua

@@ -67,6 +67,7 @@ local default_config = {
 --- @field config ClaudeCode.Config The current plugin configuration.
 --- @field server table|nil The WebSocket server instance.
 --- @field port number|nil The port the server is running on.
+--- @field auth_token string|nil The authentication token for the current session.
 --- @field initialized boolean Whether the plugin has been initialized.
 --- @field queued_mentions table[] Array of queued @ mentions waiting for connection.
 --- @field connection_timer table|nil Timer for connection timeout.
@@ -76,6 +77,7 @@ M.state = {
   config = vim.deepcopy(default_config),
   server = nil,
   port = nil,
+  auth_token = nil,
   initialized = false,
   queued_mentions = {},
   connection_timer = nil,
@@ -357,26 +359,70 @@ function M.start(show_startup_notification)
   end
 
   local server = require("claudecode.server.init")
-  local success, result = server.start(M.state.config)
+  local lockfile = require("claudecode.lockfile")
+
+  -- Generate auth token first so we can pass it to the server
+  local auth_token
+  local auth_success, auth_result = pcall(function()
+    return lockfile.generate_auth_token()
+  end)
+
+  if not auth_success then
+    local error_msg = "Failed to generate authentication token: " .. (auth_result or "unknown error")
+    logger.error("init", error_msg)
+    return false, error_msg
+  end
+
+  auth_token = auth_result
+
+  -- Validate the generated auth token
+  if not auth_token or type(auth_token) ~= "string" or #auth_token < 10 then
+    local error_msg = "Invalid authentication token generated"
+    logger.error("init", error_msg)
+    return false, error_msg
+  end
+
+  local success, result = server.start(M.state.config, auth_token)
 
   if not success then
-    logger.error("init", "Failed to start Claude Code integration: " .. result)
-    return false, result
+    local error_msg = "Failed to start Claude Code server: " .. (result or "unknown error")
+    if result and result:find("auth") then
+      error_msg = error_msg .. " (authentication related)"
+    end
+    logger.error("init", error_msg)
+    return false, error_msg
   end
 
   M.state.server = server
   M.state.port = tonumber(result)
+  M.state.auth_token = auth_token
 
-  local lockfile = require("claudecode.lockfile")
-  local lock_success, lock_result = lockfile.create(M.state.port)
+  local lock_success, lock_result, returned_auth_token = lockfile.create(M.state.port, auth_token)
 
   if not lock_success then
     server.stop()
     M.state.server = nil
     M.state.port = nil
+    M.state.auth_token = nil
+
+    local error_msg = "Failed to create lock file: " .. (lock_result or "unknown error")
+    if lock_result and lock_result:find("auth") then
+      error_msg = error_msg .. " (authentication token issue)"
+    end
+    logger.error("init", error_msg)
+    return false, error_msg
+  end
 
-    logger.error("init", "Failed to create lock file: " .. lock_result)
-    return false, lock_result
+  -- Verify that the auth token in the lock file matches what we generated
+  if returned_auth_token ~= auth_token then
+    server.stop()
+    M.state.server = nil
+    M.state.port = nil
+    M.state.auth_token = nil
+
+    local error_msg = "Authentication token mismatch between server and lock file"
+    logger.error("init", error_msg)
+    return false, error_msg
   end
 
   if M.state.config.track_selection then
@@ -422,6 +468,7 @@ function M.stop()
 
   M.state.server = nil
   M.state.port = nil
+  M.state.auth_token = nil
 
   -- Clear any queued @ mentions when server stops
   clear_mention_queue()