Skip to main content

thincan_file_transfer/
alloc_encode.rs

1use crate::{Atlas, FileAckValue, FileChunkValue, FileOfferValue, FileReqValue, schema};
2use alloc::vec::Vec;
3
4#[cfg(feature = "std")]
5use core::cell::RefCell;
6
7#[cfg(feature = "std")]
8thread_local! {
9    static CAPNP_SCRATCH: RefCell<Vec<capnp::Word>> = const { RefCell::new(Vec::new()) };
10}
11
12fn with_capnp_scratch_bytes<R>(min_len: usize, f: impl FnOnce(&mut [u8]) -> R) -> R {
13    let min_words = min_len.div_ceil(8);
14
15    #[cfg(feature = "std")]
16    {
17        CAPNP_SCRATCH.with(|cell| {
18            let mut words = cell.borrow_mut();
19            if words.len() < min_words {
20                words.resize(min_words, capnp::word(0, 0, 0, 0, 0, 0, 0, 0));
21            }
22            let bytes = unsafe {
23                core::slice::from_raw_parts_mut(words.as_mut_ptr() as *mut u8, words.len() * 8)
24            };
25            f(&mut bytes[..min_len])
26        })
27    }
28
29    #[cfg(not(feature = "std"))]
30    {
31        let mut words: Vec<capnp::Word> = Vec::new();
32        words.resize(min_words, capnp::word(0, 0, 0, 0, 0, 0, 0, 0));
33        let bytes = unsafe {
34            core::slice::from_raw_parts_mut(words.as_mut_ptr() as *mut u8, words.len() * 8)
35        };
36        f(&mut bytes[..min_len])
37    }
38}
39
40impl<A> thincan::EncodeCapnp<<A as Atlas>::FileReq> for FileReqValue<A>
41where
42    A: Atlas,
43{
44    fn max_encoded_len(&self) -> usize {
45        crate::file_req_max_encoded_len()
46    }
47
48    fn encode(&self, out: &mut [u8]) -> Result<usize, thincan::Error> {
49        let max_len = self.max_encoded_len();
50        if out.len() < max_len {
51            return Err(thincan::Error {
52                kind: thincan::ErrorKind::Other,
53            });
54        }
55        with_capnp_scratch_bytes(max_len, |scratch| {
56            let mut message =
57                capnp::message::Builder::new(capnp::message::SingleSegmentAllocator::new(scratch));
58            let mut root: schema::file_req::Builder = message.init_root();
59            root.set_transfer_id(self.transfer_id);
60            root.set_total_len(self.total_len);
61            root.set_file_metadata(&[]);
62            root.set_sender_max_chunk_size(0);
63            root.set_file_hash_algo(schema::FileHashAlgo::Sha256);
64            root.set_file_hash(&[]);
65
66            let segments = message.get_segments_for_output();
67            if segments.len() != 1 {
68                return Err(thincan::Error {
69                    kind: thincan::ErrorKind::Other,
70                });
71            }
72            let bytes = segments[0];
73            if bytes.len() > max_len {
74                return Err(thincan::Error {
75                    kind: thincan::ErrorKind::Other,
76                });
77            }
78            out[..bytes.len()].copy_from_slice(bytes);
79            Ok(bytes.len())
80        })
81    }
82}
83
84impl<'a, A> thincan::EncodeCapnp<<A as Atlas>::FileReq> for FileOfferValue<'a, A>
85where
86    A: Atlas,
87{
88    fn max_encoded_len(&self) -> usize {
89        crate::file_offer_max_encoded_len(self.file_metadata.len(), self.file_hash.len())
90    }
91
92    fn encode(&self, out: &mut [u8]) -> Result<usize, thincan::Error> {
93        let max_len = self.max_encoded_len();
94        if out.len() < max_len {
95            return Err(thincan::Error {
96                kind: thincan::ErrorKind::Other,
97            });
98        }
99        with_capnp_scratch_bytes(max_len, |scratch| {
100            let mut message =
101                capnp::message::Builder::new(capnp::message::SingleSegmentAllocator::new(scratch));
102            let mut root: schema::file_req::Builder = message.init_root();
103            root.set_transfer_id(self.transfer_id);
104            root.set_total_len(self.total_len);
105            root.set_file_metadata(self.file_metadata);
106            root.set_sender_max_chunk_size(self.sender_max_chunk_size);
107            root.set_file_hash_algo(self.file_hash_algo);
108            root.set_file_hash(self.file_hash);
109
110            let segments = message.get_segments_for_output();
111            if segments.len() != 1 {
112                return Err(thincan::Error {
113                    kind: thincan::ErrorKind::Other,
114                });
115            }
116            let bytes = segments[0];
117            if bytes.len() > max_len {
118                return Err(thincan::Error {
119                    kind: thincan::ErrorKind::Other,
120                });
121            }
122            out[..bytes.len()].copy_from_slice(bytes);
123            Ok(bytes.len())
124        })
125    }
126}
127
128impl<'a, A> thincan::EncodeCapnp<<A as Atlas>::FileChunk> for FileChunkValue<'a, A>
129where
130    A: Atlas,
131{
132    fn max_encoded_len(&self) -> usize {
133        crate::file_chunk_max_encoded_len(self.data.len())
134    }
135
136    fn encode(&self, out: &mut [u8]) -> Result<usize, thincan::Error> {
137        let max_len = self.max_encoded_len();
138        if out.len() < max_len {
139            return Err(thincan::Error {
140                kind: thincan::ErrorKind::Other,
141            });
142        }
143        with_capnp_scratch_bytes(max_len, |scratch| {
144            let mut message =
145                capnp::message::Builder::new(capnp::message::SingleSegmentAllocator::new(scratch));
146            let mut root: schema::file_chunk::Builder = message.init_root();
147            root.set_transfer_id(self.transfer_id);
148            root.set_offset(self.offset);
149            root.set_data(self.data);
150
151            let segments = message.get_segments_for_output();
152            if segments.len() != 1 {
153                return Err(thincan::Error {
154                    kind: thincan::ErrorKind::Other,
155                });
156            }
157            let bytes = segments[0];
158            if bytes.len() > max_len {
159                return Err(thincan::Error {
160                    kind: thincan::ErrorKind::Other,
161                });
162            }
163            out[..bytes.len()].copy_from_slice(bytes);
164            Ok(bytes.len())
165        })
166    }
167}
168
169impl<A> thincan::EncodeCapnp<<A as Atlas>::FileAck> for FileAckValue<A>
170where
171    A: Atlas,
172{
173    fn max_encoded_len(&self) -> usize {
174        crate::file_ack_max_encoded_len()
175    }
176
177    fn encode(&self, out: &mut [u8]) -> Result<usize, thincan::Error> {
178        let max_len = self.max_encoded_len();
179        if out.len() < max_len {
180            return Err(thincan::Error {
181                kind: thincan::ErrorKind::Other,
182            });
183        }
184        with_capnp_scratch_bytes(max_len, |scratch| {
185            let mut message =
186                capnp::message::Builder::new(capnp::message::SingleSegmentAllocator::new(scratch));
187            let mut root: schema::file_ack::Builder = message.init_root();
188            root.set_transfer_id(self.transfer_id);
189            root.set_kind(self.kind);
190            root.set_next_offset(self.next_offset);
191            root.set_chunk_size(self.chunk_size);
192            root.set_error(self.error);
193
194            let segments = message.get_segments_for_output();
195            if segments.len() != 1 {
196                return Err(thincan::Error {
197                    kind: thincan::ErrorKind::Other,
198                });
199            }
200            let bytes = segments[0];
201            if bytes.len() > max_len {
202                return Err(thincan::Error {
203                    kind: thincan::ErrorKind::Other,
204                });
205            }
206            out[..bytes.len()].copy_from_slice(bytes);
207            Ok(bytes.len())
208        })
209    }
210}