feat: Added ftfp server implementation
This commit is contained in:
@@ -0,0 +1,353 @@
|
||||
use std::io::{BufRead, BufReader, BufWriter, Cursor, Read, Write};
|
||||
use std::net::SocketAddr;
|
||||
|
||||
use byteorder::{NetworkEndian, ReadBytesExt, WriteBytesExt};
|
||||
use num_enum::{IntoPrimitive, TryFromPrimitive};
|
||||
use tokio::net::UdpSocket;
|
||||
|
||||
use crate::tftp::Error;
|
||||
|
||||
// 2 bytes op code + 2 bytes block number + max 512 bytes data
|
||||
pub(crate) const PACKET_SIZE: usize = 2 + 2 + 512;
|
||||
|
||||
fn decode_null_terminated<R>(reader: &mut BufReader<R>) -> Result<String, Error>
|
||||
where
|
||||
R: Read + ?Sized,
|
||||
{
|
||||
let mut text = Vec::new();
|
||||
reader.read_until(0, &mut text)?;
|
||||
|
||||
// Make sure we ended with a null terminator
|
||||
match text.pop() {
|
||||
Some(0) => {}
|
||||
None | Some(_) => return Err(Error::UnterminatedString),
|
||||
}
|
||||
|
||||
Ok(String::from_utf8(text)?)
|
||||
}
|
||||
|
||||
fn encode_null_terminated<W>(writer: &mut BufWriter<W>, text: impl AsRef<str>) -> Result<(), Error>
|
||||
where
|
||||
W: Write + ?Sized,
|
||||
{
|
||||
let text = text.as_ref().as_bytes();
|
||||
writer.write_all(text)?;
|
||||
|
||||
// Add null terminator
|
||||
writer.write_u8(0)?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum Mode {
|
||||
NetAscii,
|
||||
Octet,
|
||||
Mail,
|
||||
}
|
||||
|
||||
impl Mode {
|
||||
fn decode<R>(reader: &mut BufReader<R>) -> Result<Self, Error>
|
||||
where
|
||||
R: Read + ?Sized,
|
||||
{
|
||||
let mut mode = decode_null_terminated(reader)?;
|
||||
mode.make_ascii_lowercase();
|
||||
let mode = mode;
|
||||
|
||||
Ok(match mode.as_str() {
|
||||
"netascii" => Self::NetAscii,
|
||||
"octet" => Self::Octet,
|
||||
"mail" => Self::Mail,
|
||||
_ => return Err(Error::InvalidMode(mode)),
|
||||
})
|
||||
}
|
||||
|
||||
fn encode<W>(&self, writer: &mut BufWriter<W>) -> Result<(), Error>
|
||||
where
|
||||
W: Write + ?Sized,
|
||||
{
|
||||
let mode = match &self {
|
||||
Mode::NetAscii => "netascii",
|
||||
Mode::Octet => "octet",
|
||||
Mode::Mail => "mail",
|
||||
};
|
||||
|
||||
encode_null_terminated(writer, mode)?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, PartialEq, Eq)]
|
||||
pub struct Request {
|
||||
filename: String,
|
||||
mode: Mode,
|
||||
}
|
||||
|
||||
impl Request {
|
||||
pub fn filename(&self) -> &str {
|
||||
&self.filename
|
||||
}
|
||||
|
||||
fn decode<R>(reader: &mut BufReader<R>) -> Result<Self, Error>
|
||||
where
|
||||
R: Read + ?Sized,
|
||||
{
|
||||
let filename = decode_null_terminated(reader)?;
|
||||
let mode = Mode::decode(reader)?;
|
||||
|
||||
Ok(Self { filename, mode })
|
||||
}
|
||||
|
||||
fn encode<W>(&self, writer: &mut BufWriter<W>) -> Result<(), Error>
|
||||
where
|
||||
W: Write + ?Sized,
|
||||
{
|
||||
encode_null_terminated(writer, &self.filename)?;
|
||||
self.mode.encode(writer)?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn mode(&self) -> Mode {
|
||||
self.mode
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct Data {
|
||||
block: u16,
|
||||
data: Vec<u8>,
|
||||
}
|
||||
|
||||
impl Data {
|
||||
fn decode<R>(reader: &mut BufReader<R>) -> Result<Self, Error>
|
||||
where
|
||||
R: Read + ?Sized,
|
||||
{
|
||||
let block = reader.read_u16::<NetworkEndian>()?;
|
||||
|
||||
let mut data = Vec::new();
|
||||
reader.read_to_end(&mut data)?;
|
||||
|
||||
Ok(Self { data, block })
|
||||
}
|
||||
|
||||
fn encode<W>(&self, writer: &mut BufWriter<W>) -> Result<(), Error>
|
||||
where
|
||||
W: Write + ?Sized,
|
||||
{
|
||||
writer.write_u16::<NetworkEndian>(self.block)?;
|
||||
writer.write_all(&self.data)?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct Ack(u16);
|
||||
|
||||
impl Ack {
|
||||
pub fn block(&self) -> u16 {
|
||||
self.0
|
||||
}
|
||||
|
||||
fn decode<R>(reader: &mut BufReader<R>) -> Result<Self, Error>
|
||||
where
|
||||
R: Read + ?Sized,
|
||||
{
|
||||
let number = reader.read_u16::<NetworkEndian>()?;
|
||||
|
||||
Ok(Self(number))
|
||||
}
|
||||
|
||||
fn encode<W>(&self, writer: &mut BufWriter<W>) -> Result<(), Error>
|
||||
where
|
||||
W: Write + ?Sized,
|
||||
{
|
||||
writer.write_u16::<NetworkEndian>(self.0)?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, TryFromPrimitive, IntoPrimitive, Clone, Copy)]
|
||||
#[repr(u16)]
|
||||
pub enum ErrorCode {
|
||||
NotDefined,
|
||||
FileNotFound,
|
||||
AccessViolation,
|
||||
DiskFull,
|
||||
Illegal,
|
||||
UnknownTransferId,
|
||||
FileAlreadyExists,
|
||||
NoSuchUser,
|
||||
}
|
||||
|
||||
impl ErrorCode {
|
||||
fn decode<R>(reader: &mut BufReader<R>) -> Result<Self, Error>
|
||||
where
|
||||
R: Read + ?Sized,
|
||||
{
|
||||
let code = reader.read_u16::<NetworkEndian>()?;
|
||||
Self::try_from_primitive(code).map_err(|err| Error::InvalidErrorCode(err.number))
|
||||
}
|
||||
|
||||
fn encode<W>(&self, writer: &mut BufWriter<W>) -> Result<(), Error>
|
||||
where
|
||||
W: Write + ?Sized,
|
||||
{
|
||||
let code: u16 = (*self).into();
|
||||
|
||||
writer.write_u16::<NetworkEndian>(code)?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct ErrorMessage {
|
||||
code: ErrorCode,
|
||||
message: String,
|
||||
}
|
||||
|
||||
impl ErrorMessage {
|
||||
fn decode<R>(reader: &mut BufReader<R>) -> Result<Self, Error>
|
||||
where
|
||||
R: Read + ?Sized,
|
||||
{
|
||||
let code = ErrorCode::decode(reader)?;
|
||||
let message = decode_null_terminated(reader)?;
|
||||
|
||||
Ok(Self { code, message })
|
||||
}
|
||||
|
||||
fn encode<W>(&self, writer: &mut BufWriter<W>) -> Result<(), Error>
|
||||
where
|
||||
W: Write + ?Sized,
|
||||
{
|
||||
self.code.encode(writer)?;
|
||||
|
||||
encode_null_terminated(writer, &self.message)?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum Packet {
|
||||
Read(Request),
|
||||
Write(Request),
|
||||
Data(Data),
|
||||
Ack(Ack),
|
||||
Error(ErrorMessage),
|
||||
}
|
||||
|
||||
impl Packet {
|
||||
pub fn data(block: u16, data: &[u8]) -> Self {
|
||||
Self::Data(Data {
|
||||
block,
|
||||
data: data.into(),
|
||||
})
|
||||
}
|
||||
|
||||
pub fn ack(block: u16) -> Self {
|
||||
Self::Ack(Ack(block))
|
||||
}
|
||||
|
||||
pub fn error(code: ErrorCode, message: impl Into<String>) -> Self {
|
||||
Self::Error(ErrorMessage {
|
||||
code,
|
||||
message: message.into(),
|
||||
})
|
||||
}
|
||||
|
||||
fn decode<R>(reader: &mut BufReader<R>) -> Result<Self, Error>
|
||||
where
|
||||
R: Read + ?Sized,
|
||||
{
|
||||
let op_code = reader.read_u16::<NetworkEndian>()?;
|
||||
|
||||
Ok(match op_code {
|
||||
1 => Self::Read(Request::decode(reader)?),
|
||||
2 => Self::Write(Request::decode(reader)?),
|
||||
3 => Self::Data(Data::decode(reader)?),
|
||||
4 => Self::Ack(Ack::decode(reader)?),
|
||||
5 => Self::Error(ErrorMessage::decode(reader)?),
|
||||
_ => return Err(Error::InvalidOpCode(op_code)),
|
||||
})
|
||||
}
|
||||
|
||||
fn encode<W>(&self, writer: &mut BufWriter<W>) -> Result<(), Error>
|
||||
where
|
||||
W: Write + ?Sized,
|
||||
{
|
||||
match &self {
|
||||
Self::Read(inner) => {
|
||||
writer.write_u16::<NetworkEndian>(1)?;
|
||||
inner.encode(writer)?;
|
||||
}
|
||||
Self::Write(inner) => {
|
||||
writer.write_u16::<NetworkEndian>(2)?;
|
||||
inner.encode(writer)?;
|
||||
}
|
||||
Self::Data(inner) => {
|
||||
writer.write_u16::<NetworkEndian>(3)?;
|
||||
inner.encode(writer)?;
|
||||
}
|
||||
Self::Ack(inner) => {
|
||||
writer.write_u16::<NetworkEndian>(4)?;
|
||||
inner.encode(writer)?;
|
||||
}
|
||||
Self::Error(inner) => {
|
||||
writer.write_u16::<NetworkEndian>(5)?;
|
||||
inner.encode(writer)?;
|
||||
}
|
||||
};
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn recv_from(socket: &UdpSocket) -> Result<(Self, SocketAddr), Error> {
|
||||
let mut buf = Vec::with_capacity(PACKET_SIZE);
|
||||
let (_, src) = socket.recv_buf_from(&mut buf).await?;
|
||||
|
||||
let packet = buf.try_into()?;
|
||||
|
||||
Ok((packet, src))
|
||||
}
|
||||
|
||||
pub async fn send_to(&self, socket: &UdpSocket, src: &SocketAddr) -> Result<(), Error> {
|
||||
let buf: Vec<_> = self.try_into()?;
|
||||
socket.send_to(&buf, src).await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl TryFrom<&Packet> for Vec<u8> {
|
||||
type Error = Error;
|
||||
|
||||
fn try_from(packet: &Packet) -> Result<Self, Self::Error> {
|
||||
let mut buf = Vec::with_capacity(PACKET_SIZE);
|
||||
let mut writer = BufWriter::new(Cursor::new(&mut buf));
|
||||
packet.encode(&mut writer)?;
|
||||
|
||||
// Normally the writer would get dropped after we return, but since it has a mutable
|
||||
// borrow in buf we have to drop it first.
|
||||
drop(writer);
|
||||
|
||||
Ok(buf)
|
||||
}
|
||||
}
|
||||
|
||||
impl TryFrom<Vec<u8>> for Packet {
|
||||
type Error = Error;
|
||||
|
||||
fn try_from(buf: Vec<u8>) -> Result<Self, <Packet as TryFrom<Vec<u8>>>::Error> {
|
||||
let mut reader = BufReader::new(Cursor::new(buf));
|
||||
|
||||
Packet::decode(&mut reader)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user