snell 0.1.0

data/lib/snell/session_pool.rb

@@ -0,0 +1,100 @@
+# frozen_string_literal: true
+
+module Snell
+  # reactor 内的连接池，每个物理会话最多有一个活跃借用者。
+  # 拨号让出执行权期间，预留项仍计入容量；只有上层确认排空且
+  # 会话健康时才能回到空闲池；此处不重试应用操作。
+  class SessionPool
+    def initialize(limits:)
+      @limits = limits
+      @sessions = []
+      @closed = false
+      @reaper = nil
+    end
+
+    def acquire(reuse: true)
+      raise ClosedError, "Client is closed" if @closed
+
+      expire_idle
+      session = reuse && @sessions.grep(Session).find { |candidate| candidate.state == :ready }
+      unless session
+        raise CapacityError, "Connection limit reached" if @sessions.size >= @limits[:max_connections]
+
+        # 拨号会让出执行权，因此先预留容量，防止其他任务超额接入。
+        reservation = Object.new
+        @sessions << reservation
+        begin
+          session = yield
+          if @closed
+            session.close
+            raise ClosedError, "Client closed during connection attempt"
+          end
+          @sessions[@sessions.index(reservation)] = session
+        ensure
+          @sessions.delete(reservation)
+        end
+      end
+      session.activate!
+      session
+    end
+
+    # 借用者判断协议是否可复用；连接池负责空闲数量和期限限制。
+    def release(session, reusable:)
+      return if @closed || !@sessions.include?(session)
+
+      if reusable && !session.closed?
+        session.ready!
+        idle = @sessions.count { |candidate| candidate.is_a?(Session) && candidate.state == :ready }
+        discard(session) if idle > @limits[:pool_max_idle]
+        schedule_expiration
+      else
+        discard(session)
+      end
+    end
+
+    def close_idle_connections
+      @sessions.grep(Session).select { |session| session.state == :ready }.each { |session| discard(session) }
+      @reaper&.stop
+    end
+
+    def close
+      @closed = true
+      @reaper&.stop
+      @sessions.grep(Session).each(&:close)
+      @sessions.clear
+    end
+
+    private
+
+    def schedule_expiration
+      return if @reaper&.alive?
+
+      # 空闲清理任务不应阻止已完成业务的 reactor 退出。
+      @reaper = Async::Task.current.async(transient: true) do
+        until @closed
+          idle = @sessions.grep(Session).select { |session| session.state == :ready }
+          break if idle.empty?
+
+          deadline = idle.map(&:last_used).min + @limits[:pool_idle_timeout]
+          sleep([deadline - Process.clock_gettime(Process::CLOCK_MONOTONIC), 0].max)
+          expire_idle
+        end
+      ensure
+        @reaper = nil
+      end
+    end
+
+    def expire_idle
+      now = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+      @sessions.grep(Session).each do |session|
+        discard(session) if session.closed? ||
+                            (session.state == :ready && now - session.last_used >= @limits[:pool_idle_timeout])
+      end
+    end
+
+    def discard(session)
+      session.close
+      @sessions.delete(session)
+    end
+  end
+end

data/lib/snell/socks_server.rb

@@ -0,0 +1,371 @@
+# frozen_string_literal: true
+
+require "socket"
+require "ipaddr"
+require "stringio"
+require "io/wait"
+
+module Snell
+  # 仅供本地使用的无认证 SOCKS5 桥接器；UDP 按 RFC 1928 编解码，
+  # 独立于 Client#open_udp 使用的 Snell 报文格式。
+  class SocksServer
+    # SOCKS UDP ASSOCIATE 允许通配地址和零端口，因此协商值
+    # 独立于要求有效目标端口的 Snell::Address。
+    Target = Data.define(:host, :port)
+    class RequestError < StandardError
+      attr_reader :reply
+
+      def initialize(reply)
+        @reply = reply
+        super("invalid SOCKS request")
+      end
+    end
+
+    attr_reader :bound_address
+
+    def initialize(client:, limits: {})
+      @client = client
+      @limits = (limits.is_a?(Limits) ? limits : Limits.new(limits)).to_h
+      @sessions = {}
+      @handshakes = 0
+      @udp_count = 0
+      @tokens = @limits.fetch(:handshake_burst).to_f
+      @token_time = monotonic
+      @stopping = false
+    end
+
+    def run(listen:)
+      raise ArgumentError, "SOCKS listen must be loopback" unless loopback?(listen.host)
+      return if @stopping
+
+      @listener = TCPServer.new(listen.host, listen.port)
+      local = @listener.local_address
+      @bound_address = Address.new(host: local.ip_address, port: local.ip_port)
+      owner = Async::Task.current
+      @accept_task = owner.async do
+        until @stopping
+          socket = accept_socket
+          unless admit?
+            socket.close
+            sleep(0)
+            next
+          end
+          @handshakes += 1
+          start_session(socket, owner)
+        end
+      rescue IOError, SystemCallError
+        raise unless @stopping
+      end
+      @accept_task.wait
+    ensure
+      @listener&.close unless @listener&.closed?
+    end
+
+    def stop(grace: @limits.fetch(:shutdown_grace))
+      @stopping = true
+      @accept_task&.stop
+      @listener&.close unless @listener&.closed?
+      deadline = monotonic + grace
+      sleep([0.01, deadline - monotonic].min) while @sessions.any? && monotonic < deadline
+      @sessions.values.compact.each(&:stop)
+      nil
+    end
+
+    private
+
+    def loopback?(host)
+      host == "localhost" || IPAddr.new(host).loopback?
+    rescue IPAddr::InvalidAddressError
+      false
+    end
+
+    def monotonic
+      Process.clock_gettime(Process::CLOCK_MONOTONIC)
+    end
+
+    def start_session(socket, owner)
+      # 每个 socket 使用独立方法作用域，并在 Async 启动子任务前登记；
+      # 循环中的局部变量可能被后续接入连接的闭包共享。
+      key = Object.new
+      @sessions[key] = nil
+      child = owner.async do
+        serve(socket)
+      ensure
+        socket.close unless socket.closed?
+        @sessions.delete(key)
+      end
+      @sessions[key] = child if @sessions.key?(key)
+    end
+
+    def admit?
+      now = monotonic
+      @tokens = [@limits.fetch(:handshake_burst), @tokens + ((now - @token_time) * @limits.fetch(:handshake_rate))].min
+      @token_time = now
+      return false if @sessions.size >= @limits.fetch(:max_connections)
+      return false if @handshakes >= @limits.fetch(:handshake_concurrency) || @tokens < 1
+
+      @tokens -= 1
+      true
+    end
+
+    def accept_socket
+      loop do
+        result = @listener.accept_nonblock(exception: false)
+        return result unless result == :wait_readable
+
+        @listener.wait_readable
+      end
+    end
+
+    # 协商结束后释放握手名额；连接与 UDP 名额持续到转发清理完成。
+    # 超时或取消时也必须关闭已取得的 resource。
+    def serve(socket)
+      resource = nil
+      command = nil
+      replied = false
+      udp_slot = false
+      begin
+        within(@limits.fetch(:handshake_timeout)) do
+          command, target = negotiate(socket)
+          case command
+          when 1
+            raise RequestError, 8 if target.port.zero?
+
+            resource = @client.connect(target.host, target.port)
+            # Snell 首回复可能等待应用数据；此处只确认请求已发送，
+            # 目标拨号失败仍可能在随后读取时表现为连接关闭。
+            reply(socket, 0)
+          when 3
+            raise RequestError, 1 if @udp_count >= @limits.fetch(:udp_associations)
+
+            @udp_count += 1
+            udp_slot = true
+            resource = @client.open_udp
+          else
+            raise RequestError, 7
+          end
+          replied = command == 1
+        end
+      ensure
+        @handshakes -= 1
+      end
+      if command == 1
+        relay(socket, resource)
+      else
+        replied = true
+        relay_udp(socket, resource)
+      end
+    rescue RequestError => e
+      reply(socket, e.reply) unless replied
+    rescue IOError, SystemCallError
+      nil
+    rescue StandardError
+      reply(socket, 1) unless replied
+    ensure
+      resource&.close
+      @udp_count -= 1 if udp_slot
+    end
+
+    def negotiate(socket)
+      version, count = exact(socket, 2).unpack("CC")
+      raise EOFError unless version == 5 && count.positive?
+
+      methods = exact(socket, count)
+      if methods.bytes.include?(0)
+        socket.write("\x05\x00".b)
+      else
+        socket.write("\x05\xff".b)
+        raise EOFError
+      end
+      version, command, reserved, type = exact(socket, 4).unpack("CCCC")
+      raise RequestError, 1 unless version == 5 && reserved.zero?
+      raise RequestError, 7 unless [1, 3].include?(command)
+
+      [command, read_address(socket, type)]
+    end
+
+    def read_address(source, type)
+      host = case type
+             when 1 then IPAddr.new_ntoh(exact(source, 4)).to_s
+             when 4 then IPAddr.new_ntoh(exact(source, 16)).to_s
+             when 3
+               size = exact(source, 1).unpack1("C")
+               raise RequestError, 8 if size.zero?
+
+               name = exact(source, size)
+               raise RequestError, 8 if !name.ascii_only? || name.match?(/[\x00-\x20\x7f]/)
+
+               name
+             else
+               raise RequestError, 8
+             end
+      Target.new(host: host, port: exact(source, 2).unpack1("n"))
+    end
+
+    def exact(source, size)
+      result = +"".b
+      while result.bytesize < size
+        part = source.read(size - result.bytesize)
+        raise EOFError unless part && !part.empty?
+
+        result << part
+      end
+      result
+    end
+
+    def reply(socket, status, address = Target.new(host: "0.0.0.0", port: 0))
+      socket.write([5, status, 0].pack("CCC") + encode_address(address))
+    rescue IOError, SystemCallError
+      nil
+    end
+
+    def encode_address(address)
+      ip = begin
+        IPAddr.new(address.host)
+      rescue IPAddr::InvalidAddressError
+        nil
+      end
+      host = if ip
+               [ip.ipv4? ? 1 : 4].pack("C") + ip.hton
+             else
+               raise RequestError, 8 unless address.host.bytesize.between?(1, 255)
+
+               [3, address.host.bytesize].pack("CC") + address.host.b
+             end
+      host + [address.port].pack("n")
+    end
+
+    def within(seconds, &)
+      Async::Task.current.with_timeout(seconds, &)
+    end
+
+    def copy(source, destination)
+      loop do
+        data = within(@limits.fetch(:idle_timeout)) { source.readpartial(@limits.fetch(:buffer_bytes)) }
+        break unless data && !data.empty?
+
+        within(@limits.fetch(:idle_timeout)) { destination.write(data) }
+      end
+    rescue EOFError
+      nil
+    ensure
+      begin
+        destination.close_write
+      rescue IOError, SystemCallError
+        nil
+      end
+    end
+
+    def relay(socket, upstream)
+      tasks = [[socket, upstream], [upstream, socket]].map do |source, destination|
+        Async::Task.current.async do
+          copy(source, destination)
+        rescue StandardError
+          socket.close unless socket.closed?
+          upstream.close
+        end
+      end
+      tasks.each(&:wait)
+    ensure
+      tasks&.each(&:stop)
+    end
+
+    # 先绑定 TCP 对端 IP，再固定第一个有效 UDP 报文的来源端口。
+    # 控制连接 EOF 或空闲期限到达时，结束两个转发子任务。
+    def relay_udp(control, association)
+      local = control.local_address
+      peer_ip = IPAddr.new(control.remote_address.ip_address)
+      socket = UDPSocket.new(local.afamily)
+      socket.bind(local.ip_address, 0)
+      bound = socket.local_address
+      reply(control, 0, Target.new(host: bound.ip_address, port: bound.ip_port))
+      peer = nil
+      last_activity = monotonic
+      failed = false
+      children = []
+      children << Async::Task.current.async do
+        received = 0
+        loop do
+          packet, sender = receive_datagram(socket)
+          received += 1
+          sleep(0) if (received % 32).zero?
+          next unless IPAddr.new(sender[3]) == peer_ip
+          next if peer && sender[1] != peer.fetch(:port)
+
+          parsed = decode_datagram(packet)
+          next unless parsed
+
+          address, payload = parsed
+          peer ||= { host: sender[3], port: sender[1] }
+          last_activity = monotonic
+          association.send_packet(payload, to: address)
+        end
+      rescue StandardError
+        failed = true
+      end
+      children << Async::Task.current.async do
+        received = 0
+        loop do
+          datagram = association.receive_packet
+          break unless datagram
+
+          received += 1
+          sleep(0) if (received % 32).zero?
+          next unless peer
+
+          packet = "\0\0\0".b + encode_address(datagram.address) + datagram.payload
+          next if packet.bytesize > 65_507
+
+          send_datagram(socket, packet, peer)
+          last_activity = monotonic
+        end
+      rescue StandardError
+        failed = true
+      end
+      loop do
+        remaining = @limits.fetch(:idle_timeout) - (monotonic - last_activity)
+        break if failed || remaining <= 0
+        next unless control.wait_readable([remaining, 1].min)
+
+        break if control.read_nonblock(1, exception: false).nil?
+      end
+    ensure
+      children&.each(&:stop)
+      socket&.close unless socket&.closed?
+    end
+
+    # RFC 1928 的 RSV/FRAG 必须为零；非法或分片 UDP 报文
+    # 在固定来源端口和向上游发送之前丢弃。
+    def decode_datagram(packet)
+      return unless packet.bytesize >= 4 && packet.byteslice(0, 3) == "\0\0\0".b
+
+      input = StringIO.new(packet)
+      input.pos = 3
+      target = read_address(input, exact(input, 1).unpack1("C"))
+      return if target.port.zero?
+
+      [Address.new(host: target.host, port: target.port), input.read]
+    rescue RequestError, EOFError, ArgumentError
+      nil
+    end
+
+    def receive_datagram(socket)
+      loop do
+        result = socket.recvfrom_nonblock(65_535, exception: false)
+        return result unless result == :wait_readable
+
+        socket.wait_readable
+      end
+    end
+
+    def send_datagram(socket, packet, peer)
+      address = Socket.sockaddr_in(peer.fetch(:port), peer.fetch(:host))
+      loop do
+        result = socket.sendmsg_nonblock(packet, 0, address, exception: false)
+        return result unless result == :wait_writable
+
+        socket.wait_writable
+      end
+    end
+  end
+end

data/lib/snell/stream.rb

@@ -0,0 +1,203 @@
+# frozen_string_literal: true
+
+require "async/semaphore"
+
+module Snell
+  # 采用 IO 风格的逻辑 TCP 流，允许一个读者与一个写者同时推进；
+  # 同方向操作串行执行，首次读取负责确认远端回复。
+  # 关闭时可排空数据以便复用，但不会重新发送失败的应用写入。
+  class Stream
+    def initialize(session, limits:, reusable:, release:)
+      @session = session
+      @limits = limits
+      @reusable = reusable
+      @release = release
+      @owner = Fiber.scheduler
+      @read_lock = Async::Semaphore.new(1)
+      @write_lock = Async::Semaphore.new(1)
+      @buffer = +"".b
+      @reply_received = false
+      @read_eof = false
+      @write_eof = false
+      @closing = false
+      @closed = false
+    end
+
+    # 指定长度时最多返回该数量的字节，EOF 返回 nil；省略长度时
+    # 收集到 EOF。超时按记录计算，不覆盖整个收集过程。
+    def read(length = nil, outbuf = nil)
+      validate_length!(length) unless length.nil?
+      check_owner!
+      check_open!
+      result = @read_lock.acquire do
+        check_open!
+        if length
+          fill_buffer while @buffer.bytesize < length && !@read_eof
+          length.zero? ? +"".b : take_buffer(length)
+        else
+          chunks = [take_buffer(@buffer.bytesize)]
+          until @read_eof
+            fill_buffer
+            chunks << take_buffer(@buffer.bytesize)
+          end
+          chunks.compact.join.b
+        end
+      end
+      outbuf&.replace(result || "")
+      result.nil? ? nil : (outbuf || result)
+    end
+
+    def readpartial(length, outbuf = nil)
+      validate_length!(length)
+      check_owner!
+      check_open!
+      result = @read_lock.acquire do
+        check_open!
+        if length.zero?
+          +"".b
+        else
+          fill_buffer while @buffer.empty? && !@read_eof
+          raise EOFError, "Stream reached EOF" if @buffer.empty?
+
+          take_buffer(length)
+        end
+      end
+      outbuf ? outbuf.replace(result) : result
+    end
+
+    def write(bytes)
+      check_owner!
+      check_open!
+      bytes = String(bytes).b
+      @write_lock.acquire do
+        # 等待锁期间可能已开始关闭；此时不能访问已交给下一个借用者的会话。
+        check_open!
+        raise ClosedError, "Stream write side is closed" if @write_eof
+
+        begin
+          @session.write_stream(bytes) unless bytes.empty?
+        rescue Error, IOError, SystemCallError
+          # 只有实际传输失败才损坏会话；拒绝过期流的操作不能关闭新租约。
+          abort_session
+          raise
+        end
+      end
+      bytes.bytesize
+    end
+
+    def close_write
+      check_owner!
+      return if @closed || @write_eof
+
+      @write_lock.acquire do
+        next if @closed || @write_eof
+
+        @write_eof = true
+        begin
+          @session.write_eof
+        rescue Error, IOError, SystemCallError
+          # reactor 检查和等待后的关闭拒绝不应损坏其他借用者的会话。
+          abort_session
+          raise
+        end
+      end
+      nil
+    end
+
+    # 排空同时受时间与字节预算限制；清理失败时关闭物理会话，
+    # 避免将未排空或已损坏的会话归还为空闲连接。
+    def close
+      check_owner!
+      return if @closed || @closing
+
+      # 排空会让出执行权；先取得关闭权，避免另一个 Fiber 重复归还新租约。
+      # 关闭期间仍需调用 close_write，因此与最终关闭状态分别记录。
+      @closing = true
+
+      begin
+        if @reusable && !@session.closed?
+          @session.draining!
+          Async::Task.current.with_timeout(@limits[:drain_timeout]) do
+            close_write
+            @read_lock.acquire do
+              discarded = @buffer.bytesize
+              @buffer.clear
+              until @read_eof
+                fill_buffer
+                discarded += @buffer.bytesize
+                @buffer.clear
+                raise CapacityError, "Stream drain limit reached" if discarded > @limits[:drain_max_bytes]
+              end
+            end
+          end
+        else
+          @session.close
+        end
+      rescue StandardError
+        @session.close
+      ensure
+        @closed = true
+        @release.call(@session, reusable: @reusable && @read_eof && @write_eof && !@session.closed?)
+      end
+      nil
+    end
+
+    def closed?
+      @closed || @session.closed?
+    end
+
+    def inspect
+      "#<#{self.class} closed=#{closed?} read_eof=#{@read_eof} write_eof=#{@write_eof}>"
+    end
+
+    private
+
+    def fill_buffer
+      record = @session.read_record
+      unless @reply_received
+        raise ProtocolError, "Missing tunnel reply" if record.empty?
+
+        reply = Protocol.decode_reply(record)
+        raise RemoteError.new(reply.message.to_s, code: reply.code) if reply.kind == 2
+        raise ProtocolError, "Unexpected tunnel reply" unless reply.kind.zero?
+
+        @reply_received = true
+        record = reply.payload
+      end
+      if record.empty? && @reply_received
+        # 只有回复头而没有载荷的隧道回复，不是逻辑 EOF 标记。
+        @read_eof = true unless defined?(reply) && reply
+      else
+        @buffer << record
+      end
+    rescue EOFError
+      @read_eof = true
+      @session.close
+    rescue Error, IOError, SystemCallError
+      abort_session
+      raise
+    end
+
+    def take_buffer(length)
+      return nil if @buffer.empty? && @read_eof && length.positive?
+
+      @buffer.slice!(0, length)
+    end
+
+    def abort_session
+      @session.close
+    end
+
+    def check_owner!
+      raise Error, "Stream belongs to another reactor" unless Fiber.scheduler.equal?(@owner)
+    end
+
+    def check_open!
+      raise ClosedError, "Stream is closed" if @closed || @closing
+    end
+
+    def validate_length!(length)
+      raise ArgumentError, "length must be nonnegative" unless length.is_a?(Integer) && length >= 0
+    end
+  end
+end