vortex_protocol/tlv/
download_persistent_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/// DownloadPersistentCachePiece represents a download persistent 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 DownloadPersistentCachePiece {
40    task_id: String,
41    piece_number: u32,
42}
43
44/// DownloadPersistentCachePiece implements the DownloadPersistentCachePiece functions.
45impl DownloadPersistentCachePiece {
46    /// new creates a new DownloadPersistentCachePiece 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 persistent 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 persistent 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 DownloadPersistentCachePiece for conversion from a byte slice.
76impl TryFrom<Bytes> for DownloadPersistentCachePiece {
77    type Error = Error;
78
79    /// try_from decodes the download persistent 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 DownloadPersistentCachePiece, got {}",
84                TASK_ID_SIZE + PIECE_NUMBER_SIZE,
85                bytes.len()
86            )));
87        }
88
89        Ok(DownloadPersistentCachePiece {
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<DownloadPersistentCachePiece> for Bytes for conversion to a byte slice.
111impl From<DownloadPersistentCachePiece> for Bytes {
112    /// from converts the download persistent cache piece request to a byte slice.
113    fn from(piece: DownloadPersistentCachePiece) -> 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_persistent_cache_piece =
131            DownloadPersistentCachePiece::new(task_id.clone(), piece_number);
132
133        assert_eq!(download_persistent_cache_piece.task_id(), task_id);
134        assert_eq!(download_persistent_cache_piece.piece_number(), piece_number);
135        assert_eq!(
136            download_persistent_cache_piece.len(),
137            TASK_ID_SIZE + PIECE_NUMBER_SIZE
138        );
139    }
140
141    #[test]
142    fn test_is_empty() {
143        let download_persistent_cache_piece_empty =
144            DownloadPersistentCachePiece::new("".to_string(), 0);
145        let download_persistent_cache_piece_non_empty =
146            DownloadPersistentCachePiece::new("a".repeat(32), 1);
147
148        assert!(download_persistent_cache_piece_empty.is_empty());
149        assert!(!download_persistent_cache_piece_non_empty.is_empty());
150    }
151
152    #[test]
153    fn test_valid_conversion() {
154        let task_id = "a".repeat(64);
155        let piece_number = 42;
156        let download_persistent_cache_piece =
157            DownloadPersistentCachePiece::new(task_id.clone(), piece_number);
158
159        let bytes: Bytes = download_persistent_cache_piece.into();
160        let download_persistent_cache_piece =
161            DownloadPersistentCachePiece::try_from(bytes).unwrap();
162
163        assert_eq!(download_persistent_cache_piece.task_id(), task_id);
164        assert_eq!(download_persistent_cache_piece.piece_number(), piece_number);
165    }
166
167    #[test]
168    fn test_invalid_conversion() {
169        let invalid_bytes =
170            Bytes::from("c993dfb0ecfbe1b4e158891bafff709e5d29d3fcd522e09b183aeb5db1db50111111111");
171        let result = DownloadPersistentCachePiece::try_from(invalid_bytes);
172        assert!(result.is_err());
173        assert!(matches!(result.unwrap_err(), Error::InvalidLength(_)));
174
175        let invalid_bytes = Bytes::from("task_id");
176        let result = DownloadPersistentCachePiece::try_from(invalid_bytes);
177        assert!(result.is_err());
178        assert!(matches!(result.unwrap_err(), Error::InvalidLength(_)));
179
180        let invalid_bytes = Bytes::from("");
181        let result = DownloadPersistentCachePiece::try_from(invalid_bytes);
182        assert!(result.is_err());
183        assert!(matches!(result.unwrap_err(), Error::InvalidLength(_)));
184    }
185}