diff options
Diffstat (limited to 'src')
| -rw-r--r-- | src/dhcp.rs | 214 | ||||
| -rw-r--r-- | src/main.rs | 7 | ||||
| -rw-r--r-- | src/tftp.rs | 365 | ||||
| -rw-r--r-- | src/wire.rs | 81 |
4 files changed, 666 insertions, 1 deletions
diff --git a/src/dhcp.rs b/src/dhcp.rs new file mode 100644 index 0000000..b53ee92 --- /dev/null +++ b/src/dhcp.rs | |||
| @@ -0,0 +1,214 @@ | |||
| 1 | use std::{ | ||
| 2 | io::{Result, Write}, | ||
| 3 | net::Ipv4Addr, | ||
| 4 | }; | ||
| 5 | |||
| 6 | use crate::wire; | ||
| 7 | |||
| 8 | const MAGIC_COOKIE: [u8; 4] = [0x63, 0x82, 0x53, 0x63]; | ||
| 9 | |||
| 10 | #[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Hash)] | ||
| 11 | pub enum BootOp { | ||
| 12 | #[default] | ||
| 13 | Request, | ||
| 14 | Reply, | ||
| 15 | } | ||
| 16 | |||
| 17 | impl BootOp { | ||
| 18 | pub const OP_REQUEST: u8 = 1; | ||
| 19 | pub const OP_REPLY: u8 = 2; | ||
| 20 | } | ||
| 21 | |||
| 22 | impl From<BootOp> for u8 { | ||
| 23 | fn from(value: BootOp) -> Self { | ||
| 24 | match value { | ||
| 25 | BootOp::Request => BootOp::OP_REQUEST, | ||
| 26 | BootOp::Reply => BootOp::OP_REPLY, | ||
| 27 | } | ||
| 28 | } | ||
| 29 | } | ||
| 30 | |||
| 31 | #[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Hash)] | ||
| 32 | pub enum HardwareType { | ||
| 33 | #[default] | ||
| 34 | Ethernet, | ||
| 35 | } | ||
| 36 | |||
| 37 | impl HardwareType { | ||
| 38 | pub const TYPE_ETHER: u8 = 1; | ||
| 39 | pub const LEN_ETHER: u8 = 6; | ||
| 40 | |||
| 41 | pub fn hardware_len(&self) -> u8 { | ||
| 42 | match self { | ||
| 43 | HardwareType::Ethernet => Self::LEN_ETHER, | ||
| 44 | } | ||
| 45 | } | ||
| 46 | } | ||
| 47 | |||
| 48 | impl From<HardwareType> for u8 { | ||
| 49 | fn from(value: HardwareType) -> Self { | ||
| 50 | match value { | ||
| 51 | HardwareType::Ethernet => HardwareType::TYPE_ETHER, | ||
| 52 | } | ||
| 53 | } | ||
| 54 | } | ||
| 55 | |||
| 56 | #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] | ||
| 57 | pub enum DhcpMessageType { | ||
| 58 | Ack, | ||
| 59 | } | ||
| 60 | |||
| 61 | impl DhcpMessageType { | ||
| 62 | pub const CODE_ACK: u8 = 5; | ||
| 63 | |||
| 64 | pub fn code(&self) -> u8 { | ||
| 65 | match self { | ||
| 66 | DhcpMessageType::Ack => Self::CODE_ACK, | ||
| 67 | } | ||
| 68 | } | ||
| 69 | } | ||
| 70 | |||
| 71 | #[derive(Debug, Clone)] | ||
| 72 | pub enum DhcpOption { | ||
| 73 | Pad, | ||
| 74 | End, | ||
| 75 | MessageType(DhcpMessageType), | ||
| 76 | ServerIdentifier(Ipv4Addr), | ||
| 77 | VendorClassIdentifier(String), | ||
| 78 | TftpServerName(String), | ||
| 79 | TftpFileName(String), | ||
| 80 | UserClassInformation(String), | ||
| 81 | ClientMachineIdentifier(String), | ||
| 82 | Unknown { code: u8, data: Vec<u8> }, | ||
| 83 | } | ||
| 84 | |||
| 85 | impl DhcpOption { | ||
| 86 | pub const CODE_PAD: u8 = 0; | ||
| 87 | pub const CODE_END: u8 = 255; | ||
| 88 | pub const CODE_DHCP_MESSAGE_TYPE: u8 = 53; | ||
| 89 | pub const CODE_DHCP_SERVER_IDENTIFIER: u8 = 54; | ||
| 90 | pub const CODE_VENDOR_CLASS_IDENTIFIER: u8 = 60; | ||
| 91 | pub const CODE_TFTP_SERVER_NAME: u8 = 66; | ||
| 92 | pub const CODE_TFTP_FILE_NAME: u8 = 67; | ||
| 93 | pub const CODE_USER_CLASS_INFORMATION: u8 = 77; | ||
| 94 | pub const CODE_CLIENT_MACHINE_IDENTIFIER: u8 = 97; | ||
| 95 | |||
| 96 | pub fn code(&self) -> u8 { | ||
| 97 | match self { | ||
| 98 | DhcpOption::Pad => Self::CODE_PAD, | ||
| 99 | DhcpOption::End => Self::CODE_END, | ||
| 100 | DhcpOption::MessageType(_) => Self::CODE_DHCP_MESSAGE_TYPE, | ||
| 101 | DhcpOption::ServerIdentifier(_) => Self::CODE_DHCP_SERVER_IDENTIFIER, | ||
| 102 | DhcpOption::VendorClassIdentifier(_) => Self::CODE_VENDOR_CLASS_IDENTIFIER, | ||
| 103 | DhcpOption::TftpServerName(_) => Self::CODE_TFTP_SERVER_NAME, | ||
| 104 | DhcpOption::TftpFileName(_) => Self::CODE_TFTP_FILE_NAME, | ||
| 105 | DhcpOption::UserClassInformation(_) => Self::CODE_USER_CLASS_INFORMATION, | ||
| 106 | DhcpOption::ClientMachineIdentifier(_) => Self::CODE_CLIENT_MACHINE_IDENTIFIER, | ||
| 107 | DhcpOption::Unknown { code, .. } => *code, | ||
| 108 | } | ||
| 109 | } | ||
| 110 | } | ||
| 111 | |||
| 112 | #[derive(Debug)] | ||
| 113 | pub struct DhcpPacket { | ||
| 114 | pub op: BootOp, | ||
| 115 | pub htype: HardwareType, | ||
| 116 | pub xid: u32, | ||
| 117 | pub secs: u16, | ||
| 118 | pub flags: u16, | ||
| 119 | pub ciaddr: Ipv4Addr, | ||
| 120 | pub yiaddr: Ipv4Addr, | ||
| 121 | pub siaddr: Ipv4Addr, | ||
| 122 | pub giaddr: Ipv4Addr, | ||
| 123 | pub chaddr: [u8; 16], | ||
| 124 | // server host name | ||
| 125 | pub sname: Option<String>, | ||
| 126 | // boot file name | ||
| 127 | pub file: Option<String>, | ||
| 128 | pub options: Vec<DhcpOption>, | ||
| 129 | } | ||
| 130 | |||
| 131 | impl Default for DhcpPacket { | ||
| 132 | fn default() -> Self { | ||
| 133 | Self { | ||
| 134 | op: Default::default(), | ||
| 135 | htype: Default::default(), | ||
| 136 | xid: Default::default(), | ||
| 137 | secs: Default::default(), | ||
| 138 | flags: Default::default(), | ||
| 139 | ciaddr: Ipv4Addr::UNSPECIFIED, | ||
| 140 | yiaddr: Ipv4Addr::UNSPECIFIED, | ||
| 141 | siaddr: Ipv4Addr::UNSPECIFIED, | ||
| 142 | giaddr: Ipv4Addr::UNSPECIFIED, | ||
| 143 | chaddr: Default::default(), | ||
| 144 | sname: Default::default(), | ||
| 145 | file: Default::default(), | ||
| 146 | options: Default::default(), | ||
| 147 | } | ||
| 148 | } | ||
| 149 | } | ||
| 150 | |||
| 151 | pub fn write_packet<W: Write>(mut writer: W, packet: &DhcpPacket) -> Result<()> { | ||
| 152 | wire::write_u8(&mut writer, u8::from(packet.op))?; | ||
| 153 | wire::write_u8(&mut writer, u8::from(packet.htype))?; | ||
| 154 | wire::write_u8(&mut writer, packet.htype.hardware_len())?; | ||
| 155 | wire::write_u8(&mut writer, 0)?; // hops | ||
| 156 | wire::write_u32(&mut writer, packet.xid)?; | ||
| 157 | wire::write_u16(&mut writer, packet.secs)?; | ||
| 158 | wire::write_u16(&mut writer, packet.flags)?; | ||
| 159 | wire::write_ipv4(&mut writer, packet.ciaddr)?; | ||
| 160 | wire::write_ipv4(&mut writer, packet.yiaddr)?; | ||
| 161 | wire::write_ipv4(&mut writer, packet.siaddr)?; | ||
| 162 | wire::write_ipv4(&mut writer, packet.giaddr)?; | ||
| 163 | wire::write(&mut writer, &packet.chaddr)?; | ||
| 164 | match &packet.sname { | ||
| 165 | Some(name) => wire::write_null_terminated_string(&mut writer, &name)?, | ||
| 166 | None => wire::write_null_terminated_string(&mut writer, "")?, | ||
| 167 | }; | ||
| 168 | match &packet.file { | ||
| 169 | Some(name) => wire::write_null_terminated_string(&mut writer, &name)?, | ||
| 170 | None => wire::write_null_terminated_string(&mut writer, "")?, | ||
| 171 | }; | ||
| 172 | wire::write(&mut writer, &MAGIC_COOKIE)?; | ||
| 173 | for option in &packet.options { | ||
| 174 | write_option(&mut writer, option)?; | ||
| 175 | } | ||
| 176 | write_option(&mut writer, &DhcpOption::End)?; | ||
| 177 | Ok(()) | ||
| 178 | } | ||
| 179 | |||
| 180 | pub fn write_option<W: Write>(mut writer: W, option: &DhcpOption) -> Result<()> { | ||
| 181 | wire::write_u8(&mut writer, option.code())?; | ||
| 182 | match option { | ||
| 183 | DhcpOption::Pad | DhcpOption::End => {} | ||
| 184 | DhcpOption::MessageType(t) => { | ||
| 185 | wire::write_u8(&mut writer, 1)?; | ||
| 186 | wire::write_u8(&mut writer, t.code())?; | ||
| 187 | } | ||
| 188 | DhcpOption::ServerIdentifier(ip) => { | ||
| 189 | wire::write_u8(&mut writer, 4)?; | ||
| 190 | wire::write_ipv4(&mut writer, *ip)?; | ||
| 191 | } | ||
| 192 | DhcpOption::VendorClassIdentifier(vendor_class) => { | ||
| 193 | write_option_len_prefixed_string(&mut writer, &vendor_class)? | ||
| 194 | } | ||
| 195 | DhcpOption::TftpServerName(name) => write_option_len_prefixed_string(&mut writer, &name)?, | ||
| 196 | DhcpOption::TftpFileName(name) => write_option_len_prefixed_string(&mut writer, &name)?, | ||
| 197 | DhcpOption::UserClassInformation(user_class) => { | ||
| 198 | write_option_len_prefixed_string(&mut writer, &user_class)? | ||
| 199 | } | ||
| 200 | DhcpOption::ClientMachineIdentifier(identifier) => { | ||
| 201 | write_option_len_prefixed_string(&mut writer, &identifier)? | ||
| 202 | } | ||
| 203 | DhcpOption::Unknown { data, .. } => { | ||
| 204 | wire::write_u8(&mut writer, u8::try_from(data.len()).unwrap())?; | ||
| 205 | wire::write(&mut writer, &data)?; | ||
| 206 | } | ||
| 207 | } | ||
| 208 | Ok(()) | ||
| 209 | } | ||
| 210 | |||
| 211 | fn write_option_len_prefixed_string<W: Write>(mut writer: W, s: &str) -> Result<()> { | ||
| 212 | wire::write_u8(&mut writer, u8::try_from(s.len()).unwrap())?; | ||
| 213 | wire::write(&mut writer, s.as_bytes()) | ||
| 214 | } | ||
diff --git a/src/main.rs b/src/main.rs index 87ea283..3d86a8e 100644 --- a/src/main.rs +++ b/src/main.rs | |||
| @@ -1,3 +1,8 @@ | |||
| 1 | #![feature(cursor_split)] | ||
| 2 | pub mod dhcp; | ||
| 3 | pub mod tftp; | ||
| 4 | pub mod wire; | ||
| 5 | |||
| 1 | use std::io::{BufRead, Cursor, Read, Result, Write}; | 6 | use std::io::{BufRead, Cursor, Read, Result, Write}; |
| 2 | use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4, UdpSocket}; | 7 | use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4, UdpSocket}; |
| 3 | 8 | ||
| @@ -15,7 +20,7 @@ const MAGIC_COOKIE: [u8; 4] = [0x63, 0x82, 0x53, 0x63]; | |||
| 15 | const BOOT_FILE_NAME: &[u8] = b"ipxe.efi"; | 20 | const BOOT_FILE_NAME: &[u8] = b"ipxe.efi"; |
| 16 | const BOOT_FILE_NAME_IPXE: &[u8] = b"test.ipxe"; | 21 | const BOOT_FILE_NAME_IPXE: &[u8] = b"test.ipxe"; |
| 17 | 22 | ||
| 18 | const LOCAL_IPV4: Ipv4Addr = Ipv4Addr::new(192, 168, 2, 184); | 23 | const LOCAL_IPV4: Ipv4Addr = Ipv4Addr::new(192, 168, 1, 100); |
| 19 | 24 | ||
| 20 | #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] | 25 | #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] |
| 21 | enum BootOp { | 26 | enum BootOp { |
diff --git a/src/tftp.rs b/src/tftp.rs new file mode 100644 index 0000000..becaa65 --- /dev/null +++ b/src/tftp.rs | |||
| @@ -0,0 +1,365 @@ | |||
| 1 | use std::{ | ||
| 2 | io::{Cursor, Read as _, Result, Write}, | ||
| 3 | net::UdpSocket, | ||
| 4 | path::{Path, PathBuf}, | ||
| 5 | str::FromStr, | ||
| 6 | }; | ||
| 7 | |||
| 8 | use crate::wire; | ||
| 9 | |||
| 10 | pub const PORT: u16 = 69; | ||
| 11 | |||
| 12 | #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] | ||
| 13 | pub struct InvalidTftpOp(u16); | ||
| 14 | |||
| 15 | impl std::fmt::Display for InvalidTftpOp { | ||
| 16 | fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { | ||
| 17 | write!(f, "invalid tftp opcode '{}'", self.0) | ||
| 18 | } | ||
| 19 | } | ||
| 20 | |||
| 21 | impl std::error::Error for InvalidTftpOp {} | ||
| 22 | #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] | ||
| 23 | pub enum TftpOp { | ||
| 24 | ReadRequest, | ||
| 25 | WriteRequest, | ||
| 26 | Data, | ||
| 27 | Ack, | ||
| 28 | Error, | ||
| 29 | Oack, | ||
| 30 | } | ||
| 31 | |||
| 32 | impl Into<u16> for TftpOp { | ||
| 33 | fn into(self) -> u16 { | ||
| 34 | match self { | ||
| 35 | TftpOp::ReadRequest => 1, | ||
| 36 | TftpOp::WriteRequest => 2, | ||
| 37 | TftpOp::Data => 3, | ||
| 38 | TftpOp::Ack => 4, | ||
| 39 | TftpOp::Error => 5, | ||
| 40 | TftpOp::Oack => 6, | ||
| 41 | } | ||
| 42 | } | ||
| 43 | } | ||
| 44 | |||
| 45 | impl TryFrom<u16> for TftpOp { | ||
| 46 | type Error = InvalidTftpOp; | ||
| 47 | |||
| 48 | fn try_from(value: u16) -> std::result::Result<Self, InvalidTftpOp> { | ||
| 49 | match value { | ||
| 50 | 1 => Ok(Self::ReadRequest), | ||
| 51 | 2 => Ok(Self::WriteRequest), | ||
| 52 | 3 => Ok(Self::Data), | ||
| 53 | 4 => Ok(Self::Ack), | ||
| 54 | 5 => Ok(Self::Error), | ||
| 55 | 6 => Ok(Self::Oack), | ||
| 56 | unknown => Err(InvalidTftpOp(unknown)), | ||
| 57 | } | ||
| 58 | } | ||
| 59 | } | ||
| 60 | |||
| 61 | #[derive(Debug)] | ||
| 62 | pub struct InvalidTftpMode(String); | ||
| 63 | |||
| 64 | impl std::fmt::Display for InvalidTftpMode { | ||
| 65 | fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { | ||
| 66 | write!(f, "invalid tftp mode '{}'", self.0) | ||
| 67 | } | ||
| 68 | } | ||
| 69 | |||
| 70 | impl std::error::Error for InvalidTftpMode {} | ||
| 71 | |||
| 72 | #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] | ||
| 73 | pub enum TftpMode { | ||
| 74 | NetAscii, | ||
| 75 | Octet, | ||
| 76 | Mail, | ||
| 77 | } | ||
| 78 | |||
| 79 | impl FromStr for TftpMode { | ||
| 80 | type Err = InvalidTftpMode; | ||
| 81 | |||
| 82 | fn from_str(s: &str) -> std::result::Result<Self, Self::Err> { | ||
| 83 | match s.to_lowercase().as_str() { | ||
| 84 | "netascii" => Ok(Self::NetAscii), | ||
| 85 | "octet" => Ok(Self::Octet), | ||
| 86 | "mail" => Ok(Self::Mail), | ||
| 87 | _ => Err(InvalidTftpMode(s.to_string())), | ||
| 88 | } | ||
| 89 | } | ||
| 90 | } | ||
| 91 | |||
| 92 | #[derive(Debug)] | ||
| 93 | pub enum TftpPacket { | ||
| 94 | Request(TftpRequestPacket), | ||
| 95 | Data(TftpDataPacket), | ||
| 96 | Ack(TftpAckPacket), | ||
| 97 | OAck(TftpOAckPacket), | ||
| 98 | Error(TftpErrorPacket), | ||
| 99 | } | ||
| 100 | |||
| 101 | impl TftpPacket { | ||
| 102 | pub fn write<W: Write>(&self, writer: W) -> Result<()> { | ||
| 103 | match self { | ||
| 104 | TftpPacket::Request(p) => p.write(writer), | ||
| 105 | TftpPacket::Data(p) => p.write(writer), | ||
| 106 | TftpPacket::Ack(p) => p.write(writer), | ||
| 107 | TftpPacket::OAck(p) => p.write(writer), | ||
| 108 | TftpPacket::Error(p) => p.write(writer), | ||
| 109 | } | ||
| 110 | } | ||
| 111 | } | ||
| 112 | |||
| 113 | #[derive(Debug)] | ||
| 114 | pub struct TftpRequestPacket { | ||
| 115 | pub filename: String, | ||
| 116 | pub mode: TftpMode, | ||
| 117 | pub tsize: Option<u64>, | ||
| 118 | pub blksize: Option<u64>, | ||
| 119 | } | ||
| 120 | |||
| 121 | impl TftpRequestPacket { | ||
| 122 | pub fn write<W: Write>(&self, mut writer: W) -> Result<()> { | ||
| 123 | todo!() | ||
| 124 | } | ||
| 125 | } | ||
| 126 | |||
| 127 | #[derive(Debug)] | ||
| 128 | pub struct TftpDataPacket { | ||
| 129 | pub block: u16, | ||
| 130 | pub data: Vec<u8>, | ||
| 131 | } | ||
| 132 | |||
| 133 | impl TftpDataPacket { | ||
| 134 | pub fn write<W: Write>(&self, mut writer: W) -> Result<()> { | ||
| 135 | wire::write_u16(&mut writer, TftpOp::Data.into())?; | ||
| 136 | wire::write_u16(&mut writer, self.block)?; | ||
| 137 | wire::write(&mut writer, &self.data)?; | ||
| 138 | Ok(()) | ||
| 139 | } | ||
| 140 | } | ||
| 141 | |||
| 142 | #[derive(Debug)] | ||
| 143 | pub struct TftpAckPacket { | ||
| 144 | pub block: u16, | ||
| 145 | } | ||
| 146 | |||
| 147 | impl TftpAckPacket { | ||
| 148 | pub fn write<W: Write>(&self, mut writer: W) -> Result<()> { | ||
| 149 | wire::write_u16(&mut writer, TftpOp::Data.into())?; | ||
| 150 | wire::write_u16(&mut writer, self.block)?; | ||
| 151 | Ok(()) | ||
| 152 | } | ||
| 153 | } | ||
| 154 | |||
| 155 | #[derive(Debug)] | ||
| 156 | pub struct TftpOAckPacket { | ||
| 157 | pub tsize: Option<u64>, | ||
| 158 | pub blksize: Option<u64>, | ||
| 159 | } | ||
| 160 | |||
| 161 | impl TftpOAckPacket { | ||
| 162 | pub fn write<W: Write>(&self, mut writer: W) -> Result<()> { | ||
| 163 | wire::write_u16(&mut writer, TftpOp::Oack.into())?; | ||
| 164 | |||
| 165 | // Only include options that were requested by the client | ||
| 166 | if let Some(blksize_val) = self.blksize { | ||
| 167 | wire::write(&mut writer, b"blksize")?; | ||
| 168 | wire::write_u8(&mut writer, 0)?; // null terminator | ||
| 169 | let blksize_str = blksize_val.to_string(); | ||
| 170 | wire::write(&mut writer, blksize_str.as_bytes())?; | ||
| 171 | wire::write_u8(&mut writer, 0)?; // null terminator | ||
| 172 | } | ||
| 173 | |||
| 174 | if let Some(tsize_val) = self.tsize { | ||
| 175 | wire::write(&mut writer, b"tsize")?; | ||
| 176 | wire::write_u8(&mut writer, 0)?; // null terminator | ||
| 177 | let tsize_str = tsize_val.to_string(); | ||
| 178 | wire::write(&mut writer, tsize_str.as_bytes())?; | ||
| 179 | wire::write_u8(&mut writer, 0)?; // null terminator | ||
| 180 | } | ||
| 181 | |||
| 182 | Ok(()) | ||
| 183 | } | ||
| 184 | } | ||
| 185 | |||
| 186 | #[derive(Debug)] | ||
| 187 | pub struct TftpErrorPacket { | ||
| 188 | pub code: u16, | ||
| 189 | pub message: String, | ||
| 190 | } | ||
| 191 | |||
| 192 | impl TftpErrorPacket { | ||
| 193 | pub fn write<W: Write>(&self, mut writer: W) -> Result<()> { | ||
| 194 | wire::write_u16(&mut writer, TftpOp::Error.into())?; | ||
| 195 | wire::write_u16(&mut writer, self.code)?; | ||
| 196 | wire::write_null_terminated_string(&mut writer, &self.message)?; | ||
| 197 | Ok(()) | ||
| 198 | } | ||
| 199 | } | ||
| 200 | |||
| 201 | pub fn parse_packet(buf: &[u8]) -> Result<TftpPacket> { | ||
| 202 | let mut cursor = Cursor::new(buf); | ||
| 203 | let op = TftpOp::try_from(wire::read_u16(&mut cursor)?).unwrap(); | ||
| 204 | |||
| 205 | match op { | ||
| 206 | TftpOp::ReadRequest => { | ||
| 207 | let filename = wire::read_null_terminated_string(&mut cursor)?; | ||
| 208 | let mode = wire::read_null_terminated_string(&mut cursor)? | ||
| 209 | .parse::<TftpMode>() | ||
| 210 | .unwrap(); | ||
| 211 | |||
| 212 | let mut tsize = None; | ||
| 213 | let mut blksize = None; | ||
| 214 | |||
| 215 | while let Ok(opt_name) = wire::read_null_terminated_string(&mut cursor) { | ||
| 216 | if opt_name.is_empty() { | ||
| 217 | break; | ||
| 218 | } | ||
| 219 | let opt_data = wire::read_null_terminated_string(&mut cursor)?; | ||
| 220 | match opt_name.as_str() { | ||
| 221 | "tsize" => tsize = Some(opt_data.parse::<u64>().unwrap()), | ||
| 222 | "blksize" => blksize = Some(opt_data.parse::<u64>().unwrap()), | ||
| 223 | _ => eprintln!("unknown tftp request option '{opt_name}'"), | ||
| 224 | } | ||
| 225 | } | ||
| 226 | |||
| 227 | Ok(TftpPacket::Request(TftpRequestPacket { | ||
| 228 | filename, | ||
| 229 | mode, | ||
| 230 | tsize, | ||
| 231 | blksize, | ||
| 232 | })) | ||
| 233 | } | ||
| 234 | TftpOp::WriteRequest => unimplemented!(), | ||
| 235 | TftpOp::Data => { | ||
| 236 | let block = wire::read_u16(&mut cursor)?; | ||
| 237 | let mut data = Vec::new(); | ||
| 238 | cursor.read_to_end(&mut data)?; | ||
| 239 | Ok(TftpPacket::Data(TftpDataPacket { block, data })) | ||
| 240 | } | ||
| 241 | TftpOp::Ack => { | ||
| 242 | let block = wire::read_u16(&mut cursor)?; | ||
| 243 | Ok(TftpPacket::Ack(TftpAckPacket { block })) | ||
| 244 | } | ||
| 245 | TftpOp::Error => { | ||
| 246 | let code = wire::read_u16(&mut cursor)?; | ||
| 247 | let message = wire::read_null_terminated_string(&mut cursor)?; | ||
| 248 | Ok(TftpPacket::Error(TftpErrorPacket { code, message })) | ||
| 249 | } | ||
| 250 | TftpOp::Oack => { | ||
| 251 | let mut tsize = None; | ||
| 252 | let mut blksize = None; | ||
| 253 | |||
| 254 | while let Ok(opt_name) = wire::read_null_terminated_string(&mut cursor) { | ||
| 255 | if opt_name.is_empty() { | ||
| 256 | break; | ||
| 257 | } | ||
| 258 | let opt_data = wire::read_null_terminated_string(&mut cursor)?; | ||
| 259 | match opt_name.as_str() { | ||
| 260 | "tsize" => tsize = Some(opt_data.parse::<u64>().unwrap()), | ||
| 261 | "blksize" => blksize = Some(opt_data.parse::<u64>().unwrap()), | ||
| 262 | _ => eprintln!("unknown tftp ack option '{opt_name}'"), | ||
| 263 | } | ||
| 264 | } | ||
| 265 | Ok(TftpPacket::OAck(TftpOAckPacket { tsize, blksize })) | ||
| 266 | } | ||
| 267 | } | ||
| 268 | } | ||
| 269 | |||
| 270 | pub fn serve(dir: &Path) -> Result<()> { | ||
| 271 | let socket = UdpSocket::bind(format!("0.0.0.0:{PORT}"))?; | ||
| 272 | |||
| 273 | // TODO: this needs to be done per addr | ||
| 274 | let mut last_blksize = 512u64; | ||
| 275 | let mut current_file = PathBuf::default(); | ||
| 276 | |||
| 277 | loop { | ||
| 278 | let mut buf = [0u8; 1500]; | ||
| 279 | let (n, addr) = socket.recv_from(&mut buf)?; | ||
| 280 | let packet = parse_packet(&buf[..n]).unwrap(); | ||
| 281 | |||
| 282 | let response = match packet { | ||
| 283 | TftpPacket::Request(req) => { | ||
| 284 | println!( | ||
| 285 | "Request options: tsize={:?}, blksize={:?}", | ||
| 286 | req.tsize, req.blksize | ||
| 287 | ); | ||
| 288 | |||
| 289 | let filepath = dir.join(req.filename); | ||
| 290 | current_file = filepath.clone(); | ||
| 291 | let meta = std::fs::metadata(&filepath).unwrap(); | ||
| 292 | let actual_file_size = meta.len(); | ||
| 293 | |||
| 294 | // Only send OACK if client sent options | ||
| 295 | if req.tsize.is_some() || req.blksize.is_some() { | ||
| 296 | if let Some(blksize) = req.blksize { | ||
| 297 | last_blksize = blksize; | ||
| 298 | } | ||
| 299 | |||
| 300 | let tsize_response = if req.tsize.is_some() { | ||
| 301 | Some(actual_file_size) | ||
| 302 | } else { | ||
| 303 | None | ||
| 304 | }; | ||
| 305 | |||
| 306 | Some(TftpPacket::OAck(TftpOAckPacket { | ||
| 307 | tsize: req.tsize, | ||
| 308 | blksize: req.blksize, | ||
| 309 | })) | ||
| 310 | } else { | ||
| 311 | // No options, send first data block directly | ||
| 312 | let contents = std::fs::read(&filepath).unwrap(); | ||
| 313 | let block_size = 512; | ||
| 314 | let first_block = if contents.len() > block_size { | ||
| 315 | contents[..block_size].to_vec() | ||
| 316 | } else { | ||
| 317 | contents | ||
| 318 | }; | ||
| 319 | |||
| 320 | Some(TftpPacket::Data(TftpDataPacket { | ||
| 321 | block: 1, | ||
| 322 | data: first_block, | ||
| 323 | })) | ||
| 324 | } | ||
| 325 | } | ||
| 326 | TftpPacket::Data(dat) => unimplemented!(), | ||
| 327 | TftpPacket::Ack(ack) => { | ||
| 328 | println!("Received ACK packet: block {}", ack.block); | ||
| 329 | |||
| 330 | let contents = std::fs::read(¤t_file).unwrap(); | ||
| 331 | let next_block = ack.block + 1; | ||
| 332 | let start_offset = (next_block - 1) as u64 * last_blksize; | ||
| 333 | let end_offset = next_block as u64 * last_blksize; | ||
| 334 | let prev_start_offset = (next_block.saturating_sub(2)) as u64 * last_blksize; | ||
| 335 | let prev_remain = contents.len() - prev_start_offset as usize; | ||
| 336 | if prev_remain as u64 >= last_blksize || ack.block == 0 { | ||
| 337 | let end = std::cmp::min(end_offset as usize, contents.len()); | ||
| 338 | let block_data = contents[start_offset as usize..end].to_vec(); | ||
| 339 | println!("sending tftp data packet with {} bytes", block_data.len()); | ||
| 340 | Some(TftpPacket::Data(TftpDataPacket { | ||
| 341 | block: next_block, | ||
| 342 | data: block_data, | ||
| 343 | })) | ||
| 344 | } else { | ||
| 345 | None | ||
| 346 | } | ||
| 347 | } | ||
| 348 | TftpPacket::OAck(ack) => todo!(), | ||
| 349 | TftpPacket::Error(err) => { | ||
| 350 | println!( | ||
| 351 | "Received ERROR packet: code {}, message: {}", | ||
| 352 | err.code, err.message | ||
| 353 | ); | ||
| 354 | None | ||
| 355 | } | ||
| 356 | }; | ||
| 357 | |||
| 358 | if let Some(response) = response { | ||
| 359 | let mut writer = Cursor::new(&mut buf[..]); | ||
| 360 | response.write(&mut writer).unwrap(); | ||
| 361 | let (response, _) = writer.split(); | ||
| 362 | socket.send_to(&response, addr).unwrap(); | ||
| 363 | } | ||
| 364 | } | ||
| 365 | } | ||
diff --git a/src/wire.rs b/src/wire.rs new file mode 100644 index 0000000..dda7690 --- /dev/null +++ b/src/wire.rs | |||
| @@ -0,0 +1,81 @@ | |||
| 1 | use std::{ | ||
| 2 | io::{BufRead, Read, Result, Write}, | ||
| 3 | net::Ipv4Addr, | ||
| 4 | }; | ||
| 5 | |||
| 6 | pub fn write<W: Write>(mut writer: W, v: &[u8]) -> Result<()> { | ||
| 7 | writer.write_all(v) | ||
| 8 | } | ||
| 9 | |||
| 10 | pub fn write_u8<W: Write>(mut writer: W, v: u8) -> Result<()> { | ||
| 11 | writer.write_all(&[v]) | ||
| 12 | } | ||
| 13 | |||
| 14 | pub fn write_u16<W: Write>(mut writer: W, v: u16) -> Result<()> { | ||
| 15 | writer.write_all(&u16::to_be_bytes(v)) | ||
| 16 | } | ||
| 17 | |||
| 18 | pub fn write_u32<W: Write>(mut writer: W, v: u32) -> Result<()> { | ||
| 19 | writer.write_all(&u32::to_be_bytes(v)) | ||
| 20 | } | ||
| 21 | |||
| 22 | pub fn write_ipv4<W: Write>(mut writer: W, v: Ipv4Addr) -> Result<()> { | ||
| 23 | writer.write_all(&v.octets()) | ||
| 24 | } | ||
| 25 | |||
| 26 | pub fn write_null_terminated_string<W: Write>(mut writer: W, v: &str) -> Result<()> { | ||
| 27 | writer.write_all(v.as_bytes())?; | ||
| 28 | writer.write_all(&[0u8]) | ||
| 29 | } | ||
| 30 | |||
| 31 | pub fn read_u8<R: Read>(mut reader: R) -> Result<u8> { | ||
| 32 | let mut buf = [0u8; 1]; | ||
| 33 | reader.read_exact(&mut buf)?; | ||
| 34 | Ok(buf[0]) | ||
| 35 | } | ||
| 36 | |||
| 37 | pub fn read_u16<R: Read>(mut reader: R) -> Result<u16> { | ||
| 38 | let mut buf = [0u8; 2]; | ||
| 39 | reader.read_exact(&mut buf)?; | ||
| 40 | Ok(u16::from_be_bytes(buf)) | ||
| 41 | } | ||
| 42 | |||
| 43 | pub fn read_u32<R: Read>(mut reader: R) -> Result<u32> { | ||
| 44 | let mut buf = [0u8; 4]; | ||
| 45 | reader.read_exact(&mut buf)?; | ||
| 46 | Ok(u32::from_be_bytes(buf)) | ||
| 47 | } | ||
| 48 | |||
| 49 | pub fn read_arr<const N: usize, R: Read>(mut reader: R) -> Result<[u8; N]> { | ||
| 50 | let mut buf = [0u8; N]; | ||
| 51 | reader.read_exact(&mut buf)?; | ||
| 52 | Ok(buf) | ||
| 53 | } | ||
| 54 | |||
| 55 | pub fn read_null_terminated_vec<R: BufRead>(mut reader: R) -> Result<Vec<u8>> { | ||
| 56 | let mut buf = Vec::default(); | ||
| 57 | reader.read_until(0, &mut buf)?; | ||
| 58 | buf.pop(); | ||
| 59 | Ok(buf) | ||
| 60 | } | ||
| 61 | |||
| 62 | pub fn read_null_terminated_string<R: BufRead>(reader: R) -> Result<String> { | ||
| 63 | let buf = read_null_terminated_vec(reader)?; | ||
| 64 | Ok(String::from_utf8(buf).unwrap()) | ||
| 65 | } | ||
| 66 | |||
| 67 | pub fn read_len8_prefixed_vec<R: BufRead>(mut reader: R) -> Result<Vec<u8>> { | ||
| 68 | let len = read_u8(&mut reader)?; | ||
| 69 | let mut buf = vec![0u8; len as usize]; | ||
| 70 | reader.read_exact(&mut buf)?; | ||
| 71 | Ok(buf) | ||
| 72 | } | ||
| 73 | |||
| 74 | pub fn read_len8_prefixed_string<R: BufRead>(reader: R) -> Result<String> { | ||
| 75 | let buf = read_len8_prefixed_vec(reader)?; | ||
| 76 | Ok(String::from_utf8(buf).unwrap()) | ||
| 77 | } | ||
| 78 | |||
| 79 | pub fn read_ipv4<R: BufRead>(reader: R) -> Result<Ipv4Addr> { | ||
| 80 | Ok(Ipv4Addr::from_octets(read_arr(reader)?)) | ||
| 81 | } | ||
