From 7ad382f2ac6d3a5bad7e20e7b14a0ae311eba4e8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jeremy=20Lain=C3=A9?= <jeremy.laine@m4x.org>
Date: Sat, 6 Jul 2024 23:30:37 +0200
Subject: [PATCH] Allow the asyncio protocol to specify the close code / reason

The asyncio wrapper for the `close` method should allow the API user to
specify the error code and the reason phrase.
---
 examples/http3_client.py        |  2 +-
 src/aioquic/asyncio/protocol.py | 17 +++++++++++++++--
 2 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/examples/http3_client.py b/examples/http3_client.py
index ecd9df8d2..36bb42023 100644
--- a/examples/http3_client.py
+++ b/examples/http3_client.py
@@ -425,7 +425,7 @@ async def main(
 
             # process http pushes
             process_http_pushes(client=client, include=include, output_dir=output_dir)
-        client._quic.close(error_code=ErrorCode.H3_NO_ERROR)
+        client.close(error_code=ErrorCode.H3_NO_ERROR)
 
 
 if __name__ == "__main__":
diff --git a/src/aioquic/asyncio/protocol.py b/src/aioquic/asyncio/protocol.py
index 3cbc18fb9..9e530ae0c 100644
--- a/src/aioquic/asyncio/protocol.py
+++ b/src/aioquic/asyncio/protocol.py
@@ -3,6 +3,7 @@
 
 from ..quic import events
 from ..quic.connection import NetworkAddress, QuicConnection
+from ..quic.packet import QuicErrorCode
 
 QuicConnectionIdHandler = Callable[[bytes], None]
 QuicStreamHandler = Callable[[asyncio.StreamReader, asyncio.StreamWriter], None]
@@ -44,11 +45,23 @@ def change_connection_id(self) -> None:
         self._quic.change_connection_id()
         self.transmit()
 
-    def close(self) -> None:
+    def close(
+        self,
+        error_code: int = QuicErrorCode.NO_ERROR,
+        reason_phrase: str = "",
+    ) -> None:
         """
         Close the connection.
+
+        :param error_code: An error code indicating why the connection is
+                           being closed.
+        :param reason_phrase: A human-readable explanation of why the
+                              connection is being closed.
         """
-        self._quic.close()
+        self._quic.close(
+            error_code=error_code,
+            reason_phrase=reason_phrase,
+        )
         self.transmit()
 
     def connect(self, addr: NetworkAddress) -> None: