vector_mcp 0.2.0 → 0.3.1

data/lib/vector_mcp/server.rb

@@ -12,6 +12,13 @@ require_relative "util" # Needed if not using Handlers::Core
 require_relative "server/registry"
 require_relative "server/capabilities"
 require_relative "server/message_handling"
+require_relative "security/auth_manager"
+require_relative "security/authorization"
+require_relative "security/middleware"
+require_relative "security/session_context"
+require_relative "security/strategies/api_key"
+require_relative "security/strategies/jwt_token"
+require_relative "security/strategies/custom"
 
 module VectorMCP
   # The `Server` class is the central component for an MCP server implementation.
@@ -62,7 +69,8 @@ module VectorMCP
     # The specific version of the Model Context Protocol this server implements.
     PROTOCOL_VERSION = "2024-11-05"
 
-    attr_reader :logger, :name, :version, :protocol_version, :tools, :resources, :prompts, :roots, :in_flight_requests
+    attr_reader :logger, :name, :version, :protocol_version, :tools, :resources, :prompts, :roots, :in_flight_requests,
+                :auth_manager, :authorization, :security_middleware
     attr_accessor :transport
 
     # Initializes a new VectorMCP server.
@@ -108,6 +116,11 @@ module VectorMCP
       # Configure sampling capabilities
       @sampling_config = configure_sampling_capabilities(options[:sampling_config] || {})
 
+      # Initialize security components
+      @auth_manager = Security::AuthManager.new
+      @authorization = Security::Authorization.new
+      @security_middleware = Security::Middleware.new(@auth_manager, @authorization)
+
       setup_default_handlers
 
       @logger.info("Server instance '#{@name}' v#{@version} (MCP Protocol: #{@protocol_version}, Gem: v#{VectorMCP::VERSION}) initialized.")
@@ -148,6 +161,133 @@ module VectorMCP
       self.transport = active_transport
       active_transport.run
     end
+
+    # --- Security Configuration ---
+
+    # Enable authentication with specified strategy and configuration
+    # @param strategy [Symbol] the authentication strategy (:api_key, :jwt, :custom)
+    # @param options [Hash] strategy-specific configuration options
+    # @return [void]
+    def enable_authentication!(strategy: :api_key, **options, &block)
+      # Clear existing strategies when switching to a new configuration
+      clear_auth_strategies unless @auth_manager.strategies.empty?
+
+      @auth_manager.enable!(default_strategy: strategy)
+
+      case strategy
+      when :api_key
+        add_api_key_auth(options[:keys] || [])
+      when :jwt
+        add_jwt_auth(options)
+      when :custom
+        handler = block || options[:handler]
+        raise ArgumentError, "Custom authentication strategy requires a handler block" unless handler
+
+        add_custom_auth(&handler)
+
+      else
+        raise ArgumentError, "Unknown authentication strategy: #{strategy}"
+      end
+
+      @logger.info("Authentication enabled with strategy: #{strategy}")
+    end
+
+    # Disable authentication (return to pass-through mode)
+    # @return [void]
+    def disable_authentication!
+      @auth_manager.disable!
+      @logger.info("Authentication disabled")
+    end
+
+    # Enable authorization with optional policy configuration block
+    # @param block [Proc] optional block for configuring authorization policies
+    # @return [void]
+    def enable_authorization!(&)
+      @authorization.enable!
+      instance_eval(&) if block_given?
+      @logger.info("Authorization enabled")
+    end
+
+    # Disable authorization (return to pass-through mode)
+    # @return [void]
+    def disable_authorization!
+      @authorization.disable!
+      @logger.info("Authorization disabled")
+    end
+
+    # Add authorization policy for tools
+    # @param block [Proc] policy block that receives (user, action, tool)
+    # @return [void]
+    def authorize_tools(&)
+      @authorization.add_policy(:tool, &)
+    end
+
+    # Add authorization policy for resources
+    # @param block [Proc] policy block that receives (user, action, resource)
+    # @return [void]
+    def authorize_resources(&)
+      @authorization.add_policy(:resource, &)
+    end
+
+    # Add authorization policy for prompts
+    # @param block [Proc] policy block that receives (user, action, prompt)
+    # @return [void]
+    def authorize_prompts(&)
+      @authorization.add_policy(:prompt, &)
+    end
+
+    # Add authorization policy for roots
+    # @param block [Proc] policy block that receives (user, action, root)
+    # @return [void]
+    def authorize_roots(&)
+      @authorization.add_policy(:root, &)
+    end
+
+    # Check if security features are enabled
+    # @return [Boolean] true if authentication or authorization is enabled
+    def security_enabled?
+      @security_middleware.security_enabled?
+    end
+
+    # Get current security status for debugging/monitoring
+    # @return [Hash] security configuration status
+    def security_status
+      @security_middleware.security_status
+    end
+
+    private
+
+    # Add API key authentication strategy
+    # @param keys [Array<String>] array of valid API keys
+    # @return [void]
+    def add_api_key_auth(keys)
+      strategy = Security::Strategies::ApiKey.new(keys: keys)
+      @auth_manager.add_strategy(:api_key, strategy)
+    end
+
+    # Add JWT authentication strategy
+    # @param options [Hash] JWT configuration options
+    # @return [void]
+    def add_jwt_auth(options)
+      strategy = Security::Strategies::JwtToken.new(**options)
+      @auth_manager.add_strategy(:jwt, strategy)
+    end
+
+    # Add custom authentication strategy
+    # @param handler [Proc] custom authentication handler block
+    # @return [void]
+    def add_custom_auth(&)
+      strategy = Security::Strategies::Custom.new(&)
+      @auth_manager.add_strategy(:custom, strategy)
+    end
+
+    # Clear all authentication strategies
+    # @return [void]
+    def clear_auth_strategies
+      @auth_manager.strategies.each_key do |strategy_name|
+        @auth_manager.remove_strategy(strategy_name)
+      end
+    end
   end
 
   module Transport

data/lib/vector_mcp/transport/sse/client_connection.rb

@@ -0,0 +1,113 @@
+# frozen_string_literal: true
+
+module VectorMCP
+  module Transport
+    class SSE
+      # Manages individual client connection state for SSE transport.
+      # Each client connection has a unique session ID, message queue, and streaming thread.
+      class ClientConnection
+        attr_reader :session_id, :message_queue, :logger
+        attr_accessor :stream_thread, :stream_io
+
+        # Initializes a new client connection.
+        #
+        # @param session_id [String] Unique identifier for this client session
+        # @param logger [Logger] Logger instance for debugging and error reporting
+        def initialize(session_id, logger)
+          @session_id = session_id
+          @logger = logger
+          @message_queue = Queue.new
+          @stream_thread = nil
+          @stream_io = nil
+          @closed = false
+          @mutex = Mutex.new
+
+          logger.debug { "Client connection created: #{session_id}" }
+        end
+
+        # Checks if the connection is closed
+        #
+        # @return [Boolean] true if connection is closed
+        def closed?
+          @mutex.synchronize { @closed }
+        end
+
+        # Closes the client connection and cleans up resources.
+        # This method is thread-safe and can be called multiple times.
+        def close
+          @mutex.synchronize do
+            return if @closed
+
+            @closed = true
+            logger.debug { "Closing client connection: #{session_id}" }
+
+            # Close the message queue to signal streaming thread to stop
+            @message_queue.close if @message_queue.respond_to?(:close)
+
+            # Close the stream I/O if it exists
+            begin
+              @stream_io&.close
+            rescue StandardError => e
+              logger.warn { "Error closing stream I/O for #{session_id}: #{e.message}" }
+            end
+
+            # Stop the streaming thread
+            if @stream_thread&.alive?
+              @stream_thread.kill
+              @stream_thread.join(1) # Wait up to 1 second for clean shutdown
+            end
+
+            logger.debug { "Client connection closed: #{session_id}" }
+          end
+        end
+
+        # Enqueues a message to be sent to this client.
+        # This method is thread-safe.
+        #
+        # @param message [Hash] The JSON-RPC message to send
+        # @return [Boolean] true if message was enqueued successfully
+        def enqueue_message(message)
+          return false if closed?
+
+          begin
+            @message_queue.push(message)
+            logger.debug { "Message enqueued for client #{session_id}: #{message.inspect}" }
+            true
+          rescue ClosedQueueError
+            logger.warn { "Attempted to enqueue message to closed queue for client #{session_id}" }
+            false
+          rescue StandardError => e
+            logger.error { "Error enqueuing message for client #{session_id}: #{e.message}" }
+            false
+          end
+        end
+
+        # Dequeues the next message from the client's message queue.
+        # This method blocks until a message is available or the queue is closed.
+        #
+        # @return [Hash, nil] The next message, or nil if queue is closed
+        def dequeue_message
+          return nil if closed?
+
+          begin
+            @message_queue.pop
+          rescue ClosedQueueError
+            nil
+          rescue StandardError => e
+            logger.error { "Error dequeuing message for client #{session_id}: #{e.message}" }
+            nil
+          end
+        end
+
+        # Gets the current queue size
+        #
+        # @return [Integer] Number of messages waiting in the queue
+        def queue_size
+          @message_queue.size
+        rescue StandardError
+          0
+        end
+      end
+    end
+  end
+end

data/lib/vector_mcp/transport/sse/message_handler.rb

@@ -0,0 +1,166 @@
+# frozen_string_literal: true
+
+require "json"
+
+module VectorMCP
+  module Transport
+    class SSE
+      # Handles JSON-RPC message processing for POST requests.
+      # Processes incoming messages and sends responses via SSE streams.
+      class MessageHandler
+        # Initializes a new message handler.
+        #
+        # @param server [VectorMCP::Server] The MCP server instance
+        # @param session [VectorMCP::Session] The server session
+        # @param logger [Logger] Logger instance for debugging
+        def initialize(server, session, logger)
+          @server = server
+          @session = session
+          @logger = logger
+        end
+
+        # Handles a POST message request from a client.
+        #
+        # @param env [Hash] Rack environment hash
+        # @param client_conn [ClientConnection] The client connection
+        # @return [Array] Rack response triplet
+        def handle_post_message(env, client_conn)
+          request_body = read_request_body(env)
+          return error_response(nil, -32_600, "Request body is empty") if request_body.nil? || request_body.empty?
+
+          message = parse_json_message(request_body, client_conn)
+          return message if message.is_a?(Array) # Error response
+
+          process_message(message, client_conn)
+        rescue VectorMCP::ProtocolError => e
+          @logger.error { "Protocol error for client #{client_conn.session_id}: #{e.message}" }
+          request_id = e.request_id || message&.dig("id")
+          enqueue_error_response(client_conn, request_id, e.code, e.message, e.details)
+          error_response(request_id, e.code, e.message, e.details)
+        rescue StandardError => e
+          @logger.error { "Unexpected error for client #{client_conn.session_id}: #{e.message}\n#{e.backtrace.join("\n")}" }
+          request_id = message&.dig("id")
+          enqueue_error_response(client_conn, request_id, -32_603, "Internal server error")
+          error_response(request_id, -32_603, "Internal server error")
+        end
+
+        private
+
+        # Reads the request body from the Rack environment.
+        #
+        # @param env [Hash] Rack environment
+        # @return [String, nil] Request body as string
+        def read_request_body(env)
+          input = env["rack.input"]
+          return nil unless input
+
+          body = input.read
+          input.rewind if input.respond_to?(:rewind)
+          body
+        end
+
+        # Parses JSON message from request body.
+        #
+        # @param body_str [String] JSON string from request body
+        # @param client_conn [ClientConnection] Client connection for error handling
+        # @return [Hash, Array] Parsed message or error response triplet
+        def parse_json_message(body_str, client_conn)
+          JSON.parse(body_str)
+        rescue JSON::ParserError => e
+          @logger.error { "JSON parse error for client #{client_conn.session_id}: #{e.message}" }
+          malformed_id = VectorMCP::Util.extract_id_from_invalid_json(body_str)
+          enqueue_error_response(client_conn, malformed_id, -32_700, "Parse error")
+          error_response(malformed_id, -32_700, "Parse error")
+        end
+
+        # Processes a valid JSON-RPC message.
+        #
+        # @param message [Hash] Parsed JSON-RPC message
+        # @param client_conn [ClientConnection] Client connection
+        # @return [Array] Rack response triplet
+        def process_message(message, client_conn)
+          # Handle the message through the server
+          response_data = @server.handle_message(message, @session, client_conn.session_id)
+
+          # If it's a request (has id), send response via SSE
+          if message["id"]
+            enqueue_success_response(client_conn, message["id"], response_data)
+          else
+            @logger.debug { "Processed notification for client #{client_conn.session_id}" }
+          end
+
+          # Always return 202 Accepted for valid POST messages
+          success_response(message["id"])
+        end
+
+        # Enqueues a successful response to the client's SSE stream.
+        #
+        # @param client_conn [ClientConnection] Client connection
+        # @param request_id [String, Integer] Original request ID
+        # @param result [Object] Response result data
+        def enqueue_success_response(client_conn, request_id, result)
+          response = {
+            jsonrpc: "2.0",
+            id: request_id,
+            result: result
+          }
+          StreamManager.enqueue_message(client_conn, response)
+        end
+
+        # Enqueues an error response to the client's SSE stream.
+        #
+        # @param client_conn [ClientConnection] Client connection
+        # @param request_id [String, Integer, nil] Original request ID
+        # @param code [Integer] Error code
+        # @param message [String] Error message
+        # @param data [Object, nil] Additional error data
+        def enqueue_error_response(client_conn, request_id, code, message, data = nil)
+          error_payload = { code: code, message: message }
+          error_payload[:data] = data if data
+
+          error_response = {
+            jsonrpc: "2.0",
+            id: request_id,
+            error: error_payload
+          }
+          StreamManager.enqueue_message(client_conn, error_response)
+        end
+
+        # Creates a successful HTTP response for the POST request.
+        #
+        # @param request_id [String, Integer, nil] Request ID
+        # @return [Array] Rack response triplet
+        def success_response(request_id)
+          body = { status: "accepted", id: request_id }.to_json
+          [202, { "Content-Type" => "application/json" }, [body]]
+        end
+
+        # Creates an error HTTP response for the POST request.
+        #
+        # @param id [String, Integer, nil] Request ID
+        # @param code [Integer] Error code
+        # @param message [String] Error message
+        # @param data [Object, nil] Additional error data
+        # @return [Array] Rack response triplet
+        def error_response(id, code, message, data = nil)
+          status = case code
+                   when -32_700, -32_600, -32_602 then 400 # Parse, Invalid Request, Invalid Params
+                   when -32_601, -32_001 then 404 # Method Not Found, Not Found
+                   else 500 # Internal Error, Server Error
+                   end
+
+          error_payload = { code: code, message: message }
+          error_payload[:data] = data if data
+
+          body = {
+            jsonrpc: "2.0",
+            id: id,
+            error: error_payload
+          }.to_json
+
+          [status, { "Content-Type" => "application/json" }, [body]]
+        end
+      end
+    end
+  end
+end

data/lib/vector_mcp/transport/sse/puma_config.rb

@@ -0,0 +1,77 @@
+# frozen_string_literal: true
+
+module VectorMCP
+  module Transport
+    class SSE
+      # Configures Puma server for production-ready SSE transport.
+      # Handles server setup, threading, and resource management.
+      class PumaConfig
+        attr_reader :host, :port, :logger
+
+        # Initializes Puma configuration.
+        #
+        # @param host [String] Host to bind to
+        # @param port [Integer] Port to listen on
+        # @param logger [Logger] Logger instance
+        def initialize(host, port, logger)
+          @host = host
+          @port = port
+          @logger = logger
+        end
+
+        # Configures a Puma server instance.
+        #
+        # @param server [Puma::Server] The Puma server to configure
+        def configure(server)
+          server.add_tcp_listener(host, port)
+
+          # Configure threading for production use
+          configure_threading(server)
+
+          # Set up server options
+          configure_server_options(server)
+
+          logger.debug { "Puma server configured for #{host}:#{port}" }
+        end
+
+        private
+
+        # Configures threading parameters for optimal performance.
+        #
+        # @param server [Puma::Server] The Puma server
+        def configure_threading(server)
+          # Set thread pool size based on CPU cores and expected load
+          min_threads = 2
+          max_threads = [4, Etc.nprocessors * 2].max
+
+          # Puma 6.x does not expose min_threads= and max_threads= as public API.
+          # Thread pool sizing should be set via Puma DSL/config before server creation.
+          # For legacy compatibility, set if possible, otherwise log a warning.
+          if server.respond_to?(:min_threads=) && server.respond_to?(:max_threads=)
+            server.min_threads = min_threads
+            server.max_threads = max_threads
+            logger.debug { "Puma configured with #{min_threads}-#{max_threads} threads" }
+          else
+            logger.warn { "Puma::Server does not support direct thread pool sizing; set threads via Puma config DSL before server creation." }
+          end
+        end
+
+        # Configures server-specific options.
+        #
+        # @param server [Puma::Server] The Puma server
+        def configure_server_options(server)
+          # Set server-specific options for SSE handling
+          server.leak_stack_on_error = false if server.respond_to?(:leak_stack_on_error=)
+
+          # Configure timeouts appropriate for SSE connections
+          # SSE connections should be long-lived, so we set generous timeouts
+          if server.respond_to?(:first_data_timeout=)
+            server.first_data_timeout = 30 # 30 seconds to send first data
+          end
+
+          logger.debug { "Puma server options configured for SSE transport" }
+        end
+      end
+    end
+  end
+end

data/lib/vector_mcp/transport/sse/stream_manager.rb

@@ -0,0 +1,92 @@
+# frozen_string_literal: true
+
+module VectorMCP
+  module Transport
+    class SSE
+      # Manages Server-Sent Events streaming for client connections.
+      # Handles creation of streaming responses and message broadcasting.
+      class StreamManager
+        class << self
+          # Creates an SSE streaming response body for a client connection.
+          #
+          # @param client_conn [ClientConnection] The client connection to stream to
+          # @param endpoint_url [String] The URL for the client to POST messages to
+          # @param logger [Logger] Logger instance for debugging
+          # @return [Enumerator] Rack-compatible streaming response body
+          def create_sse_stream(client_conn, endpoint_url, logger)
+            Enumerator.new do |yielder|
+              # Send initial endpoint event
+              yielder << format_sse_event("endpoint", endpoint_url)
+              logger.debug { "Sent endpoint event to client #{client_conn.session_id}: #{endpoint_url}" }
+
+              # Start streaming thread for this client
+              client_conn.stream_thread = Thread.new do
+                stream_messages_to_client(client_conn, yielder, logger)
+              end
+
+              # Keep the connection alive by yielding from the streaming thread
+              client_conn.stream_thread.join
+            rescue StandardError => e
+              logger.error { "Error in SSE stream for client #{client_conn.session_id}: #{e.message}\n#{e.backtrace.join("\n")}" }
+            ensure
+              logger.debug { "SSE stream ended for client #{client_conn.session_id}" }
+              client_conn.close
+            end
+          end
+
+          # Enqueues a message to a specific client connection.
+          #
+          # @param client_conn [ClientConnection] The target client connection
+          # @param message [Hash] The JSON-RPC message to send
+          # @return [Boolean] true if message was enqueued successfully
+          def enqueue_message(client_conn, message)
+            return false unless client_conn && !client_conn.closed?
+
+            client_conn.enqueue_message(message)
+          end
+
+          private
+
+          # Streams messages from a client's queue to the SSE connection.
+          # This method runs in a dedicated thread per client.
+          #
+          # @param client_conn [ClientConnection] The client connection
+          # @param yielder [Enumerator::Yielder] The response yielder
+          # @param logger [Logger] Logger instance
+          def stream_messages_to_client(client_conn, yielder, logger)
+            logger.debug { "Starting message streaming thread for client #{client_conn.session_id}" }
+
+            loop do
+              message = client_conn.dequeue_message
+              break if message.nil? # Queue closed or connection closed
+
+              begin
+                json_message = message.to_json
+                sse_data = format_sse_event("message", json_message)
+                yielder << sse_data
+
+                logger.debug { "Streamed message to client #{client_conn.session_id}: #{json_message}" }
+              rescue StandardError => e
+                logger.error { "Error streaming message to client #{client_conn.session_id}: #{e.message}" }
+                break
+              end
+            end
+
+            logger.debug { "Message streaming thread ended for client #{client_conn.session_id}" }
+          rescue StandardError => e
+            logger.error { "Fatal error in streaming thread for client #{client_conn.session_id}: #{e.message}\n#{e.backtrace.join("\n")}" }
+          end
+
+          # Formats data as a Server-Sent Event.
+          #
+          # @param event [String] The event type
+          # @param data [String] The event data
+          # @return [String] Properly formatted SSE event
+          def format_sse_event(event, data)
+            "event: #{event}\ndata: #{data}\n\n"
+          end
+        end
+      end
+    end
+  end
+end