From e7e4411f444480822ea1826071d534d9b75b4d8d Mon Sep 17 00:00:00 2001
From: Dominik Csapak <d.csapak@proxmox.com>
Date: Tue, 12 Apr 2022 13:04:14 +0200
Subject: [PATCH] proxmox-compression: add streaming zstd encoder

similar to our DeflateEncoder, takes a Stream and implements it itself,
so that we can use it as an adapter for async api calls

Signed-off-by: Dominik Csapak <d.csapak@proxmox.com>
Signed-off-by: Wolfgang Bumiller <w.bumiller@proxmox.com>
---
 proxmox-compression/Cargo.toml  |   1 +
 proxmox-compression/src/lib.rs  |   1 +
 proxmox-compression/src/zstd.rs | 126 ++++++++++++++++++++++++++++++++
 3 files changed, 128 insertions(+)
 create mode 100644 proxmox-compression/src/zstd.rs

diff --git a/proxmox-compression/Cargo.toml b/proxmox-compression/Cargo.toml
index c3f7f497..5ca67b2a 100644
--- a/proxmox-compression/Cargo.toml
+++ b/proxmox-compression/Cargo.toml
@@ -18,6 +18,7 @@ futures = "0.3"
 tokio = { version = "1.6", features = [ "fs", "io-util"] }
 walkdir = "2"
 tar = "0.4"
+zstd = { version = "0.6", features = []}
 
 proxmox-time = { path = "../proxmox-time", version = "1" }
 proxmox-io = { path = "../proxmox-io", version = "1", features = [ "tokio" ] }
diff --git a/proxmox-compression/src/lib.rs b/proxmox-compression/src/lib.rs
index e9dd1138..1fcfb977 100644
--- a/proxmox-compression/src/lib.rs
+++ b/proxmox-compression/src/lib.rs
@@ -3,3 +3,4 @@ pub use compression::*;
 
 pub mod tar;
 pub mod zip;
+pub mod zstd;
diff --git a/proxmox-compression/src/zstd.rs b/proxmox-compression/src/zstd.rs
new file mode 100644
index 00000000..0b480f6e
--- /dev/null
+++ b/proxmox-compression/src/zstd.rs
@@ -0,0 +1,126 @@
+//! zstd helper
+use std::io;
+use std::pin::Pin;
+use std::task::{Context, Poll};
+
+use anyhow::{format_err, Error};
+use bytes::Bytes;
+use futures::ready;
+use futures::stream::Stream;
+use zstd::stream::raw::{Encoder, Operation, OutBuffer};
+
+use proxmox_io::ByteBuffer;
+
+const BUFFER_SIZE: usize = 8192;
+
+#[derive(Eq, PartialEq)]
+enum EncoderState {
+    Reading,
+    Writing,
+    Finishing,
+    Finished,
+}
+
+/// An async ZstdEncoder that implements [Stream] for another [Stream]
+///
+/// Useful for on-the-fly zstd compression in streaming api calls
+pub struct ZstdEncoder<'a, T> {
+    inner: T,
+    compressor: Encoder<'a>,
+    buffer: ByteBuffer,
+    input_buffer: Bytes,
+    state: EncoderState,
+}
+
+impl<'a, T> ZstdEncoder<'a, T> {
+    /// Returns a new [ZstdEncoder] with default level 3
+    pub fn new(inner: T) -> Result<Self, io::Error> {
+        Self::with_quality(inner, 3)
+    }
+
+    /// Returns a new [ZstdEncoder] with the given level
+    pub fn with_quality(inner: T, level: i32) -> Result<Self, io::Error> {
+        Ok(Self {
+            inner,
+            compressor: Encoder::new(level)?,
+            buffer: ByteBuffer::with_capacity(BUFFER_SIZE),
+            input_buffer: Bytes::new(),
+            state: EncoderState::Reading,
+        })
+    }
+
+    /// Returns the wrapped [Stream]
+    pub fn into_inner(self) -> T {
+        self.inner
+    }
+
+    fn encode(&mut self, inbuf: &[u8]) -> Result<zstd::stream::raw::Status, io::Error> {
+        let res = self
+            .compressor
+            .run_on_buffers(inbuf, self.buffer.get_free_mut_slice())?;
+        self.buffer.add_size(res.bytes_written);
+
+        Ok(res)
+    }
+
+    fn finish(&mut self) -> Result<usize, io::Error> {
+        let mut outbuf = OutBuffer::around(self.buffer.get_free_mut_slice());
+        let res = self.compressor.finish(&mut outbuf, true);
+        let size = outbuf.pos;
+        drop(outbuf);
+        self.buffer.add_size(size);
+        res
+    }
+}
+
+impl<'a, T, O> Stream for ZstdEncoder<'a, T>
+where
+    T: Stream<Item = Result<O, Error>> + Unpin,
+    O: Into<Bytes>,
+{
+    type Item = Result<Bytes, Error>;
+
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        let this = self.get_mut();
+
+        loop {
+            match this.state {
+                EncoderState::Reading => {
+                    if let Some(res) = ready!(Pin::new(&mut this.inner).poll_next(cx)) {
+                        let buf = res?;
+                        this.input_buffer = buf.into();
+                        this.state = EncoderState::Writing;
+                    } else {
+                        this.state = EncoderState::Finishing;
+                    }
+                }
+                EncoderState::Writing => {
+                    if this.input_buffer.is_empty() {
+                        return Poll::Ready(Some(Err(format_err!("empty input during write"))));
+                    }
+                    let mut buf = this.input_buffer.split_off(0);
+                    let status = this.encode(&buf[..])?;
+                    this.input_buffer = buf.split_off(status.bytes_read);
+                    if this.input_buffer.is_empty() {
+                        this.state = EncoderState::Reading;
+                    }
+                    if this.buffer.is_full() {
+                        let bytes = this.buffer.remove_data(this.buffer.len()).to_vec();
+                        return Poll::Ready(Some(Ok(bytes.into())));
+                    }
+                }
+                EncoderState::Finishing => {
+                    let remaining = this.finish()?;
+                    if remaining == 0 {
+                        this.state = EncoderState::Finished;
+                    }
+                    if !this.buffer.is_empty() {
+                        let bytes = this.buffer.remove_data(this.buffer.len()).to_vec();
+                        return Poll::Ready(Some(Ok(bytes.into())));
+                    }
+                }
+                EncoderState::Finished => return Poll::Ready(None),
+            }
+        }
+    }
+}