drawbridge_hash/
writer.rs

1// SPDX-FileCopyrightText: 2022 Profian Inc. <opensource@profian.com>
2// SPDX-License-Identifier: AGPL-3.0-only
3
4use std::io;
5use std::pin::Pin;
6use std::task::{Context, Poll};
7
8use futures::AsyncWrite;
9use sha2::digest::Digest;
10use sha2::{Sha224, Sha256, Sha384, Sha512};
11
12pub(super) enum Inner {
13    Sha224(Sha224),
14    Sha256(Sha256),
15    Sha384(Sha384),
16    Sha512(Sha512),
17}
18
19pub struct Writer<T> {
20    pub(super) writer: T,
21    pub(super) inner: Inner,
22}
23
24impl<T: AsyncWrite + Unpin> AsyncWrite for Writer<T> {
25    fn poll_write(
26        mut self: Pin<&mut Self>,
27        cx: &mut Context<'_>,
28        buf: &[u8],
29    ) -> Poll<io::Result<usize>> {
30        Pin::new(&mut self.writer).poll_write(cx, buf).map_ok(|n| {
31            match &mut self.inner {
32                Inner::Sha224(h) => h.update(&buf[..n]),
33                Inner::Sha256(h) => h.update(&buf[..n]),
34                Inner::Sha384(h) => h.update(&buf[..n]),
35                Inner::Sha512(h) => h.update(&buf[..n]),
36            };
37            n
38        })
39    }
40
41    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
42        Pin::new(&mut self.writer).poll_flush(cx)
43    }
44
45    fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
46        Pin::new(&mut self.writer).poll_close(cx)
47    }
48}