<commit_before>require 'stringio'


module SpiderGazelle
    class Connection
        Hijack = Struct.new(:socket, :env)

        REQUEST_METHOD = 'REQUEST_METHOD'.freeze # NOTE:: duplicate in gazelle.rb
        HEAD = 'HEAD'.freeze

        RACK = 'rack'.freeze            # used for filtering headers
        CLOSE = 'close'.freeze
        CONNECTION = 'Connection'.freeze
        CONTENT_LENGTH = 'Content-Length'.freeze
        TRANSFER_ENCODING = 'Transfer-Encoding'.freeze
        CHUNKED = 'chunked'.freeze
        COLON_SPACE = ': '.freeze
        EOF = "0\r\n\r\n".freeze
        CRLF = "\r\n".freeze
        ZERO = '0'.freeze

        HTTP_STATUS_CODES = ::Rack::Utils::HTTP_STATUS_CODES
        HTTP_STATUS_DEFAULT = proc { 'CUSTOM' }

        HTTP_11_400 = "HTTP/1.1 400 Bad Request\r\n\r\n".freeze


        def self.on_progress(data, socket); end
        DUMMY_PROGRESS = self.method(:on_progress)


        # For Gazelle
        attr_reader :state, :parsing
        # For Request
        attr_reader :tls, :port, :loop, :socket, :async_callback


        def initialize(gazelle, loop, socket, port, state, app, queue)
            # A single parser instance per-connection (supports pipelining)
            @state = state
            @pending = []

            # Work callback for thread pool processing
            @request = nil
            @work = method(:work)

            # Called after the work on the thread pool is complete
            @send_response = method(:send_response)
            @send_error = method(:send_error)

            # Used to chain promises (ensures requests are processed in order)
            @process_next = method(:process_next)
            @write_chunk = method(:write_chunk)
            @current_worker = queue # keep track of work queue head to prevent unintentional GC
            @queue_worker = queue   # start queue with an existing resolved promise (::Libuv::Q::ResolvedPromise.new(@loop, true))
            
            # Socket for writing the response
            @socket = socket
            @app = app
            @port = port
            @tls = @socket.tls?
            @loop = loop
            @gazelle = gazelle
            @async_callback = method(:deferred_callback)

            # Remove connection if the socket closes
            socket.finally &method(:unlink)
        end

        # Lazy eval the IP
        def remote_ip
            @remote_ip ||= @socket.peername[0]
        end

        # Creates a new request state object
        def start_parsing
            @parsing = Request.new(self, @app)
        end

        # Chains the work in a promise queue
        def finished_parsing
            if !@state.keep_alive?
                @parsing.keep_alive = false
                @socket.stop_read   # we don't want to do any more work than we need to
            end
            @parsing.upgrade = @state.upgrade?
            @pending.push @parsing
            @queue_worker = @queue_worker.then @process_next
        end

        # The parser encountered an error
        def parsing_error
            # Grab the error
            send_error(@state.error)

            # We no longer care for any further requests from this client
            # however we will finish processing any valid pipelined requests before shutting down
            @socket.stop_read
            @queue_worker = @queue_worker.then do
                # TODO:: send response (400 bad request)
                @socket.write HTTP_11_400
                @socket.shutdown
            end
        end

        # Schedule send
        def response(data)
            @loop.schedule
        end


        protected


        # --------------
        # State handlers:
        # --------------


        # Called when an error occurs at any point while responding
        def send_error(reason)
            # Close the socket as this is fatal (file read error, gazelle error etc)
            @socket.close

            # Log the error in a worker thread
            @loop.work do
                msg = "connection error: #{reason.message}\n#{reason.backtrace.join("\n") if reason.backtrace}\n"
                puts msg
                @gazelle.logger.error msg
            end
        end

        # We use promise chaining to move the requests forward
        # This provides an elegant way to handle persistent and pipelined connections
        def process_next(result)
            @request = @pending.shift
            @current_worker = @loop.work @work
            @current_worker.then @send_response, @send_error   # resolves the promise with a promise
        end

        # returns the response as the result of the work
        # We support the unofficial rack async api (multi-call version for chunked responses)
        def work
            @request.execute!
        end

        # Unlinks the connection from the rack app
        # This occurs when requested and when the socket closes
        def unlink
            if not @gazelle.nil?
                @socket.progress &DUMMY_PROGRESS        # unlink the progress callback (prevent funny business)
                @gazelle.discard(self)
                @gazelle = nil
                @state = nil
            end
        end


        # ----------------------
        # Core response handlers:
        # ----------------------


        def send_response(result)
            # As we have come back from another thread the socket may have closed
            # This check is an optimisation, the call to write and shutdown would fail safely

            if @request.hijacked
                unlink              # unlink the management of the socket
                        
                # Pass the hijack response to the captor using the promise
                # This forwards the socket and environment as well as moving
                #  continued execution onto the event loop.
                @request.hijacked.resolve(Hijack.new(@socket, @request.env))

            elsif !@socket.closed
                if @request.deferred
                    # Wait for the response using this promise
                    promise = @request.deferred.promise

                    # Process any responses that might have made it here first
                    if @deferred_responses
                        @deferred_responses.each &method(:respond_with)
                        @deferred_responses = nil
                    end

                    return promise

                # NOTE:: Somehow getting to here with a nil request... needs investigation
                elsif not result.nil?
                    # clear any cached responses just in case
                    # could be set by error in the rack application
                    @deferred_responses = nil if @deferred_responses

                    status, headers, body = result
                    send_body = @request.env[REQUEST_METHOD] != HEAD

                    # If a file, stream the body in a non-blocking fashion
                    if body.respond_to? :to_path
                        if headers[CONTENT_LENGTH]
                            type = :raw
                        else
                            type = :http
                            headers[TRANSFER_ENCODING] = CHUNKED
                        end

                        write_headers(status, headers)

                        if send_body
                            file = @loop.file(body.to_path, File::RDONLY)
                            file.progress do    # File is open and available for reading
                                file.send_file(@socket, type).finally do
                                    file.close
                                    if @request.keep_alive == false
                                        @socket.shutdown
                                    end
                                end
                            end
                            return file
                        end
                    else
                        # Optimize the response
                        begin
                            body_size = body.size
                            if body_size < 2
                                if body_size == 1
                                    headers[CONTENT_LENGTH] = body[0].bytesize
                                else
                                    headers[CONTENT_LENGTH] = ZERO
                                end
                            end
                        rescue # just in case
                        end

                        if send_body
                            write_response(status, headers, body)
                        else
                            write_headers(status, headers)
                            @socket.shutdown if @request.keep_alive == false
                        end
                    end
                end
            end

            # continue processing (don't wait for write to complete)
            # if the write fails it will close the socket
            nil
        end

        def write_response(status, headers, body)
            if headers[CONTENT_LENGTH]
                headers[CONTENT_LENGTH] = headers[CONTENT_LENGTH].to_s
                write_headers(status, headers)

                # Stream the response (pass directly into @socket.write)
                body.each &@socket.method(:write)

                if @request.deferred
                    @request.deferred.resolve(true)
                    @request.deferred = nil # prevent data being sent after completed
                end

                @socket.shutdown if @request.keep_alive == false
            else
                headers[TRANSFER_ENCODING] = CHUNKED
                write_headers(status, headers)

                # Stream the response
                body.each &@write_chunk

                if @request.deferred.nil?
                    @socket.write EOF
                    @socket.shutdown if @request.keep_alive == false
                else
                    @async_state = :chunked
                end
            end
        end

        def write_headers(status, headers)
            headers[CONNECTION] = CLOSE if @request.keep_alive == false

            header = "HTTP/1.1 #{status} #{fetch_code(status)}\r\n"
            headers.each do |key, value|
                next if key.start_with? RACK

                header << key
                header << COLON_SPACE
                header << value
                header << CRLF
            end
            header << CRLF
            @socket.write header
        end

        def write_chunk(part)
            chunk = part.bytesize.to_s(16) << CRLF << part << CRLF
            @socket.write chunk
        end

        def fetch_code(status)
            HTTP_STATUS_CODES.fetch(status, &HTTP_STATUS_DEFAULT)
        end


        # ------------------------
        # Async response functions:
        # ------------------------


        # Callback from a response that was marked async<commit_msg>Fixed response headers not being sent properly when more than one header of the same kind was set by the app.<commit_after>require 'stringio'


module SpiderGazelle
    class Connection
        Hijack = Struct.new(:socket, :env)

        REQUEST_METHOD = 'REQUEST_METHOD'.freeze # NOTE:: duplicate in gazelle.rb
        HEAD = 'HEAD'.freeze

        RACK = 'rack'.freeze            # used for filtering headers
        CLOSE = 'close'.freeze
        CONNECTION = 'Connection'.freeze
        CONTENT_LENGTH = 'Content-Length'.freeze
        TRANSFER_ENCODING = 'Transfer-Encoding'.freeze
        CHUNKED = 'chunked'.freeze
        COLON_SPACE = ': '.freeze
        EOF = "0\r\n\r\n".freeze
        CRLF = "\r\n".freeze
        NEWLINE = "\n".freeze
        ZERO = '0'.freeze

        HTTP_STATUS_CODES = ::Rack::Utils::HTTP_STATUS_CODES
        HTTP_STATUS_DEFAULT = proc { 'CUSTOM' }

        HTTP_11_400 = "HTTP/1.1 400 Bad Request\r\n\r\n".freeze


        def self.on_progress(data, socket); end
        DUMMY_PROGRESS = self.method(:on_progress)


        # For Gazelle
        attr_reader :state, :parsing
        # For Request
        attr_reader :tls, :port, :loop, :socket, :async_callback


        def initialize(gazelle, loop, socket, port, state, app, queue)
            # A single parser instance per-connection (supports pipelining)
            @state = state
            @pending = []

            # Work callback for thread pool processing
            @request = nil
            @work = method(:work)

            # Called after the work on the thread pool is complete
            @send_response = method(:send_response)
            @send_error = method(:send_error)

            # Used to chain promises (ensures requests are processed in order)
            @process_next = method(:process_next)
            @write_chunk = method(:write_chunk)
            @current_worker = queue # keep track of work queue head to prevent unintentional GC
            @queue_worker = queue   # start queue with an existing resolved promise (::Libuv::Q::ResolvedPromise.new(@loop, true))
            
            # Socket for writing the response
            @socket = socket
            @app = app
            @port = port
            @tls = @socket.tls?
            @loop = loop
            @gazelle = gazelle
            @async_callback = method(:deferred_callback)

            # Remove connection if the socket closes
            socket.finally &method(:unlink)
        end

        # Lazy eval the IP
        def remote_ip
            @remote_ip ||= @socket.peername[0]
        end

        # Creates a new request state object
        def start_parsing
            @parsing = Request.new(self, @app)
        end

        # Chains the work in a promise queue
        def finished_parsing
            if !@state.keep_alive?
                @parsing.keep_alive = false
                @socket.stop_read   # we don't want to do any more work than we need to
            end
            @parsing.upgrade = @state.upgrade?
            @pending.push @parsing
            @queue_worker = @queue_worker.then @process_next
        end

        # The parser encountered an error
        def parsing_error
            # Grab the error
            send_error(@state.error)

            # We no longer care for any further requests from this client
            # however we will finish processing any valid pipelined requests before shutting down
            @socket.stop_read
            @queue_worker = @queue_worker.then do
                # TODO:: send response (400 bad request)
                @socket.write HTTP_11_400
                @socket.shutdown
            end
        end

        # Schedule send
        def response(data)
            @loop.schedule
        end


        protected


        # --------------
        # State handlers:
        # --------------


        # Called when an error occurs at any point while responding
        def send_error(reason)
            # Close the socket as this is fatal (file read error, gazelle error etc)
            @socket.close

            # Log the error in a worker thread
            @loop.work do
                msg = "connection error: #{reason.message}\n#{reason.backtrace.join("\n") if reason.backtrace}\n"
                puts msg
                @gazelle.logger.error msg
            end
        end

        # We use promise chaining to move the requests forward
        # This provides an elegant way to handle persistent and pipelined connections
        def process_next(result)
            @request = @pending.shift
            @current_worker = @loop.work @work
            @current_worker.then @send_response, @send_error   # resolves the promise with a promise
        end

        # returns the response as the result of the work
        # We support the unofficial rack async api (multi-call version for chunked responses)
        def work
            @request.execute!
        end

        # Unlinks the connection from the rack app
        # This occurs when requested and when the socket closes
        def unlink
            if not @gazelle.nil?
                @socket.progress &DUMMY_PROGRESS        # unlink the progress callback (prevent funny business)
                @gazelle.discard(self)
                @gazelle = nil
                @state = nil
            end
        end


        # ----------------------
        # Core response handlers:
        # ----------------------


        def send_response(result)
            # As we have come back from another thread the socket may have closed
            # This check is an optimisation, the call to write and shutdown would fail safely

            if @request.hijacked
                unlink              # unlink the management of the socket
                        
                # Pass the hijack response to the captor using the promise
                # This forwards the socket and environment as well as moving
                #  continued execution onto the event loop.
                @request.hijacked.resolve(Hijack.new(@socket, @request.env))

            elsif !@socket.closed
                if @request.deferred
                    # Wait for the response using this promise
                    promise = @request.deferred.promise

                    # Process any responses that might have made it here first
                    if @deferred_responses
                        @deferred_responses.each &method(:respond_with)
                        @deferred_responses = nil
                    end

                    return promise

                # NOTE:: Somehow getting to here with a nil request... needs investigation
                elsif not result.nil?
                    # clear any cached responses just in case
                    # could be set by error in the rack application
                    @deferred_responses = nil if @deferred_responses

                    status, headers, body = result

                    send_body = @request.env[REQUEST_METHOD] != HEAD

                    # If a file, stream the body in a non-blocking fashion
                    if body.respond_to? :to_path
                        if headers[CONTENT_LENGTH]
                            type = :raw
                        else
                            type = :http
                            headers[TRANSFER_ENCODING] = CHUNKED
                        end

                        write_headers(status, headers)

                        if send_body
                            file = @loop.file(body.to_path, File::RDONLY)
                            file.progress do    # File is open and available for reading
                                file.send_file(@socket, type).finally do
                                    file.close
                                    if @request.keep_alive == false
                                        @socket.shutdown
                                    end
                                end
                            end
                            return file
                        end
                    else
                        # Optimize the response
                        begin
                            body_size = body.size
                            if body_size < 2
                                if body_size == 1
                                    headers[CONTENT_LENGTH] = body[0].bytesize
                                else
                                    headers[CONTENT_LENGTH] = ZERO
                                end
                            end
                        rescue # just in case
                        end

                        if send_body
                            write_response(status, headers, body)
                        else
                            write_headers(status, headers)
                            @socket.shutdown if @request.keep_alive == false
                        end
                    end
                end
            end

            # continue processing (don't wait for write to complete)
            # if the write fails it will close the socket
            nil
        end

        def write_response(status, headers, body)
            if headers[CONTENT_LENGTH]
                headers[CONTENT_LENGTH] = headers[CONTENT_LENGTH].to_s
                write_headers(status, headers)

                # Stream the response (pass directly into @socket.write)
                body.each &@socket.method(:write)

                if @request.deferred
                    @request.deferred.resolve(true)
                    @request.deferred = nil # prevent data being sent after completed
                end

                @socket.shutdown if @request.keep_alive == false
            else
                headers[TRANSFER_ENCODING] = CHUNKED
                write_headers(status, headers)

                # Stream the response
                body.each &@write_chunk

                if @request.deferred.nil?
                    @socket.write EOF
                    @socket.shutdown if @request.keep_alive == false
                else
                    @async_state = :chunked
                end
            end
        end

        def write_headers(status, headers)
            headers[CONNECTION] = CLOSE if @request.keep_alive == false

            header = "HTTP/1.1 #{status} #{fetch_code(status)}\r\n"
            headers.each do |key, value|
                next if key.start_with? RACK

                value.split(NEWLINE).each do |unique_value|
                    header << key
                    header << COLON_SPACE
                    header << unique_value
                    header << CRLF
                end
            end
            header << CRLF
            @socket.write header
        end

        def write_chunk(part)
            chunk = part.bytesize.to_s(16) << CRLF << part << CRLF
            @socket.write chunk
        end

        def fetch_code(status)
            HTTP_STATUS_CODES.fetch(status, &HTTP_STATUS_DEFAULT)
        end


        # ------------------------
        # Async response functions:
        # ------------------------


        # Callback from a response that was marked async
{'fraction_non_alphanumeric': 0.07047311248765599, 'fraction_numerical': 0.002513690636502379, 'mean_word_length': 1.2350521669341894, 'pattern_counts': {'":': 0, '<': 37, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 14, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2247702', 'n_tokens_mistral': 5443, 'n_tokens_neox': 5061, 'n_words': 1647}