From 9f9c35b248a5b03e34d6cb852d7504a54cd9d730 Mon Sep 17 00:00:00 2001 From: int 80h Date: Mon, 27 Apr 2020 15:37:54 -0400 Subject: [PATCH] Made the send functions methods of Connection struct --- src/cgi.rs | 7 +++---- src/main.rs | 27 ++++++++------------------- src/util.rs | 29 +++++++++++++++++++---------- 3 files changed, 30 insertions(+), 33 deletions(-) diff --git a/src/cgi.rs b/src/cgi.rs index c41d744..1d4955f 100644 --- a/src/cgi.rs +++ b/src/cgi.rs @@ -9,9 +9,8 @@ use url::Url; use crate::config; use crate::status; use crate::util; -use crate::Connection; -pub async fn cgi(con: Connection, path: PathBuf, url: Url) -> Result<(), io::Error> { +pub async fn cgi(mut con: util::Connection, path: PathBuf, url: Url) -> Result<(), io::Error> { let mut envs = HashMap::new(); envs.insert("GEMINI_URL", url.as_str()); envs.insert("SERVER_NAME", url.host_str().unwrap()); @@ -33,12 +32,12 @@ pub async fn cgi(con: Connection, path: PathBuf, url: Url) -> Result<(), io::Err .output() .unwrap(); if !cmd.status.success() { - util::send_status(con.stream, status::Status::CGIError, "CGI Error!").await?; + con.send_status(status::Status::CGIError, "CGI Error!").await?; return Ok(()); } let cmd = String::from_utf8(cmd.stdout).unwrap(); // if cmd.starts_with("20") { - util::send_raw(con.stream, cmd).await?; + con.send_raw(cmd).await?; //util::send_body(stream, status::Status::Success, "text/gemini", Some(cmd)).await?; // } return Ok(()); diff --git a/src/main.rs b/src/main.rs index cfde722..4f3a537 100644 --- a/src/main.rs +++ b/src/main.rs @@ -33,14 +33,6 @@ mod status; mod tls; mod util; -pub struct Connection { - stream: TlsStream, - peer_addr: SocketAddr, - hostname: String, - dir: String, - cgi: String, -} - fn get_content(path: PathBuf, u: url::Url) -> Result { let meta = fs::metadata(&path).expect("Unable to read metadata"); if meta.is_file() { @@ -60,7 +52,7 @@ fn get_content(path: PathBuf, u: url::Url) -> Result { return Ok(list); } -async fn handle_connection(mut con: Connection) -> Result<(), io::Error> { +async fn handle_connection(mut con: util::Connection) -> Result<(), io::Error> { let now: DateTime = Utc::now(); println!("{} New Connection: {}", now, con.peer_addr); let mut buffer = [0; 512]; @@ -71,13 +63,12 @@ async fn handle_connection(mut con: Connection) -> Result<(), io::Error> { let url = Url::parse(&request).unwrap(); if Some(con.hostname.as_str()) != url.host_str() { - util::send_status(con.stream, status::Status::PermanentFailure, "Url doesn't match certificate!").await?; + con.send_status(status::Status::PermanentFailure, "Url doesn't match certificate!").await?; return Ok(()); } if url.scheme() != "gemini" { - util::send_status( - con.stream, + con.send_status( status::Status::ProxyRequestRefused, "Not a gemini scheme!\r\n", ) @@ -86,7 +77,7 @@ async fn handle_connection(mut con: Connection) -> Result<(), io::Error> { } if url.path().to_string().contains("..") { - util::send_status(con.stream, status::Status::PermanentFailure, "Not in path!").await?; + con.send_status(status::Status::PermanentFailure, "Not in path!").await?; return Ok(()); } @@ -96,7 +87,7 @@ async fn handle_connection(mut con: Connection) -> Result<(), io::Error> { } if !path.exists() { - util::send_status(con.stream, status::Status::NotFound, "Not found!\r\n").await?; + con.send_status(status::Status::NotFound, "Not found!\r\n").await?; return Ok(()); } @@ -105,8 +96,7 @@ async fn handle_connection(mut con: Connection) -> Result<(), io::Error> { if meta.is_dir() { if !url.path().ends_with("/") { - util::send_status( - con.stream, + con.send_status( status::Status::RedirectPermanent, format!("{}/\r\n", url).as_str(), ) @@ -125,8 +115,7 @@ async fn handle_connection(mut con: Connection) -> Result<(), io::Error> { } let content = get_content(path, url)?; - util::send_body( - con.stream, + con.send_body( status::Status::Success, "text/gemini", Some(content), @@ -179,7 +168,7 @@ fn main() -> io::Result<()> { } } - let con = Connection { + let con = util::Connection { stream, peer_addr, hostname, diff --git a/src/util.rs b/src/util.rs index 5ea171f..e92f8e4 100644 --- a/src/util.rs +++ b/src/util.rs @@ -1,37 +1,46 @@ use std::io; +use std::net::SocketAddr; use tokio::net::TcpStream; use tokio::prelude::*; use tokio_rustls::server::TlsStream; use crate::status; +pub struct Connection { + pub stream: TlsStream, + pub peer_addr: SocketAddr, + pub hostname: String, + pub dir: String, + pub cgi: String, +} + +impl Connection { pub async fn send_status( - stream: TlsStream, + &mut self, stat: status::Status, meta: &str, ) -> Result<(), io::Error> { - send_body(stream, stat, meta, None).await?; + self.send_body(stat, meta, None).await?; Ok(()) } pub async fn send_body( - mut stream: TlsStream, + &mut self, stat: status::Status, meta: &str, body: Option, ) -> Result<(), io::Error> { let mut s = format!("{}\t{}\r\n", stat as u8, meta); - stream.write_all(s.as_bytes()).await?; - stream.flush().await?; if let Some(b) = body { - s = format!("{}", b); + s += &b; } - send_raw(stream, s).await?; + self.send_raw(s).await?; Ok(()) } -pub async fn send_raw(mut stream: TlsStream, body: String) -> Result<(), io::Error> { - stream.write_all(body.as_bytes()).await?; - stream.flush().await?; +pub async fn send_raw(&mut self, body: String) -> Result<(), io::Error> { + self.stream.write_all(body.as_bytes()).await?; + self.stream.flush().await?; Ok(()) } +}