aws_smithy_eventstream/buf/
crc.rs

Help
1/*
2 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3 * SPDX-License-Identifier: Apache-2.0
4 */
5
6//! Utilities for calculating CRC-32 while reading from a [`Buf`] or writing to a [`BufMut`].
7
8use bytes::buf::UninitSlice;
9use bytes::{Buf, BufMut};
10use crc32fast::Hasher;
11
12/// Implementation of [`Buf`] that calculates a CRC-32 checksum of the data
13/// being read from an underlying `Buf` instance.
14pub(crate) struct CrcBuf<'a, B>
15where
16    B: Buf,
17{
18    buffer: &'a mut B,
19    crc: Hasher,
20}
21
22impl<'a, B> CrcBuf<'a, B>
23where
24    B: Buf,
25{
26    /// Creates a new `CrcBuf` by wrapping the given `buffer`.
27    pub(crate) fn new(buffer: &'a mut B) -> Self {
28        CrcBuf {
29            buffer,
30            crc: Hasher::new(),
31        }
32    }
33
34    /// Consumes the `CrcBuf` and returns the calculated checksum.
35    pub(crate) fn into_crc(self) -> u32 {
36        self.crc.finalize()
37    }
38}
39
40impl<'a, B> Buf for CrcBuf<'a, B>
41where
42    B: Buf,
43{
44    fn remaining(&self) -> usize {
45        self.buffer.remaining()
46    }
47
48    fn chunk(&self) -> &[u8] {
49        self.buffer.chunk()
50    }
51
52    fn advance(&mut self, cnt: usize) {
53        let chunk = self.buffer.chunk();
54        self.crc.update(&chunk[0..cnt]);
55        self.buffer.advance(cnt);
56    }
57}
58
59#[cfg(test)]
60mod crc_buf_tests {
61    use super::CrcBuf;
62    use bytes::Buf;
63
64    #[test]
65    fn crc_no_data_read() {
66        let mut data: &[u8] = &[];
67        let buf = CrcBuf::new(&mut data);
68        assert_eq!(0, buf.into_crc());
69    }
70
71    #[test]
72    fn crc_data_read() {
73        let mut data: &[u8] = &[0, 0, 0, 5, 0, 10u8];
74        let mut buf = CrcBuf::new(&mut data);
75        assert_eq!(5, buf.get_i32());
76        assert_eq!(0x512E2B93, buf.into_crc());
77
78        let mut data: &[u8] = &[0, 0, 0, 5, 0, 10u8];
79        let mut buf = CrcBuf::new(&mut data);
80        assert_eq!(5, buf.get_i32());
81        assert_eq!(10, buf.get_i16());
82        assert_eq!(0x57DC8A56, buf.into_crc());
83    }
84
85    #[test]
86    fn chunk_called_multiple_times_before_advance() {
87        let mut data: &[u8] = &[0, 0, 0, 5, 0, 10u8];
88        let mut buf = CrcBuf::new(&mut data);
89        for _ in 0..3 {
90            buf.chunk();
91        }
92        buf.advance(4);
93        assert_eq!(10, buf.get_i16());
94        assert_eq!(0x57DC8A56, buf.into_crc());
95    }
96}
97
98/// Implementation of [`BufMut`] that calculates a CRC-32 checksum of the data
99/// being written to an underlying `Buf` instance, with a function to then write
100/// the calculated CRC-32 to the buffer.
101pub(crate) struct CrcBufMut<'a> {
102    buffer: &'a mut dyn BufMut,
103    crc: Hasher,
104}
105
106impl<'a> CrcBufMut<'a> {
107    /// Creates a new `CrcBufMut` by wrapping the given `buffer`.
108    pub(crate) fn new(buffer: &'a mut dyn BufMut) -> Self {
109        CrcBufMut {
110            buffer,
111            crc: Hasher::new(),
112        }
113    }
114
115    /// Puts the calculated CRC-32 to the buffer as a Big Endian 32-bit integer.
116    /// This can be called multiple times, and each successive call will include
117    /// the previously written checksum in its new checksum.
118    pub(crate) fn put_crc(&mut self) {
119        self.put_u32(self.crc.clone().finalize());
120    }
121}
122
123unsafe impl<'a> BufMut for CrcBufMut<'a> {
124    fn remaining_mut(&self) -> usize {
125        self.buffer.remaining_mut()
126    }
127
128    unsafe fn advance_mut(&mut self, cnt: usize) {
129        // Safety: There is no guarantee the bytes being advanced are initialized, which is why
130        // this trait method is unsafe. The best we can do is assume they have been initialized
131        // by the caller before `advance_mut` was called.
132        let written = std::slice::from_raw_parts_mut(self.chunk_mut().as_mut_ptr(), cnt);
133        self.crc.update(written);
134        self.buffer.advance_mut(cnt);
135    }
136
137    fn chunk_mut(&mut self) -> &mut UninitSlice {
138        self.buffer.chunk_mut()
139    }
140}
141
142#[cfg(test)]
143mod crc_buf_mut_tests {
144    use super::CrcBufMut;
145    use bytes::BufMut;
146
147    #[test]
148    fn crc_no_bytes_written() {
149        let mut buffer: Vec<u8> = Vec::new();
150        let mut crc_buf = CrcBufMut::new(&mut buffer);
151        crc_buf.put_crc();
152        crc_buf.put_crc();
153        assert_eq!(vec![0, 0, 0, 0u8, 0x21, 0x44, 0xDF, 0x1C], buffer);
154    }
155
156    #[test]
157    fn crc_bytes_written() {
158        let mut buffer: Vec<u8> = Vec::new();
159        let mut crc_buf = CrcBufMut::new(&mut buffer);
160        crc_buf.put_u32(5);
161        crc_buf.put_crc();
162        assert_eq!(vec![0, 0, 0, 5, 0x51, 0x2E, 0x2B, 0x93], buffer);
163    }
164}
aws_smithy_eventstream/buf/crc.rs

aws_smithy_eventstream/buf/
crc.rs
