use std::collections::HashMap;
use std::io::{Read, Write, BufWriter};
use std::net::TcpStream;
use std::pin::{Pin, pin};
use std::future::Future;
use std::task::{Context, Poll, Waker};
use std::hint;
use std::thread;
pub trait BodyStream: Send {
fn next_chunk(&mut self) -> Option<Vec<u8>>;
}
impl BodyStream for Vec<u8> {
fn next_chunk(&mut self) -> Option<Vec<u8>> {
if self.is_empty() {
None
} else {
Some(std::mem::take(self))
}
}
}
pub trait IntoBytes {
fn into_bytes(self) -> Vec<u8>;
}
impl IntoBytes for String {
fn into_bytes(self) -> Vec<u8> { self.into_bytes() }
}
impl IntoBytes for &str {
fn into_bytes(self) -> Vec<u8> { self.as_bytes().to_vec() }
}
impl IntoBytes for Vec<u8> {
fn into_bytes(self) -> Vec<u8> { self }
}
pub struct Req {
pub method: String,
pub path: String,
pub body: String,
pub headers: HashMap<String, String>,
pub stream: TcpStream,
}
impl Req {
pub fn form(&self) -> HashMap<String, String> {
let mut map = HashMap::new();
for pair in self.body.split('&') {
let mut kv = pair.split('=');
if let (Some(k), Some(v)) = (kv.next(), kv.next()) {
let decoded_v = v.replace('+', " ");
map.insert(k.to_string(), decoded_v);
}
}
map
}
pub fn save_to_file(&mut self, path: &str) -> std::io::Result<u64> {
let file = std::fs::File::create(path)?;
let mut writer = BufWriter::new(file);
let mut buffer = [0; 65536]; let mut total = 0;
while let Ok(n) = self.stream.read(&mut buffer) {
if n == 0 { break; }
writer.write_all(&buffer[..n])?;
total += n as u64;
}
writer.flush()?;
Ok(total)
}
pub fn upgrade_websocket(&self) -> Option<TcpStream> {
let key = self.headers.get("sec-websocket-key")?;
let accept_key = crate::utils::derive_websocket_accept(key);
let response = format!(
"HTTP/1.1 101 Switching Protocols\r\n\
Upgrade: websocket\r\n\
Connection: Upgrade\r\n\
Sec-WebSocket-Accept: {}\r\n\r\n",
accept_key
);
let mut s = self.stream.try_clone().ok()?;
let _ = s.write_all(response.as_bytes());
Some(s)
}
}
pub struct Params(pub HashMap<String, String>);
#[derive(Copy, Clone)]
#[repr(u16)]
pub enum StatusCode {
Continue = 100,
SwitchingProtocols = 101,
Processing = 102,
EarlyHints = 103,
Ok = 200,
Created = 201,
Accepted = 202,
NonAuthoritativeInformation = 203,
NoContent = 204,
ResetContent = 205,
PartialContent = 206,
MultiStatus = 207,
AlreadyReported = 208,
IMUsed = 226,
MultipleChoices = 300,
MovedPermanently = 301,
Found = 302,
SeeOther = 303,
NotModified = 304,
UseProxy = 305,
TemporaryRedirect = 307,
PermanentRedirect = 308,
BadRequest = 400,
Unauthorized = 401,
PaymentRequired = 402,
Forbidden = 403,
NotFound = 404,
MethodNotAllowed = 405,
NotAcceptable = 406,
ProxyAuthenticationRequired = 407,
RequestTimeout = 408,
Conflict = 409,
Gone = 410,
LengthRequired = 411,
PreconditionFailed = 412,
PayloadTooLarge = 413,
URITooLong = 414,
UnsupportedMediaType = 415,
RangeNotSatisfiable = 416,
ExpectationFailed = 417,
ImATeapot = 418,
MisdirectedRequest = 421,
UnprocessableEntity = 422,
Locked = 423,
FailedDependency = 424,
UpgradeRequired = 426,
PreconditionRequired = 428,
TooManyRequests = 429,
RequestHeaderFieldsTooLarge = 431,
UnavailableForLegalReasons = 451,
InternalServerError = 500,
NotImplemented = 501,
BadGateway = 502,
ServiceUnavailable = 503,
GatewayTimeout = 504,
HTTPVersionNotSupported = 505,
VariantAlsoNegotiates = 506,
InsufficientStorage = 507,
LoopDetected = 508,
NotExtended = 510,
NetworkAuthenticationRequired = 511,
}
pub struct Reply {
pub status: u16,
pub headers: HashMap<String, String>,
pub body: Box<dyn BodyStream>,
}
impl Reply {
pub fn new(status: StatusCode) -> Self {
Self { status: status as u16, headers: HashMap::new(), body: Box::new(Vec::new()) }
}
pub fn header(mut self, key: &str, value: &str) -> Self {
self.headers.insert(key.to_string(), value.to_string());
self
}
pub fn body<T: IntoBytes>(mut self, data: T) -> Self {
self.body = Box::new(data.into_bytes());
self
}
}
pub type BoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
pub type Handler = Box<dyn Fn(Req, Params) -> BoxFuture<'static, Reply> + Send + Sync>;
pub type Middleware = Box<dyn Fn(&str) -> Option<Reply> + Send + Sync>;
pub fn block_on<F: Future>(future: F) -> F::Output {
let mut future = pin!(future);
let waker = Waker::noop();
let mut cx = Context::from_waker(waker);
let mut spins = 0u64;
loop {
match future.as_mut().poll(&mut cx) {
Poll::Ready(v) => return v,
Poll::Pending => {
if spins < 150_000 {
hint::spin_loop();
spins += 1;
} else {
thread::yield_now();
spins = 0;
}
}
}
}
}