From 8495c5ab0d6e86e8e6e2aad5a972f02b126b5f64 Mon Sep 17 00:00:00 2001
From: Christine Dodrill <me@christine.website>
Date: Mon, 27 Jul 2020 20:11:57 -0400
Subject: [PATCH] pass gemini torture tests

---
 src/server/mod.rs | 33 +++++++++++++++++++++------------
 1 file changed, 21 insertions(+), 12 deletions(-)

diff --git a/src/server/mod.rs b/src/server/mod.rs
index 4b3bac5..1ee016b 100644
--- a/src/server/mod.rs
+++ b/src/server/mod.rs
@@ -55,8 +55,9 @@ where
         let h = h.clone();
         let acceptor = acceptor.clone();
         let addr = stream.peer_addr().unwrap();
+        let port = port.clone();
 
-        task::spawn(handle_request(h, stream, acceptor, addr));
+        task::spawn(handle_request(h, stream, acceptor, addr, port));
     }
 
     Ok(())
@@ -68,12 +69,25 @@ async fn handle_request(
     stream: TcpStream,
     acceptor: Arc<TlsAcceptor>,
     addr: SocketAddr,
+    port: u16,
 ) -> Result {
     // Perform handshake.
     let mut stream = acceptor.clone().accept(stream).await?;
 
     match parse_request(&mut stream).await {
         Ok(url) => {
+            if let Some(u_port) = url.port() {
+                if port != u_port {
+                    let _ = respond(&mut stream, "53", &["Cannot proxy"]).await;
+                    return Ok(());
+                }
+            }
+
+            if url.scheme() != "gemini" {
+                let _ = respond(&mut stream, "53", &["Cannot proxy outside geminispace"]).await;
+                Err(RequestParsingError::InvalidScheme(url.scheme().to_string()))?
+            }
+
             let req = Request {
                 url: url,
                 certs: None,
@@ -82,7 +96,7 @@ async fn handle_request(
         }
         Err(e) => {
             respond(&mut stream, "59", &["Invalid request."]).await?;
-            return Err(e);
+            log::error!("error from {}: {:?}", addr, e);
         }
     }
     Ok(())
@@ -128,9 +142,6 @@ async fn parse_request<R: Read + Unpin>(mut stream: R) -> Result<Url> {
     };
 
     // Validate the URL. TODO: Check the hostname and port.
-    if url.scheme() != "gemini" {
-        Err(RequestParsingError::InvalidScheme(url.scheme().to_string()))?
-    }
     Ok(url)
 }
 
@@ -141,18 +152,16 @@ where
     let u = req.url.clone();
     match h.handle(req).await {
         Ok(resp) => {
-            stream
+            let _ = stream
                 .write(format!("{} {}\r\n", resp.status as u8, resp.meta).as_bytes())
-                .await
-                .unwrap();
-            stream.write(&resp.body).await.unwrap();
+                .await;
+            let _ = stream.write(&resp.body).await;
             log::info!("{}: {} {} {:?}", addr, u, resp.meta, resp.status);
         }
         Err(why) => {
-            stream
+            let _ = stream
                 .write(format!("{} {:?}\r\n", StatusCode::PermanentFailure as u8, why).as_bytes())
-                .await
-                .unwrap();
+                .await;
             log::error!("{}: {}: {:?}", addr, u, why);
         }
     };