shellcoder/
ops.rs

1//! All operations available for writing shellcodes.
2
3use core::fmt;
4#[cfg(feature = "std")]
5use std::io;
6
7use crate::prelude::*;
8
9#[cfg(feature = "serde")]
10pub trait WithOrWithoutSerde: Serialize + for<'de> Deserialize<'de> {}
11#[cfg(feature = "serde")]
12impl<T> WithOrWithoutSerde for T where T: Serialize + for<'de> Deserialize<'de> {}
13
14#[cfg(not(feature = "serde"))]
15pub trait WithOrWithoutSerde {}
16#[cfg(not(feature = "serde"))]
17impl<T> WithOrWithoutSerde for T {}
18
19/// An operation that moves the cursor ahead by n bytes.
20/// The gap will be filled by zeroes.
21#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
22#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
23pub struct Advance(usize);
24
25impl Advance {
26    /// Instantiates a new [`Advance`] to move the cursor ahead by n bytes.
27    #[inline]
28    #[must_use]
29    pub const fn new(n: usize) -> Self {
30        Self(n)
31    }
32}
33
34impl Op for Advance {
35    #[cfg(feature = "std")]
36    #[inline]
37    fn write_to_io(&self, stream: &mut dyn io::Write) -> Result<usize> {
38        Fill::new(self.0, 0).write_to_io(stream)
39    }
40
41    #[inline]
42    fn write_to(&self, out: impl AsMut<[u8]>) -> Result<usize> {
43        Fill::new(self.0, 0).write_to(out)
44    }
45}
46
47/// An operation that fills with a value.
48#[derive(Clone, Copy, Debug, PartialEq, Eq)]
49#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
50pub struct Fill(usize, u8);
51
52impl Fill {
53    /// Instantiates a new [`Fill`].
54    #[inline]
55    #[must_use]
56    pub const fn new(len: usize, chr: u8) -> Self {
57        Self(len, chr)
58    }
59}
60
61impl Op for Fill {
62    #[cfg(feature = "std")]
63    #[inline]
64    fn write_to_io(&self, stream: &mut dyn io::Write) -> Result<usize> {
65        use core::slice;
66        let rchr = slice::from_ref(&self.1);
67        for _ in 0..self.0 {
68            stream.write_all(rchr)?;
69        }
70        Ok(self.0)
71    }
72
73    #[inline]
74    fn write_to(&self, mut out: impl AsMut<[u8]>) -> Result<usize> {
75        out.as_mut()
76            .get_mut(..self.0)
77            .ok_or_else(|| Error::buffer_too_small(self.0))?
78            .fill(self.1);
79        Ok(self.0)
80    }
81}
82
83/// An integer that is encodable.
84pub trait EncodableInteger:
85    Copy + Clone + Sized + fmt::Debug + PartialEq + Eq + Send + Sync + WithOrWithoutSerde
86{
87    /// Returns the number of bytes needed to encode the integer.
88    fn n(self) -> usize;
89
90    /// Writes in big endian.
91    ///
92    /// # Errors
93    ///
94    /// An I/O error may be raised here.
95    #[cfg(feature = "std")]
96    fn write_be_io(self, stream: &mut dyn io::Write) -> Result<()>;
97
98    /// Writes in little endian.
99    ///
100    /// # Errors
101    ///
102    /// An I/O error may be raised here.
103    #[cfg(feature = "std")]
104    fn write_le_io(self, stream: &mut dyn io::Write) -> Result<()>;
105
106    /// Writes in big endian.
107    ///
108    /// # Errors
109    ///
110    /// [`Error::OutputBufferTooSmall`] is raised if `out` cannot contain the encoded
111    /// integer.
112    fn write_be(self, out: impl AsMut<[u8]>) -> Result<()>;
113
114    /// Writes in little endian.
115    ///
116    /// # Errors
117    ///
118    /// [`Error::OutputBufferTooSmall`] is raised if `out` cannot contain the encoded
119    /// integer.
120    fn write_le(self, out: impl AsMut<[u8]>) -> Result<()>;
121}
122
123/// Implements [`EncodableInteger`] for a given type.
124macro_rules! impl_encodable_integer_for {
125    ($i:ident) => {
126        impl EncodableInteger for $i {
127            #[inline]
128            #[must_use]
129            fn n(self) -> usize {
130                ($i::BITS >> 3).try_into().expect("unreachable")
131            }
132
133            #[cfg(feature = "std")]
134            #[inline]
135            fn write_be_io(self, stream: &mut dyn io::Write) -> Result<()> {
136                stream.write_all(&self.to_be_bytes()).map_err(Error::from)
137            }
138
139            #[cfg(feature = "std")]
140            #[inline]
141            fn write_le_io(self, stream: &mut dyn io::Write) -> Result<()> {
142                stream.write_all(&self.to_le_bytes()).map_err(Error::from)
143            }
144
145            #[inline]
146            fn write_be(self, mut out: impl AsMut<[u8]>) -> Result<()> {
147                let n = self.n();
148                let out = out
149                    .as_mut()
150                    .get_mut(..n)
151                    .ok_or(Error::buffer_too_small(n))?;
152                // SAFETY:
153                //
154                // Length of `out` has been checked previously.
155                unsafe {
156                    out.as_mut_ptr().copy_from(self.to_be_bytes().as_ptr(), n);
157                }
158                Ok(())
159            }
160
161            #[inline]
162            fn write_le(self, mut out: impl AsMut<[u8]>) -> Result<()> {
163                let n = self.n();
164                let out = out
165                    .as_mut()
166                    .get_mut(..n)
167                    .ok_or(Error::buffer_too_small(n))?;
168                // SAFETY:
169                //
170                // Length of `out` has been checked previously.
171                unsafe {
172                    out.as_mut_ptr().copy_from(self.to_le_bytes().as_ptr(), n);
173                }
174                Ok(())
175            }
176        }
177    };
178}
179
180impl_encodable_integer_for!(u8);
181impl_encodable_integer_for!(u16);
182impl_encodable_integer_for!(u32);
183impl_encodable_integer_for!(u64);
184
185/// An operation that writes an integer.
186/// The cursor will be moved ahead by n bytes, n depending on the integer's
187/// encoded size.
188#[derive(Clone, Copy, Debug, PartialEq, Eq)]
189#[non_exhaustive]
190pub enum WriteInteger<I>
191where
192    I: EncodableInteger + WithOrWithoutSerde,
193{
194    /// The integer's value, to encode in big-endian.
195    BigEndian(I),
196
197    /// The integer's value, to encode in little-endian.
198    LittleEndian(I),
199}
200
201impl<I> WriteInteger<I>
202where
203    I: EncodableInteger,
204{
205    /// Instantiates a new [`WriteInteger`] to write a big-endian encoded integer.
206    #[inline]
207    #[must_use]
208    pub const fn new_be(value: I) -> Self {
209        Self::BigEndian(value)
210    }
211
212    /// Instantiates a new [`WriteInteger`] to write a little-endian encoded integer.
213    #[inline]
214    #[must_use]
215    pub const fn new_le(value: I) -> Self {
216        Self::LittleEndian(value)
217    }
218}
219
220impl<I> Op for WriteInteger<I>
221where
222    I: EncodableInteger,
223{
224    #[cfg(feature = "std")]
225    #[inline]
226    fn write_to_io(&self, stream: &mut dyn io::Write) -> Result<usize> {
227        match self {
228            Self::BigEndian(n) => n.write_be_io(stream).map(|()| n.n()),
229            Self::LittleEndian(n) => n.write_le_io(stream).map(|()| n.n()),
230        }
231    }
232
233    #[inline]
234    fn write_to(&self, out: impl AsMut<[u8]>) -> Result<usize> {
235        match self {
236            Self::BigEndian(n) => n.write_be(out).map(|()| n.n()),
237            Self::LittleEndian(n) => n.write_le(out).map(|()| n.n()),
238        }
239    }
240}
241
242/// An operation that writes a buffer.
243/// The cursor will be moved ahead by the length in bytes of the given buffer.
244#[derive(Clone, Copy, Debug, PartialEq, Eq)]
245#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
246pub struct WriteBuffer<'buf>(&'buf [u8]);
247
248impl<'buf> WriteBuffer<'buf> {
249    /// Instantiates a new [`WriteBuffer`].
250    #[inline]
251    #[must_use]
252    pub fn new(buffer: &'buf (impl AsRef<[u8]> + 'buf)) -> Self {
253        Self(buffer.as_ref())
254    }
255}
256
257impl Op for WriteBuffer<'_> {
258    #[cfg(feature = "std")]
259    #[inline]
260    fn write_to_io(&self, stream: &mut dyn io::Write) -> Result<usize> {
261        stream
262            .write_all(self.0)
263            .map(|()| self.0.len())
264            .map_err(Error::from)
265    }
266
267    #[inline]
268    fn write_to(&self, mut out: impl AsMut<[u8]>) -> Result<usize> {
269        let n = self.0.len();
270        let out_slice = out
271            .as_mut()
272            .get_mut(..n)
273            .ok_or_else(|| Error::buffer_too_small(n))?;
274        // SAFETY:
275        //
276        // Length of `out` has been checked previously.
277        unsafe {
278            out_slice.as_mut_ptr().copy_from(self.0.as_ptr(), n);
279        }
280        Ok(n)
281    }
282}
283
284#[cfg(test)]
285mod tests {
286    mod advance {
287        use crate::ops::Advance;
288
289        use crate::prelude::*;
290
291        #[cfg(feature = "std")]
292        #[test]
293        fn test_io() -> Result<()> {
294            {
295                let mut stream = Vec::new();
296                let advance = Advance::new(0);
297                assert_eq!(advance.write_to_io(&mut stream).unwrap(), 0);
298                assert!(stream.is_empty());
299            }
300            {
301                let mut stream = Vec::new();
302                let advance = Advance::new(42);
303                assert_eq!(advance.write_to_io(&mut stream).unwrap(), 42);
304                assert_eq!(stream.len(), 42);
305                assert_eq!(
306                    stream.as_slice(),
307                    &[
308                        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
309                        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
310                    ]
311                );
312            }
313            {
314                let mut stream = vec![42u8; 2];
315                let advance = Advance::new(42);
316                assert_eq!(advance.write_to_io(&mut stream).unwrap(), 42);
317                assert_eq!(stream.len(), 44);
318                assert_eq!(
319                    stream.as_slice(),
320                    &[
321                        42, 42, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
322                        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
323                    ]
324                );
325            }
326            Ok(())
327        }
328
329        #[cfg(feature = "std")]
330        #[test]
331        fn test() -> Result<()> {
332            {
333                let mut stream = vec![0u8; 10];
334                let advance = Advance::new(10);
335                assert_eq!(advance.write_to(&mut stream).unwrap(), 10);
336                assert_eq!(stream.len(), 10);
337                assert_eq!(stream.as_slice(), &[0, 0, 0, 0, 0, 0, 0, 0, 0, 0,]);
338            }
339            {
340                let mut stream = vec![0u8; 9];
341                let advance = Advance::new(10);
342                let err = advance.write_to(&mut stream).unwrap_err();
343                let io_err = err.io().unwrap();
344                assert_eq!(io_err.kind(), std::io::ErrorKind::WriteZero);
345                assert_eq!(stream.len(), 9);
346                assert_eq!(stream.as_slice(), &[0, 0, 0, 0, 0, 0, 0, 0, 0,]);
347            }
348            Ok(())
349        }
350    }
351
352    mod fill {
353        use crate::ops::Fill;
354
355        use crate::prelude::*;
356
357        #[cfg(feature = "std")]
358        #[test]
359        fn test_io() -> Result<()> {
360            {
361                let mut stream = Vec::new();
362                let fill = Fill::new(0, 0x41);
363                assert_eq!(fill.write_to_io(&mut stream).unwrap(), 0);
364                assert_eq!(stream.len(), 0);
365            }
366            {
367                let mut stream = Vec::new();
368                let fill = Fill::new(42, 0x41);
369                assert_eq!(fill.write_to_io(&mut stream).unwrap(), 42);
370                assert_eq!(stream.len(), 42);
371                assert_eq!(
372                    String::from_utf8(stream).as_deref(),
373                    Ok("AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA")
374                );
375            }
376            {
377                let mut stream = vec![0x42u8; 4];
378                let fill = Fill::new(42, 0x41);
379                assert_eq!(fill.write_to_io(&mut stream).unwrap(), 42);
380                assert_eq!(stream.len(), 46);
381                assert_eq!(
382                    String::from_utf8(stream).as_deref(),
383                    Ok("BBBBAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA")
384                );
385            }
386            Ok(())
387        }
388
389        #[cfg(feature = "std")]
390        #[test]
391        fn test() -> Result<()> {
392            {
393                let mut stream = vec![0u8; 10];
394                let fill = Fill::new(10, 0x41);
395                assert_eq!(fill.write_to(&mut stream).unwrap(), 10);
396                assert_eq!(stream.len(), 10);
397                assert_eq!(String::from_utf8(stream).as_deref(), Ok("AAAAAAAAAA"));
398            }
399
400            {
401                let mut stream = vec![0u8; 9];
402                let fill = Fill::new(10, 0x41);
403                let err = fill.write_to(&mut stream).unwrap_err();
404                let io_err = err.io().unwrap();
405                assert_eq!(io_err.kind(), std::io::ErrorKind::WriteZero);
406                assert_eq!(stream.len(), 9);
407                assert_eq!(
408                    stream.as_slice(),
409                    &[0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41,]
410                );
411            }
412            Ok(())
413        }
414    }
415
416    mod integers {
417        use crate::ops::{EncodableInteger, WriteInteger};
418
419        use crate::prelude::*;
420
421        #[test]
422        fn test_constructors() -> Result<()> {
423            assert_eq!(WriteInteger::new_be(1u8), WriteInteger::BigEndian(1u8));
424            assert_eq!(<_ as EncodableInteger>::n(1u8), 1);
425            assert_eq!(WriteInteger::new_le(1u8), WriteInteger::LittleEndian(1u8));
426            assert_eq!(WriteInteger::new_be(1u16), WriteInteger::BigEndian(1u16));
427            assert_eq!(<_ as EncodableInteger>::n(1u16), 2);
428            assert_eq!(WriteInteger::new_le(1u16), WriteInteger::LittleEndian(1u16));
429            assert_eq!(WriteInteger::new_be(1u32), WriteInteger::BigEndian(1u32));
430            assert_eq!(<_ as EncodableInteger>::n(1u32), 4);
431            assert_eq!(WriteInteger::new_le(1u32), WriteInteger::LittleEndian(1u32));
432            assert_eq!(WriteInteger::new_be(1u64), WriteInteger::BigEndian(1u64));
433            assert_eq!(<_ as EncodableInteger>::n(1u64), 8);
434            assert_eq!(WriteInteger::new_le(1u64), WriteInteger::LittleEndian(1u64));
435            Ok(())
436        }
437
438        #[cfg(feature = "std")]
439        #[test]
440        fn test_encodable_integer_io() -> Result<()> {
441            {
442                let mut stream = Vec::new();
443                assert_eq!(
444                    WriteInteger::new_be(1u8).write_to_io(&mut stream).unwrap(),
445                    1
446                );
447                assert_eq!(stream.len(), 1);
448                assert_eq!(
449                    WriteInteger::new_le(2u8).write_to_io(&mut stream).unwrap(),
450                    1
451                );
452                assert_eq!(stream.len(), 2);
453                assert_eq!(stream.as_slice(), &[1, 2]);
454            }
455
456            {
457                let mut stream = vec![2u8; 1];
458                assert_eq!(
459                    WriteInteger::new_be(1u8).write_to_io(&mut stream).unwrap(),
460                    1
461                );
462                assert_eq!(stream.len(), 2);
463                assert_eq!(stream.as_slice(), &[2u8, 1u8]);
464            }
465
466            {
467                let mut stream = Vec::new();
468                assert_eq!(
469                    WriteInteger::new_be(0xdeadu16)
470                        .write_to_io(&mut stream)
471                        .unwrap(),
472                    2
473                );
474                assert_eq!(stream.len(), 2);
475                assert_eq!(stream.as_slice(), &[0xde, 0xad]);
476            }
477
478            {
479                let mut stream = Vec::new();
480                assert_eq!(
481                    WriteInteger::new_le(0xdeadu16)
482                        .write_to_io(&mut stream)
483                        .unwrap(),
484                    2
485                );
486                assert_eq!(stream.len(), 2);
487                assert_eq!(stream.as_slice(), &[0xad, 0xde]);
488            }
489
490            {
491                let mut stream = Vec::new();
492                assert_eq!(
493                    WriteInteger::new_be(0xdeadbeefu32)
494                        .write_to_io(&mut stream)
495                        .unwrap(),
496                    4
497                );
498                assert_eq!(stream.len(), 4);
499                assert_eq!(stream.as_slice(), &[0xde, 0xad, 0xbe, 0xef]);
500            }
501
502            {
503                let mut stream = Vec::new();
504                assert_eq!(
505                    WriteInteger::new_le(0xdeadbeefu32)
506                        .write_to_io(&mut stream)
507                        .unwrap(),
508                    4
509                );
510                assert_eq!(stream.len(), 4);
511                assert_eq!(stream.as_slice(), &[0xef, 0xbe, 0xad, 0xde]);
512            }
513
514            {
515                let mut stream = Vec::new();
516                assert_eq!(
517                    WriteInteger::new_be(0xdeadbeefcafebabeu64)
518                        .write_to_io(&mut stream)
519                        .unwrap(),
520                    8
521                );
522                assert_eq!(stream.len(), 8);
523                assert_eq!(
524                    stream.as_slice(),
525                    &[0xde, 0xad, 0xbe, 0xef, 0xca, 0xfe, 0xba, 0xbe]
526                );
527            }
528
529            {
530                let mut stream = Vec::new();
531                assert_eq!(
532                    WriteInteger::new_le(0xdeadbeefcafebabeu64)
533                        .write_to_io(&mut stream)
534                        .unwrap(),
535                    8
536                );
537                assert_eq!(stream.len(), 8);
538                assert_eq!(
539                    stream.as_slice(),
540                    &[0xbe, 0xba, 0xfe, 0xca, 0xef, 0xbe, 0xad, 0xde]
541                );
542            }
543            Ok(())
544        }
545    }
546}
shellcoder/ops.rs

shellcoder/
ops.rs
