1  import json
   2  import os.path
   3  import socket
   4  import socketserver
   5  import threading
   6  from contextlib import closing, contextmanager
   7  from http.server import SimpleHTTPRequestHandler
   8  from typing import Callable, Generator
   9  from urllib.request import urlopen
  10  
  11  import h11
  12  
  13  
  14  @contextmanager
  15  def socket_server(
  16      handler: Callable[..., socketserver.BaseRequestHandler]
  17  ) -> Generator[socketserver.TCPServer, None, None]:
  18      httpd = socketserver.TCPServer(("127.0.0.1", 0), handler)
  19      thread = threading.Thread(
  20          target=httpd.serve_forever, kwargs={"poll_interval": 0.01}
  21      )
  22      thread.daemon = True
  23      try:
  24          thread.start()
  25          yield httpd
  26      finally:
  27          httpd.shutdown()
  28  
  29  
  30  test_file_path = os.path.join(os.path.dirname(__file__), "data/test-file")
  31  with open(test_file_path, "rb") as f:
  32      test_file_data = f.read()
  33  
  34  
  35  class SingleMindedRequestHandler(SimpleHTTPRequestHandler):
  36      def translate_path(self, path: str) -> str:
  37          return test_file_path
  38  
  39  
  40  def test_h11_as_client() -> None:
  41      with socket_server(SingleMindedRequestHandler) as httpd:
  42          with closing(socket.create_connection(httpd.server_address)) as s:
  43              c = h11.Connection(h11.CLIENT)
  44  
  45              s.sendall(
  46                  c.send(  # type: ignore[arg-type]
  47                      h11.Request(
  48                          method="GET", target="/foo", headers=[("Host", "localhost")]
  49                      )
  50                  )
  51              )
  52              s.sendall(c.send(h11.EndOfMessage()))  # type: ignore[arg-type]
  53  
  54              data = bytearray()
  55              while True:
  56                  event = c.next_event()
  57                  print(event)
  58                  if event is h11.NEED_DATA:
  59                      # Use a small read buffer to make things more challenging
  60                      # and exercise more paths :-)
  61                      c.receive_data(s.recv(10))
  62                      continue
  63                  if type(event) is h11.Response:
  64                      assert event.status_code == 200
  65                  if type(event) is h11.Data:
  66                      data += event.data
  67                  if type(event) is h11.EndOfMessage:
  68                      break
  69              assert bytes(data) == test_file_data
  70  
  71  
  72  class H11RequestHandler(socketserver.BaseRequestHandler):
  73      def handle(self) -> None:
  74          with closing(self.request) as s:
  75              c = h11.Connection(h11.SERVER)
  76              request = None
  77              while True:
  78                  event = c.next_event()
  79                  if event is h11.NEED_DATA:
  80                      # Use a small read buffer to make things more challenging
  81                      # and exercise more paths :-)
  82                      c.receive_data(s.recv(10))
  83                      continue
  84                  if type(event) is h11.Request:
  85                      request = event
  86                  if type(event) is h11.EndOfMessage:
  87                      break
  88              assert request is not None
  89              info = json.dumps(
  90                  {
  91                      "method": request.method.decode("ascii"),
  92                      "target": request.target.decode("ascii"),
  93                      "headers": {
  94                          name.decode("ascii"): value.decode("ascii")
  95                          for (name, value) in request.headers
  96                      },
  97                  }
  98              )
  99              s.sendall(c.send(h11.Response(status_code=200, headers=[])))  # type: ignore[arg-type]
 100              s.sendall(c.send(h11.Data(data=info.encode("ascii"))))
 101              s.sendall(c.send(h11.EndOfMessage()))
 102  
 103  
 104  def test_h11_as_server() -> None:
 105      with socket_server(H11RequestHandler) as httpd:
 106          host, port = httpd.server_address
 107          url = "http://{}:{}/some-path".format(host, port)
 108          with closing(urlopen(url)) as f:
 109              assert f.getcode() == 200
 110              data = f.read()
 111      info = json.loads(data.decode("ascii"))
 112      print(info)
 113      assert info["method"] == "GET"
 114      assert info["target"] == "/some-path"
 115      assert "urllib" in info["headers"]["user-agent"]
 116