Don't hardcode opcode values everywhere

3 files changed, 23 insertions, 12 deletions

diff --git a/src/lib.rs b/src/lib.rs
index b1ff232..9d47cff 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -11,6 +11,16 @@ use std::io;
 use std::io::prelude::*;
 use std::time::Duration;
 
+#[repr(u16)]
+pub enum Opcodes {
+    RRQ   = 0x01,
+    WRQ   = 0x02,
+    DATA  = 0x03,
+    ACK   = 0x04,
+    ERROR = 0x05,
+    OACK  = 0x06,
+}
+
 pub struct TftpOptions {
     blksize: usize,
     timeout: u8,
@@ -63,7 +73,7 @@ impl Tftp {
         let opcode = u16::from_be_bytes([buf[0], buf[1]]);
         let block_nr = u16::from_be_bytes([buf[2], buf[3]]);
 
-        if opcode == 4 && block_nr == expected_block {
+        if opcode == Opcodes::ACK as u16 && block_nr == expected_block {
             return Ok(true)
         }
 
@@ -80,7 +90,7 @@ impl Tftp {
         }
 
         let mut buf = Vec::with_capacity(512);
-        buf.extend([0x00, 0x06].iter());  // opcode
+        buf.extend((Opcodes::OACK as u16).to_be_bytes().iter());
 
         for (key, val) in options {
             buf.extend(key.bytes());
@@ -186,7 +196,8 @@ impl Tftp {
     }
 
     pub fn send_error(&self, socket: &UdpSocket, code: u16, msg: &str) -> Result<(), io::Error> {
-        let mut buf = vec![0x00, 0x05];  // opcode
+        let mut buf = Vec::with_capacity(512);
+        buf.extend((Opcodes::ERROR as u16).to_be_bytes().iter());
         buf.extend(code.to_be_bytes().iter());
         buf.extend(msg.as_bytes());
 
@@ -195,11 +206,11 @@ impl Tftp {
     }
 
     pub fn send_ack(&self, sock: &UdpSocket, block_nr: u16) -> Result<(), io::Error> {
-        let mut buf = vec![0x00, 0x04];  // opcode
+        let mut buf = Vec::with_capacity(4);
+        buf.extend((Opcodes::ACK as u16).to_be_bytes().iter());
         buf.extend(block_nr.to_be_bytes().iter());
 
         sock.send(&buf)?;
-
         Ok(())
     }
 
@@ -218,7 +229,7 @@ impl Tftp {
             };
 
             let mut sendbuf = Vec::with_capacity(4 + len);
-            sendbuf.extend([0x00, 0x03].iter());  // opcode
+            sendbuf.extend((Opcodes::DATA as u16).to_be_bytes().iter());
             sendbuf.extend(block_nr.to_be_bytes().iter());
             sendbuf.extend(filebuf[0..len].iter());
 
@@ -276,7 +287,7 @@ impl Tftp {
             }
 
             let _opcode = match u16::from_be_bytes([buf[0], buf[1]]) {
-                3 /* DATA */ => (),
+                opc if opc == Opcodes::DATA as u16 => (),
                 _ => return Err(io::Error::new(io::ErrorKind::Other, "unexpected opcode")),
             };
             if u16::from_be_bytes([buf[2], buf[3]]) != block_nr {
diff --git a/src/tftpc.rs b/src/tftpc.rs
index 15cadc3..8815bb7 100644
--- a/src/tftpc.rs
+++ b/src/tftpc.rs
@@ -78,7 +78,7 @@ impl Tftpc {
         let mut remote = None;
         for _ in 1 .. 3 {
             sock.send_to(&buf, self.conf.remote)?;
-            remote = self.wait_for_response(&sock, 4, 0);
+            remote = self.wait_for_response(&sock, rtftp::Opcodes::ACK as u16, 0);
             if let Some(_) = remote {
                 break;
             }
@@ -121,7 +121,7 @@ impl Tftpc {
         let mut remote = None;
         for _ in 1 .. 3 {
             sock.send_to(&buf, self.conf.remote)?;
-            remote = self.wait_for_response(&sock, 3, 1);
+            remote = self.wait_for_response(&sock, rtftp::Opcodes::DATA as u16, 1);
             if let Some(_) = remote {
                 break;
             }
diff --git a/src/tftpd.rs b/src/tftpd.rs
index 0b06310..48e374e 100644
--- a/src/tftpd.rs
+++ b/src/tftpd.rs
@@ -168,7 +168,7 @@ impl Tftpd {
         socket.connect(cl)?;
 
         let _opcode = match u16::from_be_bytes([buf[0], buf[1]]) {
-            1 /* RRQ */ => {
+            o if o == rtftp::Opcodes::RRQ  as u16 => {
                 if self.conf.wo {
                     self.tftp.send_error(&socket, 4, "reading not allowed")?;
                     return Err(io::Error::new(io::ErrorKind::Other, "unallowed mode"));
@@ -176,7 +176,7 @@ impl Tftpd {
                     self.handle_rrq(&socket, &cl, &buf[2..])?;
                 }
             },
-            2 /* WRQ */ => {
+            o if o == rtftp::Opcodes::WRQ as u16 => {
                 if self.conf.ro {
                     self.tftp.send_error(&socket, 4, "writing not allowed")?;
                     return Err(io::Error::new(io::ErrorKind::Other, "unallowed mode"));
@@ -184,7 +184,7 @@ impl Tftpd {
                     self.handle_wrq(&socket, &cl, &buf[2..])?;
                 }
             },
-            5 /* ERROR */ => println!("Received ERROR from {}", cl),
+            o if o == rtftp::Opcodes::ERROR as u16 => println!("Received ERROR from {}", cl),
             _ => {
                 self.tftp.send_error(&socket, 4, "Unexpected opcode")?;
                 return Err(io::Error::new(io::ErrorKind::Other, "unexpected opcode"));