add problem 11
This commit is contained in:
timeshifter
Generated
+242
@@ -0,0 +1,242 @@
|
||||
# This file is automatically @generated by Cargo.
|
||||
# It is not intended for manual editing.
|
||||
version = 4
|
||||
|
||||
[[package]]
|
||||
name = "anyhow"
|
||||
version = "1.0.103"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2a4385e2e34eb35d6b3efe798b9eb88096925d87726c0798709bf56d9ed84af3"
|
||||
|
||||
[[package]]
|
||||
name = "bitflags"
|
||||
version = "2.13.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b588b76d00fde79687d7646a9b5bdf3cc0f655e0bbd080335a95d7e96f3587da"
|
||||
|
||||
[[package]]
|
||||
name = "bytes"
|
||||
version = "1.12.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "fc652a48c352aef3ea3aed32080501cf3ef6ed5da78602a020c991775b0aff04"
|
||||
|
||||
[[package]]
|
||||
name = "cfg-if"
|
||||
version = "1.0.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801"
|
||||
|
||||
[[package]]
|
||||
name = "either"
|
||||
version = "1.16.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "91622ff5e7162018101f2fea40d6ebf4a78bbe5a49736a2020649edf9693679e"
|
||||
|
||||
[[package]]
|
||||
name = "errno"
|
||||
version = "0.3.14"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb"
|
||||
dependencies = [
|
||||
"libc",
|
||||
"windows-sys",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "itertools"
|
||||
version = "0.15.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8b4baf93f58d4425749ca49a51c50ebab072c5df6994d08fed93541c331481dc"
|
||||
dependencies = [
|
||||
"either",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "libc"
|
||||
version = "0.2.186"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66"
|
||||
|
||||
[[package]]
|
||||
name = "lock_api"
|
||||
version = "0.4.14"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965"
|
||||
dependencies = [
|
||||
"scopeguard",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "mio"
|
||||
version = "1.2.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "30d65c71f1ce40ab09135ce117d742b9f8a19ff91a41a8b57ed50bc2de59c427"
|
||||
dependencies = [
|
||||
"libc",
|
||||
"wasi",
|
||||
"windows-sys",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "parking_lot"
|
||||
version = "0.12.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "93857453250e3077bd71ff98b6a65ea6621a19bb0f559a85248955ac12c45a1a"
|
||||
dependencies = [
|
||||
"lock_api",
|
||||
"parking_lot_core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "parking_lot_core"
|
||||
version = "0.9.12"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1"
|
||||
dependencies = [
|
||||
"cfg-if",
|
||||
"libc",
|
||||
"redox_syscall",
|
||||
"smallvec",
|
||||
"windows-link",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pin-project-lite"
|
||||
version = "0.2.17"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd"
|
||||
|
||||
[[package]]
|
||||
name = "problem11"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"bytes",
|
||||
"itertools",
|
||||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "proc-macro2"
|
||||
version = "1.0.106"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934"
|
||||
dependencies = [
|
||||
"unicode-ident",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "quote"
|
||||
version = "1.0.46"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "dfbc457d0c7a0759a614551b11a6409e5951f6c7537be1f1b7682b9ae9230368"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "redox_syscall"
|
||||
version = "0.5.18"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d"
|
||||
dependencies = [
|
||||
"bitflags",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "scopeguard"
|
||||
version = "1.2.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49"
|
||||
|
||||
[[package]]
|
||||
name = "signal-hook-registry"
|
||||
version = "1.4.8"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c4db69cba1110affc0e9f7bcd48bbf87b3f4fc7c61fc9155afd4c469eb3d6c1b"
|
||||
dependencies = [
|
||||
"errno",
|
||||
"libc",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "smallvec"
|
||||
version = "1.15.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8ed6a63f02c8539c91a8685a86f4099661ba3da017932f6ebbea6de3f0fa7c90"
|
||||
|
||||
[[package]]
|
||||
name = "socket2"
|
||||
version = "0.6.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c3d1e2c7f27f8d4cb10542a02c49005dbd6e93095799d6f3be745fae9f8fedd4"
|
||||
dependencies = [
|
||||
"libc",
|
||||
"windows-sys",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "syn"
|
||||
version = "2.0.119"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"unicode-ident",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tokio"
|
||||
version = "1.53.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d988bcd52dbe076d3d46903332f58c912b87a2c49b1428419a5845154762ffee"
|
||||
dependencies = [
|
||||
"bytes",
|
||||
"libc",
|
||||
"mio",
|
||||
"parking_lot",
|
||||
"pin-project-lite",
|
||||
"signal-hook-registry",
|
||||
"socket2",
|
||||
"tokio-macros",
|
||||
"windows-sys",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tokio-macros"
|
||||
version = "2.7.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "6328af13490e73a9b4694030fafd93f8c8c6a9dede33e821c3fc63eddf8042ba"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "unicode-ident"
|
||||
version = "1.0.24"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75"
|
||||
|
||||
[[package]]
|
||||
name = "wasi"
|
||||
version = "0.11.1+wasi-snapshot-preview1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b"
|
||||
|
||||
[[package]]
|
||||
name = "windows-link"
|
||||
version = "0.2.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5"
|
||||
|
||||
[[package]]
|
||||
name = "windows-sys"
|
||||
version = "0.61.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc"
|
||||
dependencies = [
|
||||
"windows-link",
|
||||
]
|
||||
@@ -0,0 +1,10 @@
|
||||
[package]
|
||||
name = "problem11"
|
||||
version = "0.1.0"
|
||||
edition = "2024"
|
||||
|
||||
[dependencies]
|
||||
anyhow = "1"
|
||||
tokio = { version = "1", features = ["full"] }
|
||||
bytes = "*"
|
||||
itertools = "*"
|
||||
@@ -0,0 +1,565 @@
|
||||
use std::{collections::HashMap, fmt::Display, sync::Arc};
|
||||
|
||||
use anyhow::{Result, anyhow, bail};
|
||||
use bytes::BytesMut;
|
||||
use itertools::Itertools;
|
||||
use tokio::{
|
||||
io::{AsyncReadExt, AsyncWriteExt},
|
||||
net::{TcpListener, TcpStream},
|
||||
sync::{Mutex, MutexGuard},
|
||||
};
|
||||
|
||||
const MAX_MSG_LEN: usize = 1 << 20;
|
||||
const UPSTREAM: &str = "pestcontrol.protohackers.com:20547";
|
||||
|
||||
type Policies = HashMap<(Action, u32, String), u32>; // action, site, species, policyid
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> Result<()> {
|
||||
let listener = TcpListener::bind("0.0.0.0:10000").await?;
|
||||
let policies = Arc::new(Mutex::new(Policies::new()));
|
||||
|
||||
loop {
|
||||
let (mut socket, _) = listener.accept().await?;
|
||||
let this_policies = policies.clone();
|
||||
let mut hello_has_been_sent = false;
|
||||
|
||||
tokio::spawn(async move {
|
||||
match handle_connection(&mut socket, this_policies, &mut hello_has_been_sent).await {
|
||||
Ok(()) => (),
|
||||
Err(e) => {
|
||||
if !hello_has_been_sent {
|
||||
socket
|
||||
.write_all(&encode_message(Message::Hello).unwrap())
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
socket
|
||||
.write_all(
|
||||
&encode_message(Message::Error {
|
||||
message: e.root_cause().to_string(),
|
||||
})
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
async fn handle_connection(
|
||||
socket: &mut TcpStream,
|
||||
policies: Arc<Mutex<Policies>>,
|
||||
hello_has_been_sent: &mut bool,
|
||||
) -> Result<()> {
|
||||
let (mut reader, mut writer) = socket.split();
|
||||
let mut buf = BytesMut::with_capacity(4096);
|
||||
|
||||
loop {
|
||||
while let Some(frame) = try_extract_frame(&mut buf)? {
|
||||
let msg = decode_message(&frame)?;
|
||||
match msg {
|
||||
Message::Hello => {
|
||||
let reply = encode_message(Message::Hello)?;
|
||||
writer.write_all(&reply).await?;
|
||||
*hello_has_been_sent = true;
|
||||
}
|
||||
Message::Error { message } => {
|
||||
bail!("client sent error: {message}");
|
||||
}
|
||||
Message::SiteVisit { site, populations } => {
|
||||
match handle_site_visit(site, populations, &policies).await {
|
||||
Ok(()) => (),
|
||||
Err(e) => {
|
||||
if e.root_cause().to_string()
|
||||
== DataError::ConflictingPopulationReport.to_string()
|
||||
{
|
||||
writer
|
||||
.write_all(&encode_message(Message::Error {
|
||||
message: e.root_cause().to_string(),
|
||||
})?)
|
||||
.await?;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
msg => bail!("unexpected client message: {msg:?}"),
|
||||
};
|
||||
}
|
||||
|
||||
if reader.read_buf(&mut buf).await? == 0 {
|
||||
if !buf.is_empty() {
|
||||
return Err(anyhow!("connection closed mid-message"));
|
||||
}
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn handle_site_visit(
|
||||
reported_site: u32,
|
||||
reported_populations: Vec<(String, u32)>,
|
||||
policies: &Arc<Mutex<Policies>>,
|
||||
) -> Result<()> {
|
||||
let mut upstream = TcpStream::connect(UPSTREAM).await?;
|
||||
|
||||
if conflicted(&reported_populations) {
|
||||
return Err(anyhow!(DataError::ConflictingPopulationReport));
|
||||
}
|
||||
|
||||
let buffer = encode_message(Message::Hello)?;
|
||||
upstream.write_all(&buffer).await?;
|
||||
|
||||
let mut read_buf = BytesMut::with_capacity(4096);
|
||||
loop {
|
||||
upstream.read_buf(&mut read_buf).await?;
|
||||
if let Some(frame) = try_extract_frame(&mut read_buf)? {
|
||||
if decode_message(&frame)? == Message::Hello {
|
||||
break;
|
||||
} else {
|
||||
bail!("expected hello, received something else");
|
||||
}
|
||||
}
|
||||
}
|
||||
let buf = encode_message(Message::DialAuthority {
|
||||
site: reported_site,
|
||||
})?;
|
||||
upstream.write_all(&buf).await?;
|
||||
|
||||
upstream.read_buf(&mut read_buf).await?;
|
||||
|
||||
{
|
||||
if let Some(frame) = try_extract_frame(&mut read_buf)? {
|
||||
let msg = decode_message(&frame)?;
|
||||
let goal_policies = match msg {
|
||||
Message::TargetPopulations { site, populations } => {
|
||||
work_out_policies(site, populations, reported_site, reported_populations)
|
||||
}
|
||||
m => {
|
||||
return Err(anyhow!(
|
||||
"unexpected upstream message after site visit: {m:?}"
|
||||
));
|
||||
}
|
||||
};
|
||||
let mut established_policies = policies.lock().await;
|
||||
for (species, action) in goal_policies {
|
||||
if established_policies.contains_key(&(
|
||||
action.clone(),
|
||||
reported_site,
|
||||
species.clone(),
|
||||
)) {
|
||||
continue;
|
||||
} else {
|
||||
let policy_id = create_policy(&species, &action, &mut upstream).await?;
|
||||
let entry = (action, reported_site, species);
|
||||
if let Some(policy_id) = policy_id {
|
||||
established_policies.insert(entry.clone(), policy_id);
|
||||
}
|
||||
remove_conflicting_policies(&mut established_policies, entry, &mut upstream)
|
||||
.await?;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn conflicted(reported_populations: &[(String, u32)]) -> bool {
|
||||
reported_populations
|
||||
.iter()
|
||||
.combinations(2)
|
||||
.any(|vec| vec[0].0 == vec[1].0 && vec[0].1 != vec[1].1)
|
||||
}
|
||||
|
||||
async fn remove_conflicting_policies(
|
||||
established_policies: &mut MutexGuard<'_, HashMap<(Action, u32, String), u32>>,
|
||||
entry: (Action, u32, String),
|
||||
upstream: &mut TcpStream,
|
||||
) -> Result<()> {
|
||||
let actions_to_remove = match entry.0 {
|
||||
Action::Conserve => vec![Action::Cull],
|
||||
Action::Cull => vec![Action::Conserve],
|
||||
Action::Maintain => vec![Action::Cull, Action::Conserve],
|
||||
};
|
||||
|
||||
for action_to_remove in actions_to_remove {
|
||||
let mut this_entry = entry.clone();
|
||||
this_entry.0 = action_to_remove;
|
||||
if let Some(policy_id) = established_policies.remove(&this_entry) {
|
||||
let msg = Message::DeletePolicy { policy: policy_id };
|
||||
let buf = encode_message(msg)?;
|
||||
upstream.write_all(&buf).await?;
|
||||
|
||||
let mut buf = BytesMut::with_capacity(4096);
|
||||
upstream.read_buf(&mut buf).await?;
|
||||
if let Some(frame) = try_extract_frame(&mut buf)? {
|
||||
let msg = decode_message(&frame)?;
|
||||
match msg {
|
||||
Message::Ok => (),
|
||||
_ => bail!("expected ok after DeletePolicy, received something else"),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn create_policy(
|
||||
species: &str,
|
||||
action: &Action,
|
||||
upstream: &mut TcpStream,
|
||||
) -> Result<Option<u32>> {
|
||||
if *action == Action::Maintain {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let msg = Message::CreatePolicy {
|
||||
species: species.to_string(),
|
||||
action: action.to_owned(),
|
||||
};
|
||||
|
||||
let bytes = encode_message(msg)?;
|
||||
upstream.write_all(&bytes).await?;
|
||||
|
||||
let mut buf = BytesMut::with_capacity(4096);
|
||||
loop {
|
||||
upstream.read_buf(&mut buf).await?;
|
||||
if let Some(frame) = try_extract_frame(&mut buf)? {
|
||||
let msg = decode_message(&frame)?;
|
||||
match msg {
|
||||
Message::PolicyResult { policy } => return Ok(Some(policy)),
|
||||
_ => bail!("expected 0x57 PolicyResult, got something else"),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn work_out_policies(
|
||||
site: u32,
|
||||
target_populations: Vec<(String, u32, u32)>,
|
||||
reported_site: u32,
|
||||
reported_populations: Vec<(String, u32)>,
|
||||
) -> Vec<(String, Action)> {
|
||||
assert!(site == reported_site);
|
||||
let mut goal_policies = vec![];
|
||||
|
||||
for (species, count) in &reported_populations {
|
||||
match target_populations.iter().find(|(s, _, _)| *species == *s) {
|
||||
Some((_, min, max)) => {
|
||||
let action = if (*min..=*max).contains(count) {
|
||||
Action::Maintain
|
||||
} else if count < min {
|
||||
Action::Conserve
|
||||
} else {
|
||||
Action::Cull
|
||||
};
|
||||
goal_policies.push((species.clone(), action))
|
||||
}
|
||||
None => {
|
||||
continue;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for (species, min, _) in target_populations {
|
||||
if min > 0 && reported_populations.iter().all(|(s, _)| *s != species) {
|
||||
goal_policies.push((species, Action::Conserve));
|
||||
}
|
||||
}
|
||||
|
||||
goal_policies
|
||||
}
|
||||
|
||||
fn encode_message(msg: Message) -> Result<Vec<u8>, ProtoError> {
|
||||
let mut w = Writer::new();
|
||||
|
||||
let msg_bytes = match msg {
|
||||
Message::Hello => {
|
||||
w.str("pestcontrol");
|
||||
w.u32(1);
|
||||
write_message(0x50, &w.buf)
|
||||
}
|
||||
Message::Error { message } => {
|
||||
w.str(&message);
|
||||
write_message(0x51, &w.buf)
|
||||
}
|
||||
Message::Ok => write_message(0x52, &w.buf),
|
||||
Message::DialAuthority { site } => {
|
||||
w.u32(site);
|
||||
write_message(0x53, &w.buf)
|
||||
}
|
||||
Message::TargetPopulations { site, populations } => {
|
||||
w.u32(site);
|
||||
w.array(&populations, |w, v| {
|
||||
w.str(&v.0);
|
||||
w.u32(v.1);
|
||||
w.u32(v.2)
|
||||
});
|
||||
write_message(0x54, &w.buf)
|
||||
}
|
||||
Message::CreatePolicy { species, action } => {
|
||||
w.str(&species);
|
||||
match action {
|
||||
Action::Conserve => w.u8(0xa0),
|
||||
Action::Cull => w.u8(0x90),
|
||||
Action::Maintain => panic!("maintain is not for sending"),
|
||||
};
|
||||
write_message(0x55, &w.buf)
|
||||
}
|
||||
Message::DeletePolicy { policy } => {
|
||||
w.u32(policy);
|
||||
write_message(0x56, &w.buf)
|
||||
}
|
||||
Message::PolicyResult { policy } => {
|
||||
w.u32(policy);
|
||||
write_message(0x57, &w.buf)
|
||||
}
|
||||
Message::SiteVisit { site, populations } => {
|
||||
w.u32(site);
|
||||
w.array(&populations, |w, v| {
|
||||
w.str(&v.0);
|
||||
w.u32(v.1)
|
||||
});
|
||||
write_message(0x58, &w.buf)
|
||||
}
|
||||
};
|
||||
|
||||
Ok(msg_bytes)
|
||||
}
|
||||
|
||||
fn decode_message(frame: &[u8]) -> Result<Message, ProtoError> {
|
||||
let sum: u32 = frame.iter().map(|&b| b as u32).sum();
|
||||
if !sum.is_multiple_of(256) {
|
||||
return Err(ProtoError::WrongChecksum);
|
||||
}
|
||||
|
||||
let type_byte = frame[0];
|
||||
let mut r = Reader::new(&frame[5..frame.len() - 1]);
|
||||
let msg = match type_byte {
|
||||
0x50 => {
|
||||
if r.str()? != "pestcontrol" || r.u32()? != 1 {
|
||||
return Err(ProtoError::Malformed);
|
||||
}
|
||||
Message::Hello
|
||||
}
|
||||
0x51 => Message::Error { message: r.str()? },
|
||||
0x52 => Message::Ok,
|
||||
0x53 => Message::DialAuthority { site: r.u32()? },
|
||||
0x54 => {
|
||||
let site = r.u32()?;
|
||||
let populations = r.array(|r| Ok((r.str()?, r.u32()?, r.u32()?)))?;
|
||||
Message::TargetPopulations { site, populations }
|
||||
}
|
||||
0x55 => {
|
||||
let species = r.str()?;
|
||||
let action = match r.u8()? {
|
||||
0x90 => Action::Cull,
|
||||
0xa0 => Action::Conserve,
|
||||
_ => {
|
||||
return Err(ProtoError::InvalidPolicy);
|
||||
}
|
||||
};
|
||||
Message::CreatePolicy { species, action }
|
||||
}
|
||||
0x56 => Message::DeletePolicy { policy: r.u32()? },
|
||||
0x57 => Message::PolicyResult { policy: r.u32()? },
|
||||
0x58 => {
|
||||
let site = r.u32()?;
|
||||
let populations = r.array(|r| Ok((r.str()?, r.u32()?)))?;
|
||||
Message::SiteVisit { site, populations }
|
||||
}
|
||||
_ => return Err(ProtoError::WrongTypeByte),
|
||||
};
|
||||
|
||||
if !r.finished() {
|
||||
return Err(ProtoError::Malformed);
|
||||
}
|
||||
|
||||
Ok(msg)
|
||||
}
|
||||
|
||||
fn try_extract_frame(buf: &mut BytesMut) -> Result<Option<BytesMut>, ProtoError> {
|
||||
if buf.len() < 5 {
|
||||
return Ok(None); // less than type byte + length
|
||||
}
|
||||
let len = u32::from_be_bytes(buf[1..5].try_into().unwrap()) as usize;
|
||||
|
||||
if !(6..MAX_MSG_LEN).contains(&len) {
|
||||
return Err(ProtoError::Malformed);
|
||||
}
|
||||
|
||||
if buf.len() < len {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
Ok(Some(buf.split_to(len)))
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct Writer {
|
||||
buf: Vec<u8>,
|
||||
}
|
||||
|
||||
impl Writer {
|
||||
fn u8(&mut self, val: u8) {
|
||||
self.buf.push(val);
|
||||
}
|
||||
|
||||
fn u32(&mut self, val: u32) {
|
||||
self.buf.extend_from_slice(&val.to_be_bytes());
|
||||
}
|
||||
|
||||
fn str(&mut self, s: &str) {
|
||||
self.u32(s.len() as u32);
|
||||
self.buf.extend_from_slice(s.as_bytes());
|
||||
}
|
||||
|
||||
fn array<T>(&mut self, items: &[T], mut elem: impl FnMut(&mut Self, &T)) {
|
||||
self.u32(items.len() as u32);
|
||||
for item in items {
|
||||
elem(self, item);
|
||||
}
|
||||
}
|
||||
|
||||
fn new() -> Self {
|
||||
Self {
|
||||
buf: Vec::with_capacity(4096),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct Reader<'a> {
|
||||
buf: &'a [u8],
|
||||
pos: usize,
|
||||
}
|
||||
|
||||
impl<'a> Reader<'a> {
|
||||
fn u8(&mut self) -> Result<u8, ProtoError> {
|
||||
Ok(self.take(1)?[0])
|
||||
}
|
||||
|
||||
fn u32(&mut self) -> Result<u32, ProtoError> {
|
||||
let b = self.take(4)?;
|
||||
Ok(u32::from_be_bytes(b.try_into().unwrap()))
|
||||
}
|
||||
|
||||
fn str(&mut self) -> Result<String, ProtoError> {
|
||||
let len = self.u32()? as usize;
|
||||
let bytes = self.take(len)?;
|
||||
// spec says ASCII, so this can be a straight cast rather than utf8 validation
|
||||
Ok(bytes.iter().map(|&b| b as char).collect())
|
||||
}
|
||||
|
||||
fn array<T>(
|
||||
&mut self,
|
||||
mut elem: impl FnMut(&mut Self) -> Result<T, ProtoError>,
|
||||
) -> Result<Vec<T>, ProtoError> {
|
||||
let len = self.u32()? as usize;
|
||||
(0..len).map(|_| elem(self)).collect()
|
||||
}
|
||||
|
||||
fn take(&mut self, n: usize) -> Result<&'a [u8], ProtoError> {
|
||||
if self.pos + n > self.buf.len() {
|
||||
return Err(ProtoError::Malformed);
|
||||
}
|
||||
let s = &self.buf[self.pos..self.pos + n];
|
||||
self.pos += n;
|
||||
Ok(s)
|
||||
}
|
||||
|
||||
fn finished(&self) -> bool {
|
||||
self.pos == self.buf.len()
|
||||
}
|
||||
|
||||
fn new(buf: &'a [u8]) -> Self {
|
||||
Self { buf, pos: 0 }
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
enum ProtoError {
|
||||
Malformed,
|
||||
WrongTypeByte,
|
||||
WrongChecksum,
|
||||
InvalidPolicy,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
enum DataError {
|
||||
ConflictingPopulationReport,
|
||||
}
|
||||
|
||||
impl Display for DataError {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
DataError::ConflictingPopulationReport => f.write_str("conflicting population report"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::error::Error for DataError {}
|
||||
|
||||
impl Display for ProtoError {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
ProtoError::Malformed => f.write_str("malformed packet"),
|
||||
ProtoError::WrongTypeByte => f.write_str("wrong type byte"),
|
||||
ProtoError::WrongChecksum => f.write_str("checksum error"),
|
||||
ProtoError::InvalidPolicy => f.write_str("invalid policy"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::error::Error for ProtoError {}
|
||||
|
||||
fn write_message(type_byte: u8, payload: &[u8]) -> Vec<u8> {
|
||||
let mut out = Vec::with_capacity(payload.len() + 6);
|
||||
out.push(type_byte);
|
||||
out.extend_from_slice(&((payload.len() + 6) as u32).to_be_bytes());
|
||||
out.extend_from_slice(payload);
|
||||
let sum: u32 = out.iter().map(|&b| b as u32).sum();
|
||||
out.push((256 - sum % 256) as u8);
|
||||
out
|
||||
}
|
||||
|
||||
#[derive(Debug, PartialEq)]
|
||||
enum Message {
|
||||
Hello,
|
||||
Error {
|
||||
message: String,
|
||||
},
|
||||
Ok,
|
||||
DialAuthority {
|
||||
site: u32,
|
||||
},
|
||||
TargetPopulations {
|
||||
site: u32,
|
||||
populations: Vec<(String, u32, u32)>,
|
||||
},
|
||||
CreatePolicy {
|
||||
species: String,
|
||||
action: Action,
|
||||
},
|
||||
DeletePolicy {
|
||||
policy: u32,
|
||||
},
|
||||
PolicyResult {
|
||||
policy: u32,
|
||||
},
|
||||
SiteVisit {
|
||||
site: u32,
|
||||
populations: Vec<(String, u32)>,
|
||||
},
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
|
||||
enum Action {
|
||||
Conserve,
|
||||
Cull,
|
||||
Maintain,
|
||||
}
|
||||
Reference in New Issue
Block a user