1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110

// SPDX-FileCopyrightText: 2022 Profian Inc. <opensource@profian.com>
// SPDX-License-Identifier: AGPL-3.0-only

use super::{Algorithm, ContentDigest};

use std::io;
use std::pin::Pin;
use std::task::{Context, Poll};

use futures::AsyncWrite;
use sha2::digest::DynDigest;

/// A hashing writer
///
/// This type wraps another writer and hashes the bytes as they are written.
#[allow(missing_debug_implementations)] // DynDigest does not implement Debug
pub struct Writer<T> {
    writer: T,
    digests: Vec<(Algorithm, Box<dyn DynDigest>)>,
}

#[allow(unsafe_code)]
unsafe impl<T> Sync for Writer<T> where T: Sync {}

#[allow(unsafe_code)]
unsafe impl<T> Send for Writer<T> where T: Send {}

impl<T> Writer<T> {
    pub(crate) fn new(writer: T, digests: impl IntoIterator<Item = Algorithm>) -> Self {
        let digests = digests.into_iter().map(|a| (a, a.hasher())).collect();
        Writer { writer, digests }
    }

    fn update(&mut self, buf: &[u8]) {
        for digest in &mut self.digests {
            digest.1.update(buf);
        }
    }

    /// Calculates the digests for all the bytes written so far.
    pub fn digests(&self) -> ContentDigest<Box<[u8]>> {
        let mut set = ContentDigest::default();

        for digest in &self.digests {
            _ = set.insert(digest.0, digest.1.clone().finalize().into());
        }

        set
    }
}

impl<T: AsyncWrite + Unpin> AsyncWrite for Writer<T> {
    fn poll_write(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<io::Result<usize>> {
        Pin::new(&mut self.writer).poll_write(cx, buf).map_ok(|n| {
            self.update(&buf[..n]);
            n
        })
    }

    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        Pin::new(&mut self.writer).poll_flush(cx)
    }

    fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        Pin::new(&mut self.writer).poll_close(cx)
    }
}

impl<T: io::Write> io::Write for Writer<T> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        let n = self.writer.write(buf)?;
        self.update(&buf[..n]);
        Ok(n)
    }

    fn flush(&mut self) -> io::Result<()> {
        self.writer.flush()
    }
}

#[cfg(test)]
mod tests {
    use futures::io::{copy, sink};

    use super::*;

    #[async_std::test]
    async fn success() {
        const HASH: &str = "sha-256=:LCa0a2j/xo/5m0U8HTBBNBNCLXBkg7+g+YpeiGJm564=:";
        let set = HASH.parse::<ContentDigest>().unwrap();

        let mut writer = set.clone().writer(sink());
        assert_eq!(copy(&mut &b"foo"[..], &mut writer).await.unwrap(), 3);
        assert_eq!(writer.digests(), set);
    }

    #[async_std::test]
    async fn failure() {
        const HASH: &str = "sha-256=:LCa0a2j/xo/5m0U8HTBBNBNCLXBkg7+g+YpeiGJm564=:";
        let set = HASH.parse::<ContentDigest>().unwrap();

        let mut writer = set.clone().writer(sink());
        assert_eq!(copy(&mut &b"bar"[..], &mut writer).await.unwrap(), 3);
        assert_ne!(writer.digests(), set);
    }
}