"""Holding a connection open across several requests.

The reply path has two shapes. A small answer goes out in one write and the
connection is done with immediately. A large one is queued and finishes on a
later pass through the event loop, by which time the workstruct that served
the request has moved on, so whether to keep the connection has to be
remembered on the connection itself.
"""

import time
import unittest
import urllib.parse

from harness import TrackerTestCase, bdecode


class KeepAlive(TrackerTestCase):
    def announce_path(self, name, port=6881):
        return "/announce?info_hash=%s&peer_id=%s&port=%d&left=0" % (
            urllib.parse.quote_from_bytes(self.hash_for(name), safe=""),
            "x" * 20,
            port,
        )

    def test_two_requests_over_one_connection(self):
        c = self.tracker.connect()
        try:
            c.send(self.announce_path("kaA"))
            status, _, first = c.reply()
            self.assertIn(b"200", status)
            c.send(self.announce_path("kaB"))
            status, _, second = c.reply()
            self.assertIn(b"200", status)
            self.assertIn(b"complete", bdecode(second))
        finally:
            c.close()

    def test_close_ends_the_connection(self):
        c = self.tracker.connect()
        try:
            c.send(self.announce_path("kaclose"), connection="close")
            c.reply()
            self.assertFalse(c.is_open(), "the tracker was asked to hang up")
        finally:
            c.close()

    def test_a_connection_survives_a_reply_too_large_for_one_write(self):
        """The case that was broken: a queued reply finished draining and the
        connection was closed regardless of what the client had asked for."""
        for i in range(400):
            self.tracker.announce(bytes([i // 256, i % 256]) + b"\0" * 18, port=7000 + i, left=0)

        c = self.tracker.connect()
        try:
            c.send("/scrape")
            status, headers, body = c.reply()
            self.assertIn(b"200", status)
            self.assertGreater(len(body), 16384, "this reply has to outgrow one write")

            c.send(self.announce_path("kaafter"))
            status, _, after = c.reply()
            self.assertIn(b"200", status)
            self.assertIn(b"complete", bdecode(after))
        finally:
            c.close()

    def test_several_requests_in_a_row(self):
        c = self.tracker.connect()
        try:
            for i in range(20):
                c.send(self.announce_path("karow%d" % i))
                status, _, body = c.reply()
                self.assertIn(b"200", status)
            self.assertEqual(self.tracker.torrent_count(), 20)
        finally:
            c.close()

    def test_without_the_header_the_connection_is_not_held(self):
        """HTTP/1.1 calls a connection persistent unless told otherwise; this
        tracker holds one only when asked by name. Pinned as it is."""
        c = self.tracker.connect()
        try:
            c.send(self.announce_path("kanone"), connection=None)
            c.reply()
            self.assertFalse(c.is_open())
        finally:
            c.close()

    def test_the_error_path_does_not_hold_the_connection(self):
        c = self.tracker.connect()
        try:
            c.send("/announce?info_hash=short")
            status, _, _ = c.reply()
            self.assertIn(b"400", status)
        finally:
            c.close()

    def test_an_error_after_a_kept_connection_is_still_an_error(self):
        """The keep-alive decision used to be taken after the handler had run,
        so an error reply inherited whatever the previous request wanted."""
        c = self.tracker.connect()
        try:
            c.send(self.announce_path("kabefore"))
            self.assertIn(b"200", c.reply()[0])
            c.send("/announce?info_hash=short")
            self.assertIn(b"400", c.reply()[0])
        finally:
            c.close()

    def test_a_tiny_split_request_does_not_overread(self):
        """A request that arrives in pieces and completes as a buffer shorter
        than the dispatch literals (\"GET /\" etc.) used to read past the
        exactly sized buffer. It must be rejected without a sanitizer hit."""
        c = self.tracker.connect()
        try:
            c.sock.sendall(b"\r\n")
            time.sleep(0.05)
            c.sock.sendall(b"\r\n")          # completes as a 4-byte buffer
            time.sleep(0.1)
        finally:
            c.close()
        # the tracker must still answer; tearDown's stop() surfaces any sanitizer abort
        answer = self.tracker.announce(self.hash_for("after-tiny"), port=6881, left=0)
        self.assertNotIn(b"failure reason", answer, "the tracker stopped serving after a tiny request")

    def test_a_split_request_answered_by_a_worker_does_not_spin(self):
        """A reply that a worker produces does not go through the send path,
        so the request stays in the buffer the read loop is reading. The loop
        used to find the same request again on every pass, queue a task for
        it, and never hand the thread back to the event loop."""
        for path in ("/stats", "/scrape"):
            c = self.tracker.connect()
            try:
                c.sock.sendall(("GET %s HTTP/1.1\r\nConnection: keep-alive\r\n" % path).encode())
                time.sleep(0.2)
                c.sock.sendall(b"\r\n")
                self.assertIn(b"200", c.reply()[0], "%s was not answered" % path)
            finally:
                c.close()

            answer = self.tracker.announce(self.hash_for("after-" + path), port=6881, left=0)
            self.assertNotIn(b"failure reason", answer, "the tracker stopped serving after %s" % path)


if __name__ == "__main__":
    unittest.main()