From ac11ae64a85b52894d04f3fc252c2606580db4b5 Mon Sep 17 00:00:00 2001
From: Adam Wick <awick@uhsure.com>
Date: Sat, 6 Nov 2021 20:56:19 -0700
Subject: [PATCH] Add a test that our little mini-firewall works as intended.

---
 src/client.rs |  3 +++
 src/lib.rs    | 24 ++++++++++++++++++++++++
 2 files changed, 27 insertions(+)

diff --git a/src/client.rs b/src/client.rs
index d06fe6f..4245297 100644
--- a/src/client.rs
+++ b/src/client.rs
@@ -4,6 +4,7 @@ use crate::messages::{
     ServerResponseStatus,
 };
 use crate::network::generic::Networklike;
+use async_std::io;
 use futures::io::{AsyncRead, AsyncWrite};
 use log::{trace, warn};
 use thiserror::Error;
@@ -22,6 +23,8 @@ pub enum SOCKSv5Error {
     UnsupportedAuthMethodChosen(AuthenticationMethod),
     #[error("Server said no: {0}")]
     ServerFailure(#[from] ServerResponseStatus),
+    #[error("Connection error: {0}")]
+    ConnectionError(#[from] io::Error),
 }
 
 pub struct SOCKSv5Client<S, N>
diff --git a/src/lib.rs b/src/lib.rs
index 753fbc1..83db050 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -103,4 +103,28 @@ mod test {
             assert!(client.is_err());
         })
     }
+
+    #[test]
+    fn firewall_blocks() {
+        task::block_on(async {
+            let mut network_stack = TestingStack::default();
+
+            // generate the server
+            let mut security_parameters = SecurityParameters::unrestricted();
+            security_parameters.allow_connection = Some(|_, _| false);
+            let default_port = network_stack.listen("localhost", 9999).await.unwrap();
+            let server =
+                SOCKSv5Server::new(network_stack.clone(), security_parameters, default_port);
+
+            let _server_task = task::spawn(async move { server.run().await });
+
+            let stream = network_stack.connect("localhost", 9999).await.unwrap();
+            let login_info = LoginInfo {
+                username_password: None,
+            };
+            let client = SOCKSv5Client::new(network_stack, stream, &login_info).await;
+
+            assert!(client.is_err());
+        })
+    }
 }