working base
This commit is contained in:
parent
e090b70dba
commit
616c17501a
11 changed files with 1142 additions and 243 deletions
265
src/main.rs
265
src/main.rs
|
|
@ -6,18 +6,19 @@ use http::HeaderLineIterator;
|
|||
use policy::{CompiledPolicies, Policy};
|
||||
|
||||
use rand::Rng;
|
||||
use regex::{bytes::Regex, bytes::RegexSet};
|
||||
use std::{
|
||||
io::{BufReader, Write}, net::SocketAddr, sync::atomic::ATOMIC_BOOL_INIT, time::Duration
|
||||
};
|
||||
use regex::bytes::Regex;
|
||||
use std::{io::Write, net::SocketAddr, time::Duration};
|
||||
use tokio::{
|
||||
io::{AsyncWriteExt, ReadBuf}, net::TcpSocket, time::timeout
|
||||
io::{AsyncWriteExt, ReadBuf},
|
||||
net::{TcpSocket, TcpStream},
|
||||
time::timeout,
|
||||
};
|
||||
|
||||
const SALT_LEN: usize = 16;
|
||||
const SECRET_LEN: usize = 32;
|
||||
const MAC_LEN: usize = 32;
|
||||
const CHALLENGE_TIMEOUT: u64 = 3600;
|
||||
const TARGET_ZEROS: u32 = 15;
|
||||
|
||||
static CHALLENGE_BODY: &str = include_str!("challenge.html");
|
||||
|
||||
|
|
@ -36,28 +37,34 @@ async fn main() {
|
|||
|
||||
let listen_addr: SocketAddr = "127.0.0.1:8504".parse().unwrap();
|
||||
let pass_addr: SocketAddr = "127.0.0.1:80".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 policy_groups = vec![vec![
|
||||
Policy {
|
||||
name: String::from("Favicon"),
|
||||
first_line_regex: String::from(r"^GET /favicon.ico"),
|
||||
action: policy::Action::Allow,
|
||||
},
|
||||
Policy {
|
||||
name: String::from("robots.txt"),
|
||||
first_line_regex: String::from(r"^GET /robots.txt"),
|
||||
action: policy::Action::Allow,
|
||||
},
|
||||
Policy {
|
||||
name: String::from("Block"),
|
||||
first_line_regex: String::from(r"^GET /block"),
|
||||
action: policy::Action::Drop,
|
||||
},
|
||||
]];
|
||||
let default_action = policy::Action::Challenge;
|
||||
let secret: [u8; SECRET_LEN] = rng.r#gen();
|
||||
|
||||
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 = &*mk_static!(
|
||||
Vec<CompiledPolicies>,
|
||||
policy_groups
|
||||
.into_iter()
|
||||
.map(CompiledPolicies::new)
|
||||
.collect()
|
||||
);
|
||||
|
||||
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();
|
||||
|
|
@ -68,16 +75,26 @@ async fn main() {
|
|||
let listener = tokio::net::TcpListener::from_std(socket.into()).unwrap();
|
||||
|
||||
let proof_regex =
|
||||
Regex::new(r"^Cookie: *(?:[^;=]+=[^;=]* *; *)*mesozoa-proof *= *([0-9a-zA-Z_-]{4})").unwrap();
|
||||
Regex::new(r"^Cookie: *(?:[^;=]+=[^;=]* *; *)*mesozoa-proof *= *([0-9a-zA-Z_-]{8})")
|
||||
.unwrap();
|
||||
let challenge_regex =
|
||||
Regex::new(r"^Cookie: *(?:[^;=]+=[^;=]* *; *)*mesozoa-challenge *= *([0-9a-zA-Z_-]{75})").unwrap();
|
||||
let ip_regex =
|
||||
Regex::new(r"^X-Forwarded-For: *([a-fA-F0-9.:]+)$").unwrap();
|
||||
let user_agent_regex =
|
||||
Regex::new(r"^User-Agent: *([a-zA-Z0-9.,:;/ _()-]+)$").unwrap();
|
||||
Regex::new(r"^Cookie: *(?:[^;=]+=[^;=]* *; *)*mesozoa-challenge *= *([0-9a-zA-Z_-]{75})")
|
||||
.unwrap();
|
||||
let ip_regex = Regex::new(r"^X-Forwarded-For: *([a-fA-F0-9.:]+)$").unwrap();
|
||||
let user_agent_regex = Regex::new(r"^User-Agent: *([a-zA-Z0-9.,:;/ _()-]+)$").unwrap();
|
||||
|
||||
let response_begin = &*mk_static!(
|
||||
String,
|
||||
format!(
|
||||
"HTTP/1.1 200\r\n\
|
||||
content-type: text/html\r\n\
|
||||
content-length: {}\r\n",
|
||||
CHALLENGE_BODY.len(),
|
||||
)
|
||||
);
|
||||
|
||||
loop {
|
||||
let Ok((mut client_stream, client_addr)) = listener.accept().await else {
|
||||
let Ok((mut client_stream, _client_addr)) = listener.accept().await else {
|
||||
continue;
|
||||
};
|
||||
//client_stream.set_nodelay(true).ok();
|
||||
|
|
@ -89,11 +106,12 @@ async fn main() {
|
|||
tokio::spawn(async move {
|
||||
let mut buf = [0u8; 1024];
|
||||
let mut buf_reader = ReadBuf::new(&mut buf);
|
||||
if let Err(_) = timeout(
|
||||
if timeout(
|
||||
Duration::from_millis(100),
|
||||
std::future::poll_fn(|cx| client_stream.poll_peek(cx, &mut buf_reader)),
|
||||
)
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
println!("peek timeout");
|
||||
return;
|
||||
|
|
@ -102,85 +120,136 @@ async fn main() {
|
|||
let mut stdout = std::io::stdout();
|
||||
stdout.write_all(&buf).unwrap();
|
||||
stdout.flush().unwrap();
|
||||
println!("");
|
||||
println!();
|
||||
|
||||
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!
|
||||
let mut req_challenge = None;
|
||||
let mut req_proof = None;
|
||||
let mut req_user_agent: &[u8] = &[];
|
||||
let mut req_ip: &[u8] = &[];
|
||||
for line in header_line_iter {
|
||||
if let Some(Some(m)) = challenge_regex.captures(line).map(|c| c.get(1)) {
|
||||
req_challenge = Some(m.as_bytes());
|
||||
}
|
||||
if let Some(Some(m)) = proof_regex.captures(line).map(|c| c.get(1)) {
|
||||
req_proof = Some(m.as_bytes());
|
||||
}
|
||||
if let Some(Some(m)) = user_agent_regex.captures(line).map(|c| c.get(1)) {
|
||||
req_user_agent = m.as_bytes();
|
||||
}
|
||||
if let Some(Some(m)) = ip_regex.captures(line).map(|c| c.get(1)) {
|
||||
req_ip = m.as_bytes();
|
||||
|
||||
let mut action = default_action;
|
||||
for policy_group in policy_groups.iter() {
|
||||
if let Some(policy) = policy_group.evaluate(first_line) {
|
||||
println!("Applying policy {}", policy.name);
|
||||
action = policy.action;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
let mut allow = false;
|
||||
|
||||
if let Some(req_challenge) = req_challenge {
|
||||
allow = challenge::verify_challenge_cookie(req_challenge, &secret, req_user_agent, req_ip);
|
||||
}
|
||||
|
||||
if allow {
|
||||
// TODO reuse connections
|
||||
let pass_socket = realm_syscall::new_tcp_socket(&pass_addr).unwrap();
|
||||
|
||||
pass_socket.set_reuse_address(true).ok();
|
||||
|
||||
let pass_socket = TcpSocket::from_std_stream(pass_socket.into());
|
||||
|
||||
let mut pass_stream = pass_socket.connect(pass_addr).await.unwrap();
|
||||
|
||||
match realm_io::bidi_zero_copy(&mut client_stream, &mut pass_stream).await {
|
||||
Ok(_) => {},
|
||||
Err(ref e) if e.kind() == tokio::io::ErrorKind::InvalidInput => {
|
||||
realm_io::bidi_copy(&mut client_stream, &mut pass_stream).await.unwrap();
|
||||
match action {
|
||||
policy::Action::Drop => {}
|
||||
policy::Action::Allow => {
|
||||
do_proxy(pass_addr, client_stream).await;
|
||||
}
|
||||
policy::Action::Challenge => {
|
||||
let mut req_challenge = None;
|
||||
let mut req_proof = None;
|
||||
let mut req_user_agent: &[u8] = &[];
|
||||
let mut req_ip: &[u8] = &[];
|
||||
for line in header_line_iter {
|
||||
if let Some(Some(m)) = challenge_regex.captures(line).map(|c| c.get(1)) {
|
||||
req_challenge = Some(m.as_bytes());
|
||||
}
|
||||
if let Some(Some(m)) = proof_regex.captures(line).map(|c| c.get(1)) {
|
||||
req_proof = Some(m.as_bytes());
|
||||
}
|
||||
if let Some(Some(m)) = user_agent_regex.captures(line).map(|c| c.get(1)) {
|
||||
req_user_agent = m.as_bytes();
|
||||
}
|
||||
if let Some(Some(m)) = ip_regex.captures(line).map(|c| c.get(1)) {
|
||||
req_ip = m.as_bytes();
|
||||
}
|
||||
}
|
||||
|
||||
let mut allow = false;
|
||||
let mut valid_challenge = false;
|
||||
|
||||
if let (Some(req_challenge), Some(req_proof)) = (req_challenge, req_proof) {
|
||||
valid_challenge = challenge::verify_challenge_cookie(
|
||||
req_challenge,
|
||||
&secret,
|
||||
req_user_agent,
|
||||
req_ip,
|
||||
);
|
||||
allow = dbg!(valid_challenge)
|
||||
&& dbg!(challenge::check_challenge(
|
||||
req_challenge,
|
||||
req_proof,
|
||||
TARGET_ZEROS
|
||||
));
|
||||
}
|
||||
|
||||
if allow {
|
||||
do_proxy(pass_addr, client_stream).await;
|
||||
} else {
|
||||
let salt: [u8; SALT_LEN] = rand::thread_rng().r#gen();
|
||||
|
||||
let timestamp = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap()
|
||||
.as_secs();
|
||||
let timestamp_bytes = timestamp.to_be_bytes();
|
||||
|
||||
let challenge_mac = challenge::compute_challenge_mac(
|
||||
&secret,
|
||||
&salt,
|
||||
timestamp_bytes,
|
||||
req_ip,
|
||||
req_user_agent,
|
||||
);
|
||||
|
||||
let challenge_cookie = challenge::format_challenge_cookie(
|
||||
&salt,
|
||||
timestamp_bytes,
|
||||
&challenge_mac,
|
||||
);
|
||||
|
||||
client_stream.writable().await.unwrap();
|
||||
client_stream
|
||||
.write_all(response_begin.as_bytes())
|
||||
.await
|
||||
.unwrap();
|
||||
if !valid_challenge {
|
||||
client_stream
|
||||
.write_all(b"set-cookie: mesozoa-challenge=")
|
||||
.await
|
||||
.unwrap();
|
||||
client_stream
|
||||
.write_all(challenge_cookie.as_bytes())
|
||||
.await
|
||||
.unwrap();
|
||||
client_stream.write_all(b"; domain=127.0.0.1; path=/; max-age=3600; samesite=strict\r\n").await.unwrap();
|
||||
}
|
||||
client_stream.write_all(b"\r\n").await.unwrap();
|
||||
client_stream
|
||||
.write_all(CHALLENGE_BODY.as_bytes())
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
Err(e) => panic!("err {}", e),
|
||||
}
|
||||
} else {
|
||||
let salt: [u8; SALT_LEN] = rand::thread_rng().r#gen();
|
||||
|
||||
let timestamp = std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap().as_secs();
|
||||
let timestamp_bytes = timestamp.to_be_bytes();
|
||||
|
||||
let challenge_mac = challenge::compute_challenge_mac(&secret, &salt, timestamp_bytes, req_ip, req_user_agent);
|
||||
|
||||
let challenge_cookie = challenge::format_challenge_cookie(&salt, timestamp_bytes, &challenge_mac);
|
||||
|
||||
let response = format!(
|
||||
"HTTP/1.1 200\r\n\
|
||||
content-type: text/html\r\n\
|
||||
content-length: {}\r\n\
|
||||
set-cookie: mesozoa-challenge={}; max-age=3600; samesite=strict\r\n\
|
||||
\r\n\
|
||||
{}",
|
||||
CHALLENGE_BODY.len(),
|
||||
challenge_cookie,
|
||||
CHALLENGE_BODY,
|
||||
);
|
||||
|
||||
client_stream.writable().await.unwrap();
|
||||
client_stream
|
||||
.write_all(response.as_bytes())
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
async fn do_proxy(pass_addr: SocketAddr, mut client_stream: TcpStream) {
|
||||
// TODO reuse connections
|
||||
let pass_socket = realm_syscall::new_tcp_socket(&pass_addr).unwrap();
|
||||
|
||||
pass_socket.set_reuse_address(true).ok();
|
||||
|
||||
let pass_socket = TcpSocket::from_std_stream(pass_socket.into());
|
||||
|
||||
let mut pass_stream = pass_socket.connect(pass_addr).await.unwrap();
|
||||
|
||||
match realm_io::bidi_zero_copy(&mut client_stream, &mut pass_stream).await {
|
||||
Ok(_) => {}
|
||||
Err(ref e) if e.kind() == tokio::io::ErrorKind::InvalidInput => {
|
||||
realm_io::bidi_copy(&mut client_stream, &mut pass_stream)
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
Err(e) => panic!("err {}", e),
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue