import: split out compression logic, so that we can share it with between import and pull calls

5 files changed, 270 insertions, 145 deletions

diff --git a/Makefile.am b/Makefile.am
index 19071d9ffc..d430fff974 100644
--- a/Makefile.am
+++ b/Makefile.am
@@ -5456,6 +5456,8 @@ systemd_pull_SOURCES = \
 	src/import/pull-job.h \
 	src/import/pull-common.c \
 	src/import/pull-common.h \
+	src/import/import-compress.c \
+	src/import/import-compress.h \
 	src/import/curl-util.c \
 	src/import/curl-util.h \
 	src/import/aufs-util.c \
diff --git a/src/import/import-compress.c b/src/import/import-compress.c
new file mode 100644
index 0000000000..629605a44c
--- /dev/null
+++ b/src/import/import-compress.c
@@ -0,0 +1,194 @@
+/*-*- Mode: C; c-basic-offset: 8; indent-tabs-mode: nil -*-*/
+
+/***
+  This file is part of systemd.
+
+  Copyright 2015 Lennart Poettering
+
+  systemd is free software; you can redistribute it and/or modify it
+  under the terms of the GNU Lesser General Public License as published by
+  the Free Software Foundation; either version 2.1 of the License, or
+  (at your option) any later version.
+
+  systemd is distributed in the hope that it will be useful, but
+  WITHOUT ANY WARRANTY; without even the implied warranty of
+  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
+  Lesser General Public License for more details.
+
+  You should have received a copy of the GNU Lesser General Public License
+  along with systemd; If not, see <http://www.gnu.org/licenses/>.
+***/
+
+#include "util.h"
+#include "import-compress.h"
+
+void import_compress_free(ImportCompress *c) {
+        assert(c);
+
+        if (c->type == IMPORT_COMPRESS_XZ)
+                lzma_end(&c->xz);
+        else if (c->type == IMPORT_COMPRESS_GZIP)
+                inflateEnd(&c->gzip);
+        else if (c->type == IMPORT_COMPRESS_BZIP2)
+                BZ2_bzDecompressEnd(&c->bzip2);
+
+        c->type = IMPORT_COMPRESS_UNKNOWN;
+}
+
+int import_uncompress_detect(ImportCompress *c, const void *data, size_t size) {
+        static const uint8_t xz_signature[] = {
+                0xfd, '7', 'z', 'X', 'Z', 0x00
+        };
+        static const uint8_t gzip_signature[] = {
+                0x1f, 0x8b
+        };
+        static const uint8_t bzip2_signature[] = {
+                'B', 'Z', 'h'
+        };
+
+        int r;
+
+        assert(c);
+
+        if (c->type != IMPORT_COMPRESS_UNKNOWN)
+                return 1;
+
+        if (size < MAX3(sizeof(xz_signature),
+                        sizeof(gzip_signature),
+                        sizeof(bzip2_signature)))
+                return 0;
+
+        assert(data);
+
+        if (memcmp(data, xz_signature, sizeof(xz_signature)) == 0) {
+                lzma_ret xzr;
+
+                xzr = lzma_stream_decoder(&c->xz, UINT64_MAX, LZMA_TELL_UNSUPPORTED_CHECK);
+                if (xzr != LZMA_OK)
+                        return -EIO;
+
+                c->type = IMPORT_COMPRESS_XZ;
+
+        } else if (memcmp(data, gzip_signature, sizeof(gzip_signature)) == 0) {
+                r = inflateInit2(&c->gzip, 15+16);
+                if (r != Z_OK)
+                        return -EIO;
+
+                c->type = IMPORT_COMPRESS_GZIP;
+
+        } else if (memcmp(data, bzip2_signature, sizeof(bzip2_signature)) == 0) {
+                r = BZ2_bzDecompressInit(&c->bzip2, 0, 0);
+                if (r != BZ_OK)
+                        return -EIO;
+
+                c->type = IMPORT_COMPRESS_BZIP2;
+        } else
+                c->type = IMPORT_COMPRESS_UNCOMPRESSED;
+
+        return 1;
+}
+
+int import_uncompress(ImportCompress *c, const void *data, size_t size, ImportCompressCallback callback, void *userdata) {
+        int r;
+
+        assert(c);
+        assert(callback);
+
+        r = import_uncompress_detect(c, data, size);
+        if (r <= 0)
+                return r;
+
+        if (size <= 0)
+                return 1;
+
+        assert(data);
+
+        switch (c->type) {
+
+        case IMPORT_COMPRESS_UNCOMPRESSED:
+                r = callback(data, size, userdata);
+                if (r < 0)
+                        return r;
+
+                break;
+
+        case IMPORT_COMPRESS_XZ:
+                c->xz.next_in = data;
+                c->xz.avail_in = size;
+
+                while (c->xz.avail_in > 0) {
+                        uint8_t buffer[16 * 1024];
+                        lzma_ret lzr;
+
+                        c->xz.next_out = buffer;
+                        c->xz.avail_out = sizeof(buffer);
+
+                        lzr = lzma_code(&c->xz, LZMA_RUN);
+                        if (lzr != LZMA_OK && lzr != LZMA_STREAM_END)
+                                return -EIO;
+
+                        r = callback(buffer, sizeof(buffer) - c->xz.avail_out, userdata);
+                        if (r < 0)
+                                return r;
+                }
+
+                break;
+
+        case IMPORT_COMPRESS_GZIP:
+                c->gzip.next_in = (void*) data;
+                c->gzip.avail_in = size;
+
+                while (c->gzip.avail_in > 0) {
+                        uint8_t buffer[16 * 1024];
+
+                        c->gzip.next_out = buffer;
+                        c->gzip.avail_out = sizeof(buffer);
+
+                        r = inflate(&c->gzip, Z_NO_FLUSH);
+                        if (r != Z_OK && r != Z_STREAM_END)
+                                return -EIO;
+
+                        r = callback(buffer, sizeof(buffer) - c->gzip.avail_out, userdata);
+                        if (r < 0)
+                                return r;
+                }
+
+                break;
+
+        case IMPORT_COMPRESS_BZIP2:
+                c->bzip2.next_in = (void*) data;
+                c->bzip2.avail_in = size;
+
+                while (c->bzip2.avail_in > 0) {
+                        uint8_t buffer[16 * 1024];
+
+                        c->bzip2.next_out = (char*) buffer;
+                        c->bzip2.avail_out = sizeof(buffer);
+
+                        r = BZ2_bzDecompress(&c->bzip2);
+                        if (r != BZ_OK && r != BZ_STREAM_END)
+                                return -EIO;
+
+                        r = callback(buffer, sizeof(buffer) - c->bzip2.avail_out, userdata);
+                        if (r < 0)
+                                return r;
+                }
+
+                break;
+
+        default:
+                assert_not_reached("Unknown compression");
+        }
+
+        return 1;
+}
+
+static const char* const import_compress_type_table[_IMPORT_COMPRESS_TYPE_MAX] = {
+        [IMPORT_COMPRESS_UNKNOWN] = "unknown",
+        [IMPORT_COMPRESS_UNCOMPRESSED] = "uncompressed",
+        [IMPORT_COMPRESS_XZ] = "xz",
+        [IMPORT_COMPRESS_GZIP] = "gzip",
+        [IMPORT_COMPRESS_BZIP2] = "bzip2",
+};
+
+DEFINE_STRING_TABLE_LOOKUP(import_compress_type, ImportCompressType);
diff --git a/src/import/import-compress.h b/src/import/import-compress.h
new file mode 100644
index 0000000000..4f97f20b45
--- /dev/null
+++ b/src/import/import-compress.h
@@ -0,0 +1,61 @@
+/*-*- Mode: C; c-basic-offset: 8; indent-tabs-mode: nil -*-*/
+
+#pragma once
+
+/***
+  This file is part of systemd.
+
+  Copyright 2015 Lennart Poettering
+
+  systemd is free software; you can redistribute it and/or modify it
+  under the terms of the GNU Lesser General Public License as published by
+  the Free Software Foundation; either version 2.1 of the License, or
+  (at your option) any later version.
+
+  systemd is distributed in the hope that it will be useful, but
+  WITHOUT ANY WARRANTY; without even the implied warranty of
+  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
+  Lesser General Public License for more details.
+
+  You should have received a copy of the GNU Lesser General Public License
+  along with systemd; If not, see <http://www.gnu.org/licenses/>.
+***/
+
+#include <sys/types.h>
+
+#include <lzma.h>
+#include <zlib.h>
+#include <bzlib.h>
+
+#include "macro.h"
+
+typedef enum ImportCompressType {
+        IMPORT_COMPRESS_UNKNOWN,
+        IMPORT_COMPRESS_UNCOMPRESSED,
+        IMPORT_COMPRESS_XZ,
+        IMPORT_COMPRESS_GZIP,
+        IMPORT_COMPRESS_BZIP2,
+        _IMPORT_COMPRESS_TYPE_MAX,
+        _IMPORT_COMPRESS_TYPE_INVALID = -1,
+} ImportCompressType;
+
+typedef struct ImportCompress {
+        ImportCompressType type;
+
+        union {
+                lzma_stream xz;
+                z_stream gzip;
+                bz_stream bzip2;
+        };
+} ImportCompress;
+
+typedef int (*ImportCompressCallback)(const void *data, size_t size, void *userdata);
+
+void import_compress_free(ImportCompress *c);
+
+int import_uncompress_detect(ImportCompress *c, const void *data, size_t size);
+
+int import_uncompress(ImportCompress *c, const void *data, size_t size, ImportCompressCallback callback, void *userdata);
+
+const char* import_compress_type_to_string(ImportCompressType t) _const_;
+ImportCompressType import_compress_type_from_string(const char *s) _pure_;
diff --git a/src/import/pull-job.c b/src/import/pull-job.c
index 165dae6619..b5237449c9 100644
--- a/src/import/pull-job.c
+++ b/src/import/pull-job.c
@@ -37,12 +37,7 @@ PullJob* pull_job_unref(PullJob *j) {
 
         safe_close(j->disk_fd);
 
-        if (j->compressed == PULL_JOB_XZ)
-                lzma_end(&j->xz);
-        else if (j->compressed == PULL_JOB_GZIP)
-                inflateEnd(&j->gzip);
-        else if (j->compressed == PULL_JOB_BZIP2)
-                BZ2_bzDecompressEnd(&j->bzip2);
+        import_compress_free(&j->compress);
 
         if (j->checksum_context)
                 gcry_md_close(j->checksum_context);
@@ -180,7 +175,8 @@ finish:
         pull_job_finish(j, r);
 }
 
-static int pull_job_write_uncompressed(PullJob *j, void *p, size_t sz) {
+static int pull_job_write_uncompressed(const void *p, size_t sz, void *userdata) {
+        PullJob *j = userdata;
         ssize_t n;
 
         assert(j);
@@ -261,88 +257,9 @@ static int pull_job_write_compressed(PullJob *j, void *p, size_t sz) {
         if (j->checksum_context)
                 gcry_md_write(j->checksum_context, p, sz);
 
-        switch (j->compressed) {
-
-        case PULL_JOB_UNCOMPRESSED:
-                r = pull_job_write_uncompressed(j, p, sz);
-                if (r < 0)
-                        return r;
-
-                break;
-
-        case PULL_JOB_XZ:
-                j->xz.next_in = p;
-                j->xz.avail_in = sz;
-
-                while (j->xz.avail_in > 0) {
-                        uint8_t buffer[16 * 1024];
-                        lzma_ret lzr;
-
-                        j->xz.next_out = buffer;
-                        j->xz.avail_out = sizeof(buffer);
-
-                        lzr = lzma_code(&j->xz, LZMA_RUN);
-                        if (lzr != LZMA_OK && lzr != LZMA_STREAM_END) {
-                                log_error("Decompression error.");
-                                return -EIO;
-                        }
-
-                        r = pull_job_write_uncompressed(j, buffer, sizeof(buffer) - j->xz.avail_out);
-                        if (r < 0)
-                                return r;
-                }
-
-                break;
-
-        case PULL_JOB_GZIP:
-                j->gzip.next_in = p;
-                j->gzip.avail_in = sz;
-
-                while (j->gzip.avail_in > 0) {
-                        uint8_t buffer[16 * 1024];
-
-                        j->gzip.next_out = buffer;
-                        j->gzip.avail_out = sizeof(buffer);
-
-                        r = inflate(&j->gzip, Z_NO_FLUSH);
-                        if (r != Z_OK && r != Z_STREAM_END) {
-                                log_error("Decompression error.");
-                                return -EIO;
-                        }
-
-                        r = pull_job_write_uncompressed(j, buffer, sizeof(buffer) - j->gzip.avail_out);
-                        if (r < 0)
-                                return r;
-                }
-
-                break;
-
-        case PULL_JOB_BZIP2:
-                j->bzip2.next_in = p;
-                j->bzip2.avail_in = sz;
-
-                while (j->bzip2.avail_in > 0) {
-                        uint8_t buffer[16 * 1024];
-
-                        j->bzip2.next_out = (char*) buffer;
-                        j->bzip2.avail_out = sizeof(buffer);
-
-                        r = BZ2_bzDecompress(&j->bzip2);
-                        if (r != BZ_OK && r != BZ_STREAM_END) {
-                                log_error("Decompression error.");
-                                return -EIO;
-                        }
-
-                        r = pull_job_write_uncompressed(j,  buffer, sizeof(buffer) - j->bzip2.avail_out);
-                        if (r < 0)
-                                return r;
-                }
-
-                break;
-
-        default:
-                assert_not_reached("Unknown compression");
-        }
+        r = import_uncompress(&j->compress, p, sz, pull_job_write_uncompressed, j);
+        if (r < 0)
+                return r;
 
         j->written_compressed += sz;
 
@@ -384,16 +301,6 @@ static int pull_job_open_disk(PullJob *j) {
 }
 
 static int pull_job_detect_compression(PullJob *j) {
-        static const uint8_t xz_signature[] = {
-                0xfd, '7', 'z', 'X', 'Z', 0x00
-        };
-        static const uint8_t gzip_signature[] = {
-                0x1f, 0x8b
-        };
-        static const uint8_t bzip2_signature[] = {
-                'B', 'Z', 'h'
-        };
-
         _cleanup_free_ uint8_t *stub = NULL;
         size_t stub_size;
 
@@ -401,47 +308,13 @@ static int pull_job_detect_compression(PullJob *j) {
 
         assert(j);
 
-        if (j->payload_size < MAX3(sizeof(xz_signature),
-                                   sizeof(gzip_signature),
-                                   sizeof(bzip2_signature)))
+        r = import_uncompress_detect(&j->compress, j->payload, j->payload_size);
+        if (r < 0)
+                return log_error_errno(r, "Failed to initialize compressor: %m");
+        if (r == 0)
                 return 0;
 
-        if (memcmp(j->payload, xz_signature, sizeof(xz_signature)) == 0)
-                j->compressed = PULL_JOB_XZ;
-        else if (memcmp(j->payload, gzip_signature, sizeof(gzip_signature)) == 0)
-                j->compressed = PULL_JOB_GZIP;
-        else if (memcmp(j->payload, bzip2_signature, sizeof(bzip2_signature)) == 0)
-                j->compressed = PULL_JOB_BZIP2;
-        else
-                j->compressed = PULL_JOB_UNCOMPRESSED;
-
-        log_debug("Stream is XZ compressed: %s", yes_no(j->compressed == PULL_JOB_XZ));
-        log_debug("Stream is GZIP compressed: %s", yes_no(j->compressed == PULL_JOB_GZIP));
-        log_debug("Stream is BZIP2 compressed: %s", yes_no(j->compressed == PULL_JOB_BZIP2));
-
-        if (j->compressed == PULL_JOB_XZ) {
-                lzma_ret xzr;
-
-                xzr = lzma_stream_decoder(&j->xz, UINT64_MAX, LZMA_TELL_UNSUPPORTED_CHECK);
-                if (xzr != LZMA_OK) {
-                        log_error("Failed to initialize XZ decoder.");
-                        return -EIO;
-                }
-        }
-        if (j->compressed == PULL_JOB_GZIP) {
-                r = inflateInit2(&j->gzip, 15+16);
-                if (r != Z_OK) {
-                        log_error("Failed to initialize gzip decoder.");
-                        return -EIO;
-                }
-        }
-        if (j->compressed == PULL_JOB_BZIP2) {
-                r = BZ2_bzDecompressInit(&j->bzip2, 0, 0);
-                if (r != BZ_OK) {
-                        log_error("Failed to initialize bzip2 decoder.");
-                        return -EIO;
-                }
-        }
+        log_debug("Stream is compressed: %s", import_compress_type_to_string(j->compress.type));
 
         r = pull_job_open_disk(j);
         if (r < 0)
diff --git a/src/import/pull-job.h b/src/import/pull-job.h
index b807bd1b41..3239aeac20 100644
--- a/src/import/pull-job.h
+++ b/src/import/pull-job.h
@@ -21,13 +21,11 @@
   along with systemd; If not, see <http://www.gnu.org/licenses/>.
 ***/
 
-#include <lzma.h>
-#include <zlib.h>
-#include <bzlib.h>
 #include <gcrypt.h>
 
 #include "macro.h"
 #include "curl-util.h"
+#include "import-compress.h"
 
 typedef struct PullJob PullJob;
 
@@ -92,10 +90,7 @@ struct PullJob {
 
         usec_t mtime;
 
-        PullJobCompression compressed;
-        lzma_stream xz;
-        z_stream gzip;
-        bz_stream bzip2;
+        ImportCompress compress;
 
         unsigned progress_percent;
         usec_t start_usec;