vector_mcp 0.4.0 → 0.6.0

data/lib/vector_mcp/middleware/anonymizer.rb

@@ -0,0 +1,186 @@
+# frozen_string_literal: true
+
+require "json"
+
+require_relative "base"
+require_relative "../token_store"
+require_relative "../util/token_sweeper"
+
+module VectorMCP
+  module Middleware
+    # Middleware that rewrites selected string fields in outbound tool
+    # results into opaque tokens and restores them on inbound tool
+    # invocations. All domain knowledge (which keys to match, token
+    # prefixes, which keys to treat as atomic blobs) is supplied by the
+    # application via the constructor.
+    #
+    # @example Wiring on a server
+    #   anonymizer = VectorMCP::Middleware::Anonymizer.new(
+    #     store:       VectorMCP::TokenStore.new,
+    #     field_rules: [
+    #       { pattern: /\bname\b/i, prefix: "NAME"  },
+    #       { pattern: /email/i,    prefix: "EMAIL" }
+    #     ],
+    #     atomic_keys: /address/i
+    #   )
+    #   anonymizer.install_on(server)
+    class Anonymizer
+      # @param store [VectorMCP::TokenStore] the backing token store.
+      # @param field_rules [Array<Hash>] an array of +{ pattern: Regexp, prefix: String }+
+      #   hashes. The pattern is matched against each Hash key whose value is a String.
+      # @param atomic_keys [Regexp, nil] optional pattern; Hash values whose parent
+      #   key matches are serialized and tokenized as a single opaque unit instead
+      #   of recursed into.
+      def initialize(store:, field_rules:, atomic_keys: nil)
+        raise ArgumentError, "store is required"       if store.nil?
+        raise ArgumentError, "field_rules is required" if field_rules.nil?
+
+        @store = store
+        @field_rules = field_rules.map { |rule| validate_rule!(rule) }.freeze
+        @atomic_keys = atomic_keys
+      end
+
+      # Tokenize sensitive string fields in an outbound payload.
+      #
+      # @param obj [Object] a parsed JSON-like Ruby structure.
+      # @return [Object] a new structure with matched values replaced by tokens.
+      def sweep_outbound(obj)
+        replace_atomic_nodes(obj).then do |shaped|
+          VectorMCP::Util::TokenSweeper.sweep(shaped) do |value, parent_key|
+            rule = rule_for(parent_key)
+            rule ? @store.tokenize(value, prefix: rule[:prefix]) : value
+          end
+        end
+      end
+
+      # Resolve tokens in an inbound payload back to their original values.
+      # Unknown token-shaped strings pass through unchanged.
+      #
+      # @param obj [Object] a parsed JSON-like Ruby structure.
+      # @return [Object] a new structure with tokens resolved to original values.
+      def sweep_inbound(obj)
+        VectorMCP::Util::TokenSweeper.sweep(obj) do |value, _parent_key|
+          if @store.token?(value)
+            resolved = @store.resolve(value)
+            resolved.nil? ? value : resolved
+          else
+            value
+          end
+        end
+      end
+
+      # Middleware hook: rewrite tool arguments before the handler runs.
+      # @param context [VectorMCP::Middleware::Context]
+      def before_tool_call(context)
+        return unless context.params.is_a?(Hash)
+
+        arguments = context.params["arguments"]
+        return unless arguments.is_a?(Hash) || arguments.is_a?(Array)
+
+        context.params = context.params.merge("arguments" => sweep_inbound(arguments))
+      end
+
+      # Middleware hook: tokenize matched fields in the tool result.
+      # @param context [VectorMCP::Middleware::Context]
+      def after_tool_call(context)
+        return if context.result.nil?
+
+        context.result = sweep_outbound(context.result)
+      end
+
+      # Register this anonymizer instance on +server+ for the tool call
+      # lifecycle. Creates a thin adapter class so the middleware manager's
+      # argumentless instantiation can still deliver the configured instance.
+      #
+      # @param server [VectorMCP::Server] the server instance.
+      # @param priority [Integer] middleware priority.
+      # @return [Class] the adapter class registered with the server.
+      def install_on(server, priority: Hook::DEFAULT_PRIORITY)
+        instance = self
+        adapter = Class.new(Base) do
+          define_method(:before_tool_call) { |context| instance.before_tool_call(context) }
+          define_method(:after_tool_call)  { |context| instance.after_tool_call(context) }
+        end
+        server.use_middleware(adapter, %i[before_tool_call after_tool_call], priority: priority)
+        adapter
+      end
+
+      private
+
+      def validate_rule!(rule)
+        unless rule.is_a?(Hash) && rule[:pattern].is_a?(Regexp) && rule[:prefix].is_a?(String)
+          raise ArgumentError,
+                "each field_rule must be a Hash with :pattern (Regexp) and :prefix (String)"
+        end
+
+        rule
+      end
+
+      def rule_for(parent_key)
+        return nil if parent_key.nil?
+
+        key_string = parent_key.to_s
+        @field_rules.find { |rule| rule[:pattern].match?(key_string) }
+      end
+
+      def atomic_match?(parent_key)
+        return false if @atomic_keys.nil? || parent_key.nil?
+
+        @atomic_keys.match?(parent_key.to_s)
+      end
+
+      # First pass: collapse Hash nodes whose parent key matches +atomic_keys+
+      # into a single tokenized string. A fresh traversal avoids entanglement
+      # with the field-rule sweep that runs afterwards.
+      def replace_atomic_nodes(obj, parent_key = nil, visited = {}.compare_by_identity)
+        case obj
+        when Hash
+          return obj if visited[obj]
+
+          if atomic_match?(parent_key)
+            @store.tokenize(canonical_json(obj), prefix: atomic_prefix_for(parent_key))
+          else
+            visited[obj] = true
+            begin
+              obj.each_with_object({}) do |(key, value), out|
+                out[key] = replace_atomic_nodes(value, key, visited)
+              end
+            ensure
+              visited.delete(obj)
+            end
+          end
+        when Array
+          return obj if visited[obj]
+
+          visited[obj] = true
+          begin
+            obj.map { |element| replace_atomic_nodes(element, parent_key, visited) }
+          ensure
+            visited.delete(obj)
+          end
+        else
+          obj
+        end
+      end
+
+      # Atomic nodes use the prefix of the first field rule whose pattern
+      # matches the enclosing key. If no field rule matches, fall back to a
+      # neutral default so the token is still well-formed.
+      def atomic_prefix_for(parent_key)
+        rule_for(parent_key)&.dig(:prefix) || "ATOM"
+      end
+
+      def canonical_json(hash)
+        JSON.generate(canonicalize(hash))
+      end
+
+      def canonicalize(obj)
+        case obj
+        when Hash  then obj.keys.map(&:to_s).sort.to_h { |k| [k, canonicalize(obj[k] || obj[k.to_sym])] }
+        when Array then obj.map { |element| canonicalize(element) }
+        else obj
+        end
+      end
+    end
+  end
+end

data/lib/vector_mcp/middleware/hook.rb

@@ -4,7 +4,7 @@ module VectorMCP
   module Middleware
     # Represents a single middleware hook with priority and execution logic
     class Hook
-      attr_reader :middleware_class, :hook_type, :priority, :conditions
+      attr_reader :middleware_class, :hook_type, :operation_type, :priority, :conditions
 
       # Default priority for middleware (lower numbers execute first)
       DEFAULT_PRIORITY = 100
@@ -21,6 +21,8 @@ module VectorMCP
 
         validate_hook_type!
         validate_middleware_class!
+
+        @operation_type = HOOK_OPERATION_TYPES[@hook_type]
       end
 
       # Execute this hook with the given context
@@ -72,12 +74,12 @@ module VectorMCP
         raise ArgumentError, "middleware_class must inherit from VectorMCP::Middleware::Base"
       end
 
+      # Whether this hook's operation_type matches the context's. Transport- and
+      # auth-level hooks (operation_type == nil) match any context.
       def matches_operation_type?(context)
-        operation_prefix = @hook_type.split("_")[1..].join("_")
-
-        return true if transport_or_auth_hook?(operation_prefix)
+        return true if @operation_type.nil?
 
-        operation_matches?(operation_prefix, context.operation_type)
+        context.operation_type == @operation_type
       end
 
       def condition_matches?(key, value, context)
@@ -91,25 +93,6 @@ module VectorMCP
         end
       end
 
-      def transport_or_auth_hook?(operation_prefix)
-        %w[request response transport_error auth auth_error].include?(operation_prefix)
-      end
-
-      def operation_matches?(operation_prefix, operation_type)
-        case operation_prefix
-        when "tool_call", "tool_error"
-          operation_type == :tool_call
-        when "resource_read", "resource_error"
-          operation_type == :resource_read
-        when "prompt_get", "prompt_error"
-          operation_type == :prompt_get
-        when "sampling_request", "sampling_response", "sampling_error"
-          operation_type == :sampling
-        else
-          true
-        end
-      end
-
       def operation_condition_matches?(key, value, context)
         included = Array(value).include?(context.operation_name)
         key == :only_operations ? included : !included

data/lib/vector_mcp/middleware.rb

@@ -12,15 +12,32 @@ module VectorMCP
   # Middleware system for pluggable hooks around MCP operations
   # Allows developers to add custom behavior without modifying core code
   module Middleware
-    # Hook types available in the system
-    HOOK_TYPES = %w[
-      before_tool_call after_tool_call on_tool_error
-      before_resource_read after_resource_read on_resource_error
-      before_prompt_get after_prompt_get on_prompt_error
-      before_sampling_request after_sampling_response on_sampling_error
-      before_request after_response on_transport_error
-      before_auth after_auth on_auth_error
-    ].freeze
+    # Maps each hook type to the operation_type it matches. `nil` means the
+    # hook is transport- or auth-level and matches any operation_type.
+    # This is the single source of truth — HOOK_TYPES is derived from it.
+    HOOK_OPERATION_TYPES = {
+      "before_tool_call" => :tool_call,
+      "after_tool_call" => :tool_call,
+      "on_tool_error" => :tool_call,
+      "before_resource_read" => :resource_read,
+      "after_resource_read" => :resource_read,
+      "on_resource_error" => :resource_read,
+      "before_prompt_get" => :prompt_get,
+      "after_prompt_get" => :prompt_get,
+      "on_prompt_error" => :prompt_get,
+      "before_sampling_request" => :sampling,
+      "after_sampling_response" => :sampling,
+      "on_sampling_error" => :sampling,
+      "before_request" => nil,
+      "after_response" => nil,
+      "on_transport_error" => nil,
+      "before_auth" => nil,
+      "after_auth" => nil,
+      "on_auth_error" => nil
+    }.freeze
+
+    # Hook types available in the system (derived from HOOK_OPERATION_TYPES)
+    HOOK_TYPES = HOOK_OPERATION_TYPES.keys.freeze
 
     # Error raised when invalid hook type is specified
     class InvalidHookTypeError < VectorMCP::Error

data/lib/vector_mcp/request_context.rb

@@ -91,6 +91,25 @@ module VectorMCP
       http_methods.include?(@method.upcase)
     end
 
+    # Return a new context with the given attributes merged in. Nested hashes
+    # (headers, params, transport_metadata) are merged key-wise; scalar
+    # attributes are replaced.
+    #
+    # @param attributes [Hash] attributes to merge
+    # @return [RequestContext] a new, merged context
+    def merge(attributes)
+      current = to_h
+      merged = current.dup
+      attributes.each do |key, value|
+        merged[key] = if value.is_a?(Hash) && current[key].is_a?(Hash)
+                        current[key].merge(value)
+                      else
+                        value
+                      end
+      end
+      RequestContext.new(**merged)
+    end
+
     # Create a minimal request context for non-HTTP transports.
     # This is useful for non-HTTP transports or testing contexts.
     #
@@ -106,6 +125,73 @@ module VectorMCP
       )
     end
 
+    # Coerce an arbitrary request representation into a RequestContext.
+    # This is the single boundary where legacy request shapes (raw Rack
+    # environments, `{headers:, params:}` hashes with symbol or string keys)
+    # are converted; everything downstream of authentication trusts this type.
+    #
+    # @param request [RequestContext, Hash] the request to coerce
+    # @param transport_type [String] transport identifier used when building from a Rack env
+    # @return [RequestContext]
+    # @raise [ArgumentError] if the input cannot be coerced
+    def self.coerce(request, transport_type = "http")
+      case request
+      when RequestContext
+        request
+      when Hash
+        return from_rack_env(request, transport_type) if request.key?("REQUEST_METHOD")
+
+        from_normalized_hash(request)
+      else
+        raise ArgumentError, "Cannot coerce #{request.class} into a VectorMCP::RequestContext"
+      end
+    end
+
+    # Build a context from a legacy `{headers:, params:, ...}` hash with either
+    # symbol or string keys.
+    # @api private
+    def self.from_normalized_hash(hash)
+      fetch = ->(key) { hash[key] || hash[key.to_s] }
+      new(
+        headers: fetch.call(:headers) || {},
+        params: fetch.call(:params) || {},
+        method: fetch.call(:method),
+        path: fetch.call(:path),
+        transport_metadata: fetch.call(:transport_metadata) || {}
+      )
+    end
+    private_class_method :from_normalized_hash
+
+    # Hash-style read access for compatibility with custom authentication
+    # handlers written against the legacy `{headers:, params:}` request hash.
+    #
+    # @param key [String, Symbol] one of :headers, :params, :method, :path, :transport_metadata
+    # @return [Object, nil]
+    def [](key)
+      hash_view[key.to_sym]
+    end
+
+    # @param key [String, Symbol] the top-level key to check
+    # @return [Boolean]
+    def key?(key)
+      hash_view.key?(key.to_sym)
+    end
+    alias has_key? key?
+    alias include? key?
+    alias member? key?
+
+    # Dig into the context like a nested hash, e.g. `context.dig(:headers, "X-API-Key")`.
+    #
+    # @param key [String, Symbol] the top-level key
+    # @param rest [Array] further keys passed to the nested value's #dig
+    # @return [Object, nil]
+    def dig(key, *rest)
+      value = self[key]
+      return value if rest.empty?
+
+      value&.dig(*rest)
+    end
+
     # Create a request context from a Rack environment.
     # This is a convenience method for HTTP-based transports.
     #
@@ -149,25 +235,22 @@ module VectorMCP
 
     private
 
-    # Normalize headers to ensure consistent format.
-    #
-    # @param headers [Hash] Raw headers hash
-    # @return [Hash] Normalized headers hash
-    def normalize_headers(headers)
-      return {} unless headers.is_a?(Hash)
-
-      headers.transform_keys(&:to_s).transform_values { |v| v.nil? ? "" : v.to_s }
+    # Memoized nested-hash view backing the hash-style read accessors.
+    def hash_view
+      @hash_view ||= to_h.freeze
     end
 
-    # Normalize parameters to ensure consistent format.
+    # Normalize a headers/params hash to string keys and string values.
     #
-    # @param params [Hash] Raw parameters hash
-    # @return [Hash] Normalized parameters hash
-    def normalize_params(params)
-      return {} unless params.is_a?(Hash)
+    # @param hash [Hash] Raw hash
+    # @return [Hash] Normalized hash
+    def normalize_string_map(hash)
+      return {} unless hash.is_a?(Hash)
 
-      params.transform_keys(&:to_s).transform_values { |v| v.nil? ? "" : v.to_s }
+      hash.transform_keys(&:to_s).transform_values { |v| v.nil? ? "" : v.to_s }
     end
+    alias normalize_headers normalize_string_map
+    alias normalize_params normalize_string_map
 
     # Normalize transport metadata to ensure consistent format.
     #

data/lib/vector_mcp/security/auth_manager.rb

@@ -1,5 +1,7 @@
 # frozen_string_literal: true
 
+require_relative "auth_result"
+
 module VectorMCP
   module Security
     # Manages authentication strategies for VectorMCP servers
@@ -40,27 +42,26 @@ module VectorMCP
       end
 
       # Authenticate a request using the specified or default strategy
-      # @param request [Hash] the request object containing headers, params, etc.
+      # @param request [VectorMCP::RequestContext, Hash] the request. Legacy hash
+      #   shapes (raw Rack env, {headers:, params:}) are coerced here — the single
+      #   normalization boundary for the authentication pipeline.
       # @param strategy [Symbol] optional strategy override
-      # @return [Object, false] authentication result or false if failed
+      # @return [AuthResult] the authentication outcome
       def authenticate(request, strategy: nil)
-        return { authenticated: true, user: nil } unless @enabled
+        return AuthResult.passthrough unless @enabled
 
         strategy_name = strategy || @default_strategy
         auth_strategy = @strategies[strategy_name]
+        return AuthResult.failure unless auth_strategy
 
-        return { authenticated: false, error: "Unknown strategy: #{strategy_name}" } unless auth_strategy
-
-        begin
-          result = auth_strategy.authenticate(request)
-          if result
-            { authenticated: true, user: result }
-          else
-            { authenticated: false, error: "Authentication failed" }
-          end
-        rescue StandardError => e
-          { authenticated: false, error: "Authentication error: #{e.message}" }
+        result = auth_strategy.authenticate(VectorMCP::RequestContext.coerce(request))
+        if result == false
+          AuthResult.failure
+        else
+          AuthResult.success(user: result, strategy: strategy_name.to_s)
         end
+      rescue StandardError
+        AuthResult.failure
       end
 
       # Check if authentication is required

data/lib/vector_mcp/security/auth_result.rb

@@ -0,0 +1,33 @@
+# frozen_string_literal: true
+
+module VectorMCP
+  module Security
+    # Value object representing the outcome of an authentication attempt.
+    # Replaces the unstructured Hash that previously flowed through the auth pipeline.
+    class AuthResult
+      attr_reader :user, :strategy, :authenticated_at
+
+      def initialize(authenticated:, user: nil, strategy: nil, authenticated_at: nil)
+        @authenticated = authenticated
+        @user = user
+        @strategy = strategy
+        @authenticated_at = authenticated_at || (Time.now if authenticated)
+        freeze
+      end
+
+      def authenticated? = @authenticated
+
+      def self.success(user:, strategy:, authenticated_at: Time.now)
+        new(authenticated: true, user: user, strategy: strategy, authenticated_at: authenticated_at)
+      end
+
+      def self.failure
+        new(authenticated: false)
+      end
+
+      def self.passthrough
+        new(authenticated: true)
+      end
+    end
+  end
+end

data/lib/vector_mcp/security/authorization.rb

@@ -10,6 +10,7 @@ module VectorMCP
       def initialize
         @policies = {}
         @enabled = false
+        @logger = VectorMCP.logger_for("authorization")
       end
 
       # Enable authorization system
@@ -45,17 +46,12 @@ module VectorMCP
 
         resource_type = determine_resource_type(resource)
         policy = @policies[resource_type]
-
-        # If no policy is defined, allow access (opt-in authorization)
         return true unless policy
 
-        begin
-          policy_result = policy.call(user, action, resource)
-          policy_result ? true : false
-        rescue StandardError
-          # Log error but deny access for safety
-          false
-        end
+        !!policy.call(user, action, resource)
+      rescue StandardError => e
+        @logger.error("Authorization policy error for #{resource_type}: #{e.message}")
+        false
       end
 
       # Check if authorization is required

data/lib/vector_mcp/security/middleware.rb

@@ -73,25 +73,6 @@ module VectorMCP
         }
       end
 
-      # Create a request object from different transport formats
-      # @param transport_request [Object] the transport-specific request
-      # @return [Hash] normalized request object
-      def normalize_request(transport_request)
-        case transport_request
-        when Hash
-          # Check if it's a Rack environment (has REQUEST_METHOD key)
-          if transport_request.key?("REQUEST_METHOD")
-            extract_from_rack_env(transport_request)
-          else
-            # Already normalized
-            transport_request
-          end
-        else
-          # Extract from transport-specific request (e.g., custom objects)
-          extract_request_data(transport_request)
-        end
-      end
-
       # Check if security is enabled
       # @return [Boolean] true if any security features are enabled
       def security_enabled?
@@ -113,34 +94,6 @@ module VectorMCP
           }
         }
       end
-
-      private
-
-      # Extract request data from transport-specific formats
-      # @param transport_request [Object] the transport request
-      # @return [Hash] extracted request data
-      def extract_request_data(transport_request)
-        # Handle Rack environment (for HTTP transports)
-        if transport_request.respond_to?(:[]) && transport_request["REQUEST_METHOD"]
-          extract_from_rack_env(transport_request)
-        else
-          # Default fallback for non-HTTP request formats
-          { headers: {}, params: {} }
-        end
-      end
-
-      # Extract data from Rack environment
-      # @param env [Hash] the Rack environment
-      # @return [Hash] extracted request data
-      def extract_from_rack_env(env)
-        {
-          headers: VectorMCP::Util.extract_headers_from_rack_env(env),
-          params: VectorMCP::Util.extract_params_from_rack_env(env),
-          method: env["REQUEST_METHOD"],
-          path: env["PATH_INFO"],
-          rack_env: env
-        }
-      end
     end
   end
 end

data/lib/vector_mcp/security/session_context.rb

@@ -113,34 +113,18 @@ module VectorMCP
         new(authenticated: false)
       end
 
-      # Create an authenticated session context from auth result
-      # @param auth_result [Hash] the authentication result
-      # @return [SessionContext] an authenticated session
+      # Create an authenticated session context from an AuthResult
+      # @param auth_result [VectorMCP::Security::AuthResult] the authentication outcome
+      # @return [SessionContext] an authenticated or anonymous session
       def self.from_auth_result(auth_result)
-        return anonymous unless auth_result&.dig(:authenticated)
-
-        user_data = auth_result[:user]
-
-        # Handle special marker for authenticated nil user
-        if user_data == :authenticated_nil_user
-          new(
-            user: nil,
-            authenticated: true,
-            auth_strategy: "custom",
-            authenticated_at: Time.now
-          )
-        else
-          # Extract strategy and authenticated_at only if user_data is a Hash
-          strategy = user_data.is_a?(Hash) ? user_data[:strategy] : nil
-          auth_time = user_data.is_a?(Hash) ? user_data[:authenticated_at] : nil
-
-          new(
-            user: user_data,
-            authenticated: true,
-            auth_strategy: strategy,
-            authenticated_at: auth_time
-          )
-        end
+        return anonymous unless auth_result&.authenticated?
+
+        new(
+          user: auth_result.user,
+          authenticated: true,
+          auth_strategy: auth_result.strategy,
+          authenticated_at: auth_result.authenticated_at
+        )
       end
     end
   end