ruby-utcp 1.1.0

data/lib/utcp/protocols/websocket.rb

@@ -0,0 +1,350 @@
+# frozen_string_literal: true
+
+require "base64"
+require "digest/sha1"
+require "openssl"
+require "socket"
+require "uri"
+
+module UTCP
+  module WebSocketURLSecurity
+    module_function
+
+    def validate!(url)
+      uri = URI.parse(url.to_s)
+      raise SecurityError, "WebSocket URL must use ws or wss" unless %w[ws wss].include?(uri.scheme)
+      raise SecurityError, "WebSocket URL must contain a host" if uri.host.nil? || uri.host.empty?
+      raise SecurityError, "WebSocket URL must not contain user information" if uri.userinfo
+      if uri.scheme == "ws" && !URLSecurity.loopback_host?(uri.host)
+        raise SecurityError, "plain WebSocket is allowed only for loopback hosts"
+      end
+      uri
+    rescue URI::InvalidURIError => error
+      raise SecurityError, "Invalid WebSocket URL: #{error.message}"
+    end
+  end
+
+  class WebSocketConnection
+    GUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
+    MAX_HEADER_SIZE = 65_536
+    MAX_MESSAGE_SIZE = 16 * 1024 * 1024
+
+    def initialize(url, headers = {}, protocol = nil, timeout = 30)
+      @uri = WebSocketURLSecurity.validate!(url)
+      @timeout = Float(timeout)
+      @read_buffer = +"".b
+      @closed = false
+      open_socket
+      handshake(headers, protocol)
+    end
+
+    def closed?
+      @closed || @socket.closed?
+    end
+
+    def send_text(value)
+      send_frame(0x1, value.to_s.encode(Encoding::UTF_8))
+    end
+
+    def send_binary(value)
+      send_frame(0x2, value.to_s.b)
+    end
+
+    def read_message
+      message = +"".b
+      message_opcode = nil
+      loop do
+        fin, opcode, payload = read_frame
+        case opcode
+        when 0x0
+          raise ToolCallError, "unexpected WebSocket continuation frame" unless message_opcode
+          message << payload
+        when 0x1, 0x2
+          raise ToolCallError, "interleaved WebSocket data frame" if message_opcode
+          message_opcode = opcode
+          message << payload
+        when 0x8
+          send_frame(0x8, payload) unless @closed
+          @closed = true
+          return nil
+        when 0x9
+          send_frame(0xA, payload)
+          next
+        when 0xA
+          next
+        else
+          raise ToolCallError, "unsupported WebSocket opcode #{opcode}"
+        end
+        raise ToolCallError, "WebSocket message exceeds #{MAX_MESSAGE_SIZE} bytes" if message.bytesize > MAX_MESSAGE_SIZE
+        return [message_opcode, message] if fin
+      end
+    end
+
+    def close
+      send_frame(0x8, [1000].pack("n")) unless closed?
+    rescue IOError, SystemCallError
+      nil
+    ensure
+      @closed = true
+      @socket.close if @socket && !@socket.closed?
+    end
+
+    private
+
+    def open_socket
+      tcp = Socket.tcp(@uri.host, websocket_port, connect_timeout: @timeout)
+      if @uri.scheme == "wss"
+        context = OpenSSL::SSL::SSLContext.new
+        context.set_params(verify_mode: OpenSSL::SSL::VERIFY_PEER)
+        ssl = OpenSSL::SSL::SSLSocket.new(tcp, context)
+        ssl.hostname = @uri.host if ssl.respond_to?(:hostname=)
+        ssl.sync_close = true
+        ssl.connect
+        @socket = ssl
+      else
+        @socket = tcp
+      end
+    rescue StandardError
+      tcp.close if tcp && !tcp.closed?
+      raise
+    end
+
+    def handshake(headers, protocol)
+      key = Base64.strict_encode64(Random.new.bytes(16))
+      path = websocket_request_target
+      host = @uri.host.include?(":") ? "[#{@uri.host}]" : @uri.host
+      default_port = @uri.scheme == "wss" ? 443 : 80
+      host = "#{host}:#{websocket_port}" unless websocket_port == default_port
+      values = {
+        "Host" => host,
+        "Upgrade" => "websocket",
+        "Connection" => "Upgrade",
+        "Sec-WebSocket-Key" => key,
+        "Sec-WebSocket-Version" => "13"
+      }
+      values["Sec-WebSocket-Protocol"] = protocol if protocol
+      reserved = %w[host upgrade connection sec-websocket-key sec-websocket-version sec-websocket-protocol]
+      Utils.stringify_keys(headers || {}).each do |name, value|
+        raise SecurityError, "WebSocket header #{name.inspect} is reserved" if reserved.include?(name.downcase)
+
+        values[name] = value.to_s
+      end
+      values.each do |name, value|
+        raise SecurityError, "WebSocket header contains CR/LF" if name.to_s.match?(/[\r\n]/) || value.match?(/[\r\n]/)
+      end
+      request = "GET #{path} HTTP/1.1\r\n" + values.map { |name, value| "#{name}: #{value}\r\n" }.join + "\r\n"
+      @socket.write(request)
+      header_text = read_headers
+      lines = header_text.split("\r\n")
+      status = lines.shift.to_s.split[1].to_i
+      raise ToolCallError, "WebSocket handshake failed with status #{status}" unless status == 101
+
+      response_headers = lines.each_with_object({}) do |line, result|
+        name, value = line.split(":", 2)
+        result[name.to_s.downcase] = value.to_s.strip
+      end
+      expected = Base64.strict_encode64(Digest::SHA1.digest(key + GUID))
+      unless secure_compare(response_headers["sec-websocket-accept"].to_s, expected)
+        raise SecurityError, "WebSocket handshake returned an invalid Sec-WebSocket-Accept"
+      end
+      unless response_headers["upgrade"].to_s.casecmp?("websocket") &&
+             response_headers["connection"].to_s.downcase.split(/\s*,\s*/).include?("upgrade")
+        raise ToolCallError, "WebSocket handshake did not confirm the protocol upgrade"
+      end
+      if protocol && response_headers["sec-websocket-protocol"] != protocol
+        raise ToolCallError, "WebSocket server did not select the requested subprotocol"
+      end
+    end
+
+    def websocket_port
+      @uri.port || (@uri.scheme == "wss" ? 443 : 80)
+    end
+
+    def websocket_request_target
+      path = @uri.path.to_s
+      path = "/" if path.empty?
+      @uri.query ? "#{path}?#{@uri.query}" : path
+    end
+
+    def read_headers
+      until (index = @read_buffer.index("\r\n\r\n"))
+        raise ToolCallError, "WebSocket handshake headers are too large" if @read_buffer.bytesize >= MAX_HEADER_SIZE
+        wait_readable
+        @read_buffer << @socket.readpartial(4096)
+      end
+      @read_buffer.slice!(0, index + 4).byteslice(0, index)
+    end
+
+    def send_frame(opcode, payload)
+      raise IOError, "WebSocket is closed" if closed?
+
+      bytes = payload.to_s.b
+      mask = Random.new.bytes(4)
+      header = [0x80 | opcode].pack("C")
+      header << if bytes.bytesize < 126
+                  [0x80 | bytes.bytesize].pack("C")
+                elsif bytes.bytesize <= 65_535
+                  [0x80 | 126, bytes.bytesize].pack("Cn")
+                else
+                  [0x80 | 127, bytes.bytesize].pack("CQ>")
+                end
+      masked = bytes.bytes.each_with_index.map { |byte, index| byte ^ mask.getbyte(index % 4) }.pack("C*")
+      @socket.write(header + mask + masked)
+    end
+
+    def read_frame
+      head = read_exact(2)
+      first, second = head.unpack("CC")
+      fin = (first & 0x80) != 0
+      opcode = first & 0x0F
+      masked = (second & 0x80) != 0
+      length = second & 0x7F
+      length = read_exact(2).unpack1("n") if length == 126
+      length = read_exact(8).unpack1("Q>") if length == 127
+      raise ToolCallError, "WebSocket frame exceeds #{MAX_MESSAGE_SIZE} bytes" if length > MAX_MESSAGE_SIZE
+
+      mask = masked ? read_exact(4) : nil
+      payload = read_exact(length)
+      if mask
+        payload = payload.bytes.each_with_index.map { |byte, index| byte ^ mask.getbyte(index % 4) }.pack("C*")
+      end
+      [fin, opcode, payload]
+    end
+
+    def read_exact(length)
+      while @read_buffer.bytesize < length
+        wait_readable
+        @read_buffer << @socket.readpartial([4096, length - @read_buffer.bytesize].max)
+      end
+      @read_buffer.slice!(0, length)
+    rescue EOFError
+      @closed = true
+      raise ToolCallError, "WebSocket connection closed unexpectedly"
+    end
+
+    def wait_readable
+      raise TimeoutError, "WebSocket read timed out" unless IO.select([@socket], nil, nil, @timeout)
+    end
+
+    def secure_compare(first, second)
+      return false unless first.bytesize == second.bytesize
+
+      first.bytes.zip(second.bytes).reduce(0) { |memo, pair| memo | (pair[0] ^ pair[1]) }.zero?
+    end
+  end
+
+  class WebSocketProtocol < HTTPProtocol
+    ConnectionEntry = Struct.new(:connection, :mutex)
+
+    def initialize(connection_factory: nil, **options)
+      super(**options)
+      @connection_factory = connection_factory || lambda do |url, headers, protocol, timeout|
+        WebSocketConnection.new(url, headers, protocol, timeout)
+      end
+      @connections = {}
+      @connections_mutex = Mutex.new
+    end
+
+    def register_manual(client, template)
+      assert_websocket_template!(template)
+      entry, transient = connection_for(template, {})
+      payload = entry.mutex.synchronize do
+        entry.connection.send_text(JSON.generate("type" => "utcp"))
+        _opcode, bytes = entry.connection.read_message
+        bytes
+      end
+      entry.connection.close if transient
+      success(template, manual_from_payload(template, payload, source: "WebSocket discovery response"))
+    rescue StandardError => error
+      client.logger.warn("Unable to register WebSocket manual #{template.name.inspect}: #{error.message}")
+      failure(template, error)
+    end
+
+    def deregister_manual(_client, template)
+      prefix = [template.url, template.protocol].join("\0") + "\0"
+      entries = @connections_mutex.synchronize do
+        keys = @connections.keys.select { |key| key.start_with?(prefix) }
+        keys.map { |key| @connections.delete(key) }
+      end
+      entries.each { |entry| entry.connection.close }
+      nil
+    end
+
+    def call_tool(_client, tool_name, tool_args, template)
+      assert_websocket_template!(template)
+      args = Utils.stringify_keys(tool_args || {})
+      entry, transient, message_args = connection_for(template, args, include_arguments: true)
+      result = entry.mutex.synchronize do
+        message = format_message(template, message_args)
+        entry.connection.send_text(message)
+        frame = entry.connection.read_message
+        raise ToolCallError, "WebSocket closed without a response" unless frame
+
+        decode_message(frame[1], template.response_format, frame[0])
+      end
+      result
+    rescue Error
+      raise
+    rescue StandardError => error
+      raise ToolCallError.new("WebSocket tool #{tool_name.inspect} failed: #{error.message}", tool_name: tool_name)
+    ensure
+      entry.connection.close if defined?(entry) && entry && defined?(transient) && transient
+    end
+
+    private
+
+    def assert_websocket_template!(template)
+      return if template.is_a?(WebSocketCallTemplate)
+
+      raise ValidationError, "WebSocket protocol requires a WebSocketCallTemplate"
+    end
+
+    def connection_for(template, arguments, include_arguments: false)
+      args = arguments.dup
+      headers = Utils.stringify_keys(template.headers || {})
+      template.header_fields.each { |field| headers[field] = args.delete(field).to_s if args.key?(field) }
+      query = {}
+      cookies = {}
+      apply_auth(template.auth, headers, query, cookies)
+      if template.auth.is_a?(OAuth2Auth)
+        headers["Authorization"] = "Bearer #{oauth_token(template.auth)}"
+      end
+      headers["Cookie"] = cookies.map { |key, value| "#{key}=#{value}" }.join("; ") unless cookies.empty?
+      url = append_query(template.url, query)
+      WebSocketURLSecurity.validate!(url)
+      key = connection_key(template, url, headers)
+      if template.keep_alive
+        entry = @connections_mutex.synchronize do
+          current = @connections[key]
+          if current.nil? || current.connection.closed?
+            current = ConnectionEntry.new(@connection_factory.call(url, headers, template.protocol, template.timeout), Mutex.new)
+            @connections[key] = current
+          end
+          current
+        end
+        include_arguments ? [entry, false, args] : [entry, false]
+      else
+        entry = ConnectionEntry.new(@connection_factory.call(url, headers, template.protocol, template.timeout), Mutex.new)
+        include_arguments ? [entry, true, args] : [entry, true]
+      end
+    end
+
+    def connection_key(template, url = template.url, headers = template.headers)
+      [url, template.protocol, JSON.generate(Utils.stringify_keys(headers || {}).sort)].join("\0")
+    end
+
+    def format_message(template, arguments)
+      value = template.message.nil? ? arguments : substitute_message_template(template.message, arguments)
+      value.is_a?(String) ? value : JSON.generate(value)
+    end
+
+    def decode_message(bytes, response_format, opcode)
+      return bytes.b if response_format == "raw" || opcode == 0x2
+
+      text = bytes.dup.force_encoding(Encoding::UTF_8)
+      response_format == "text" ? text : decode_json_or_text(text)
+    end
+  end
+  WebsocketCommunicationProtocol = WebSocketProtocol
+  WebSocketCommunicationProtocol = WebSocketProtocol
+end

data/lib/utcp/registry.rb

@@ -0,0 +1,67 @@
+# frozen_string_literal: true
+
+require "thread"
+
+module UTCP
+  @registry_mutex = Mutex.new
+  @call_template_classes = {}
+  @auth_classes = {}
+  @protocols = {}
+
+  class << self
+    def register_call_template(type, klass)
+      validate_plugin_type!(type)
+      unless klass.respond_to?(:new)
+        raise ArgumentError, "call template implementation must be a class-like object"
+      end
+      @registry_mutex.synchronize { @call_template_classes[type.to_s] = klass }
+      klass
+    end
+
+    def call_template_class(type)
+      @registry_mutex.synchronize { @call_template_classes[type.to_s] }
+    end
+
+    def call_template_types
+      @registry_mutex.synchronize { @call_template_classes.keys.sort }
+    end
+
+    def register_auth(type, klass)
+      validate_plugin_type!(type)
+      @registry_mutex.synchronize { @auth_classes[type.to_s] = klass }
+      klass
+    end
+
+    def auth_class(type)
+      @registry_mutex.synchronize { @auth_classes[type.to_s] }
+    end
+
+    def register_protocol(type, implementation)
+      validate_plugin_type!(type)
+      protocol = implementation.is_a?(Class) ? implementation.new : implementation
+      required = %i[register_manual deregister_manual call_tool]
+      missing = required.reject { |method_name| protocol.respond_to?(method_name) }
+      raise ArgumentError, "protocol is missing: #{missing.join(', ')}" unless missing.empty?
+
+      @registry_mutex.synchronize { @protocols[type.to_s] = protocol }
+      protocol
+    end
+
+    def protocol(type)
+      @registry_mutex.synchronize { @protocols[type.to_s] }
+    end
+
+    def protocol_types
+      @registry_mutex.synchronize { @protocols.keys.sort }
+    end
+
+    private
+
+    def validate_plugin_type!(type)
+      unless type.to_s.match?(/\A[a-zA-Z0-9_]+\z/)
+        raise ArgumentError, "plugin type must contain only letters, numbers, and underscores"
+      end
+    end
+  end
+end
+

data/lib/utcp/repository.rb

@@ -0,0 +1,137 @@
+# frozen_string_literal: true
+
+module UTCP
+  class InMemoryToolRepository
+    attr_reader :tool_repository_type
+
+    def initialize
+      @tool_repository_type = "in_memory"
+      @tools = {}
+      @manuals = {}
+      @templates = {}
+      @mutex = Mutex.new
+    end
+
+    def save_manual(manual_call_template, manual)
+      @mutex.synchronize do
+        old_manual = @manuals[manual_call_template.name]
+        old_manual&.tools&.each { |tool| @tools.delete(tool.name) }
+
+        @templates[manual_call_template.name] = manual_call_template
+        @manuals[manual_call_template.name] = manual
+        manual.tools.each { |tool| @tools[tool.name] = tool }
+      end
+      nil
+    end
+
+    def remove_manual(manual_name)
+      @mutex.synchronize do
+        manual = @manuals.delete(manual_name.to_s)
+        return false unless manual
+
+        manual.tools.each { |tool| @tools.delete(tool.name) }
+        @templates.delete(manual_name.to_s)
+        true
+      end
+    end
+
+    def remove_tool(tool_name)
+      @mutex.synchronize do
+        tool = @tools.delete(tool_name.to_s)
+        return false unless tool
+
+        @manuals.each_value { |manual| manual.tools.delete_if { |candidate| candidate.name == tool.name } }
+        true
+      end
+    end
+
+    def get_tool(tool_name)
+      @mutex.synchronize { @tools[tool_name.to_s] }
+    end
+
+    def get_tools
+      @mutex.synchronize { @tools.values.dup }
+    end
+
+    def get_tools_by_manual(manual_name)
+      @mutex.synchronize do
+        manual = @manuals[manual_name.to_s]
+        manual&.tools&.dup
+      end
+    end
+
+    def get_manual(manual_name)
+      @mutex.synchronize { @manuals[manual_name.to_s] }
+    end
+
+    def get_manuals
+      @mutex.synchronize { @manuals.values.dup }
+    end
+
+    def get_manual_call_template(manual_name)
+      @mutex.synchronize { @templates[manual_name.to_s] }
+    end
+
+    def get_manual_call_templates
+      @mutex.synchronize { @templates.values.dup }
+    end
+
+    def to_h
+      { "tool_repository_type" => tool_repository_type }
+    end
+  end
+  InMemToolRepository = InMemoryToolRepository
+
+  class TagSearchStrategy
+    attr_reader :tool_search_strategy_type, :description_weight, :tag_weight
+
+    def initialize(description_weight: 1, tag_weight: 3)
+      @tool_search_strategy_type = "tag_and_description_word_match"
+      @description_weight = Float(description_weight)
+      @tag_weight = Float(tag_weight)
+    end
+
+    def search_tools(tool_repository:, query:, limit: 10, any_of_tags_required: nil)
+      raise ArgumentError, "limit must be non-negative" if limit.negative?
+
+      query_text = query.to_s.downcase
+      query_words = query_text.scan(/[[:alnum:]_]+/).uniq
+      required_tags = Array(any_of_tags_required).map { |tag| tag.to_s.downcase }
+      tools = tool_repository.get_tools
+      unless required_tags.empty?
+        tools = tools.select do |tool|
+          (tool.tags.map(&:downcase) & required_tags).any?
+        end
+      end
+
+      scored = tools.each_with_index.map do |tool, index|
+        score = tool.tags.sum do |tag|
+          normalized = tag.downcase
+          if query_text.include?(normalized) || (normalized.scan(/[[:alnum:]_]+/) & query_words).any?
+            tag_weight
+          else
+            0
+          end
+        end
+        description_words = tool.description.downcase.scan(/[[:alnum:]_]+/).uniq
+        score += (description_words & query_words).count { |word| word.length > 2 } * description_weight
+        name_words = tool.name.downcase.scan(/[[:alnum:]_]+/).uniq
+        score += (name_words & query_words).length * description_weight
+        [tool, score, index]
+      end
+
+      results = scored.sort_by { |_tool, score, index| [-score, index] }.map(&:first)
+      limit.zero? ? results : results.first(limit)
+    end
+
+    def to_h
+      {
+        "tool_search_strategy_type" => tool_search_strategy_type,
+        "description_weight" => description_weight,
+        "tag_weight" => tag_weight
+      }
+    end
+  end
+  TagAndDescriptionWordMatchStrategy = TagSearchStrategy
+end
+

data/lib/utcp/serializer.rb

@@ -0,0 +1,71 @@
+# frozen_string_literal: true
+
+module UTCP
+  class Serializer
+    def initialize(model_class)
+      @model_class = model_class
+    end
+
+    def to_dict(object)
+      unless object.respond_to?(:to_h)
+        raise SerializerValidationError, "#{object.class} cannot be serialized"
+      end
+
+      object.to_h
+    end
+    alias to_h to_dict
+
+    def validate_dict(value)
+      @model_class.from_h(value)
+    rescue Error
+      raise
+    rescue StandardError => error
+      raise SerializerValidationError, error.message
+    end
+    alias from_h validate_dict
+
+    def copy(object)
+      validate_dict(to_dict(object))
+    end
+  end
+
+  class AuthSerializer < Serializer
+    def initialize
+      super(Auth)
+    end
+  end
+
+  class CallTemplateSerializer < Serializer
+    def initialize
+      super(CallTemplate)
+    end
+  end
+
+  class JsonSchemaSerializer < Serializer
+    def initialize
+      super(JsonSchema)
+    end
+  end
+  JSONSchemaSerializer = JsonSchemaSerializer
+
+  class ToolSerializer < Serializer
+    def initialize
+      super(Tool)
+    end
+  end
+
+  class ManualSerializer < Serializer
+    def initialize
+      super(Manual)
+    end
+  end
+  UtcpManualSerializer = ManualSerializer
+
+  class ClientConfigSerializer < Serializer
+    def initialize
+      super(ClientConfig)
+    end
+  end
+  UtcpClientConfigSerializer = ClientConfigSerializer
+end
+