#![deny(clippy::all)]
#![warn(clippy::pedantic)]
use casper_utils::csl;
use kstring::KString;
use std::{
self,
collections::HashMap,
fmt::Write,
net::UdpSocket,
time::Duration,
};
use tftp_client::{
download,
upload,
};
use thiserror::Error;
use tracing::debug;
pub const FLASH_SECTOR_SIZE: u32 = 0x10000;
pub const DEFAULT_TIMEOUT: Duration = Duration::from_millis(500);
pub const MAX_TIMEOUT: Duration = Duration::from_secs(5);
#[derive(Debug, Error)]
pub enum Error {
#[error(transparent)]
Tftp(#[from] tftp_client::Error),
#[error("Some part of the received payload was incomplete")]
Incomplete,
#[error("While trying to parse a string from a response, we received invalid UTF8")]
Utf8(#[from] std::str::Utf8Error),
#[error("No metadata returned when we requested metadata")]
MissingMetadata,
#[error(transparent)]
Csl(#[from] csl::Error),
}
fn retrying_download(
filename: &str,
socket: &UdpSocket,
timeout: Duration,
max_timeout: Duration,
retries: usize,
) -> Result<Vec<u8>, Error> {
let mut local_retries = 0;
let mut this_timeout = timeout;
loop {
if local_retries == retries {
return Err(Error::Tftp(tftp_client::Error::Timeout));
}
let res = download(filename, socket, timeout, max_timeout, retries);
match res {
Ok(v) => return Ok(v),
Err(tftp_client::Error::Protocol { code, msg }) => {
debug!("Protocol error: {:?} {msg}", code);
std::thread::sleep(this_timeout);
local_retries += 1;
this_timeout += this_timeout / 2;
if this_timeout > MAX_TIMEOUT {
this_timeout = MAX_TIMEOUT;
}
continue;
}
Err(e) => {
return Err(Error::Tftp(e));
}
}
}
}
fn retrying_upload(
filename: &str,
data: &[u8],
socket: &UdpSocket,
timeout: Duration,
max_timeout: Duration,
retries: usize,
) -> Result<(), Error> {
let mut local_retries = 0;
let mut this_timeout = timeout;
loop {
if local_retries == retries {
return Err(Error::Tftp(tftp_client::Error::Timeout));
}
let res = upload(filename, data, socket, timeout, max_timeout, retries);
match res {
Ok(()) => return Ok(()),
Err(tftp_client::Error::Protocol { code, msg }) => {
debug!("Protocol error: {:?} {msg}", code);
local_retries += 1;
std::thread::sleep(this_timeout);
local_retries += 1;
this_timeout += this_timeout / 2;
if this_timeout > MAX_TIMEOUT {
this_timeout = MAX_TIMEOUT;
}
continue;
}
Err(e) => {
return Err(Error::Tftp(e));
}
}
}
}
pub fn temp(socket: &UdpSocket, retries: usize) -> Result<f32, Error> {
let bytes = retrying_download("/temp", socket, DEFAULT_TIMEOUT, MAX_TIMEOUT, retries)?;
let four_bytes = bytes.get(..4).ok_or(Error::Incomplete)?;
Ok(f32::from_be_bytes(
four_bytes.try_into().map_err(|_| Error::Incomplete)?,
))
}
pub fn help(socket: &UdpSocket, retries: usize) -> Result<String, Error> {
let bytes = retrying_download("/help", socket, DEFAULT_TIMEOUT, MAX_TIMEOUT, retries)?;
Ok(std::str::from_utf8(&bytes)?.to_string())
}
pub fn listdev(socket: &UdpSocket, retries: usize) -> Result<HashMap<String, (u32, u32)>, Error> {
let bytes = retrying_download("/listdev", socket, DEFAULT_TIMEOUT, MAX_TIMEOUT, retries)?;
let csl = csl::from_bytes(&bytes)?;
csl.into_iter()
.map(|(k, v)| {
let addr = u32::from_be_bytes(v[..4].try_into().map_err(|_| Error::Incomplete)?);
let length = u32::from_be_bytes(v[4..].try_into().map_err(|_| Error::Incomplete)?);
Ok((k, (addr, length)))
})
.collect()
}
pub fn read_device(
device: &str,
offset: usize,
n: usize,
socket: &UdpSocket,
retries: usize,
) -> Result<Vec<u8>, Error> {
let filename = format!("/dev/{device}.{offset:x}.{n:x}");
let bytes = retrying_download(&filename, socket, DEFAULT_TIMEOUT, MAX_TIMEOUT, retries)?;
if n != 0 && bytes.len() != n * 4 {
Err(Error::Incomplete)
} else {
Ok(bytes)
}
}
pub fn write_device(
device: &str,
offset: usize,
data: &[u8],
socket: &UdpSocket,
retries: usize,
) -> Result<(), Error> {
let filename = format!("/dev/{device}.{offset:x}");
retrying_upload(
&filename,
data,
socket,
DEFAULT_TIMEOUT,
MAX_TIMEOUT,
retries,
)
}
pub fn read_flash(
offset: usize,
n: usize,
socket: &UdpSocket,
retries: usize,
) -> Result<Vec<u8>, Error> {
let filename = format!("/flash.{offset:x}.{n:x}");
let bytes = retrying_download(&filename, socket, DEFAULT_TIMEOUT, MAX_TIMEOUT, retries)?;
Ok(bytes)
}
pub fn write_flash(
offset: usize,
data: &[u8],
socket: &UdpSocket,
retries: usize,
) -> Result<(), Error> {
let filename = format!("/flash.{offset:x}");
retrying_upload(
&filename,
data,
socket,
DEFAULT_TIMEOUT,
MAX_TIMEOUT,
retries,
)
}
pub fn progdev(addr: u32, socket: &UdpSocket) -> Result<(), Error> {
match upload(
"/progdev",
&addr.to_be_bytes(),
socket,
DEFAULT_TIMEOUT,
MAX_TIMEOUT,
0,
) {
Ok(()) | Err(_) => (),
}
std::thread::sleep(Duration::from_secs(10));
Ok(())
}
pub fn get_metadata(
socket: &UdpSocket,
user_flash_loc: u32,
retries: usize,
) -> Result<HashMap<KString, String>, Error> {
let mut dict_str = String::new();
let mut chunks = 0;
let chunk_size = 1024 / 4;
loop {
if chunks > 128 {
return Err(Error::MissingMetadata);
}
let raw = read_flash(
(user_flash_loc / 4 + chunks * chunk_size) as usize,
chunk_size as usize,
socket,
retries,
)?;
dict_str.push_str(std::str::from_utf8(&raw)?);
match dict_str.find("?end") {
Some(idx) => {
dict_str = dict_str.split_at(idx).0.to_string();
break;
}
None => chunks += 1,
}
}
Ok(dict_str
.split('?')
.filter_map(|kv| kv.split_once('\t'))
.map(|(k, v)| (k.to_string().into(), v.to_string()))
.collect())
}
#[allow(clippy::implicit_hasher)]
pub fn set_metadata(
data: &HashMap<KString, String>,
socket: &UdpSocket,
user_flash_loc: u32,
retries: usize,
) -> Result<(), Error> {
let mut dict_str = data.iter().fold(String::new(), |mut output, (k, v)| {
let _ = write!(output, "?{k}\t{v}");
output
});
dict_str.push_str("?end");
let mut bytes = dict_str.as_bytes().to_vec();
if bytes.len() % 1024 != 0 {
bytes.append(&mut vec![b'0'; 1024 - bytes.len() % 1024]);
}
write_flash((user_flash_loc / 4) as usize, &bytes, socket, retries)
}