vortex_protocol/tlv/
download_persistent_piece.rs

1/*
2 *     Copyright 2025 The Dragonfly Authors
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 *      http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17use crate::error::{Error, Result};
18use bytes::{BufMut, Bytes, BytesMut};
19use std::convert::TryFrom;
20
21/// TASK_ID_SIZE is the size of the task ID in bytes.
22pub const TASK_ID_SIZE: usize = 64;
23
24/// PIECE_NUMBER_SIZE is the size of the piece number in bytes.
25pub const PIECE_NUMBER_SIZE: usize = 4;
26
27/// DownloadPersistentPiece represents a download persistent piece request.
28/// Value Format:
29///   - Task ID (64 bytes): SHA-256 hash of the task ID.
30///   - Piece Number (4 bytes): Piece number to download.
31///
32/// ```text
33/// -----------------------------------------------
34/// | Task ID (64 bytes) | Piece Number (4 bytes) |
35/// -----------------------------------------------
36/// ```
37#[derive(Debug, Clone)]
38pub struct DownloadPersistentPiece {
39    task_id: String,
40    piece_number: u32,
41}
42
43/// DownloadPersistentPiece implements the DownloadPersistentPiece functions.
44impl DownloadPersistentPiece {
45    /// new creates a new DownloadPersistentPiece request.
46    pub fn new(task_id: String, piece_number: u32) -> Self {
47        Self {
48            task_id,
49            piece_number,
50        }
51    }
52
53    /// task_id returns the task ID.
54    pub fn task_id(&self) -> &str {
55        &self.task_id
56    }
57
58    /// piece_number returns the piece number.
59    pub fn piece_number(&self) -> u32 {
60        self.piece_number
61    }
62
63    /// len returns the length of the download persistent piece request.
64    pub fn len(&self) -> usize {
65        TASK_ID_SIZE + PIECE_NUMBER_SIZE
66    }
67
68    /// is_empty returns whether the download persistent piece request is empty.
69    pub fn is_empty(&self) -> bool {
70        self.task_id.is_empty()
71    }
72}
73
74/// Implement TryFrom<Bytes> for DownloadPersistentPiece for conversion from a byte slice.
75impl TryFrom<Bytes> for DownloadPersistentPiece {
76    type Error = Error;
77
78    /// try_from decodes the download persistent piece request from the byte slice.
79    fn try_from(bytes: Bytes) -> Result<Self> {
80        if bytes.len() != TASK_ID_SIZE + PIECE_NUMBER_SIZE {
81            return Err(Error::InvalidLength(format!(
82                "expected {} bytes for DownloadPersistentPiece, got {}",
83                TASK_ID_SIZE + PIECE_NUMBER_SIZE,
84                bytes.len()
85            )));
86        }
87
88        Ok(DownloadPersistentPiece {
89            task_id: String::from_utf8(
90                bytes
91                    .get(..TASK_ID_SIZE)
92                    .ok_or(Error::InvalidPacket(
93                        "insufficient bytes for task id".to_string(),
94                    ))?
95                    .to_vec(),
96            )?,
97            piece_number: u32::from_be_bytes(
98                bytes
99                    .get(TASK_ID_SIZE..TASK_ID_SIZE + PIECE_NUMBER_SIZE)
100                    .ok_or(Error::InvalidPacket(
101                        "insufficient bytes for piece number".to_string(),
102                    ))?
103                    .try_into()?,
104            ),
105        })
106    }
107}
108
109/// Implement From<DownloadPersistentPiece> for Bytes for conversion to a byte slice.
110impl From<DownloadPersistentPiece> for Bytes {
111    /// from converts the download persistent piece request to a byte slice.
112    fn from(piece: DownloadPersistentPiece) -> Self {
113        let mut bytes = BytesMut::with_capacity(piece.len());
114        bytes.extend_from_slice(piece.task_id.as_bytes());
115        bytes.put_u32(piece.piece_number);
116        bytes.freeze()
117    }
118}
119
120#[cfg(test)]
121mod tests {
122    use super::*;
123    use bytes::Bytes;
124
125    #[test]
126    fn test_new() {
127        let task_id = "a".repeat(64);
128        let piece_number = 42;
129        let download_persistent_piece = DownloadPersistentPiece::new(task_id.clone(), piece_number);
130
131        assert_eq!(download_persistent_piece.task_id(), task_id);
132        assert_eq!(download_persistent_piece.piece_number(), piece_number);
133        assert_eq!(
134            download_persistent_piece.len(),
135            TASK_ID_SIZE + PIECE_NUMBER_SIZE
136        );
137    }
138
139    #[test]
140    fn test_is_empty() {
141        let download_persistent_piece_empty = DownloadPersistentPiece::new("".to_string(), 0);
142        let download_persistent_piece_non_empty = DownloadPersistentPiece::new("a".repeat(32), 1);
143
144        assert!(download_persistent_piece_empty.is_empty());
145        assert!(!download_persistent_piece_non_empty.is_empty());
146    }
147
148    #[test]
149    fn test_valid_conversion() {
150        let task_id = "a".repeat(64);
151        let piece_number = 42;
152        let download_persistent_piece = DownloadPersistentPiece::new(task_id.clone(), piece_number);
153
154        let bytes: Bytes = download_persistent_piece.into();
155        let download_persistent_piece = DownloadPersistentPiece::try_from(bytes).unwrap();
156
157        assert_eq!(download_persistent_piece.task_id(), task_id);
158        assert_eq!(download_persistent_piece.piece_number(), piece_number);
159    }
160
161    #[test]
162    fn test_invalid_conversion() {
163        let invalid_bytes =
164            Bytes::from("c993dfb0ecfbe1b4e158891bafff709e5d29d3fcd522e09b183aeb5db1db50111111111");
165        let result = DownloadPersistentPiece::try_from(invalid_bytes);
166        assert!(result.is_err());
167        assert!(matches!(result.unwrap_err(), Error::InvalidLength(_)));
168
169        let invalid_bytes = Bytes::from("task_id");
170        let result = DownloadPersistentPiece::try_from(invalid_bytes);
171        assert!(result.is_err());
172        assert!(matches!(result.unwrap_err(), Error::InvalidLength(_)));
173
174        let invalid_bytes = Bytes::from("");
175        let result = DownloadPersistentPiece::try_from(invalid_bytes);
176        assert!(result.is_err());
177        assert!(matches!(result.unwrap_err(), Error::InvalidLength(_)));
178    }
179}
vortex_protocol/tlv/download_persistent_piece.rs

vortex_protocol/tlv/
download_persistent_piece.rs
