Merge pull request #371 from runloopai/adam/zstd

Add zstd decompression support for overlaybd-apply

Author: BigVan

src/overlaybd/CMakeLists.txt

@@ -1,6 +1,7 @@
 add_subdirectory(registryfs)
 add_subdirectory(lsmt)
 add_subdirectory(zfile)
+add_subdirectory(zstd)
 add_subdirectory(cache)
 add_subdirectory(tar)
 add_subdirectory(gzip)
@@ -16,6 +17,7 @@ target_link_libraries(overlaybd_lib INTERFACE
     registryfs_lib
     lsmt_lib
     zfile_lib
+    zstd_lib
     cache_lib
     tar_lib
     gzip_lib

src/overlaybd/zstd/CMakeLists.txt

@@ -0,0 +1,9 @@
+file(GLOB SOURCE_ZSTD "*.cpp")
+
+add_library(zstd_lib STATIC ${SOURCE_ZSTD})
+
+target_include_directories(zstd_lib PUBLIC
+    ${PHOTON_INCLUDE_DIR}
+)
+
+target_link_libraries(zstd_lib photon_static ${LIBZSTD})

src/overlaybd/zstd/zstdfile.cpp

@@ -0,0 +1,105 @@
+/*
+   Copyright The Overlaybd Authors
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
+*/
+
+#include "zstdfile.h"
+
+#include <vector>
+
+#include <fcntl.h>
+#include <zstd.h>
+#include <photon/common/alog.h>
+#include <photon/fs/filesystem.h>
+#include <photon/fs/virtual-file.h>
+
+class ZStdAdaptorFile : public photon::fs::VirtualReadOnlyFile {
+public:
+    ZStdAdaptorFile(photon::fs::IFile* file, bool ownership)
+        : m_file(file)
+        , m_ownership(ownership)
+        , m_stream(ZSTD_createDStream())
+        , m_buffer(ZSTD_DStreamInSize())
+        , m_input({m_buffer.data(), 0, 0}) {
+    }
+
+    ~ZStdAdaptorFile() {
+        ZSTD_freeDStream(m_stream);
+        if (m_ownership) {
+            delete m_file;
+        }
+    }
+
+    ssize_t read(void *buf, size_t count) override {
+        ssize_t bytes_read = 0;
+        ZSTD_outBuffer output = { buf, count, 0 };
+        while (output.pos < output.size) {
+            // Output buffer is not filled, read more.
+            if (m_input.pos == m_input.size) {
+                // Reached the end of the input buffer, read more compressed data.
+                ssize_t read_compressed_bytes = m_file->read(m_buffer.data(), m_buffer.size());
+                if (read_compressed_bytes == 0) {
+                    // EOF reached.
+                    return bytes_read;
+                }
+                if (read_compressed_bytes < 0 || read_compressed_bytes > (ssize_t) m_buffer.size()) {
+                    // Error reading file.
+                    LOG_ERRNO_RETURN(EIO, -1, "read buffer error");
+                }
+                m_input.size = read_compressed_bytes;
+                m_input.pos = 0;
+            }
+            const size_t prev_pos = output.pos;
+            const size_t ret = ZSTD_decompressStream(m_stream, &output, &m_input);
+            if (ZSTD_isError(ret)) {
+                LOG_ERRNO_RETURN(EIO, -1, "failed to decompress zstd frame");
+            }
+            bytes_read += output.pos - prev_pos;
+            if (ret == 0) {
+                // End of this ZSTD frame, set up for the next one (if there is one).
+                ZSTD_initDStream(m_stream);
+            }
+        }
+        return bytes_read;
+    }
+
+    int fstat(struct stat *buf) override {
+        return m_file->fstat(buf);
+    }
+
+    UNIMPLEMENTED_POINTER(photon::fs::IFileSystem *filesystem() override);
+    UNIMPLEMENTED(off_t lseek(off_t offset, int whence) override);
+    UNIMPLEMENTED(ssize_t readv(const struct iovec *iov, int iovcnt) override);
+    UNIMPLEMENTED(ssize_t preadv(const struct iovec *iov, int iovcnt, off_t offset) override);
+private:
+    photon::fs::IFile* m_file;
+    bool m_ownership; // whether we own m_file and should delete it.
+    ZSTD_DStream* m_stream;
+
+    std::vector<uint8_t> m_buffer;
+    ZSTD_inBuffer m_input;
+};
+
+photon::fs::IFile *open_zstdfile_adaptor(photon::fs::IFile* file, bool ownership) {
+    return new ZStdAdaptorFile(file, ownership);
+}
+
+const uint8_t ZSTD_MAGIC_HEADER[4] = {0x28, 0xB5, 0x2F, 0xFD};
+
+bool is_zstdfile(photon::fs::IFile* file) {
+    char buf[4] = {0};
+    ssize_t readn = file->read(buf, 4);
+    file->lseek(0, 0);
+    return readn == 4 && memcmp(buf, ZSTD_MAGIC_HEADER, 4) == 0;
+}

src/overlaybd/zstd/zstdfile.h

@@ -0,0 +1,25 @@
+/*
+   Copyright The Overlaybd Authors
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
+*/
+
+#pragma once
+
+#include <photon/fs/virtual-file.h>
+
+// Opens a ZSTD compressed file and provides a view of the decompressed data.
+// Takes ownership of `file` when `ownership` is true (default).
+photon::fs::IFile* open_zstdfile_adaptor(photon::fs::IFile* file, bool ownership = true);
+
+bool is_zstdfile(photon::fs::IFile* file);

src/tools/overlaybd-apply.cpp

@@ -25,6 +25,7 @@
 #include "../overlaybd/tar/libtar.h"
 #include "../overlaybd/gzindex/gzfile.h"
 #include "../overlaybd/gzip/gz.h"
+#include "../overlaybd/zstd/zstdfile.h"
 #include <errno.h>
 #include <fcntl.h>
 #include <inttypes.h>
@@ -100,7 +101,6 @@ int main(int argc, char **argv) {
     photon::init(photon::INIT_EVENT_DEFAULT, photon::INIT_IO_DEFAULT);
     DEFER({photon::fini();});
 
-
     ImageService *imgservice = nullptr;
     photon::fs::IFile *imgfile = nullptr;
     if (raw) {
@@ -141,6 +141,8 @@ int main(int argc, char **argv) {
             tarf->lseek(0, 0);
         }
         src_file = open_gzfile_adaptor(input_path.c_str());
+    } else if (is_zstdfile(tarf)) {
+        src_file = open_zstdfile_adaptor(tarf, /*ownership=*/false);
     } else {
         src_file = tarf;
     }