h2 proto

Author: qwj

pproxy/proto.py

@@ -283,43 +283,51 @@ async def guess(self, reader, **kw):
         header = await reader.read_w(4)
         reader.rollback(header)
         return header in (b'GET ', b'HEAD', b'POST', b'PUT ', b'DELE', b'CONN', b'OPTI', b'TRAC', b'PATC')
-    async def accept(self, reader, user, writer, users, authtable, httpget=None, **kw):
+    async def accept(self, reader, user, writer, **kw):
         lines = await reader.read_until(b'\r\n\r\n')
         headers = lines[:-4].decode().split('\r\n')
         method, path, ver = HTTP_LINE.match(headers.pop(0)).groups()
         lines = '\r\n'.join(i for i in headers if not i.startswith('Proxy-'))
         headers = dict(i.split(': ', 1) for i in headers if ': ' in i)
+        async def reply(code, message, body=None, wait=False):
+            writer.write(message)
+            if body:
+                writer.write(body)
+            if wait:
+                await writer.drain()
+        return await self.http_accept(user, method, path, None, ver, lines, headers['Host'], headers.get('Proxy-Authorization'), reply, **kw)
+    async def http_accept(self, user, method, path, authority, ver, lines, host, pauth, reply, authtable, users, httpget, **kw):
         url = urllib.parse.urlparse(path)
-        if method == 'GET' and not url.hostname and httpget:
-            for path, text in httpget.items():
+        if method == 'GET' and not url.hostname:
+            for path, text in (httpget.items() if httpget else ()):
                 if path == url.path:
                     user = next(filter(lambda x: x.decode()==url.query, users), None) if users else True
                     if user:
                         if users:
                             authtable.set_authed(user)
                         if type(text) is str:
-                            text = (text % dict(host=headers["Host"])).encode()
-                        writer.write(f'{ver} 200 OK\r\nConnection: close\r\nContent-Type: text/plain\r\nCache-Control: max-age=900\r\nContent-Length: {len(text)}\r\n\r\n'.encode() + text)
-                        await writer.drain()
+                            text = (text % dict(host=host)).encode()
+                        await reply(200, f'{ver} 200 OK\r\nConnection: close\r\nContent-Type: text/plain\r\nCache-Control: max-age=900\r\nContent-Length: {len(text)}\r\n\r\n'.encode(), text, True)
                         raise Exception('Connection closed')
             raise Exception(f'404 {method} {url.path}')
         if users:
-            pauth = headers.get('Proxy-Authorization', None)
             user = authtable.authed()
             if not user:
                 user = next(filter(lambda i: ('Basic '+base64.b64encode(i).decode()) == pauth, users), None)
                 if user is None:
-                    writer.write(f'{ver} 407 Proxy Authentication Required\r\nConnection: close\r\nProxy-Authenticate: Basic realm="simple"\r\n\r\n'.encode())
+                    await reply(407, f'{ver} 407 Proxy Authentication Required\r\nConnection: close\r\nProxy-Authenticate: Basic realm="simple"\r\n\r\n'.encode(), wait=True)
                     raise Exception('Unauthorized HTTP')
             authtable.set_authed(user)
         if method == 'CONNECT':
-            host_name, port = netloc_split(path)
-            return user, host_name, port, f'{ver} 200 OK\r\nConnection: close\r\n\r\n'.encode()
+            host_name, port = netloc_split(authority or path)
+            return user, host_name, port, lambda writer: reply(200, f'{ver} 200 OK\r\nConnection: close\r\n\r\n'.encode())
         else:
-            url = urllib.parse.urlparse(path)
-            host_name, port = netloc_split(url.netloc or headers.get("Host"), default_port=80)
+            host_name, port = netloc_split(url.netloc or host, default_port=80)
             newpath = url._replace(netloc='', scheme='').geturl()
-            return user, host_name, port, b'', f'{method} {newpath} {ver}\r\n{lines}\r\n\r\n'.encode()
+            async def connected(writer):
+                writer.write(f'{method} {newpath} {ver}\r\n{lines}\r\n\r\n'.encode())
+                return True
+            return user, host_name, port, connected
     async def connect(self, reader_remote, writer_remote, rauth, host_name, port, myhost, **kw):
         writer_remote.write(f'CONNECT {host_name}:{port} HTTP/1.1\r\nHost: {myhost}'.encode() + (b'\r\nProxy-Authorization: Basic '+base64.b64encode(rauth) if rauth else b'') + b'\r\n\r\n')
         await reader_remote.read_until(b'\r\n\r\n')
@@ -370,6 +378,29 @@ def write(data, o=writer_remote.write):
                 return o(data)
         writer_remote.write = write
 
+class H2(HTTP):
+    async def guess(self, reader, **kw):
+        return True
+    async def accept(self, reader, user, writer, **kw):
+        if not writer.headers.done():
+            await writer.headers
+        headers = writer.headers.result()
+        headers = {i.decode().lower():j.decode() for i,j in headers}
+        lines = '\r\n'.join(i for i in headers if not i.startswith('proxy-') and not i.startswith(':'))
+        async def reply(code, message, body=None, wait=False):
+            writer.send_headers(((':status', str(code)),))
+            if body:
+                writer.write(body)
+            if wait:
+                await writer.drain()
+        return await self.http_accept(user, headers[':method'], headers[':path'], headers[':authority'], '2.0', lines, '', headers.get('proxy-authorization'), reply, **kw)
+    async def connect(self, reader_remote, writer_remote, rauth, host_name, port, myhost, **kw):
+        headers = [(':method', 'CONNECT'), (':scheme', 'https'), (':path', '/'),
+                   (':authority', f'{host_name}:{port}')]
+        if rauth:
+            headers.append(('proxy-authorization', 'Basic '+base64.b64encode(rauth)))
+        writer_remote.send_headers(headers)
+
 class SSH(BaseProtocol):
     async def connect(self, reader_remote, writer_remote, rauth, host_name, port, myhost, **kw):
         pass
@@ -521,8 +552,8 @@ async def accept(protos, reader, **kw):
             raise Exception('Connection closed')
         if user:
             ret = await proto.accept(reader, user, **kw)
-            while len(ret) < 5:
-                ret += (b'',)
+            while len(ret) < 4:
+                ret += (None,)
             return (proto,) + ret
     raise Exception(f'Unsupported protocol')
 
@@ -533,7 +564,7 @@ def udp_accept(protos, data, **kw):
             return (proto,) + ret
     raise Exception(f'Unsupported protocol {data[:10]}')
 
-MAPPINGS = dict(direct=Direct, http=HTTP, httponly=HTTPOnly, ssh=SSH, socks5=Socks5, socks4=Socks4, socks=Socks5, ss=SS, ssr=SSR, redir=Redir, pf=Pf, tunnel=Tunnel, echo=Echo, ws=WS, trojan=Trojan, ssl='', secure='', quic='')
+MAPPINGS = dict(direct=Direct, http=HTTP, httponly=HTTPOnly, ssh=SSH, socks5=Socks5, socks4=Socks4, socks=Socks5, ss=SS, ssr=SSR, redir=Redir, pf=Pf, tunnel=Tunnel, echo=Echo, ws=WS, trojan=Trojan, h2=H2, ssl='', secure='', quic='')
 MAPPINGS['in'] = ''
 
 def get_protos(rawprotos):
@@ -553,7 +584,7 @@ def get_protos(rawprotos):
 def sslwrap(reader, writer, sslcontext, server_side=False, server_hostname=None, verbose=None):
     if sslcontext is None:
         return reader, writer
-    ssl_reader = type(reader)()
+    ssl_reader = asyncio.StreamReader()
     class Protocol(asyncio.Protocol):
         def data_received(self, data):
             ssl_reader.feed_data(data)
@@ -576,13 +607,14 @@ def close(self):
         def _force_close(self, exc):
             if not self.closed:
                 (verbose or print)(f'{exc} from {writer.get_extra_info("peername")[0]}')
+            ssl._app_transport._closed = True
             self.close()
         def abort(self):
             self.close()
     ssl.connection_made(Transport())
     async def channel():
         try:
-            while not reader.at_eof() and not ssl._app_transport.closed:
+            while not reader.at_eof() and not ssl._app_transport._closed:
                 data = await reader.read(65536)
                 if not data:
                     break
@@ -600,7 +632,8 @@ def write(self, data):
         def drain(self):
             return writer.drain()
         def is_closing(self):
-            return ssl._app_transport.closed
+            return ssl._app_transport._closed
         def close(self):
-            ssl._app_transport.close()
+            if not ssl._app_transport._closed:
+                ssl._app_transport.close()
     return ssl_reader, Writer()

pproxy/server.py

@@ -71,7 +71,7 @@ async def stream_handler(reader, writer, unix, lbind, protos, rserver, cipher, s
             remote_text = f'{remote_ip}:{remote_port}'
         local_addr = None if server_ip in ('127.0.0.1', '::1', None) else (server_ip, 0)
         reader_cipher, _ = await prepare_ciphers(cipher, reader, writer, server_side=False)
-        lproto, user, host_name, port, lbuf, rbuf = await proto.accept(protos, reader=reader, writer=writer, authtable=AuthTable(remote_ip, authtime), reader_cipher=reader_cipher, sock=writer.get_extra_info('socket'), **kwargs)
+        lproto, user, host_name, port, client_connected = await proto.accept(protos, reader=reader, writer=writer, authtable=AuthTable(remote_ip, authtime), reader_cipher=reader_cipher, sock=writer.get_extra_info('socket'), **kwargs)
         if host_name == 'echo':
             asyncio.ensure_future(lproto.channel(reader, writer, DUMMY, DUMMY))
         elif host_name == 'empty':
@@ -87,13 +87,12 @@ async def stream_handler(reader, writer, unix, lbind, protos, rserver, cipher, s
                 raise Exception(f'Connection timeout {roption.bind}')
             try:
                 reader_remote, writer_remote = await roption.prepare_connection(reader_remote, writer_remote, host_name, port)
-                writer.write(lbuf)
-                writer_remote.write(rbuf)
+                use_http = (await client_connected(writer_remote)) if client_connected else None
             except Exception:
                 writer_remote.close()
                 raise Exception('Unknown remote protocol')
             m = modstat(user, remote_ip, host_name)
-            lchannel = lproto.http_channel if rbuf else lproto.channel
+            lchannel = lproto.http_channel if use_http else lproto.channel
             asyncio.ensure_future(lproto.channel(reader_remote, writer, m(2+roption.direct), m(4+roption.direct)))
             asyncio.ensure_future(lchannel(reader, writer_remote, m(roption.direct), roption.connection_change))
     except Exception as ex:
@@ -304,6 +303,126 @@ def start_server(self, args, stream_handler=stream_handler):
         else:
             return asyncio.start_server(handler, host=self.host_name, port=self.port, reuse_port=args.get('ruport'))
 
+class ProxyH2(ProxySimple):
+    def __init__(self, sslserver, sslclient, **kw):
+        super().__init__(sslserver=None, sslclient=None, **kw)
+        self.handshake = None
+        self.h2sslserver = sslserver
+        self.h2sslclient = sslclient
+    async def handler(self, reader, writer, client_side=True, stream_handler=None, **kw):
+        import h2.connection, h2.config, h2.events
+        reader, writer = proto.sslwrap(reader, writer, self.h2sslclient if client_side else self.h2sslserver, not client_side, None)
+        config = h2.config.H2Configuration(client_side=client_side)
+        conn = h2.connection.H2Connection(config=config)
+        streams = {}
+        conn.initiate_connection()
+        writer.write(conn.data_to_send())
+        while not reader.at_eof() and not writer.is_closing():
+            try:
+                data = await reader.read(65636)
+                if not data:
+                    break
+                events = conn.receive_data(data)
+            except Exception:
+                pass
+            writer.write(conn.data_to_send())
+            for event in events:
+                if isinstance(event, h2.events.RequestReceived) and not client_side:
+                    if event.stream_id not in streams:
+                        stream_reader, stream_writer = self.get_stream(conn, writer, event.stream_id)
+                        streams[event.stream_id] = (stream_reader, stream_writer)
+                        asyncio.ensure_future(stream_handler(stream_reader, stream_writer))
+                    else:
+                        stream_reader, stream_writer = streams[event.stream_id]
+                    stream_writer.headers.set_result(event.headers)
+                elif isinstance(event, h2.events.SettingsAcknowledged) and client_side:
+                    self.handshake.set_result((conn, streams, writer))
+                elif isinstance(event, h2.events.DataReceived):
+                    stream_reader, stream_writer = streams[event.stream_id]
+                    stream_reader.feed_data(event.data)
+                    conn.acknowledge_received_data(len(event.data), event.stream_id)
+                    writer.write(conn.data_to_send())
+                elif isinstance(event, h2.events.StreamEnded) or isinstance(event, h2.events.StreamReset):
+                    stream_reader, stream_writer = streams[event.stream_id]
+                    stream_reader.feed_eof()
+                    if not stream_writer.closed:
+                        stream_writer.close()
+                elif isinstance(event, h2.events.ConnectionTerminated):
+                    break
+                elif isinstance(event, h2.events.WindowUpdated):
+                    if event.stream_id in streams:
+                        stream_reader, stream_writer = streams[event.stream_id]
+                        stream_writer.window_update()
+        writer.write(conn.data_to_send())
+        writer.close()
+    def get_stream(self, conn, writer, stream_id):
+        reader = asyncio.StreamReader()
+        write_buffer = bytearray()
+        write_wait = asyncio.Event()
+        write_full = asyncio.Event()
+        class StreamWriter():
+            def __init__(self):
+                self.closed = False
+                self.headers = asyncio.get_event_loop().create_future()
+            def get_extra_info(self, key):
+                return writer.get_extra_info(key)
+            def write(self, data):
+                write_buffer.extend(data)
+                write_wait.set()
+            def drain(self):
+                writer.write(conn.data_to_send())
+                return writer.drain()
+            def is_closing(self):
+                return self.closed
+            def close(self):
+                self.closed = True
+                write_wait.set()
+            def window_update(self):
+                write_full.set()
+            def send_headers(self, headers):
+                conn.send_headers(stream_id, headers)
+                writer.write(conn.data_to_send())
+        stream_writer = StreamWriter()
+        async def write_job():
+            while not stream_writer.closed:
+                while len(write_buffer) > 0:
+                    while conn.local_flow_control_window(stream_id) <= 0:
+                        write_full.clear()
+                        await write_full.wait()
+                        if stream_writer.closed:
+                            break
+                    chunk_size = min(conn.local_flow_control_window(stream_id), len(write_buffer), conn.max_outbound_frame_size)
+                    conn.send_data(stream_id, write_buffer[:chunk_size])
+                    writer.write(conn.data_to_send())
+                    del write_buffer[:chunk_size]
+                if not stream_writer.closed:
+                    write_wait.clear()
+                    await write_wait.wait()
+            conn.send_data(stream_id, b'', end_stream=True)
+            writer.write(conn.data_to_send())
+        asyncio.ensure_future(write_job())
+        return reader, stream_writer
+    async def wait_h2_connection(self, local_addr, family):
+        if self.handshake is not None:
+            if not self.handshake.done():
+                await self.handshake
+        else:
+            self.handshake = asyncio.get_event_loop().create_future()
+            reader, writer = await super().wait_open_connection(None, None, local_addr, family)
+            asyncio.ensure_future(self.handler(reader, writer))
+            await self.handshake
+        return self.handshake.result()
+    async def wait_open_connection(self, host, port, local_addr, family):
+        conn, streams, writer = await self.wait_h2_connection(local_addr, family)
+        stream_id = conn.get_next_available_stream_id()
+        conn._begin_new_stream(stream_id, stream_id%2)
+        stream_reader, stream_writer = self.get_stream(conn, writer, stream_id)
+        streams[stream_id] = (stream_reader, stream_writer)
+        return stream_reader, stream_writer
+    def start_server(self, args, stream_handler=stream_handler):
+        handler = functools.partial(stream_handler, **vars(self), **args)
+        return super().start_server(args, functools.partial(self.handler, client_side=False, stream_handler=handler))
+
 class ProxyQUIC(ProxySimple):
     def __init__(self, quicserver, quicclient, **kw):
         super().__init__(**kw)
@@ -544,6 +663,8 @@ def proxies_by_uri(uri_jumps):
         jump = proxy_by_uri(uri, jump)
     return jump
 
+sslcontexts = []
+
 def proxy_by_uri(uri, jump):
     scheme, _, uri = uri.partition('://')
     url = urllib.parse.urlparse('s://'+uri)
@@ -558,17 +679,25 @@ def proxy_by_uri(uri, jump):
         if 'ssl' in rawprotos:
             sslclient.check_hostname = False
             sslclient.verify_mode = ssl.CERT_NONE
+        sslcontexts.append(sslserver)
+        sslcontexts.append(sslclient)
     else:
         sslserver = sslclient = None
     if 'quic' in rawprotos:
         try:
             import ssl, aioquic.quic.configuration
         except Exception:
             raise Exception('Missing library: "pip3 install aioquic"')
-        import logging
         quicserver = aioquic.quic.configuration.QuicConfiguration(is_client=False)
         quicclient = aioquic.quic.configuration.QuicConfiguration()
         quicclient.verify_mode = ssl.CERT_NONE
+        sslcontexts.append(quicserver)
+        sslcontexts.append(quicclient)
+    if 'h2' in rawprotos:
+        try:
+            import h2
+        except Exception:
+            raise Exception('Missing library: "pip3 install h2"')
     protonames = [i.name for i in protos]
     urlpath, _, plugins = url.path.partition(',')
     urlpath, _, lbind = urlpath.partition('@')
@@ -611,6 +740,8 @@ def proxy_by_uri(uri, jump):
                       host_name=host_name, port=port, unix=not loc, lbind=lbind, sslclient=sslclient, sslserver=sslserver)
         if 'quic' in rawprotos:
             proxy = ProxyQUIC(quicserver, quicclient, **params)
+        elif 'h2' in protonames:
+            proxy = ProxyH2(**params)
         elif 'ssh' in protonames:
             proxy = ProxySSH(**params)
         else:
@@ -646,7 +777,7 @@ async def test_url(url, rserver):
         print(headers.decode()[:-4])
         print(f'--------------------------------')
         body = bytearray()
-        while 1:
+        while not reader.at_eof():
             s = await reader.read(65536)
             if not s:
                 break
@@ -677,15 +808,8 @@ def main():
     args = parser.parse_args()
     if args.sslfile:
         sslfile = args.sslfile.split(',')
-        for option in args.listen:
-            if option.sslclient:
-                option.sslclient.load_cert_chain(*sslfile)
-                option.sslserver.load_cert_chain(*sslfile)
-        for option in args.listen+args.ulisten+args.rserver+args.urserver:
-            if isinstance(option, ProxyQUIC):
-                option.quicserver.load_cert_chain(*sslfile)
-            if isinstance(option, ProxyBackward) and isinstance(option.backward, ProxyQUIC):
-                option.backward.quicserver.load_cert_chain(*sslfile)
+        for context in sslcontexts:
+            context.load_cert_chain(*sslfile)
     elif any(map(lambda o: o.sslclient or isinstance(o, ProxyQUIC), args.listen+args.ulisten)):
         print('You must specify --ssl to listen in ssl mode')
         return