Skip to content

Commit

Permalink
implement From trait for tlv struct
Browse files Browse the repository at this point in the history 
Signed-off-by: cormick <[email protected]>
  • Loading branch information
@CormickKneey
CormickKneey committed Jan 16, 2025
1 parent d1d8ace commit 08b6d72
Show file tree
Hide file tree
Showing 4 changed files with 104 additions and 68 deletions.
25 changes: 14 additions & 11 deletions src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -79,7 +79,7 @@ impl Vortex {
)));
}

let mut rng = thread_rng();
let mut rng: ThreadRng = thread_rng();
let header = Header {
packet_id: rng.gen(),
tag,
Expand All @@ -88,16 +88,16 @@ impl Vortex {

match tag {
tlv::Tag::DownloadPiece => {
let download_piece = tlv::download_piece::DownloadPiece::from_bytes(value)?;
let download_piece = tlv::download_piece::DownloadPiece::try_from(value)?;
Ok(Vortex::DownloadPiece(header, download_piece))
}
tlv::Tag::PieceContent => {
let piece_content = tlv::piece_content::PieceContent::from_bytes(value)?;
let piece_content = tlv::piece_content::PieceContent::try_from(value)?;
Ok(Vortex::PieceContent(header, piece_content))
}
tlv::Tag::Reserved(_) => Ok(Vortex::Reserved(header)),
tlv::Tag::Error => {
let err = tlv::error::Error::from_bytes(value)?;
let err = tlv::error::Error::try_from(value)?;
Ok(Vortex::Error(header, err))
}
}
Expand Down Expand Up @@ -179,17 +179,16 @@ impl Vortex {

match tag {
tlv::Tag::DownloadPiece => {
let download_piece =
tlv::download_piece::DownloadPiece::from_bytes(value.freeze())?;
let download_piece = tlv::download_piece::DownloadPiece::try_from(value.freeze())?;
Ok(Vortex::DownloadPiece(header, download_piece))
}
tlv::Tag::PieceContent => {
let piece_content = tlv::piece_content::PieceContent::from_bytes(value.freeze())?;
let piece_content = tlv::piece_content::PieceContent::try_from(value.freeze())?;
Ok(Vortex::PieceContent(header, piece_content))
}
tlv::Tag::Reserved(_) => Ok(Vortex::Reserved(header)),
tlv::Tag::Error => {
let error = tlv::error::Error::from_bytes(value.freeze())?;
let error = tlv::error::Error::try_from(value.freeze())?;
Ok(Vortex::Error(header, error))
}
}
Expand All @@ -198,10 +197,14 @@ impl Vortex {
/// to_bytes converts the Vortex packet to a byte slice.
pub fn to_bytes(&self) -> Bytes {
let (header, value) = match self {
Vortex::DownloadPiece(header, download_piece) => (header, download_piece.to_bytes()),
Vortex::PieceContent(header, piece_content) => (header, piece_content.to_bytes()),
Vortex::DownloadPiece(header, download_piece) => {
(header, Into::<Bytes>::into(download_piece.clone()))
}
Vortex::PieceContent(header, piece_content) => {
(header, Into::<Bytes>::into(piece_content.clone()))
}
Vortex::Reserved(header) => (header, Bytes::new()),
Vortex::Error(header, err) => (header, err.to_bytes()),
Vortex::Error(header, err) => (header, Into::<Bytes>::into(err.clone())),
};

let mut bytes = BytesMut::with_capacity(HEADER_SIZE + value.len());
Expand Down
56 changes: 36 additions & 20 deletions src/tlv/download_piece.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@

use crate::error::{Error, Result};
use bytes::{BufMut, Bytes, BytesMut};
use std::convert::TryFrom;

/// SEPARATOR is the separator character used in the download piece request, separating the task ID
/// and piece number. It is a hyphen character '-'.
Expand Down Expand Up @@ -72,9 +73,13 @@ impl DownloadPiece {
pub fn is_empty(&self) -> bool {
self.task_id.is_empty()
}
}

/// Implement TryFrom<Bytes> for DownloadPiece.
impl TryFrom<Bytes> for DownloadPiece {
type Error = Error;

/// from_bytes creates a download piece request from a byte slice.
pub fn from_bytes(bytes: Bytes) -> Result<Self> {
fn try_from(bytes: Bytes) -> Result<Self> {
let mut parts = bytes.splitn(2, |&b| b == SEPARATOR);
let task_id = std::str::from_utf8(
parts
Expand All @@ -96,13 +101,15 @@ impl DownloadPiece {
piece_number,
})
}
}

/// to_bytes converts the download piece request to a byte slice.
pub fn to_bytes(&self) -> Bytes {
/// Implement From<DownloadPiece> for Bytes.
impl From<DownloadPiece> for Bytes {
fn from(piece: DownloadPiece) -> Self {
let mut bytes = BytesMut::with_capacity(DOWNLOAD_PIECE_SIZE);
bytes.extend_from_slice(self.task_id.as_bytes());
bytes.extend_from_slice(piece.task_id.as_bytes());
bytes.put_u8(SEPARATOR);
bytes.extend_from_slice(self.piece_number.to_string().as_bytes());
bytes.extend_from_slice(piece.piece_number.to_string().as_bytes());
bytes.freeze()
}
}
Expand Down Expand Up @@ -133,30 +140,39 @@ mod tests {
}

#[test]
fn test_to_bytes_and_from_bytes() {
fn test_valid_conversion() {
let task_id = "a".repeat(32);
let piece_number = 42;
let download_piece = DownloadPiece::new(task_id.clone(), piece_number);

let bytes = download_piece.to_bytes();
let download_piece_decoded = DownloadPiece::from_bytes(bytes).unwrap();
// Test From<DownloadPiece> for Bytes
let bytes: Bytes = download_piece.into();

// Test TryFrom<Bytes> for DownloadPiece
let download_piece_decoded = DownloadPiece::try_from(bytes).unwrap();

assert_eq!(download_piece_decoded.task_id(), task_id);
assert_eq!(download_piece_decoded.piece_number(), piece_number);
}

#[test]
fn test_from_bytes_invalid_input() {
// Test missing separator.
fn test_invalid_conversion() {
// Test missing separator
let invalid_bytes = Bytes::from("invalid_input_without_separator");
assert!(DownloadPiece::from_bytes(invalid_bytes).is_err());

// Test missing piece number.
let invalid_bytes = Bytes::from(format!("{}-", "a".repeat(32)));
assert!(DownloadPiece::from_bytes(invalid_bytes).is_err());

// Test invalid piece number format.
let invalid_bytes = Bytes::from(format!("{}-invalid", "a".repeat(32)));
assert!(DownloadPiece::from_bytes(invalid_bytes).is_err());
let result = DownloadPiece::try_from(invalid_bytes);
assert!(result.is_err());
assert!(matches!(result.unwrap_err(), Error::InvalidPacket(_)));

// Test missing piece number
let invalid_bytes = Bytes::from("task_id-");
let result = DownloadPiece::try_from(invalid_bytes);
assert!(result.is_err());
assert!(matches!(result.unwrap_err(), Error::ParseIntError(_)));

// Test invalid piece number
let invalid_bytes = Bytes::from("task_id-invalid");
let result = DownloadPiece::try_from(invalid_bytes);
assert!(result.is_err());
assert!(matches!(result.unwrap_err(), Error::ParseIntError(_)));
}
}
38 changes: 22 additions & 16 deletions src/tlv/error.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -116,9 +116,24 @@ impl Error {
pub fn is_empty(&self) -> bool {
self.len() == 0
}
}


impl From<Error> for Bytes {
// Converts a Error request to a byte slice.
fn from(err: Error) -> Self {
let mut bytes = BytesMut::with_capacity(err.len());
bytes.put_u8(err.code.into());
bytes.put_u8(SEPARATOR);
bytes.extend_from_slice(err.message.as_bytes());
bytes.freeze()
}
}

impl TryFrom<Bytes> for Error {
type Error = VortexError;

/// from_bytes creates a error request from a byte slice.
pub fn from_bytes(bytes: Bytes) -> VortexResult<Self> {
fn try_from(bytes: Bytes) -> VortexResult<Self> {
let mut parts = bytes.splitn(2, |&b| b == SEPARATOR);

let code = parts
Expand Down Expand Up @@ -149,15 +164,6 @@ impl Error {

Ok(Error { code, message })
}

/// to_bytes converts the error request to a byte slice.
pub fn to_bytes(&self) -> Bytes {
let mut bytes = BytesMut::with_capacity(self.len());
bytes.put_u8(self.code.into());
bytes.put_u8(SEPARATOR);
bytes.extend_from_slice(self.message.as_bytes());
bytes.freeze()
}
}

#[cfg(test)]
Expand Down Expand Up @@ -188,8 +194,8 @@ mod tests {
let message = "Resource not found".to_string();
let error = Error::new(code, message.clone());

let bytes = error.to_bytes();
let error_decoded = Error::from_bytes(bytes).unwrap();
let bytes: Bytes = error.into();
let error_decoded = Error::try_from(bytes).unwrap();

assert_eq!(error_decoded.code(), code);
assert_eq!(error_decoded.message(), message);
Expand All @@ -199,14 +205,14 @@ mod tests {
fn test_from_bytes_invalid_input() {
// Test missing separator
let invalid_bytes = Bytes::from("invalid_input_without_separator");
assert!(Error::from_bytes(invalid_bytes).is_err());
assert!(Error::try_from(invalid_bytes).is_err());

// Test missing error message
let invalid_bytes = Bytes::from(format!("{}:", 1));
assert!(Error::from_bytes(invalid_bytes).is_err());
assert!(Error::try_from(invalid_bytes).is_err());

// Test invalid error code format
let invalid_bytes = Bytes::from(format!("{}:{}", 256, "Invalid code"));
assert!(Error::from_bytes(invalid_bytes).is_err());
assert!(Error::try_from(invalid_bytes).is_err());
}
}
53 changes: 32 additions & 21 deletions src/tlv/piece_content.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@

use crate::error::{Error, Result};
use bytes::Bytes;
use std::convert::TryFrom;

/// MAX_PIECE_SIZE is the maximum size of a piece content (4 GiB).
const MAX_PIECE_SIZE: usize = crate::MAX_VALUE_SIZE;
Expand Down Expand Up @@ -49,15 +50,21 @@ impl PieceContent {
pub fn is_empty(&self) -> bool {
self.0.is_empty()
}
}

/// Implement TryFrom<Bytes> for PieceContent.
impl TryFrom<Bytes> for PieceContent {
type Error = Error;

/// from_bytes creates a piece content from a byte slice.
pub fn from_bytes(bytes: Bytes) -> Result<Self> {
fn try_from(bytes: Bytes) -> Result<Self> {
Self::new(bytes)
}
}

/// to_bytes converts the piece content to a byte slice.
pub fn to_bytes(&self) -> Bytes {
self.0.clone()
/// Implement From<PieceContent> for Bytes.
impl From<PieceContent> for Bytes {
fn from(piece: PieceContent) -> Self {
piece.0
}
}

Expand All @@ -82,26 +89,30 @@ mod tests {
}

#[test]
fn test_to_bytes_and_from_bytes() {
fn test_valid_conversion() {
let content = vec![1, 2, 3, 4];
let piece_content = PieceContent::new(content.clone().into()).unwrap();
let bytes = piece_content.to_bytes();
let result = PieceContent::from_bytes(bytes);
assert!(result.is_ok());
assert_eq!(result.unwrap().to_bytes(), content);
let bytes = Bytes::from(content.clone());

// Test TryFrom<Bytes>
let piece_content = PieceContent::try_from(bytes.clone()).unwrap();

// Test From<PieceContent> for Bytes
let bytes_back: Bytes = piece_content.clone().into();
assert_eq!(bytes_back, content);
}

#[test]
fn test_from_bytes_invalid_input() {
// Test empty input
let result = PieceContent::from_bytes(Bytes::new());
assert!(result.is_ok());
assert!(result.unwrap().is_empty());

// Test oversize input
fn test_invalid_conversion() {
// Create bytes larger than MAX_PIECE_SIZE
let large_content = vec![0; MAX_PIECE_SIZE + 1];
let result: std::result::Result<PieceContent, Error> =
PieceContent::from_bytes(large_content.into());
assert!(matches!(result, Err(Error::InvalidLength(_))));
let bytes = Bytes::from(large_content);

// Test conversion fails
let result = PieceContent::try_from(bytes);
assert!(result.is_err());
match result {
Err(Error::InvalidLength(_)) => (),
_ => panic!("Expected InvalidLength error"),
}
}
}

0 comments on commit 08b6d72

Please sign in to comment.