vortex_protocol/tlv/
download_cache_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/// DownloadCachePiece represents a download cache piece request.
28///
29/// Value Format:
30///   - Task ID (64 bytes): SHA-256 hash of the task ID.
31///   - Piece Number (4 bytes): Piece number to download.
32///
33/// ```text
34/// -----------------------------------------------
35/// | Task ID (64 bytes) | Piece Number (4 bytes) |
36/// -----------------------------------------------
37/// ```
38#[derive(Debug, Clone)]
39pub struct DownloadCachePiece {
40    task_id: String,
41    piece_number: u32,
42}
43
44/// DownloadCachePiece implements the DownloadCachePiece functions.
45impl DownloadCachePiece {
46    /// new creates a new DownloadPiece request.
47    pub fn new(task_id: String, piece_number: u32) -> Self {
48        Self {
49            task_id,
50            piece_number,
51        }
52    }
53
54    /// task_id returns the task ID.
55    pub fn task_id(&self) -> &str {
56        &self.task_id
57    }
58
59    /// piece_number returns the piece number.
60    pub fn piece_number(&self) -> u32 {
61        self.piece_number
62    }
63
64    /// len returns the length of the download cache piece request.
65    pub fn len(&self) -> usize {
66        TASK_ID_SIZE + PIECE_NUMBER_SIZE
67    }
68
69    /// is_empty returns whether the download cache piece request is empty.
70    pub fn is_empty(&self) -> bool {
71        self.task_id.is_empty()
72    }
73}
74
75/// Implement TryFrom<Bytes> for DownloadCachePiece for conversion from a byte slice.
76impl TryFrom<Bytes> for DownloadCachePiece {
77    type Error = Error;
78
79    /// try_from decodes the download cache piece request from the byte slice.
80    fn try_from(bytes: Bytes) -> Result<Self> {
81        if bytes.len() != TASK_ID_SIZE + PIECE_NUMBER_SIZE {
82            return Err(Error::InvalidLength(format!(
83                "expected {} bytes for DownloadCachePiece, got {}",
84                TASK_ID_SIZE + PIECE_NUMBER_SIZE,
85                bytes.len()
86            )));
87        }
88
89        Ok(DownloadCachePiece {
90            task_id: String::from_utf8(
91                bytes
92                    .get(..TASK_ID_SIZE)
93                    .ok_or(Error::InvalidPacket(
94                        "insufficient bytes for task id".to_string(),
95                    ))?
96                    .to_vec(),
97            )?,
98            piece_number: u32::from_be_bytes(
99                bytes
100                    .get(TASK_ID_SIZE..TASK_ID_SIZE + PIECE_NUMBER_SIZE)
101                    .ok_or(Error::InvalidPacket(
102                        "insufficient bytes for piece number".to_string(),
103                    ))?
104                    .try_into()?,
105            ),
106        })
107    }
108}
109
110/// Implement From<DownloadCachePiece> for Bytes for conversion to a byte slice.
111impl From<DownloadCachePiece> for Bytes {
112    /// from converts the download piece request to a byte slice.
113    fn from(piece: DownloadCachePiece) -> Self {
114        let mut bytes = BytesMut::with_capacity(piece.len());
115        bytes.extend_from_slice(piece.task_id.as_bytes());
116        bytes.put_u32(piece.piece_number);
117        bytes.freeze()
118    }
119}
120
121#[cfg(test)]
122mod tests {
123    use super::*;
124    use bytes::Bytes;
125
126    #[test]
127    fn test_new() {
128        let task_id = "a".repeat(64);
129        let piece_number = 42;
130        let download_cache_piece = DownloadCachePiece::new(task_id.clone(), piece_number);
131
132        assert_eq!(download_cache_piece.task_id(), task_id);
133        assert_eq!(download_cache_piece.piece_number(), piece_number);
134        assert_eq!(download_cache_piece.len(), TASK_ID_SIZE + PIECE_NUMBER_SIZE);
135    }
136
137    #[test]
138    fn test_is_empty() {
139        let download_cache_piece_empty = DownloadCachePiece::new("".to_string(), 0);
140        let download_cache_piece_non_empty = DownloadCachePiece::new("a".repeat(32), 1);
141
142        assert!(download_cache_piece_empty.is_empty());
143        assert!(!download_cache_piece_non_empty.is_empty());
144    }
145
146    #[test]
147    fn test_valid_conversion() {
148        let task_id = "a".repeat(64);
149        let piece_number = 42;
150        let download_cache_piece = DownloadCachePiece::new(task_id.clone(), piece_number);
151
152        let bytes: Bytes = download_cache_piece.into();
153        let download_cache_piece = DownloadCachePiece::try_from(bytes).unwrap();
154
155        assert_eq!(download_cache_piece.task_id(), task_id);
156        assert_eq!(download_cache_piece.piece_number(), piece_number);
157    }
158
159    #[test]
160    fn test_invalid_conversion() {
161        let invalid_bytes =
162            Bytes::from("c993dfb0ecfbe1b4e158891bafff709e5d29d3fcd522e09b183aeb5db1db50111111111");
163        let result = DownloadCachePiece::try_from(invalid_bytes);
164        assert!(result.is_err());
165        assert!(matches!(result.unwrap_err(), Error::InvalidLength(_)));
166
167        let invalid_bytes = Bytes::from("task_id");
168        let result = DownloadCachePiece::try_from(invalid_bytes);
169        assert!(result.is_err());
170        assert!(matches!(result.unwrap_err(), Error::InvalidLength(_)));
171
172        let invalid_bytes = Bytes::from("");
173        let result = DownloadCachePiece::try_from(invalid_bytes);
174        assert!(result.is_err());
175        assert!(matches!(result.unwrap_err(), Error::InvalidLength(_)));
176    }
177}
vortex_protocol/tlv/download_cache_piece.rs

vortex_protocol/tlv/
download_cache_piece.rs
