Internalize json_rpc_handler along with an XSS vulnerability fix

lib/json_rpc_handler.rb

@@ -0,0 +1,151 @@
+# frozen_string_literal: true
+
+require "json"
+
+module JsonRpcHandler
+  class Version
+    V1_0 = "1.0"
+    V2_0 = "2.0"
+  end
+
+  class ErrorCode
+    INVALID_REQUEST = -32600
+    METHOD_NOT_FOUND = -32601
+    INVALID_PARAMS = -32602
+    INTERNAL_ERROR = -32603
+    PARSE_ERROR = -32700
+  end
+
+  DEFAULT_ALLOWED_ID_CHARACTERS = /\A[a-zA-Z0-9_-]+\z/
+
+  extend self
+
+  def handle(request, id_validation_pattern: DEFAULT_ALLOWED_ID_CHARACTERS, &method_finder)
+    if request.is_a?(Array)
+      return error_response(id: :unknown_id, id_validation_pattern:, error: {
+        code: ErrorCode::INVALID_REQUEST,
+        message: "Invalid Request",
+        data: "Request is an empty array",
+      }) if request.empty?
+
+      # Handle batch requests
+      responses = request.map { |req| process_request(req, id_validation_pattern:, &method_finder) }.compact
+
+      # A single item is hoisted out of the array
+      return responses.first if responses.one?
+
+      # An empty array yields nil
+      responses if responses.any?
+    elsif request.is_a?(Hash)
+      # Handle single request
+      process_request(request, id_validation_pattern:, &method_finder)
+    else
+      error_response(id: :unknown_id, id_validation_pattern:, error: {
+        code: ErrorCode::INVALID_REQUEST,
+        message: "Invalid Request",
+        data: "Request must be an array or a hash",
+      })
+    end
+  end
+
+  def handle_json(request_json, id_validation_pattern: DEFAULT_ALLOWED_ID_CHARACTERS, &method_finder)
+    begin
+      request = JSON.parse(request_json, symbolize_names: true)
+      response = handle(request, id_validation_pattern:, &method_finder)
+    rescue JSON::ParserError
+      response = error_response(id: :unknown_id, id_validation_pattern:, error: {
+        code: ErrorCode::PARSE_ERROR,
+        message: "Parse error",
+        data: "Invalid JSON",
+      })
+    end
+
+    response&.to_json
+  end
+
+  def process_request(request, id_validation_pattern:, &method_finder)
+    id = request[:id]
+
+    error = if !valid_version?(request[:jsonrpc])
+      "JSON-RPC version must be 2.0"
+    elsif !valid_id?(request[:id], id_validation_pattern)
+      "Request ID must match validation pattern, or be an integer or null"
+    elsif !valid_method_name?(request[:method])
+      'Method name must be a string and not start with "rpc."'
+    end
+
+    return error_response(id: :unknown_id, id_validation_pattern:, error: {
+      code: ErrorCode::INVALID_REQUEST,
+      message: "Invalid Request",
+      data: error,
+    }) if error
+
+    method_name = request[:method]
+    params = request[:params]
+
+    unless valid_params?(params)
+      return error_response(id:, id_validation_pattern:, error: {
+        code: ErrorCode::INVALID_PARAMS,
+        message: "Invalid params",
+        data: "Method parameters must be an array or an object or null",
+      })
+    end
+
+    begin
+      method = method_finder.call(method_name)
+
+      if method.nil?
+        return error_response(id:, id_validation_pattern:, error: {
+          code: ErrorCode::METHOD_NOT_FOUND,
+          message: "Method not found",
+          data: method_name,
+        })
+      end
+
+      result = method.call(params)
+
+      success_response(id:, result:)
+    rescue StandardError => e
+      error_response(id:, id_validation_pattern:, error: {
+        code: ErrorCode::INTERNAL_ERROR,
+        message: "Internal error",
+        data: e.message,
+      })
+    end
+  end
+
+  def valid_version?(version)
+    version == Version::V2_0
+  end
+
+  def valid_id?(id, pattern = nil)
+    return true if id.nil? || id.is_a?(Integer)
+    return false unless id.is_a?(String)
+
+    pattern ? id.match?(pattern) : true
+  end
+
+  def valid_method_name?(method)
+    method.is_a?(String) && !method.start_with?("rpc.")
+  end
+
+  def valid_params?(params)
+    params.nil? || params.is_a?(Array) || params.is_a?(Hash)
+  end
+
+  def success_response(id:, result:)
+    {
+      jsonrpc: Version::V2_0,
+      id:,
+      result:,
+    } unless id.nil?
+  end
+
+  def error_response(id:, id_validation_pattern:, error:)
+    {
+      jsonrpc: Version::V2_0,
+      id: valid_id?(id, id_validation_pattern) ? id : nil,
+      error: error.compact,
+    } unless id.nil?
+  end
+end

lib/mcp.rb

@@ -1,5 +1,6 @@
 # frozen_string_literal: true
 
+require_relative "json_rpc_handler"
 require_relative "mcp/configuration"
 require_relative "mcp/content"
 require_relative "mcp/instrumentation"

lib/mcp/server.rb

@@ -1,6 +1,6 @@
 # frozen_string_literal: true
 
-require "json_rpc_handler"
+require_relative "../json_rpc_handler"
 require_relative "instrumentation"
 require_relative "methods"
 

mcp.gemspec

@@ -29,6 +29,5 @@ Gem::Specification.new do |spec|
   spec.executables   = spec.files.grep(%r{^exe/}) { |f| File.basename(f) }
   spec.require_paths = ["lib"]
 
-  spec.add_dependency("json_rpc_handler", "~> 0.1")
   spec.add_dependency("json-schema", ">= 4.1")
 end