ClickHouse/utils/self-extracting-executable/decompressor.cpp

456 lines
13 KiB
C++
Raw Normal View History

2022-03-29 19:30:09 +00:00
#include <zstd.h>
#include <sys/mman.h>
2022-07-21 15:43:00 +00:00
#if defined(OS_DARWIN) || defined(OS_FREEBSD)
# include <sys/mount.h>
2022-06-13 10:14:21 +00:00
#else
2022-07-21 15:43:00 +00:00
# include <sys/statfs.h>
2022-06-13 10:14:21 +00:00
#endif
2022-03-29 19:30:09 +00:00
#include <fcntl.h>
#include <sys/wait.h>
2022-03-29 19:30:09 +00:00
#include <unistd.h>
2022-06-28 11:29:07 +00:00
#include <cstdlib>
#include <cstdio>
#include <cstring>
#include <iostream>
#include <filesystem>
#if (defined(OS_DARWIN) || defined(OS_FREEBSD)) && defined(__GNUC__)
2022-07-21 06:08:07 +00:00
# include <machine/endian.h>
2022-07-25 01:13:00 +00:00
#else
2022-07-21 06:08:07 +00:00
# include <endian.h>
#endif
2022-07-15 14:48:34 +00:00
2022-07-21 06:08:07 +00:00
#if defined OS_DARWIN
# include <mach-o/dyld.h>
2022-07-21 06:08:07 +00:00
# include <libkern/OSByteOrder.h>
// define 64 bit macros
# define le64toh(x) OSSwapLittleToHostInt64(x)
#endif
2022-03-29 19:30:09 +00:00
#if defined(OS_FREEBSD)
# include <sys/sysctl.h>
#endif
#include "types.h"
2022-04-07 12:21:14 +00:00
2022-03-29 19:30:09 +00:00
/// decompress part
int doDecompress(char * input, char * output, off_t & in_offset, off_t & out_offset,
off_t input_size, off_t output_size, ZSTD_DCtx* dctx)
{
size_t decompressed_size = ZSTD_decompressDCtx(dctx, output + out_offset, output_size, input + in_offset, input_size);
if (ZSTD_isError(decompressed_size))
{
std::cerr << "Error (ZSTD):" << decompressed_size << " " << ZSTD_getErrorName(decompressed_size) << std::endl;
2022-03-29 19:30:09 +00:00
return 1;
}
return 0;
}
/// decompress data from in_fd into out_fd
int decompress(char * input, char * output, off_t start, off_t end, size_t max_number_of_forks=10)
2022-04-07 12:21:14 +00:00
{
off_t in_pointer = start, out_pointer = 0;
off_t size = 0;
off_t max_block_size = 1ull<<27;
off_t decompressed_size = 0;
size_t number_of_forks = 0;
2022-04-07 12:21:14 +00:00
/// Create context
ZSTD_DCtx * dctx = ZSTD_createDCtx();
if (dctx == nullptr)
{
std::cerr << "Error (ZSTD): failed to create decompression context" << std::endl;
2022-04-07 12:21:14 +00:00
return 1;
}
pid_t pid;
bool error_happened = false;
2022-04-07 12:21:14 +00:00
/// Decompress data
while (in_pointer < end && !error_happened)
2022-04-07 12:21:14 +00:00
{
size = ZSTD_findFrameCompressedSize(input + in_pointer, max_block_size);
if (ZSTD_isError(size))
{
std::cerr << "Error (ZSTD): " << size << " " << ZSTD_getErrorName(size) << std::endl;
error_happened = true;
break;
}
2022-04-07 12:21:14 +00:00
decompressed_size = ZSTD_getFrameContentSize(input + in_pointer, max_block_size);
if (ZSTD_isError(decompressed_size))
{
std::cerr << "Error (ZSTD): " << decompressed_size << " " << ZSTD_getErrorName(decompressed_size) << std::endl;
error_happened = true;
break;
}
pid = fork();
if (-1 == pid)
{
perror("fork");
/// If fork failed just decompress data in main process.
if (0 != doDecompress(input, output, in_pointer, out_pointer, size, decompressed_size, dctx))
{
error_happened = true;
break;
}
in_pointer += size;
out_pointer += decompressed_size;
}
else if (pid == 0)
{
/// Decompress data in child process.
if (0 != doDecompress(input, output, in_pointer, out_pointer, size, decompressed_size, dctx))
2022-08-20 15:09:20 +00:00
_exit(1);
_exit(0);
}
else
{
++number_of_forks;
while (number_of_forks >= max_number_of_forks)
{
/// Wait any fork
int status;
waitpid(0, &status, 0);
/// If error happened, stop processing
if (WEXITSTATUS(status) != 0)
{
error_happened = true;
break;
}
--number_of_forks;
}
in_pointer += size;
out_pointer += decompressed_size;
}
}
/// wait for all working decompressions
while (number_of_forks > 0)
{
/// Wait any fork
int status;
waitpid(0, &status, 0);
if (WIFEXITED(status))
{
if (WEXITSTATUS(status) != 0)
error_happened = true;
}
else
{
error_happened = true;
if (WIFSIGNALED(status))
{
if (WCOREDUMP(status))
std::cerr << "Error: child process core dumped with signal " << WTERMSIG(status) << std::endl;
else
std::cerr << "Error: child process was terminated with signal " << WTERMSIG(status) << std::endl;
}
}
if (WEXITSTATUS(status) != 0)
error_happened = true;
--number_of_forks;
2022-04-07 12:21:14 +00:00
}
ZSTD_freeDCtx(dctx);
/// If error happen end of processed part will not reach end
if (in_pointer < end || error_happened)
return 1;
2022-04-07 12:21:14 +00:00
return 0;
}
/// Read data about files and decomrpess them.
int decompressFiles(int input_fd, char * path, char * name, bool & have_compressed_analoge, bool & has_exec, char * decompressed_suffix, uint64_t * decompressed_umask)
2022-03-29 19:30:09 +00:00
{
/// Read data about output file.
/// Compressed data will replace data in file
struct stat info_in;
2022-04-07 12:21:14 +00:00
if (0 != fstat(input_fd, &info_in))
{
perror("fstat");
2022-04-07 12:21:14 +00:00
return 1;
}
2022-03-29 19:30:09 +00:00
2022-04-07 12:21:14 +00:00
/// mmap input file
2022-04-07 12:31:41 +00:00
char * input = static_cast<char*>(mmap(nullptr, info_in.st_size, PROT_READ, MAP_PRIVATE, input_fd, 0));
2022-04-07 12:21:14 +00:00
if (input == MAP_FAILED)
2022-03-29 19:30:09 +00:00
{
perror("mmap");
2022-03-29 19:30:09 +00:00
return 1;
}
2022-04-07 12:21:14 +00:00
/// Read metadata from end of file
MetaData metadata = *reinterpret_cast<MetaData*>(input + info_in.st_size - sizeof(MetaData));
2022-03-29 19:30:09 +00:00
2022-04-07 12:21:14 +00:00
/// Prepare to read information about files and decompress them
off_t files_pointer = le64toh(metadata.start_of_files_data);
2022-04-07 12:21:14 +00:00
size_t decompressed_full_size = 0;
/// Read files metadata and check if decompression is possible
off_t check_pointer = le64toh(metadata.start_of_files_data);
for (size_t i = 0; i < le64toh(metadata.number_of_files); ++i)
2022-04-07 12:21:14 +00:00
{
FileData data = *reinterpret_cast<FileData*>(input + check_pointer);
decompressed_full_size += le64toh(data.uncompressed_size);
check_pointer += sizeof(FileData) + le64toh(data.name_length);
2022-04-07 12:21:14 +00:00
}
/// Check free space
struct statfs fs_info;
if (0 != fstatfs(input_fd, &fs_info))
2022-03-29 19:30:09 +00:00
{
perror("fstatfs");
2022-04-07 12:21:14 +00:00
if (0 != munmap(input, info_in.st_size))
perror("munmap");
2022-03-29 19:30:09 +00:00
return 1;
2022-04-07 12:21:14 +00:00
}
if (fs_info.f_blocks * info_in.st_blksize < decompressed_full_size)
{
std::cerr << "Not enough space for decompression. Have " << fs_info.f_blocks * info_in.st_blksize << ", need " << decompressed_full_size << std::endl;
2022-04-07 12:21:14 +00:00
return 1;
}
2022-03-29 19:30:09 +00:00
2022-04-07 12:21:14 +00:00
FileData file_info;
/// Decompress files with appropriate file names
for (size_t i = 0; i < le64toh(metadata.number_of_files); ++i)
2022-03-29 19:30:09 +00:00
{
2022-04-07 12:21:14 +00:00
/// Read information about file
file_info = *reinterpret_cast<FileData*>(input + files_pointer);
files_pointer += sizeof(FileData);
2022-03-29 19:30:09 +00:00
2022-08-22 12:14:29 +00:00
/// for output filename matching compressed allow additional 13 + 7 symbols for ".decompressed.XXXXXX" suffix
2022-08-20 21:02:36 +00:00
size_t file_name_len = file_info.exec ? strlen(name) + 13 + 7 + 1 : le64toh(file_info.name_length);
size_t file_path_len = path ? strlen(path) + 1 + file_name_len : file_name_len;
char file_name[file_path_len];
memset(file_name, '\0', file_path_len);
if (path)
2022-04-07 12:21:14 +00:00
{
strcat(file_name, path);
strcat(file_name, "/");
2022-04-07 12:21:14 +00:00
}
bool same_name = false;
if (file_info.exec)
{
has_exec = true;
strcat(file_name, name);
}
else
{
if (strcmp(name, input + files_pointer) == 0)
same_name = true;
strcat(file_name, input + files_pointer);
}
files_pointer += le64toh(file_info.name_length);
if (file_info.exec || same_name)
2022-04-07 12:21:14 +00:00
{
strcat(file_name, ".decompressed.XXXXXX");
int fd = mkstemp(file_name);
2022-08-01 15:06:54 +00:00
if (fd == -1)
{
perror("mkstemp");
return 1;
}
close(fd);
strncpy(decompressed_suffix, file_name + strlen(file_name) - 6, 6);
2022-08-03 07:34:51 +00:00
*decompressed_umask = le64toh(file_info.umask);
have_compressed_analoge = true;
2022-04-07 12:21:14 +00:00
}
int output_fd = open(file_name, O_RDWR | O_CREAT, le64toh(file_info.umask));
2022-04-07 12:21:14 +00:00
if (output_fd == -1)
2022-03-29 19:30:09 +00:00
{
perror("open");
2022-04-07 12:21:14 +00:00
if (0 != munmap(input, info_in.st_size))
perror("munmap");
return 1;
2022-04-07 12:21:14 +00:00
}
/// Prepare output file
if (0 != ftruncate(output_fd, le64toh(file_info.uncompressed_size)))
2022-04-07 12:21:14 +00:00
{
perror("ftruncate");
2022-04-07 12:21:14 +00:00
if (0 != munmap(input, info_in.st_size))
perror("munmap");
2022-04-07 12:21:14 +00:00
return 1;
}
char * output = static_cast<char*>(
mmap(nullptr,
le64toh(file_info.uncompressed_size),
PROT_READ | PROT_WRITE, MAP_SHARED,
output_fd,
2022-04-07 12:21:14 +00:00
0)
);
if (output == MAP_FAILED)
{
perror("mmap");
2022-04-07 12:21:14 +00:00
if (0 != munmap(input, info_in.st_size))
perror("munmap");
2022-03-29 19:30:09 +00:00
return 1;
}
2022-04-07 12:21:14 +00:00
/// Decompress data into file
if (0 != decompress(input, output, le64toh(file_info.start), le64toh(file_info.end)))
2022-04-07 12:21:14 +00:00
{
if (0 != munmap(input, info_in.st_size))
perror("munmap");
if (0 != munmap(output, le64toh(file_info.uncompressed_size)))
perror("munmap");
2022-04-07 12:21:14 +00:00
return 1;
}
if (0 != fsync(output_fd))
perror("fsync");
2022-04-07 12:21:14 +00:00
if (0 != close(output_fd))
perror("close");
2022-04-07 12:21:14 +00:00
}
2022-04-07 12:21:14 +00:00
if (0 != munmap(input, info_in.st_size))
perror("munmap");
2022-03-29 19:30:09 +00:00
return 0;
}
2022-03-24 19:45:54 +00:00
#if defined(OS_DARWIN)
int read_exe_path(char *exe, size_t buf_sz)
{
uint32_t size = buf_sz;
char apple[size];
if (_NSGetExecutablePath(apple, &size) != 0)
return 1;
if (realpath(apple, exe) == nullptr)
return 1;
return 0;
}
#elif defined(OS_FREEBSD)
int read_exe_path(char *exe, size_t buf_sz)
{
int name[] = { CTL_KERN, KERN_PROC, KERN_PROC_PATHNAME, -1 };
size_t length = buf_sz;
int error = sysctl(name, 4, exe, &length, nullptr, 0);
if (error < 0 || length <= 1)
return 1;
return 0;
}
#else
int read_exe_path(char *exe, size_t/* buf_sz*/)
{
if (realpath("/proc/self/exe", exe) == nullptr)
return 1;
return 0;
}
#endif
2022-08-01 00:07:33 +00:00
int main(int/* argc*/, char* argv[])
2022-03-24 19:45:54 +00:00
{
char self[4096] = {0};
if (read_exe_path(self, 4096) == -1)
{
perror("read_exe_path");
return 1;
}
char file_path[strlen(self) + 1];
strcpy(file_path, self);
char * path = nullptr;
char * name = strrchr(file_path, '/');
if (name)
{
path = file_path;
*name = 0;
++name;
}
else
name = file_path;
int input_fd = open(self, O_RDONLY);
2022-03-29 19:30:09 +00:00
if (input_fd == -1)
{
perror("open");
return 1;
2022-03-29 19:30:09 +00:00
}
bool have_compressed_analoge = false;
bool has_exec = false;
2022-08-03 07:34:51 +00:00
char decompressed_suffix[7] = {0};
uint64_t decompressed_umask = 0;
2022-04-07 12:21:14 +00:00
/// Decompress all files
if (0 != decompressFiles(input_fd, path, name, have_compressed_analoge, has_exec, decompressed_suffix, &decompressed_umask))
2022-03-29 19:30:09 +00:00
{
printf("Error happened during decompression.\n");
2022-04-07 12:21:14 +00:00
if (0 != close(input_fd))
perror("close");
2022-03-29 19:30:09 +00:00
return 1;
}
2022-04-07 12:21:14 +00:00
if (0 != close(input_fd))
perror("close");
2022-04-07 12:21:14 +00:00
if (unlink(self))
{
perror("unlink");
return 1;
}
2022-08-01 18:32:32 +00:00
if (!have_compressed_analoge)
printf("No target executable - decompression only was performed.\n");
else
{
const char * const decompressed_name_fmt = "%s.decompressed.%s";
int decompressed_name_len = snprintf(nullptr, 0, decompressed_name_fmt, self, decompressed_suffix);
char decompressed_name[decompressed_name_len + 1];
(void)snprintf(decompressed_name, decompressed_name_len + 1, decompressed_name_fmt, self, decompressed_suffix);
std::error_code ec;
std::filesystem::copy_file(static_cast<char *>(decompressed_name), static_cast<char *>(self), ec);
if (ec)
{
std::cerr << ec.message() << std::endl;
return 1;
}
2022-08-01 15:06:54 +00:00
if (chmod(self, decompressed_umask))
{
perror("chmod");
return 1;
}
if (unlink(decompressed_name))
{
perror("unlink");
return 1;
}
2022-08-01 00:07:33 +00:00
2022-08-21 22:06:11 +00:00
if (has_exec)
{
execv(self, argv);
/// This part of code will be reached only if error happened
perror("execv");
return 1;
}
printf("No target executable - decompression only was performed.\n");
}
2022-03-24 19:45:54 +00:00
}