/* * Copyright (c) 2016-present, Facebook, Inc. * All rights reserved. * * This source code is licensed under both the BSD-style license (found in the * LICENSE file in the root directory of this source tree) and the GPLv2 (found * in the COPYING file in the root directory of this source tree). */ /* Compression level or 0 to disable */ #define DO_ZLIB 9 /* Compression level or 0 to disable */ #define DO_ZSTD 0 /* Buffer size */ #define BUFFER_SIZE 4096 #include #include #include #include #include #if DO_ZSTD #include #endif #if DO_ZLIB #include #endif /* Device name to pass to register_chrdev(). */ #define DEVICE_NAME "zstd_compress_test" /* Dynamically allocated device major number */ static int device_major; /* * We reuse the same state, and thus can compress only one file at a time. */ static bool device_is_open; static void *workspace = NULL; /* * Input buffer used to put data coming from userspace. */ static uint8_t buffer_in[BUFFER_SIZE]; static uint8_t buffer_out[BUFFER_SIZE]; static uint64_t uncompressed_len; static uint64_t compressed_len; #if DO_ZSTD static ZSTD_CStream *state; static ZSTD_inBuffer input = { .src = buffer_in, .size = sizeof(buffer_in), .pos = sizeof(buffer_in), }; static ZSTD_outBuffer output = { .dst = buffer_out, .size = sizeof(buffer_out), .pos = sizeof(buffer_out), }; #endif /* DO_ZSTD */ #if DO_ZLIB static z_stream state = { .next_in = buffer_in, .avail_in = 0, .total_in = 0, .next_out = buffer_out, .avail_out = sizeof(buffer_out), .total_out = 0, .msg = NULL, .state = NULL, .workspace = NULL, }; #endif /* DO_ZLIB */ static int zstd_compress_test_open(struct inode *i, struct file *f) { if (device_is_open) return -EBUSY; device_is_open = true; uncompressed_len = compressed_len = 0; #if DO_ZSTD if (ZSTD_isError(ZSTD_resetCStream(state, 0))) return -EIO; #endif #if DO_ZLIB if (zlib_deflateReset(&state) != Z_OK) return -EIO; #endif printk(KERN_INFO DEVICE_NAME ": opened\n"); return 0; } static int zstd_compress_test_release(struct inode *i, struct file *f) { device_is_open = false; #if DO_ZSTD do { size_t ret; output.pos = 0; ret = ZSTD_endStream(state, &output); if (ZSTD_isError(ret)) { printk(KERN_INFO DEVICE_NAME ": zstd end error %u\n", ZSTD_getErrorCode(ret)); return -EIO; } compressed_len += output.pos; } while (output.pos != output.size); #endif #if DO_ZLIB for (;;) { int ret; state.next_out = buffer_out; state.avail_out = sizeof(buffer_out); ret = zlib_deflate(&state, Z_FINISH); compressed_len += sizeof(buffer_out) - state.avail_out; if (ret == Z_STREAM_END) break; if (ret != Z_OK) { printk(KERN_INFO DEVICE_NAME ": zlib end error %d: %s\n", ret, state.msg); return -EIO; } } #endif printk(KERN_INFO DEVICE_NAME ": uncompressed_len = %llu\n", uncompressed_len); printk(KERN_INFO DEVICE_NAME ": compressed_len = %llu\n", compressed_len); printk(KERN_INFO DEVICE_NAME ": closed\n"); return 0; } /* * Hash the data given to us from userspace. */ static ssize_t zstd_compress_test_write(struct file *file, const char __user *buf, size_t size, loff_t *pos) { size_t remaining = size; while (remaining > 0) { size_t const copy_size = min(remaining, sizeof(buffer_in)); if (copy_from_user(buffer_in, buf, copy_size)) return -EFAULT; buf += copy_size; remaining -= copy_size; uncompressed_len += copy_size; #if DO_ZSTD input.pos = 0; input.size = copy_size; while (input.pos != input.size) { size_t ret; output.pos = 0; ret = ZSTD_compressStream(state, &output, &input); if (ZSTD_isError(ret)) { printk(KERN_INFO DEVICE_NAME ": zstd compress error %u\n", ZSTD_getErrorCode(ret)); return -EIO; } compressed_len += output.pos; } #endif #if DO_ZLIB state.next_in = buffer_in; state.avail_in = copy_size; while (state.avail_in > 0) { int ret; state.next_out = buffer_out; state.avail_out = sizeof(buffer_out); ret = zlib_deflate(&state, Z_NO_FLUSH); compressed_len += sizeof(buffer_out) - state.avail_out; if (ret != Z_OK) { printk(KERN_INFO DEVICE_NAME ": zlib end error %d: %s\n", ret, state.msg); return -EIO; } } #endif } return size; } /* register the character device. */ static int __init zstd_compress_test_init(void) { static const struct file_operations fileops = { .owner = THIS_MODULE, .open = &zstd_compress_test_open, .release = &zstd_compress_test_release, .write = &zstd_compress_test_write }; size_t workspace_size = 0; #if DO_ZSTD ZSTD_parameters params; #endif device_major = register_chrdev(0, DEVICE_NAME, &fileops); if (device_major < 0) { return device_major; } #if DO_ZSTD params = ZSTD_getParams(DO_ZSTD, 0, 0); workspace_size = ZSTD_CStreamWorkspaceBound(params.cParams); if (!(workspace = vmalloc(workspace_size))) goto fail; if (!(state = ZSTD_initCStream(params, 0, workspace, workspace_size))) goto fail; #endif #if DO_ZLIB workspace_size = zlib_deflate_workspacesize(MAX_WBITS, MAX_MEM_LEVEL); if (!(workspace = vmalloc(workspace_size))) goto fail; state.workspace = workspace; if (zlib_deflateInit(&state, DO_ZLIB) != Z_OK) goto fail; #endif printk(KERN_INFO DEVICE_NAME ": module loaded\n"); printk(KERN_INFO DEVICE_NAME ": compression requires %zu bytes of memory\n", workspace_size); printk(KERN_INFO DEVICE_NAME ": Create a device node with " "'mknod " DEVICE_NAME " c %d 0' and write data " "to it.\n", device_major); return 0; fail: printk(KERN_INFO DEVICE_NAME ": failed to load module\n"); if (workspace) { vfree(workspace); workspace = NULL; } return -ENOMEM; } static void __exit zstd_compress_test_exit(void) { unregister_chrdev(device_major, DEVICE_NAME); #if DO_ZLIB zlib_deflateEnd(&state); #endif if (workspace) { vfree(workspace); workspace = NULL; } printk(KERN_INFO DEVICE_NAME ": module unloaded\n"); } module_init(zstd_compress_test_init); module_exit(zstd_compress_test_exit); MODULE_DESCRIPTION("Zstd compression tester"); MODULE_VERSION("1.0"); MODULE_LICENSE("Dual BSD/GPL");