use async_std::net::{ToSocketAddrs, UdpSocket};
use futures::Future;
use presence::{unhex, Presence};
use sha2::Digest;
use smol_timeout::TimeoutExt;
use std::{
collections::VecDeque,
fmt::{Debug, Display},
io,
net::SocketAddr,
pin::Pin,
time::Duration,
};
#[derive(Clone, Debug)]
pub struct Response {
pub message: Message,
pub from: SocketAddr,
}
#[derive(Clone)]
pub enum Message {
Icu(Presence, SocketAddr),
Cht(Presence, VenueId, Vec<u8>),
Err(String),
Gon(VenueId),
Rlx(String),
Bye(Presence),
}
#[derive(Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct VenueId([u8; 32]);
impl From<[u8; 32]> for VenueId {
fn from(value: [u8; 32]) -> Self {
Self(value)
}
}
impl From<&[u8; 32]> for VenueId {
fn from(value: &[u8; 32]) -> Self {
Self(value.clone())
}
}
impl Default for VenueId {
fn default() -> Self {
let empty_string_sha256_hash = b"\xe3\xb0\xc4\x42\x98\xfc\x1c\x14\x9a\xfb\xf4\xc8\x99\x6f\xb9\x24\x27\xae\x41\xe4\x64\x9b\x93\x4c\xa4\x95\x99\x1b\x78\x52\xb8\x55";
Self(empty_string_sha256_hash.clone())
}
}
impl From<VenueId> for [u8; 32] {
fn from(value: VenueId) -> Self {
value.0
}
}
impl From<&VenueId> for [u8; 32] {
fn from(value: &VenueId) -> Self {
value.0.clone()
}
}
impl VenueId {
pub fn new(name: impl AsRef<[u8]>) -> Self {
Self(
sha2::Sha256::digest(name.as_ref())[0..32]
.try_into()
.expect("sha2 venueid"),
)
}
}
impl Display for VenueId {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
self.0.iter().try_fold((), |_, b| write!(f, "{b:02X}"))
}
}
impl Debug for VenueId {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
if f.alternate() {
f.debug_tuple("VenueId").field(&self.0).finish()
} else {
Display::fmt(&self, f)
}
}
}
impl AsRef<[u8]> for VenueId {
fn as_ref(&self) -> &[u8] {
&self.0
}
}
impl std::ops::Deref for VenueId {
type Target = [u8];
fn deref(&self) -> &Self::Target {
&self.0
}
}
impl Message {
pub fn to_bytes(&self) -> Vec<u8> {
match self {
Self::Rlx(arg0) => format!("RLX {arg0}\n").into(),
Self::Cht(presence, venue_id, chat) => "chat"
.bytes()
.chain(serialize_chat_message_2(presence, venue_id, chat))
.collect(),
Self::Err(arg0) => format!("ERR {arg0}\n").into(),
Self::Bye(arg0) => format!("BYE {arg0}\n").into(),
Self::Icu(arg0, arg1) => format!("ICU {arg0} on {arg1}\n").into(),
Self::Gon(arg0) => format!("GON {arg0}\n").into(),
}
}
}
impl Debug for Message {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
if f.alternate() {
match self {
Self::Cht(presence, vid, arg0) => f
.debug_tuple("Cht")
.field(presence)
.field(vid)
.field(arg0)
.finish(),
Self::Rlx(arg0) => f.debug_tuple("Rlx").field(arg0).finish(),
Self::Err(arg0) => f.debug_tuple("Err").field(arg0).finish(),
Self::Bye(arg0) => f.debug_tuple("Bye").field(arg0).finish(),
Self::Icu(arg0, arg1) => f.debug_tuple("Icu").field(arg0).field(arg1).finish(),
Self::Gon(arg0) => f.debug_tuple("Gon").field(arg0).finish(),
}
} else {
match self {
Self::Rlx(arg0) => {
write!(f, "RLX {:?}", arg0)
}
Self::Cht(presence, vid, chat) => {
write!(f, "CHT {presence} {vid} ")?;
match std::str::from_utf8(chat) {
Ok(s) => Debug::fmt(&s, f),
Err(_) => chat.fmt(f),
}
}
Self::Err(arg0) => {
write!(f, "ERR {:?}", arg0)
}
Self::Bye(arg0) => {
write!(f, "BYE {arg0}")
}
Self::Icu(arg0, arg1) => write!(f, "ICU {arg0} on {arg1}"),
Self::Gon(arg0) => write!(f, "GON {arg0}",),
}
}
}
}
impl Message {
pub fn from_slice<T: AsRef<[u8]>>(origin: T) -> std::result::Result<Self, T> {
let value = (&origin).as_ref();
let (tag, mut rest) = if value.len() >= 4 {
value.split_at(4)
} else {
(&[][..], value)
};
if rest.last() == Some(&b'\n') {
rest = &rest[0..rest.len() - 1];
}
Ok(match tag {
b"ERR " => Message::Err(String::from_utf8_lossy(rest).to_string()),
b"BYE " => std::str::from_utf8(rest)
.ok()
.and_then(|value| Presence::from_str(value, time_std(), u16::MAX, u16::MAX).ok())
.flatten()
.map(Message::Bye)
.ok_or(origin)?,
b"RLX " => String::from_utf8(rest.to_vec())
.ok()
.map(Message::Rlx)
.ok_or(origin)?,
b"ICU " => std::str::from_utf8(rest)
.ok()
.and_then(|str| str.split_once(" on "))
.and_then(|(presence, socket)| {
let presence =
Presence::from_str(presence, time_std(), u16::MAX, u16::MAX).ok()??;
let socket = socket.parse().ok()?;
Some(Message::Icu(presence, socket))
})
.ok_or(origin)?,
b"GON " => std::str::from_utf8(rest)
.ok()
.and_then(|str| unhex(str).ok())
.map(VenueId::from)
.map(Message::Gon)
.ok_or(origin)?,
b"chat" => parse_chat_message_2(rest)
.map(|(presence, venue, chat)| Message::Cht(presence, venue, chat.to_vec()))
.ok_or(origin)?,
_ => return Err(origin),
})
}
}
pub fn parse_chat_message_2(value: &[u8]) -> Option<(Presence, VenueId, &[u8])> {
if value.len() < 40 {
return None;
}
let (presence, value) = value.split_at(8);
let presence = Presence::from_bytes(
presence.try_into().expect("8 venue_id bytes"),
time_std(),
u16::MAX,
u16::MAX,
)?;
let (venue_id, value) = value.split_at(32);
let venue_id: [u8; 32] = venue_id.try_into().expect("32 venue_id bytes");
let venue_id = VenueId::from(venue_id);
Some((presence, venue_id, value))
}
pub fn serialize_chat_message_2<'a>(
presence: &Presence,
venue_id: &'a VenueId,
chat: &'a [u8],
) -> impl Iterator<Item = u8> + 'a {
presence
.to_bytes()
.into_iter()
.chain(venue_id.into_iter().cloned())
.chain(chat.into_iter().cloned())
}
fn time_std() -> i64 {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|future| future.as_secs() as i64)
.unwrap_or_else(|past| past.duration().as_secs() as i64)
}
#[derive(Debug)]
pub struct Client<S, B> {
pub server: S,
pub bind: B,
pub bound: Option<SocketAddr>,
pub servers: Vec<SocketAddr>,
pub commands: VecDeque<Vec<u8>>,
state: ClientState,
}
impl<S, B> Client<S, B>
where
S: ToSocketAddrs + Clone + 'static,
B: ToSocketAddrs + Clone + 'static,
{
pub fn new(server: S, bind: B) -> Self {
Self {
server: server.clone(),
bind: bind.clone(),
bound: None,
servers: vec![],
commands: vec![].into(),
state: ClientState::Pending(Box::pin(ClientState::new(server, bind))),
}
}
}
impl<S, B> Client<S, B> {
pub fn chat(&mut self, venue_id: &VenueId, message: impl AsRef<[u8]>) {
let chat = message.as_ref();
self.commands.push_back(
"chat"
.bytes()
.chain(serialize_chat_message_2(&Presence::now(), &venue_id, chat))
.collect(),
)
}
pub fn quit(&mut self) {
self.commands.push_back(b"quit".to_vec())
}
pub fn gone(&mut self, venue: Option<&VenueId>) {
let venue_id = match venue {
Some(venue) => venue.as_ref(),
None => b"",
};
self.commands
.push_back(b"gone".iter().chain(venue_id).cloned().collect())
}
pub async fn response(&mut self, mut timeout: Duration) -> Result<Vec<Response>> {
loop {
break match self.state {
ClientState::Pending(ref mut fut) => {
self.state = fut.as_mut().await?;
continue;
}
ClientState::Ready(ref socket, ref servers) => {
self.bound = socket.local_addr().ok();
self.servers = servers.clone();
while let Some(cmd) = self.commands.pop_front() {
let _size = socket
.send_to(&cmd, &servers[..])
.await
.map_err(|e| Error::Sending(cmd, e))?;
}
let mut buff = [0u8; 1024];
let mut responses = vec![];
while let Some((size, address)) = socket
.recv_from(&mut buff)
.timeout(timeout)
.await
.transpose()
.map_err(|e| Error::Receiving(e))?
{
timeout = Duration::from_millis(1);
responses.push(Response {
from: address,
message: Message::from_slice(&buff[0..size]).map_err(|e| {
if let Ok(str) = String::from_utf8(e.to_vec()) {
Error::Message(str)
} else {
Error::MessageBin(e.to_vec())
}
})?,
});
}
Ok(responses)
}
};
}
}
}
enum ClientState {
Pending(Pin<Box<dyn Future<Output = Result<ClientState>>>>),
Ready(UdpSocket, Vec<SocketAddr>),
}
impl Debug for ClientState {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Pending(_) => f.debug_tuple("Pending").finish(),
Self::Ready(arg0, arg1) => f.debug_tuple("Ready").field(arg0).field(arg1).finish(),
}
}
}
impl ClientState {
async fn new(server: impl ToSocketAddrs, bind: impl ToSocketAddrs) -> Result<ClientState> {
let socket = UdpSocket::bind(bind).await?;
let servers = server.to_socket_addrs().await?.collect();
Ok(ClientState::Ready(socket, servers))
}
}
pub type CommandResult = std::result::Result<(), ()>;
pub type Result<T> = std::result::Result<T, Error>;
#[derive(Debug, thiserror::Error)]
pub enum Error {
#[error("Failed to set the client up - {0}")]
Setup(#[from] std::io::Error),
#[error("Failed to parse server message. {0:?}")]
MessageBin(Vec<u8>),
#[error("Failed to parse server message: {0}")]
Message(String),
#[error("Failed to send UDP message - {1}: {0:?}")]
Sending(Vec<u8>, io::Error),
#[error("Failed to receive UDP message - {0}")]
Receiving(io::Error),
}