From 788b7921674fc10cb07a8ca076c9b0257be98674 Mon Sep 17 00:00:00 2001
From: yggverse <yggverse@project>
Date: Mon, 3 Feb 2025 01:08:31 +0200
Subject: [PATCH] return connection in result

---
 src/client.rs                     |  2 +-
 src/client/connection.rs          | 19 ++++++-----
 src/client/connection/response.rs | 57 ++++++++++++++++---------------
 3 files changed, 42 insertions(+), 36 deletions(-)

diff --git a/src/client.rs b/src/client.rs
index bcce693..f222393 100644
--- a/src/client.rs
+++ b/src/client.rs
@@ -60,7 +60,7 @@ impl Client {
         priority: Priority,
         cancellable: Cancellable,
         certificate: Option<TlsCertificate>,
-        callback: impl FnOnce(Result<Response, Error>) + 'static,
+        callback: impl FnOnce(Result<(Response, Connection), Error>) + 'static,
     ) {
         // Begin new connection
         // * [NetworkAddress](https://docs.gtk.org/gio/class.NetworkAddress.html) required for valid
diff --git a/src/client/connection.rs b/src/client/connection.rs
index 9af51ed..b2cee84 100644
--- a/src/client/connection.rs
+++ b/src/client/connection.rs
@@ -58,7 +58,7 @@ impl Connection {
         request: Request,
         priority: Priority,
         cancellable: Cancellable,
-        callback: impl FnOnce(Result<Response, Error>) + 'static,
+        callback: impl FnOnce(Result<(Response, Self), Error>) + 'static,
     ) {
         let output_stream = self.stream().output_stream();
         output_stream.clone().write_async(
@@ -67,14 +67,17 @@ impl Connection {
             Some(&cancellable.clone()),
             move |result| match result {
                 Ok(_) => match request {
-                    Request::Gemini { .. } => {
-                        Response::from_connection_async(self, priority, cancellable, |result| {
+                    Request::Gemini { .. } => Response::from_connection_async(
+                        self,
+                        priority,
+                        cancellable,
+                        |result, connection| {
                             callback(match result {
-                                Ok(response) => Ok(response),
+                                Ok(response) => Ok((response, connection)),
                                 Err(e) => Err(Error::Response(e)),
                             })
-                        })
-                    }
+                        },
+                    ),
                     Request::Titan { data, .. } => output_stream.write_bytes_async(
                         &data,
                         priority,
@@ -84,9 +87,9 @@ impl Connection {
                                 self,
                                 priority,
                                 cancellable,
-                                |result| {
+                                |result, connection| {
                                     callback(match result {
-                                        Ok(response) => Ok(response),
+                                        Ok(response) => Ok((response, connection)),
                                         Err(e) => Err(Error::Response(e)),
                                     })
                                 },
diff --git a/src/client/connection/response.rs b/src/client/connection/response.rs
index 9f2b901..adaf209 100644
--- a/src/client/connection/response.rs
+++ b/src/client/connection/response.rs
@@ -36,7 +36,7 @@ impl Response {
         connection: Connection,
         priority: Priority,
         cancellable: Cancellable,
-        callback: impl FnOnce(Result<Self, Error>) + 'static,
+        callback: impl FnOnce(Result<Self, Error>, Connection) + 'static,
     ) {
         from_stream_async(
             Vec::with_capacity(HEADER_LEN),
@@ -44,35 +44,38 @@ impl Response {
             cancellable,
             priority,
             |result| {
-                callback(match result {
-                    Ok(buffer) => match buffer.first() {
-                        Some(byte) => match byte {
-                            1 => match Input::from_utf8(&buffer) {
-                                Ok(input) => Ok(Self::Input(input)),
-                                Err(e) => Err(Error::Input(e)),
+                callback(
+                    match result {
+                        Ok(buffer) => match buffer.first() {
+                            Some(byte) => match byte {
+                                1 => match Input::from_utf8(&buffer) {
+                                    Ok(input) => Ok(Self::Input(input)),
+                                    Err(e) => Err(Error::Input(e)),
+                                },
+                                2 => match Success::from_utf8(&buffer) {
+                                    Ok(success) => Ok(Self::Success(success)),
+                                    Err(e) => Err(Error::Success(e)),
+                                },
+                                3 => match Redirect::from_utf8(&buffer) {
+                                    Ok(redirect) => Ok(Self::Redirect(redirect)),
+                                    Err(e) => Err(Error::Redirect(e)),
+                                },
+                                4 | 5 => match Failure::from_utf8(&buffer) {
+                                    Ok(failure) => Ok(Self::Failure(failure)),
+                                    Err(e) => Err(Error::Failure(e)),
+                                },
+                                6 => match Certificate::from_utf8(&buffer) {
+                                    Ok(certificate) => Ok(Self::Certificate(certificate)),
+                                    Err(e) => Err(Error::Certificate(e)),
+                                },
+                                b => Err(Error::Code(*b)),
                             },
-                            2 => match Success::from_utf8(&buffer) {
-                                Ok(success) => Ok(Self::Success(success)),
-                                Err(e) => Err(Error::Success(e)),
-                            },
-                            3 => match Redirect::from_utf8(&buffer) {
-                                Ok(redirect) => Ok(Self::Redirect(redirect)),
-                                Err(e) => Err(Error::Redirect(e)),
-                            },
-                            4 | 5 => match Failure::from_utf8(&buffer) {
-                                Ok(failure) => Ok(Self::Failure(failure)),
-                                Err(e) => Err(Error::Failure(e)),
-                            },
-                            6 => match Certificate::from_utf8(&buffer) {
-                                Ok(certificate) => Ok(Self::Certificate(certificate)),
-                                Err(e) => Err(Error::Certificate(e)),
-                            },
-                            b => Err(Error::Code(*b)),
+                            None => Err(Error::Protocol),
                         },
-                        None => Err(Error::Protocol),
+                        Err(e) => Err(e),
                     },
-                    Err(e) => Err(e),
-                })
+                    connection,
+                )
             },
         );
     }