233 lines
8.9 KiB
C++
233 lines
8.9 KiB
C++
#include "sam_io.h"
|
||
|
||
#include <htslib/sam.h>
|
||
#include <libdeflate.h>
|
||
#include <stdint.h>
|
||
#include <stdio.h>
|
||
#include <zlib.h>
|
||
#include <spdlog/spdlog.h>
|
||
|
||
#include "const_val.h"
|
||
|
||
namespace nsgv {
|
||
extern bool gIsBigEndian;
|
||
};
|
||
|
||
static const uint8_t g_magic[19] = "\037\213\010\4\0\0\0\0\0\377\6\0\102\103\2\0\0\0";
|
||
|
||
size_t fileSize(FILE *fp) {
|
||
spdlog::info("test file start");
|
||
size_t bufSize = 64L * 1024 * 1024;
|
||
uint8_t *testBuf = (uint8_t *)malloc(bufSize);
|
||
size_t readState = 0;
|
||
size_t fileSize = 0;
|
||
while ((readState = fread(testBuf, 1, bufSize, fp)) > 0) {
|
||
fileSize += readState;
|
||
}
|
||
spdlog::info("test file size: {}", fileSize);
|
||
return fileSize;
|
||
}
|
||
|
||
static inline void packInt16(uint8_t* buffer, uint16_t value) {
|
||
buffer[0] = value;
|
||
buffer[1] = value >> 8;
|
||
}
|
||
|
||
static inline void packInt32(uint8_t* buffer, uint32_t value) {
|
||
buffer[0] = value;
|
||
buffer[1] = value >> 8;
|
||
buffer[2] = value >> 16;
|
||
buffer[3] = value >> 24;
|
||
}
|
||
|
||
// dlen本身长度一定得够存放压缩后的数据，否则会返回错误
|
||
int bgzfCompress(void* _dst, size_t* dlen, const void* src, size_t slen, int level) {
|
||
// fprintf(stderr, "bgzfCompress, slen: %zu\n", slen);
|
||
if (slen == 0) {
|
||
// EOF block
|
||
if (*dlen < 28)
|
||
return -1;
|
||
memcpy(_dst, "\037\213\010\4\0\0\0\0\0\377\6\0\102\103\2\0\033\0\3\0\0\0\0\0\0\0\0\0", 28);
|
||
*dlen = 28;
|
||
return 0;
|
||
}
|
||
|
||
uint8_t* dst = (uint8_t*)_dst;
|
||
|
||
if (level == 0) {
|
||
// Uncompressed data
|
||
if (*dlen < slen + 5 + BLOCK_HEADER_LENGTH + BLOCK_FOOTER_LENGTH)
|
||
return -1;
|
||
dst[BLOCK_HEADER_LENGTH] = 1; // BFINAL=1, BTYPE=00; see RFC1951
|
||
u16_to_le(slen, &dst[BLOCK_HEADER_LENGTH + 1]); // length
|
||
u16_to_le(~slen, &dst[BLOCK_HEADER_LENGTH + 3]); // ones-complement length
|
||
memcpy(dst + BLOCK_HEADER_LENGTH + 5, src, slen);
|
||
*dlen = slen + 5 + BLOCK_HEADER_LENGTH + BLOCK_FOOTER_LENGTH;
|
||
|
||
} else {
|
||
level = level > 0 ? level : 6; // libdeflate doesn't honour -1 as default
|
||
// NB levels go up to 12 here.
|
||
int lvl_map[] = {0, 1, 2, 3, 5, 6, 7, 8, 10, 12};
|
||
level = lvl_map[level > 9 ? 9 : level];
|
||
struct libdeflate_compressor* z = libdeflate_alloc_compressor(level);
|
||
if (!z)
|
||
return -1;
|
||
|
||
// Raw deflate
|
||
size_t clen = libdeflate_deflate_compress(z, src, slen, dst + BLOCK_HEADER_LENGTH, *dlen - BLOCK_HEADER_LENGTH - BLOCK_FOOTER_LENGTH);
|
||
|
||
if (clen <= 0) {
|
||
hts_log_error("Call to libdeflate_deflate_compress failed");
|
||
libdeflate_free_compressor(z);
|
||
return -1;
|
||
}
|
||
|
||
*dlen = clen + BLOCK_HEADER_LENGTH + BLOCK_FOOTER_LENGTH;
|
||
|
||
libdeflate_free_compressor(z);
|
||
}
|
||
|
||
// write the header
|
||
memcpy(dst, g_magic, BLOCK_HEADER_LENGTH); // the last two bytes are a place holder for the length of the block
|
||
packInt16(&dst[16], *dlen - 1); // write the compressed length; -1 to fit 2 bytes
|
||
|
||
// write the footer
|
||
uint32_t crc = libdeflate_crc32(0, src, slen);
|
||
packInt32((uint8_t*)&dst[*dlen - 8], crc);
|
||
packInt32((uint8_t*)&dst[*dlen - 4], slen);
|
||
return 0;
|
||
}
|
||
|
||
int bgzfUncompress(uint8_t *dst, size_t *dlen, const uint8_t *src, size_t slen, uint32_t expected_crc) {
|
||
struct libdeflate_decompressor *z = libdeflate_alloc_decompressor();
|
||
if (!z) {
|
||
hts_log_error("Call to libdeflate_alloc_decompressor failed");
|
||
return -1;
|
||
}
|
||
|
||
int ret = libdeflate_deflate_decompress(z, src, slen, dst, *dlen, dlen);
|
||
libdeflate_free_decompressor(z);
|
||
|
||
if (ret != LIBDEFLATE_SUCCESS) {
|
||
hts_log_error("Inflate operation failed: %d", ret);
|
||
return -1;
|
||
}
|
||
|
||
uint32_t crc = libdeflate_crc32(0, (unsigned char *)dst, *dlen);
|
||
#ifdef FUZZING_BUILD_MODE_UNSAFE_FOR_PRODUCTION
|
||
// Pretend the CRC was OK so the fuzzer doesn't have to get it right
|
||
crc = expected_crc;
|
||
#endif
|
||
if (crc != expected_crc) {
|
||
hts_log_error("CRC32 checksum mismatch");
|
||
return -2;
|
||
}
|
||
|
||
return 0;
|
||
}
|
||
|
||
// 读取并解压一个gz block
|
||
size_t readUncompressOneBlock(FILE *fpr, uint8_t *fBuf, DataBuffer *uDataPtr) {
|
||
DataBuffer &uData = *uDataPtr;
|
||
size_t readState = 0;
|
||
int blockLen = 0; // 每个gz block的大小
|
||
size_t dlen = SINGLE_BLOCK_SIZE; // gz block解压之后的真正大小
|
||
readState = fread(fBuf, 1, BLOCK_HEADER_LENGTH, fpr); // 先读取一个gz block的头部
|
||
blockLen = unpackInt16(&fBuf[16]) + 1; // block的字节数
|
||
readState = fread(&fBuf[BLOCK_HEADER_LENGTH], 1, blockLen - BLOCK_HEADER_LENGTH, fpr);
|
||
// 解压gz block
|
||
uint32_t crc = le_to_u32(fBuf + blockLen - 8);
|
||
size_t newDataSize = uData.maxLen;
|
||
while (uData.curLen + SINGLE_BLOCK_SIZE > newDataSize) newDataSize *= 2;
|
||
uData.ReAllocMem(newDataSize); // 需要重新开辟空间
|
||
int ret = bgzfUncompress(&uData.data[uData.curLen], &dlen, (Bytef *)fBuf + BLOCK_HEADER_LENGTH,
|
||
blockLen - BLOCK_HEADER_LENGTH, crc);
|
||
if (ret < 0) {
|
||
spdlog::error("Error decompressing gz block");
|
||
exit(0);
|
||
}
|
||
uData.curLen += dlen;
|
||
spdlog::info("gz block size: {}, uncompressed size: {}", blockLen, dlen);
|
||
return readState;
|
||
}
|
||
|
||
// 解析sam头部信息
|
||
void parseSamHeader(FILE *fpr, HeaderBuf &hdrBuf) {
|
||
const int kMaxBlockSize = SINGLE_BLOCK_SIZE;
|
||
sam_hdr_t *header = NULL;
|
||
uint8_t *fBuf = (uint8_t*)malloc(kMaxBlockSize);
|
||
DataBuffer uData; // 用来放解压缩后的数据，待解析
|
||
int magicLen;
|
||
int32_t i, nameLen, numNames = 0;
|
||
|
||
header = sam_hdr_init(); // 初始化header
|
||
uData.AllocMem(kMaxBlockSize); // 初始化解压数据的buffer
|
||
readUncompressOneBlock(fpr, fBuf, &uData); // 读取第一个gz block
|
||
|
||
// 解析header
|
||
magicLen = 4;
|
||
if (memcmp(uData.data, "BAM\1", magicLen)) {
|
||
spdlog::error("Invalid BAM binary header");
|
||
return;
|
||
}
|
||
uData.readPos += magicLen; // 在解压数据中读取了magicLen个字节
|
||
header->l_text = le_to_u32(uData.data + uData.readPos); uData.readPos += 4; // 解析l_text的长度
|
||
header->text = (char *)malloc(header->l_text + 1);
|
||
header->text[header->l_text] = 0; // 在text末尾添加'\0'
|
||
|
||
while (uData.readPos + header->l_text > uData.curLen) { // header内容超过了一个block，继续读
|
||
readUncompressOneBlock(fpr, fBuf, &uData); // 读取包含header内容的gz block
|
||
}
|
||
|
||
memcpy(header->text, &uData.data[uData.readPos], header->l_text);
|
||
uData.readPos += header->l_text;
|
||
|
||
memcpy(&header->n_targets, &uData.data[uData.readPos], 4); // n_targets
|
||
uData.readPos += 4;
|
||
spdlog::info("num target: {}", header->n_targets);
|
||
if (nsgv::gIsBigEndian) ed_swap_4p(&header->n_targets);
|
||
if (header->n_targets > 0) {
|
||
header->target_name = (char **)calloc(header->n_targets, sizeof(char *));
|
||
header->target_len = (uint32_t *)calloc(header->n_targets, sizeof(uint32_t));
|
||
} else {
|
||
header->target_name = NULL;
|
||
header->target_len = NULL;
|
||
}
|
||
|
||
for (int i = 0; i != header->n_targets; ++i) {
|
||
while (uData.readPos + 4 > uData.curLen) readUncompressOneBlock(fpr, fBuf, &uData);
|
||
memcpy(&nameLen, &uData.data[uData.readPos], 4);
|
||
uData.readPos += 4;
|
||
if (nsgv::gIsBigEndian) ed_swap_4p(&nameLen);
|
||
header->target_name[i] = (char *)malloc(nameLen);
|
||
++numNames;
|
||
while (uData.readPos + nameLen > uData.curLen) readUncompressOneBlock(fpr, fBuf, &uData);
|
||
memcpy(header->target_name[i], &uData.data[uData.readPos], nameLen);
|
||
uData.readPos += nameLen;
|
||
if (header->target_name[i][nameLen - 1] != '\0') {
|
||
/* Fix missing NUL-termination. Is this being too nice?
|
||
We could alternatively bail out with an error. */
|
||
char *newName = (char*)realloc(header->target_name[i], nameLen + 1);
|
||
header->target_name[i] = newName;
|
||
header->target_name[i][nameLen] = '\0';
|
||
}
|
||
while (uData.readPos + 4 > uData.curLen) readUncompressOneBlock(fpr, fBuf, &uData);
|
||
memcpy(&header->target_len[i], &uData.data[uData.readPos], 4);
|
||
uData.readPos += 4;
|
||
if (nsgv::gIsBigEndian) ed_swap_4p(&header->target_len[i]);
|
||
|
||
// spdlog::info("nameLen {}, {}, {}", nameLen, header->target_len[i], header->target_name[i]);
|
||
}
|
||
// spdlog::info("test res: {} {} {} {}", header->l_text, dlen, header->n_targets, header->text);
|
||
|
||
spdlog::info("udata: readPos: {}, curLen: {}, maxLen: {}", uData.readPos, uData.curLen, uData.maxLen);
|
||
|
||
hdrBuf.header = header;
|
||
if (uData.readPos < uData.curLen) { // 如果还有数据，则将其保存在hdrBuf中，用以后续的bam解析
|
||
hdrBuf.data = (uint8_t *)malloc(uData.curLen - uData.readPos);
|
||
memcpy(hdrBuf.data, &uData.data[uData.readPos], uData.curLen - uData.readPos);
|
||
hdrBuf.dataLen = uData.curLen - uData.readPos;
|
||
}
|
||
|
||
free(fBuf);
|
||
} |