1  from typing import Any, Callable, Generator, List
   2  
   3  import pytest
   4  
   5  from .._events import (
   6      ConnectionClosed,
   7      Data,
   8      EndOfMessage,
   9      Event,
  10      InformationalResponse,
  11      Request,
  12      Response,
  13  )
  14  from .._headers import Headers, normalize_and_validate
  15  from .._readers import (
  16      _obsolete_line_fold,
  17      ChunkedReader,
  18      ContentLengthReader,
  19      Http10Reader,
  20      READERS,
  21  )
  22  from .._receivebuffer import ReceiveBuffer
  23  from .._state import (
  24      CLIENT,
  25      CLOSED,
  26      DONE,
  27      IDLE,
  28      MIGHT_SWITCH_PROTOCOL,
  29      MUST_CLOSE,
  30      SEND_BODY,
  31      SEND_RESPONSE,
  32      SERVER,
  33      SWITCHED_PROTOCOL,
  34  )
  35  from .._util import LocalProtocolError
  36  from .._writers import (
  37      ChunkedWriter,
  38      ContentLengthWriter,
  39      Http10Writer,
  40      write_any_response,
  41      write_headers,
  42      write_request,
  43      WRITERS,
  44  )
  45  from .helpers import normalize_data_events
  46  
  47  SIMPLE_CASES = [
  48      (
  49          (CLIENT, IDLE),
  50          Request(
  51              method="GET",
  52              target="/a",
  53              headers=[("Host", "foo"), ("Connection", "close")],
  54          ),
  55          b"GET /a HTTP/1.1\r\nHost: foo\r\nConnection: close\r\n\r\n",
  56      ),
  57      (
  58          (SERVER, SEND_RESPONSE),
  59          Response(status_code=200, headers=[("Connection", "close")], reason=b"OK"),
  60          b"HTTP/1.1 200 OK\r\nConnection: close\r\n\r\n",
  61      ),
  62      (
  63          (SERVER, SEND_RESPONSE),
  64          Response(status_code=200, headers=[], reason=b"OK"),  # type: ignore[arg-type]
  65          b"HTTP/1.1 200 OK\r\n\r\n",
  66      ),
  67      (
  68          (SERVER, SEND_RESPONSE),
  69          InformationalResponse(
  70              status_code=101, headers=[("Upgrade", "websocket")], reason=b"Upgrade"
  71          ),
  72          b"HTTP/1.1 101 Upgrade\r\nUpgrade: websocket\r\n\r\n",
  73      ),
  74      (
  75          (SERVER, SEND_RESPONSE),
  76          InformationalResponse(status_code=101, headers=[], reason=b"Upgrade"),  # type: ignore[arg-type]
  77          b"HTTP/1.1 101 Upgrade\r\n\r\n",
  78      ),
  79  ]
  80  
  81  
  82  def dowrite(writer: Callable[..., None], obj: Any) -> bytes:
  83      got_list: List[bytes] = []
  84      writer(obj, got_list.append)
  85      return b"".join(got_list)
  86  
  87  
  88  def tw(writer: Any, obj: Any, expected: Any) -> None:
  89      got = dowrite(writer, obj)
  90      assert got == expected
  91  
  92  
  93  def makebuf(data: bytes) -> ReceiveBuffer:
  94      buf = ReceiveBuffer()
  95      buf += data
  96      return buf
  97  
  98  
  99  def tr(reader: Any, data: bytes, expected: Any) -> None:
 100      def check(got: Any) -> None:
 101          assert got == expected
 102          # Headers should always be returned as bytes, not e.g. bytearray
 103          # https://github.com/python-hyper/wsproto/pull/54#issuecomment-377709478
 104          for name, value in getattr(got, "headers", []):
 105              assert type(name) is bytes
 106              assert type(value) is bytes
 107  
 108      # Simple: consume whole thing
 109      buf = makebuf(data)
 110      check(reader(buf))
 111      assert not buf
 112  
 113      # Incrementally growing buffer
 114      buf = ReceiveBuffer()
 115      for i in range(len(data)):
 116          assert reader(buf) is None
 117          buf += data[i : i + 1]
 118      check(reader(buf))
 119  
 120      # Trailing data
 121      buf = makebuf(data)
 122      buf += b"trailing"
 123      check(reader(buf))
 124      assert bytes(buf) == b"trailing"
 125  
 126  
 127  def test_writers_simple() -> None:
 128      for ((role, state), event, binary) in SIMPLE_CASES:
 129          tw(WRITERS[role, state], event, binary)
 130  
 131  
 132  def test_readers_simple() -> None:
 133      for ((role, state), event, binary) in SIMPLE_CASES:
 134          tr(READERS[role, state], binary, event)
 135  
 136  
 137  def test_writers_unusual() -> None:
 138      # Simple test of the write_headers utility routine
 139      tw(
 140          write_headers,
 141          normalize_and_validate([("foo", "bar"), ("baz", "quux")]),
 142          b"foo: bar\r\nbaz: quux\r\n\r\n",
 143      )
 144      tw(write_headers, Headers([]), b"\r\n")
 145  
 146      # We understand HTTP/1.0, but we don't speak it
 147      with pytest.raises(LocalProtocolError):
 148          tw(
 149              write_request,
 150              Request(
 151                  method="GET",
 152                  target="/",
 153                  headers=[("Host", "foo"), ("Connection", "close")],
 154                  http_version="1.0",
 155              ),
 156              None,
 157          )
 158      with pytest.raises(LocalProtocolError):
 159          tw(
 160              write_any_response,
 161              Response(
 162                  status_code=200, headers=[("Connection", "close")], http_version="1.0"
 163              ),
 164              None,
 165          )
 166  
 167  
 168  def test_readers_unusual() -> None:
 169      # Reading HTTP/1.0
 170      tr(
 171          READERS[CLIENT, IDLE],
 172          b"HEAD /foo HTTP/1.0\r\nSome: header\r\n\r\n",
 173          Request(
 174              method="HEAD",
 175              target="/foo",
 176              headers=[("Some", "header")],
 177              http_version="1.0",
 178          ),
 179      )
 180  
 181      # check no-headers, since it's only legal with HTTP/1.0
 182      tr(
 183          READERS[CLIENT, IDLE],
 184          b"HEAD /foo HTTP/1.0\r\n\r\n",
 185          Request(method="HEAD", target="/foo", headers=[], http_version="1.0"),  # type: ignore[arg-type]
 186      )
 187  
 188      tr(
 189          READERS[SERVER, SEND_RESPONSE],
 190          b"HTTP/1.0 200 OK\r\nSome: header\r\n\r\n",
 191          Response(
 192              status_code=200,
 193              headers=[("Some", "header")],
 194              http_version="1.0",
 195              reason=b"OK",
 196          ),
 197      )
 198  
 199      # single-character header values (actually disallowed by the ABNF in RFC
 200      # 7230 -- this is a bug in the standard that we originally copied...)
 201      tr(
 202          READERS[SERVER, SEND_RESPONSE],
 203          b"HTTP/1.0 200 OK\r\n" b"Foo: a a a a a \r\n\r\n",
 204          Response(
 205              status_code=200,
 206              headers=[("Foo", "a a a a a")],
 207              http_version="1.0",
 208              reason=b"OK",
 209          ),
 210      )
 211  
 212      # Empty headers -- also legal
 213      tr(
 214          READERS[SERVER, SEND_RESPONSE],
 215          b"HTTP/1.0 200 OK\r\n" b"Foo:\r\n\r\n",
 216          Response(
 217              status_code=200, headers=[("Foo", "")], http_version="1.0", reason=b"OK"
 218          ),
 219      )
 220  
 221      tr(
 222          READERS[SERVER, SEND_RESPONSE],
 223          b"HTTP/1.0 200 OK\r\n" b"Foo: \t \t \r\n\r\n",
 224          Response(
 225              status_code=200, headers=[("Foo", "")], http_version="1.0", reason=b"OK"
 226          ),
 227      )
 228  
 229      # Tolerate broken servers that leave off the response code
 230      tr(
 231          READERS[SERVER, SEND_RESPONSE],
 232          b"HTTP/1.0 200\r\n" b"Foo: bar\r\n\r\n",
 233          Response(
 234              status_code=200, headers=[("Foo", "bar")], http_version="1.0", reason=b""
 235          ),
 236      )
 237  
 238      # Tolerate headers line endings (\r\n and \n)
 239      #    \n\r\b between headers and body
 240      tr(
 241          READERS[SERVER, SEND_RESPONSE],
 242          b"HTTP/1.1 200 OK\r\nSomeHeader: val\n\r\n",
 243          Response(
 244              status_code=200,
 245              headers=[("SomeHeader", "val")],
 246              http_version="1.1",
 247              reason="OK",
 248          ),
 249      )
 250  
 251      #   delimited only with \n
 252      tr(
 253          READERS[SERVER, SEND_RESPONSE],
 254          b"HTTP/1.1 200 OK\nSomeHeader1: val1\nSomeHeader2: val2\n\n",
 255          Response(
 256              status_code=200,
 257              headers=[("SomeHeader1", "val1"), ("SomeHeader2", "val2")],
 258              http_version="1.1",
 259              reason="OK",
 260          ),
 261      )
 262  
 263      #   mixed \r\n and \n
 264      tr(
 265          READERS[SERVER, SEND_RESPONSE],
 266          b"HTTP/1.1 200 OK\r\nSomeHeader1: val1\nSomeHeader2: val2\n\r\n",
 267          Response(
 268              status_code=200,
 269              headers=[("SomeHeader1", "val1"), ("SomeHeader2", "val2")],
 270              http_version="1.1",
 271              reason="OK",
 272          ),
 273      )
 274  
 275      # obsolete line folding
 276      tr(
 277          READERS[CLIENT, IDLE],
 278          b"HEAD /foo HTTP/1.1\r\n"
 279          b"Host: example.com\r\n"
 280          b"Some: multi-line\r\n"
 281          b" header\r\n"
 282          b"\tnonsense\r\n"
 283          b"    \t   \t\tI guess\r\n"
 284          b"Connection: close\r\n"
 285          b"More-nonsense: in the\r\n"
 286          b"    last header  \r\n\r\n",
 287          Request(
 288              method="HEAD",
 289              target="/foo",
 290              headers=[
 291                  ("Host", "example.com"),
 292                  ("Some", "multi-line header nonsense I guess"),
 293                  ("Connection", "close"),
 294                  ("More-nonsense", "in the last header"),
 295              ],
 296          ),
 297      )
 298  
 299      with pytest.raises(LocalProtocolError):
 300          tr(
 301              READERS[CLIENT, IDLE],
 302              b"HEAD /foo HTTP/1.1\r\n" b"  folded: line\r\n\r\n",
 303              None,
 304          )
 305  
 306      with pytest.raises(LocalProtocolError):
 307          tr(
 308              READERS[CLIENT, IDLE],
 309              b"HEAD /foo HTTP/1.1\r\n" b"foo  : line\r\n\r\n",
 310              None,
 311          )
 312      with pytest.raises(LocalProtocolError):
 313          tr(
 314              READERS[CLIENT, IDLE],
 315              b"HEAD /foo HTTP/1.1\r\n" b"foo\t: line\r\n\r\n",
 316              None,
 317          )
 318      with pytest.raises(LocalProtocolError):
 319          tr(
 320              READERS[CLIENT, IDLE],
 321              b"HEAD /foo HTTP/1.1\r\n" b"foo\t: line\r\n\r\n",
 322              None,
 323          )
 324      with pytest.raises(LocalProtocolError):
 325          tr(READERS[CLIENT, IDLE], b"HEAD /foo HTTP/1.1\r\n" b": line\r\n\r\n", None)
 326  
 327  
 328  def test__obsolete_line_fold_bytes() -> None:
 329      # _obsolete_line_fold has a defensive cast to bytearray, which is
 330      # necessary to protect against O(n^2) behavior in case anyone ever passes
 331      # in regular bytestrings... but right now we never pass in regular
 332      # bytestrings. so this test just exists to get some coverage on that
 333      # defensive cast.
 334      assert list(_obsolete_line_fold([b"aaa", b"bbb", b"  ccc", b"ddd"])) == [
 335          b"aaa",
 336          bytearray(b"bbb ccc"),
 337          b"ddd",
 338      ]
 339  
 340  
 341  def _run_reader_iter(
 342      reader: Any, buf: bytes, do_eof: bool
 343  ) -> Generator[Any, None, None]:
 344      while True:
 345          event = reader(buf)
 346          if event is None:
 347              break
 348          yield event
 349          # body readers have undefined behavior after returning EndOfMessage,
 350          # because this changes the state so they don't get called again
 351          if type(event) is EndOfMessage:
 352              break
 353      if do_eof:
 354          assert not buf
 355          yield reader.read_eof()
 356  
 357  
 358  def _run_reader(*args: Any) -> List[Event]:
 359      events = list(_run_reader_iter(*args))
 360      return normalize_data_events(events)
 361  
 362  
 363  def t_body_reader(thunk: Any, data: bytes, expected: Any, do_eof: bool = False) -> None:
 364      # Simple: consume whole thing
 365      print("Test 1")
 366      buf = makebuf(data)
 367      assert _run_reader(thunk(), buf, do_eof) == expected
 368  
 369      # Incrementally growing buffer
 370      print("Test 2")
 371      reader = thunk()
 372      buf = ReceiveBuffer()
 373      events = []
 374      for i in range(len(data)):
 375          events += _run_reader(reader, buf, False)
 376          buf += data[i : i + 1]
 377      events += _run_reader(reader, buf, do_eof)
 378      assert normalize_data_events(events) == expected
 379  
 380      is_complete = any(type(event) is EndOfMessage for event in expected)
 381      if is_complete and not do_eof:
 382          buf = makebuf(data + b"trailing")
 383          assert _run_reader(thunk(), buf, False) == expected
 384  
 385  
 386  def test_ContentLengthReader() -> None:
 387      t_body_reader(lambda: ContentLengthReader(0), b"", [EndOfMessage()])
 388  
 389      t_body_reader(
 390          lambda: ContentLengthReader(10),
 391          b"0123456789",
 392          [Data(data=b"0123456789"), EndOfMessage()],
 393      )
 394  
 395  
 396  def test_Http10Reader() -> None:
 397      t_body_reader(Http10Reader, b"", [EndOfMessage()], do_eof=True)
 398      t_body_reader(Http10Reader, b"asdf", [Data(data=b"asdf")], do_eof=False)
 399      t_body_reader(
 400          Http10Reader, b"asdf", [Data(data=b"asdf"), EndOfMessage()], do_eof=True
 401      )
 402  
 403  
 404  def test_ChunkedReader() -> None:
 405      t_body_reader(ChunkedReader, b"0\r\n\r\n", [EndOfMessage()])
 406  
 407      t_body_reader(
 408          ChunkedReader,
 409          b"0\r\nSome: header\r\n\r\n",
 410          [EndOfMessage(headers=[("Some", "header")])],
 411      )
 412  
 413      t_body_reader(
 414          ChunkedReader,
 415          b"5\r\n01234\r\n"
 416          + b"10\r\n0123456789abcdef\r\n"
 417          + b"0\r\n"
 418          + b"Some: header\r\n\r\n",
 419          [
 420              Data(data=b"012340123456789abcdef"),
 421              EndOfMessage(headers=[("Some", "header")]),
 422          ],
 423      )
 424  
 425      t_body_reader(
 426          ChunkedReader,
 427          b"5\r\n01234\r\n" + b"10\r\n0123456789abcdef\r\n" + b"0\r\n\r\n",
 428          [Data(data=b"012340123456789abcdef"), EndOfMessage()],
 429      )
 430  
 431      # handles upper and lowercase hex
 432      t_body_reader(
 433          ChunkedReader,
 434          b"aA\r\n" + b"x" * 0xAA + b"\r\n" + b"0\r\n\r\n",
 435          [Data(data=b"x" * 0xAA), EndOfMessage()],
 436      )
 437  
 438      # refuses arbitrarily long chunk integers
 439      with pytest.raises(LocalProtocolError):
 440          # Technically this is legal HTTP/1.1, but we refuse to process chunk
 441          # sizes that don't fit into 20 characters of hex
 442          t_body_reader(ChunkedReader, b"9" * 100 + b"\r\nxxx", [Data(data=b"xxx")])
 443  
 444      # refuses garbage in the chunk count
 445      with pytest.raises(LocalProtocolError):
 446          t_body_reader(ChunkedReader, b"10\x00\r\nxxx", None)
 447  
 448      # handles (and discards) "chunk extensions" omg wtf
 449      t_body_reader(
 450          ChunkedReader,
 451          b"5; hello=there\r\n"
 452          + b"xxxxx"
 453          + b"\r\n"
 454          + b'0; random="junk"; some=more; canbe=lonnnnngg\r\n\r\n',
 455          [Data(data=b"xxxxx"), EndOfMessage()],
 456      )
 457  
 458      t_body_reader(
 459          ChunkedReader,
 460          b"5   	 \r\n01234\r\n" + b"0\r\n\r\n",
 461          [Data(data=b"01234"), EndOfMessage()],
 462      )
 463  
 464  
 465  def test_ContentLengthWriter() -> None:
 466      w = ContentLengthWriter(5)
 467      assert dowrite(w, Data(data=b"123")) == b"123"
 468      assert dowrite(w, Data(data=b"45")) == b"45"
 469      assert dowrite(w, EndOfMessage()) == b""
 470  
 471      w = ContentLengthWriter(5)
 472      with pytest.raises(LocalProtocolError):
 473          dowrite(w, Data(data=b"123456"))
 474  
 475      w = ContentLengthWriter(5)
 476      dowrite(w, Data(data=b"123"))
 477      with pytest.raises(LocalProtocolError):
 478          dowrite(w, Data(data=b"456"))
 479  
 480      w = ContentLengthWriter(5)
 481      dowrite(w, Data(data=b"123"))
 482      with pytest.raises(LocalProtocolError):
 483          dowrite(w, EndOfMessage())
 484  
 485      w = ContentLengthWriter(5)
 486      dowrite(w, Data(data=b"123")) == b"123"
 487      dowrite(w, Data(data=b"45")) == b"45"
 488      with pytest.raises(LocalProtocolError):
 489          dowrite(w, EndOfMessage(headers=[("Etag", "asdf")]))
 490  
 491  
 492  def test_ChunkedWriter() -> None:
 493      w = ChunkedWriter()
 494      assert dowrite(w, Data(data=b"aaa")) == b"3\r\naaa\r\n"
 495      assert dowrite(w, Data(data=b"a" * 20)) == b"14\r\n" + b"a" * 20 + b"\r\n"
 496  
 497      assert dowrite(w, Data(data=b"")) == b""
 498  
 499      assert dowrite(w, EndOfMessage()) == b"0\r\n\r\n"
 500  
 501      assert (
 502          dowrite(w, EndOfMessage(headers=[("Etag", "asdf"), ("a", "b")]))
 503          == b"0\r\nEtag: asdf\r\na: b\r\n\r\n"
 504      )
 505  
 506  
 507  def test_Http10Writer() -> None:
 508      w = Http10Writer()
 509      assert dowrite(w, Data(data=b"1234")) == b"1234"
 510      assert dowrite(w, EndOfMessage()) == b""
 511  
 512      with pytest.raises(LocalProtocolError):
 513          dowrite(w, EndOfMessage(headers=[("Etag", "asdf")]))
 514  
 515  
 516  def test_reject_garbage_after_request_line() -> None:
 517      with pytest.raises(LocalProtocolError):
 518          tr(READERS[SERVER, SEND_RESPONSE], b"HTTP/1.0 200 OK\x00xxxx\r\n\r\n", None)
 519  
 520  
 521  def test_reject_garbage_after_response_line() -> None:
 522      with pytest.raises(LocalProtocolError):
 523          tr(
 524              READERS[CLIENT, IDLE],
 525              b"HEAD /foo HTTP/1.1 xxxxxx\r\n" b"Host: a\r\n\r\n",
 526              None,
 527          )
 528  
 529  
 530  def test_reject_garbage_in_header_line() -> None:
 531      with pytest.raises(LocalProtocolError):
 532          tr(
 533              READERS[CLIENT, IDLE],
 534              b"HEAD /foo HTTP/1.1\r\n" b"Host: foo\x00bar\r\n\r\n",
 535              None,
 536          )
 537  
 538  
 539  def test_reject_non_vchar_in_path() -> None:
 540      for bad_char in b"\x00\x20\x7f\xee":
 541          message = bytearray(b"HEAD /")
 542          message.append(bad_char)
 543          message.extend(b" HTTP/1.1\r\nHost: foobar\r\n\r\n")
 544          with pytest.raises(LocalProtocolError):
 545              tr(READERS[CLIENT, IDLE], message, None)
 546  
 547  
 548  # https://github.com/python-hyper/h11/issues/57
 549  def test_allow_some_garbage_in_cookies() -> None:
 550      tr(
 551          READERS[CLIENT, IDLE],
 552          b"HEAD /foo HTTP/1.1\r\n"
 553          b"Host: foo\r\n"
 554          b"Set-Cookie: ___utmvafIumyLc=kUd\x01UpAt; path=/; Max-Age=900\r\n"
 555          b"\r\n",
 556          Request(
 557              method="HEAD",
 558              target="/foo",
 559              headers=[
 560                  ("Host", "foo"),
 561                  ("Set-Cookie", "___utmvafIumyLc=kUd\x01UpAt; path=/; Max-Age=900"),
 562              ],
 563          ),
 564      )
 565  
 566  
 567  def test_host_comes_first() -> None:
 568      tw(
 569          write_headers,
 570          normalize_and_validate([("foo", "bar"), ("Host", "example.com")]),
 571          b"Host: example.com\r\nfoo: bar\r\n\r\n",
 572      )
 573