<commit_before>from __future__ import (absolute_import, print_function, division)

import sys
import traceback
import six
import struct
import threading
import Queue

from netlib import tcp
from netlib.exceptions import HttpException, HttpReadDisconnect, NetlibException
from netlib.http import http1, Headers, CONTENT_MISSING
from netlib.tcp import Address, ssl_read_select

import h2
from h2.connection import H2Connection
from h2.events import *
from hyperframe import frame

from .base import Layer, Kill
from .. import utils
from ..exceptions import HttpProtocolException, ProtocolException
from ..models import (
    HTTPFlow,
    HTTPRequest,
    HTTPResponse,
    make_error_response,
    make_connect_response,
    Error,
    expect_continue_response
)


class _HttpLayer(Layer):
    def read_request(self):
        raise NotImplementedError()

    def read_request_body(self, request):
        raise NotImplementedError()

    def send_request(self, request):
        raise NotImplementedError()

    def read_response(self, request):
        response = self.read_response_headers()
        response.data.content = b"".join(
            self.read_response_body(request, response)
        )
        return response

    def read_response_headers(self):
        raise NotImplementedError()

    def read_response_body(self, request, response):
        raise NotImplementedError()
        yield "this is a generator"  # pragma: no cover

    def send_response(self, response):
        if response.content == CONTENT_MISSING:
            raise HttpException("Cannot assemble flow with CONTENT_MISSING")
        self.send_response_headers(response)
        self.send_response_body(response, [response.content])

    def send_response_headers(self, response):
        raise NotImplementedError()

    def send_response_body(self, response, chunks):
        raise NotImplementedError()

    def check_close_connection(self, flow):
        raise NotImplementedError()


class Http1Layer(_HttpLayer):
    def __init__(self, ctx, mode):
        super(Http1Layer, self).__init__(ctx)
        self.mode = mode

    def read_request(self):
        req = http1.read_request(self.client_conn.rfile, body_size_limit=self.config.body_size_limit)
        return HTTPRequest.wrap(req)

    def read_request_body(self, request):
        expected_size = http1.expected_http_body_size(request)
        return http1.read_body(self.client_conn.rfile, expected_size, self.config.body_size_limit)

    def send_request(self, request):
        self.server_conn.wfile.write(http1.assemble_request(request))
        self.server_conn.wfile.flush()

    def read_response_headers(self):
        resp = http1.read_response_head(self.server_conn.rfile)
        return HTTPResponse.wrap(resp)

    def read_response_body(self, request, response):
        expected_size = http1.expected_http_body_size(request, response)
        return http1.read_body(self.server_conn.rfile, expected_size, self.config.body_size_limit)

    def send_response_headers(self, response):
        raw = http1.assemble_response_head(response)
        self.client_conn.wfile.write(raw)
        self.client_conn.wfile.flush()

    def send_response_body(self, response, chunks):
        for chunk in http1.assemble_body(response.headers, chunks):
            self.client_conn.wfile.write(chunk)
            self.client_conn.wfile.flush()

    def check_close_connection(self, flow):
        request_close = http1.connection_close(
            flow.request.http_version,
            flow.request.headers
        )
        response_close = http1.connection_close(
            flow.response.http_version,
            flow.response.headers
        )
        read_until_eof = http1.expected_http_body_size(flow.request, flow.response) == -1
        close_connection = request_close or response_close or read_until_eof
        if flow.request.form_in == "authority" and flow.response.status_code == 200:
            # Workaround for https://github.com/mitmproxy/mitmproxy/issues/313:
            # Charles Proxy sends a CONNECT response with HTTP/1.0
            # and no Content-Length header

            return False
        return close_connection

    def __call__(self):
        layer = HttpLayer(self, self.mode)
        layer()

class SafeH2Connection(H2Connection):
    def __init__(self, conn, *args, **kwargs):
        super(SafeH2Connection, self).__init__(*args, **kwargs)
        self.conn = conn
        self.lock = threading.RLock()

    def safe_close_connection(self, error_code):
        with self.lock:
            self.close_connection(error_code)
            self.conn.send(self.data_to_send())

    def safe_increment_flow_control(self, stream_id, length):
        with self.lock:
            self.increment_flow_control_window(length)
            self.conn.send(self.data_to_send())
        with self.lock:
            if stream_id in self.streams and not self.streams[stream_id].closed:
                self.increment_flow_control_window(length, stream_id=stream_id)
                self.conn.send(self.data_to_send())

    def safe_reset_stream(self, stream_id, error_code):
        with self.lock:
            self.reset_stream(stream_id, error_code)
            self.conn.send(self.h2.data_to_send())

    def safe_acknowledge_settings(self, event):
        with self.conn.h2.lock:
            self.conn.h2.acknowledge_settings(event)
            self.conn.send(self.data_to_send())

    def safe_update_settings(self, new_settings):
        with self.conn.h2.lock:
            self.update_settings(new_settings)
            self.conn.send(self.data_to_send())

    def safe_send_headers(self, stream_id, headers):
        with self.lock:
            self.send_headers(stream_id, headers)
            self.conn.send(self.data_to_send())

    def safe_send_body(self, stream_id, chunks):
        for chunk in chunks:
            max_outbound_frame_size = self.max_outbound_frame_size
            for i in xrange(0, len(chunk), max_outbound_frame_size):
                frame_chunk = chunk[i:i+max_outbound_frame_size]
                with self.lock:
                    self.send_data(stream_id, frame_chunk)
                    self.conn.send(self.data_to_send())
        with self.lock:
            self.end_stream(stream_id)
            self.conn.send(self.data_to_send())

class Http2Layer(Layer):
    def __init__(self, ctx, mode):
        super(Http2Layer, self).__init__(ctx)
        self.mode = mode
        self.streams = dict()
        self.server_to_client_stream_ids = dict([(0, 0)])
        self.client_conn.h2 = SafeH2Connection(self.client_conn, client_side=False)

        # make sure that we only pass actual SSL.Connection objects in here,
        # because otherwise ssl_read_select fails!
        self.active_conns = [self.client_conn.connection]

        if self.server_conn:
            self._initiate_server_conn()

    def _initiate_server_conn(self):
        self.server_conn.h2 = SafeH2Connection(self.server_conn, client_side=True)
        self.server_conn.h2.initiate_connection()
        self.server_conn.h2.update_settings({frame.SettingsFrame.ENABLE_PUSH: False})
        self.server_conn.send(self.server_conn.h2.data_to_send())
        self.active_conns.append(self.server_conn.connection)

    def connect(self):
        self.ctx.connect()<commit_msg>improve flow control<commit_after>from __future__ import (absolute_import, print_function, division)

import sys
import traceback
import six
import struct
import threading
import Queue

from netlib import tcp
from netlib.exceptions import HttpException, HttpReadDisconnect, NetlibException
from netlib.http import http1, Headers, CONTENT_MISSING
from netlib.tcp import Address, ssl_read_select

import h2
from h2.connection import H2Connection
from h2.events import *
from hyperframe import frame

from .base import Layer, Kill
from .. import utils
from ..exceptions import HttpProtocolException, ProtocolException
from ..models import (
    HTTPFlow,
    HTTPRequest,
    HTTPResponse,
    make_error_response,
    make_connect_response,
    Error,
    expect_continue_response
)


class _HttpLayer(Layer):
    def read_request(self):
        raise NotImplementedError()

    def read_request_body(self, request):
        raise NotImplementedError()

    def send_request(self, request):
        raise NotImplementedError()

    def read_response(self, request):
        response = self.read_response_headers()
        response.data.content = b"".join(
            self.read_response_body(request, response)
        )
        return response

    def read_response_headers(self):
        raise NotImplementedError()

    def read_response_body(self, request, response):
        raise NotImplementedError()
        yield "this is a generator"  # pragma: no cover

    def send_response(self, response):
        if response.content == CONTENT_MISSING:
            raise HttpException("Cannot assemble flow with CONTENT_MISSING")
        self.send_response_headers(response)
        self.send_response_body(response, [response.content])

    def send_response_headers(self, response):
        raise NotImplementedError()

    def send_response_body(self, response, chunks):
        raise NotImplementedError()

    def check_close_connection(self, flow):
        raise NotImplementedError()


class Http1Layer(_HttpLayer):
    def __init__(self, ctx, mode):
        super(Http1Layer, self).__init__(ctx)
        self.mode = mode

    def read_request(self):
        req = http1.read_request(self.client_conn.rfile, body_size_limit=self.config.body_size_limit)
        return HTTPRequest.wrap(req)

    def read_request_body(self, request):
        expected_size = http1.expected_http_body_size(request)
        return http1.read_body(self.client_conn.rfile, expected_size, self.config.body_size_limit)

    def send_request(self, request):
        self.server_conn.wfile.write(http1.assemble_request(request))
        self.server_conn.wfile.flush()

    def read_response_headers(self):
        resp = http1.read_response_head(self.server_conn.rfile)
        return HTTPResponse.wrap(resp)

    def read_response_body(self, request, response):
        expected_size = http1.expected_http_body_size(request, response)
        return http1.read_body(self.server_conn.rfile, expected_size, self.config.body_size_limit)

    def send_response_headers(self, response):
        raw = http1.assemble_response_head(response)
        self.client_conn.wfile.write(raw)
        self.client_conn.wfile.flush()

    def send_response_body(self, response, chunks):
        for chunk in http1.assemble_body(response.headers, chunks):
            self.client_conn.wfile.write(chunk)
            self.client_conn.wfile.flush()

    def check_close_connection(self, flow):
        request_close = http1.connection_close(
            flow.request.http_version,
            flow.request.headers
        )
        response_close = http1.connection_close(
            flow.response.http_version,
            flow.response.headers
        )
        read_until_eof = http1.expected_http_body_size(flow.request, flow.response) == -1
        close_connection = request_close or response_close or read_until_eof
        if flow.request.form_in == "authority" and flow.response.status_code == 200:
            # Workaround for https://github.com/mitmproxy/mitmproxy/issues/313:
            # Charles Proxy sends a CONNECT response with HTTP/1.0
            # and no Content-Length header

            return False
        return close_connection

    def __call__(self):
        layer = HttpLayer(self, self.mode)
        layer()

class SafeH2Connection(H2Connection):
    def __init__(self, conn, *args, **kwargs):
        super(SafeH2Connection, self).__init__(*args, **kwargs)
        self.conn = conn
        self.lock = threading.RLock()

    def safe_close_connection(self, error_code):
        with self.lock:
            self.close_connection(error_code)
            self.conn.send(self.data_to_send())

    def safe_increment_flow_control(self, stream_id, length):
        if length == 0:
            return

        with self.lock:
            self.increment_flow_control_window(length)
            self.conn.send(self.data_to_send())
        with self.lock:
            if stream_id in self.streams and not self.streams[stream_id].closed:
                self.increment_flow_control_window(length, stream_id=stream_id)
                self.conn.send(self.data_to_send())

    def safe_reset_stream(self, stream_id, error_code):
        with self.lock:
            self.reset_stream(stream_id, error_code)
            self.conn.send(self.data_to_send())

    def safe_acknowledge_settings(self, event):
        with self.conn.h2.lock:
            self.conn.h2.acknowledge_settings(event)
            self.conn.send(self.data_to_send())

    def safe_update_settings(self, new_settings):
        with self.conn.h2.lock:
            self.update_settings(new_settings)
            self.conn.send(self.data_to_send())

    def safe_send_headers(self, stream_id, headers):
        with self.lock:
            self.send_headers(stream_id, headers)
            self.conn.send(self.data_to_send())

    def safe_send_body(self, stream_id, chunks):
        for chunk in chunks:
            max_outbound_frame_size = self.max_outbound_frame_size
            for i in xrange(0, len(chunk), max_outbound_frame_size):
                frame_chunk = chunk[i:i+max_outbound_frame_size]

                self.lock.acquire()
                while True:
                    if self.local_flow_control_window(stream_id) < len(frame_chunk):
                        self.lock.release()
                        time.sleep(0)
                    else:
                        break
                self.send_data(stream_id, frame_chunk)
                self.conn.send(self.data_to_send())
                self.lock.release()
        with self.lock:
            self.end_stream(stream_id)
            self.conn.send(self.data_to_send())

class Http2Layer(Layer):
    def __init__(self, ctx, mode):
        super(Http2Layer, self).__init__(ctx)
        self.mode = mode
        self.streams = dict()
        self.server_to_client_stream_ids = dict([(0, 0)])
        self.client_conn.h2 = SafeH2Connection(self.client_conn, client_side=False)

        # make sure that we only pass actual SSL.Connection objects in here,
        # because otherwise ssl_read_select fails!
        self.active_conns = [self.client_conn.connection]

        if self.server_conn:
            self._initiate_server_conn()

    def _initiate_server_conn(self):
        self.server_conn.h2 = SafeH2Connection(self.server_conn, client_side=True)
        self.server_conn.h2.initiate_connection()
        self.server_conn.h2.update_settings({frame.SettingsFrame.ENABLE_PUSH: False})
        self.server_conn.send(self.server_conn.h2.data_to_send())
        self.active_conns.append(self.server_conn.connection)

    def connect(self):
        self.ctx.connect()
{'fraction_non_alphanumeric': 0.08666755283796358, 'fraction_numerical': 0.006313970490495813, 'mean_word_length': 2.858205128205128, 'pattern_counts': {'":': 0, '<': 4, '<?xml version=': 0, '>': 3, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7360581', 'n_tokens_mistral': 4508, 'n_tokens_neox': 4387, 'n_words': 854}