[Fix] Premature end for exported snapshots (#3015)

* [Fix] Premature end for exported snapshots

* add snapshot validity test

Author: ruseinov

CHANGELOG.md

@@ -35,6 +35,9 @@
 
 ### Fixed
 
+- [#3006](https://github.com/ChainSafe/forest/issues/3006): Fix `premature end`
+  error when exporting a snapshot.
+
 ## Forest v0.9.0 "Fellowship"
 
 Notable updates:

scripts/tests/calibnet_export_check.sh

@@ -9,8 +9,14 @@ source "$(dirname "$0")/harness.sh"
 
 forest_init
 
+echo "Cleaning up the initial snapshot"
+rm -rf ./*.car.*
+
 echo "Exporting zstd compressed snapshot"
 $FOREST_CLI_PATH snapshot export
 
-echo "Verifing snapshot checksum"
-sha256sum -c ./*.sha256sum
+echo "Testing snapshot validity"
+zstd --test ./*.car.zst
+
+echo "Verifying snapshot checksum"
+sha256sum --check ./*.sha256sum

src/chain/store/chain_store.rs

@@ -29,12 +29,13 @@ use crate::utils::{
     misc::Either,
 };
 use ahash::{HashMap, HashMapExt, HashSet};
-use anyhow::Result;
+use anyhow::{Context, Result};
 use async_compression::futures::write::ZstdEncoder;
 use bls_signatures::Serialize as SerializeBls;
 use cid::Cid;
 use digest::Digest;
 use futures::{io::BufWriter, AsyncWrite};
+use futures_util::AsyncWriteExt;
 use fvm_ipld_amt::Amtv0 as Amt;
 use fvm_ipld_blockstore::Blockstore;
 use fvm_ipld_car::CarHeader;
@@ -624,6 +625,9 @@ where
         );
 
         let mut writer = writer.lock().await;
+        writer.flush().await.context("failed to flush")?;
+        writer.close().await.context("failed to close")?;
+
         let digest = match &mut *writer {
             Either::Left(left) => left.get_mut().finalize().await,
             Either::Right(right) => right.finalize().await,

src/utils/io/writer_checksum.rs

@@ -4,7 +4,8 @@ use std::{pin::Pin, task::Poll};
 
 use async_trait::async_trait;
 use digest::{Digest, Output};
-use futures::{io::BufWriter, AsyncWrite, AsyncWriteExt};
+use futures::{io::BufWriter, AsyncWrite};
+use futures_util::AsyncWriteExt;
 use pin_project_lite::pin_project;
 
 pin_project! {
@@ -25,7 +26,7 @@ pub trait Checksum<D: Digest> {
     async fn finalize(&mut self) -> std::io::Result<Option<Output<D>>>;
 }
 
-impl<D: Digest, W: AsyncWrite + Unpin> AsyncWrite for AsyncWriterWithChecksum<D, W> {
+impl<D: Digest, W: AsyncWriteExt + Unpin> AsyncWrite for AsyncWriterWithChecksum<D, W> {
     fn poll_write(
         mut self: std::pin::Pin<&mut Self>,
         cx: &mut std::task::Context<'_>,
@@ -59,9 +60,10 @@ impl<D: Digest, W: AsyncWrite + Unpin> AsyncWrite for AsyncWriterWithChecksum<D,
 }
 
 #[async_trait]
-impl<D: Digest + Send, W: AsyncWrite + Send + Unpin> Checksum<D> for AsyncWriterWithChecksum<D, W> {
+impl<D: Digest + Send, W: AsyncWriteExt + Send + Unpin> Checksum<D>
+    for AsyncWriterWithChecksum<D, W>
+{
     async fn finalize(&mut self) -> std::io::Result<Option<Output<D>>> {
-        self.inner.flush().await?;
         if let Some(hasher) = &mut self.hasher {
             let hasher = std::mem::replace(hasher, D::new());
             Ok(Some(hasher.finalize()))
@@ -133,6 +135,9 @@ mod test {
             temp_file_writer.write_all(&bytes).await?;
         }
 
+        temp_file_writer.flush().await?;
+        temp_file_writer.close().await?;
+
         let checksum = temp_file_writer.finalize().await?;
 
         let file_hash = {

src/utils/misc/either.rs

@@ -1,12 +1,15 @@
 // Copyright 2019-2023 ChainSafe Systems
 // SPDX-License-Identifier: Apache-2.0, MIT
 
+use std::io::IoSlice;
 use std::{
     pin::Pin,
     task::{Context, Poll},
 };
 
-use futures::AsyncRead;
+use futures::{AsyncRead, AsyncWrite};
+use futures_util::io::{Close, Flush, IntoSink, Write, WriteAll, WriteVectored};
+use futures_util::AsyncWriteExt;
 
 pub enum Either<L, R> {
     Left(L),
@@ -25,3 +28,30 @@ impl<L: AsyncRead + Unpin, R: AsyncRead + Unpin> AsyncRead for Either<L, R> {
         }
     }
 }
+
+impl<L: AsyncWrite + Unpin, R: AsyncWrite + Unpin> AsyncWrite for Either<L, R> {
+    fn poll_write(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &[u8],
+    ) -> Poll<std::io::Result<usize>> {
+        match Pin::into_inner(self) {
+            Self::Left(left) => Pin::new(left).poll_write(cx, buf),
+            Self::Right(right) => Pin::new(right).poll_write(cx, buf),
+        }
+    }
+
+    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
+        match Pin::into_inner(self) {
+            Self::Left(left) => Pin::new(left).poll_flush(cx),
+            Self::Right(right) => Pin::new(right).poll_flush(cx),
+        }
+    }
+
+    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
+        match Pin::into_inner(self) {
+            Self::Left(left) => Pin::new(left).poll_close(cx),
+            Self::Right(right) => Pin::new(right).poll_close(cx),
+        }
+    }
+}