Initial commit

src/main.rs

@@ -0,0 +1,95 @@
+mod http;
+mod policy;
+
+use http::HeaderLineIterator;
+use policy::{CompiledPolicies, Policy};
+
+use regex::bytes::Regex;
+use std::{
+	io::{BufReader, Write},
+	net::SocketAddr,
+	time::Duration,
+};
+use tokio::{io::{ReadBuf, AsyncWriteExt}, time::timeout};
+
+static CHALLENGE_BODY: &str = include_str!("challenge.html");
+
+macro_rules! mk_static {
+	($t:ty, $val:expr) => {{
+		static STATIC_CELL: static_cell::StaticCell<$t> = static_cell::StaticCell::new();
+		#[deny(unused_attributes)]
+		let x = STATIC_CELL.uninit().write(($val));
+		x
+	}};
+}
+
+#[tokio::main]
+async fn main() {
+	let listen_addr = "127.0.0.1:8000".parse().unwrap();
+	let policy_groups = &[&[Policy {
+		name: String::from("Block"),
+		filter: policy::Filter::FirstLineMatch(String::from("GET /block")),
+		action: policy::Action::Drop,
+		priority: 0,
+	}]];
+
+	let challenge_response = &*mk_static!(String, format!("HTTP/1.1 200\r\ncontent-type: text/html\r\ncontent-length: {}\r\n\r\n{}", CHALLENGE_BODY.len(), CHALLENGE_BODY));
+
+	let policy_groups: Vec<CompiledPolicies> = policy_groups.into_iter().map(|policies| CompiledPolicies::new(*policies)).collect();
+
+	let socket = realm_syscall::new_tcp_socket(&listen_addr).unwrap();
+
+	socket.set_reuse_address(true).ok();
+
+	socket.bind(&listen_addr.into()).unwrap();
+	socket.listen(1024).unwrap();
+
+	let listener = tokio::net::TcpListener::from_std(socket.into()).unwrap();
+
+	let cookie_regex = Regex::new(r"^Cookie: *(?:[^;=]+=[^;=]* *; *)*mesozoa *= *([0-9a-zA-Z]{4})").unwrap();
+
+
+			loop {
+				let Ok((mut client_stream, client_addr)) = listener.accept().await else {
+					continue;
+				};
+				//client_stream.set_nodelay(true).ok();
+		
+				let cookie_regex = cookie_regex.clone();
+				tokio::spawn(async move {
+					let mut buf = [0u8; 1024];
+					let mut buf_reader = ReadBuf::new(&mut buf);
+					if let Err(_) = timeout(
+						Duration::from_millis(100),
+						std::future::poll_fn(|cx| client_stream.poll_peek(cx, &mut buf_reader)),
+					)
+					.await
+					{
+						println!("peek timeout");
+						return;
+					}
+		
+					let mut header_line_iter = HeaderLineIterator::new(&buf);
+					let Some(first_line) = header_line_iter.next() else {
+						println!("Not HTTP, or too long line");
+						return;
+					};
+					// TODO matching
+					// for test we will challenge everything!
+					if let Some(captures) = header_line_iter.find_map(|line| cookie_regex.captures(line)) {
+						if let Some(cookie) = captures.get(1) {
+							let mut stdout = std::io::stdout();
+							stdout.write_all(cookie.as_bytes()).unwrap();
+							stdout.flush().unwrap();
+							println!("");
+						} else {
+							println!("cookie header, but no cookie")
+						}
+					} else {
+						println!("no cookie");
+					}
+					client_stream.writable().await.unwrap();
+					client_stream.write_all(challenge_response.as_bytes()).await.unwrap();
+				});
+			}
+}