diff --git a/java/common/src/main/resources/org/apache/tsfile/i18n/messages.properties b/java/common/src/main/resources/org/apache/tsfile/i18n/messages.properties index 2836b8c1a8..3700a46f61 100644 --- a/java/common/src/main/resources/org/apache/tsfile/i18n/messages.properties +++ b/java/common/src/main/resources/org/apache/tsfile/i18n/messages.properties @@ -636,6 +636,7 @@ log.read.block_reader_fill_measurements_error = Cannot fill measurements # TableResultSet — failed to close tsBlockReader log.read.table_result_set_close_error = Failed to close tsBlockReader +log.read.tree_result_set_close_error = Failed to close tree query dataset # DeviceMetaIterator — failed to load device meta data log.read.device_meta_iterator_error = Failed to load device meta data @@ -1117,6 +1118,44 @@ error.encrypt.encrypt_no_constructor = Get constructor for encryptor failed: %1$ # EncryptUtils.getEncrypt — instantiation failed (encryptType arg) error.encrypt.encrypt_instantiation_failed = New encryptor instance failed: %1$s +# EncryptionProviderRegistry — provider registration and lookup +error.encrypt.invalid_provider_id = Encryption provider id must not be null or empty +error.encrypt.duplicate_provider = Encryption provider id is already registered: %1$s +error.encrypt.provider_not_found = Encryption provider is not registered: %1$s +error.encrypt.provider_returned_null = Encryption provider returned null: %1$s + +# EncryptParameter — use after sensitive state was destroyed +error.encrypt.parameter_destroyed = Encryption parameters have been destroyed +error.encrypt.invalid_page_body_overhead = Invalid page body overhead %1$s for encryption profile %2$s +error.encrypt.page_body_too_short = Encrypted page body size %1$s is smaller than AEAD overhead %2$s +error.encrypt.page_context_invalid_parameter = Page crypto context requires page AEAD parameters +error.encrypt.page_aead_encrypt_unsupported = Encryptor %1$s does not implement authenticated page encryption +error.encrypt.page_aead_decrypt_unsupported = Decryptor %1$s does not implement authenticated page decryption +error.encrypt.page_context_invalid_file_id = Page crypto context requires a 16-byte file crypto id +error.encrypt.page_context_invalid_chunk_ordinal = Page crypto context requires a nonnegative chunk ordinal +error.encrypt.page_read_requires_chunk_ordinal = Page-AEAD files require readPage(header, type, pageIndex, chunkOrdinal) +error.encrypt.chunk_ordinal_exhausted = File chunk ordinal space is exhausted +error.write.invalid_chunk_ordinal = Encrypted chunks require a file-unique chunk ordinal +error.write.unencrypted_chunk_has_ordinal = Unencrypted chunks must not have a chunk ordinal +error.write.aead_append_requires_file_header = Cannot append encrypted data to a file without an encryption header +error.write.encrypted_chunk_copy_unsupported = Encrypted chunks must be decrypted and re-encrypted before copying +error.encrypt.page_context_invalid_sizes = Invalid page crypto context values, uncompressed size: %1$s, compressed plaintext size: %2$s, page index: %3$s +error.encrypt.page_output_size_mismatch = AEAD page output size mismatch, expected %1$s bytes but got %2$s +error.encrypt.page_plaintext_size_mismatch = AEAD page plaintext size mismatch, expected %1$s bytes but got %2$s +error.chunk.merge_page_aead_unsupported = Page-AEAD encrypted chunks must be decrypted and rewritten before merging + +# FileEncryptionHeader — malformed or unsupported file encryption metadata +error.file.encryption_header_invalid_size = Invalid file encryption header size: %1$s +error.file.encryption_header_unsupported_version = Unsupported file encryption header version: %1$s +error.file.encryption_header_malformed = Malformed file encryption header +error.file.encryption_header_trailing_bytes = File encryption header contains trailing bytes +error.file.encryption_header_invalid_component_size = Invalid file encryption header component size: %1$s +error.file.encryption_header_truncated = File encryption header is truncated, expected %1$s bytes but read %2$s +error.file.encryption_header_invalid_field = Invalid file encryption header field: %1$s +error.write.encryption_header_must_precede_data = File encryption header must be written before TsFile data +error.write.encrypted_chunk_context_mismatch = Encrypted chunks must use the destination file encryption context +error.write.encryption_context_immutable = The file encryption context cannot change after its header is written + # === compress === # ICompressor.getCompressor / IUnCompressor.getUnCompressor — null type @@ -1573,3 +1612,15 @@ error.utils.table_point_count_tool_file_not_found = TsFile does not exist: %1$s error.utils.table_point_count_tool_incomplete_file = TsFile is incomplete: %1$s error.utils.table_point_count_tool_copy_failed = Failed to copy the TsFile data and metadata prefix info.utils.table_point_count_tool_result = Processed %1$s: %2$s +error.write.aead_raw_page_copy_unsupported = Raw page copying is not supported for page AEAD; decrypt and re-encrypt with the target file context +error.encrypt.file_context_already_used = Encryption context already belongs to a new file; generate a new DEK and file identity +error.encrypt.provider_discovery_failed = Encryption provider discovery failed +error.encrypt.page_aead_operations_missing = Page AEAD provider must supply an encryptor and a decryptor +error.encrypt.page_aead_requires_page_context = Page AEAD requires the page-aware encryption or decryption API +error.encrypt.page_aead_decryption_failed = Page AEAD decryption failed +error.encrypt.chunk_context_invalid = Chunk position authentication requires a file identity, a nonnegative marker offset and a chunk ordinal +error.encrypt.chunk_authentication_unsupported = Provider %1$s does not support Chunk position authentication +error.encrypt.invalid_chunk_authentication = The encrypted Chunk position authentication proof is missing or has an invalid size +error.encrypt.invalid_chunk_authentication_size = Invalid chunk authentication size: %1$d +error.encrypt.chunk_authentication_failed = Chunk position authentication failed +error.encrypt.invalid_page_body_lengths = Invalid authenticated page body length encoding diff --git a/java/common/src/main/resources/org/apache/tsfile/i18n/messages_zh.properties b/java/common/src/main/resources/org/apache/tsfile/i18n/messages_zh.properties index 502df532e1..e86263ab0f 100644 --- a/java/common/src/main/resources/org/apache/tsfile/i18n/messages_zh.properties +++ b/java/common/src/main/resources/org/apache/tsfile/i18n/messages_zh.properties @@ -636,6 +636,7 @@ log.read.block_reader_fill_measurements_error = 无法填充 measurements # TableResultSet — failed to close tsBlockReader log.read.table_result_set_close_error = 关闭 tsBlockReader 失败 +log.read.tree_result_set_close_error = 关闭树模型查询结果集失败 # DeviceMetaIterator — failed to load device meta data log.read.device_meta_iterator_error = 加载设备 meta data 失败 @@ -1117,6 +1118,44 @@ error.encrypt.encrypt_no_constructor = 获取 encryptor 构造函数失败: %1$s # EncryptUtils.getEncrypt — instantiation failed (encryptType arg) error.encrypt.encrypt_instantiation_failed = 实例化 encryptor 失败: %1$s +# EncryptionProviderRegistry — provider registration and lookup +error.encrypt.invalid_provider_id = encryption provider id 不能为空 +error.encrypt.duplicate_provider = encryption provider id 已注册: %1$s +error.encrypt.provider_not_found = encryption provider 未注册: %1$s +error.encrypt.provider_returned_null = encryption provider 返回 null: %1$s + +# EncryptParameter — use after sensitive state was destroyed +error.encrypt.parameter_destroyed = encryption parameters 已销毁 +error.encrypt.invalid_page_body_overhead = encryption profile %2$s 的 page body 额外开销 %1$s 无效 +error.encrypt.page_body_too_short = encrypted page body 大小 %1$s 小于 AEAD 额外开销 %2$s +error.encrypt.page_context_invalid_parameter = page crypto context 需要 page AEAD 参数 +error.encrypt.page_aead_encrypt_unsupported = 加密器 %1$s 未实现经过认证的 Page 加密 +error.encrypt.page_aead_decrypt_unsupported = 解密器 %1$s 未实现经过认证的 Page 解密 +error.encrypt.page_context_invalid_file_id = page crypto context 需要 16 字节的 file crypto id +error.encrypt.page_context_invalid_chunk_ordinal = page crypto context 需要非负的 chunk ordinal +error.encrypt.page_read_requires_chunk_ordinal = Page-AEAD 文件需要调用 readPage(header, type, pageIndex, chunkOrdinal) +error.encrypt.chunk_ordinal_exhausted = 文件 Chunk 序号已耗尽 +error.write.invalid_chunk_ordinal = 加密 Chunk 必须具有文件内唯一的 Chunk 序号 +error.write.unencrypted_chunk_has_ordinal = 未加密 Chunk 不得携带 Chunk 序号 +error.write.aead_append_requires_file_header = 无加密头的文件不能追加加密数据 +error.write.encrypted_chunk_copy_unsupported = 复制加密 Chunk 前必须先解密并重新加密 +error.encrypt.page_context_invalid_sizes = page crypto context 值无效，uncompressed size: %1$s，compressed plaintext size: %2$s，page index: %3$s +error.encrypt.page_output_size_mismatch = AEAD page 输出大小不匹配，预期 %1$s 字节，实际 %2$s 字节 +error.encrypt.page_plaintext_size_mismatch = AEAD page plaintext 大小不匹配，预期 %1$s 字节，实际 %2$s 字节 +error.chunk.merge_page_aead_unsupported = Page-AEAD encrypted chunk 合并前必须先解密并重写 + +# FileEncryptionHeader — malformed or unsupported file encryption metadata +error.file.encryption_header_invalid_size = file encryption header 大小无效: %1$s +error.file.encryption_header_unsupported_version = 不支持的 file encryption header 版本: %1$s +error.file.encryption_header_malformed = file encryption header 格式错误 +error.file.encryption_header_trailing_bytes = file encryption header 包含多余字节 +error.file.encryption_header_invalid_component_size = file encryption header 组件大小无效: %1$s +error.file.encryption_header_truncated = file encryption header 已截断，预期 %1$s 字节，实际读取 %2$s 字节 +error.file.encryption_header_invalid_field = file encryption header 字段无效: %1$s +error.write.encryption_header_must_precede_data = file encryption header 必须写在 TsFile data 之前 +error.write.encrypted_chunk_context_mismatch = encrypted Chunk 必须使用目标文件的加密上下文 +error.write.encryption_context_immutable = 文件加密头写入后不能更改文件加密上下文 + # === compress === # ICompressor.getCompressor / IUnCompressor.getUnCompressor — null type @@ -1573,3 +1612,15 @@ error.utils.table_point_count_tool_file_not_found = TsFile 不存在：%1$s error.utils.table_point_count_tool_incomplete_file = TsFile 不完整：%1$s error.utils.table_point_count_tool_copy_failed = 复制 TsFile 数据和元数据前缀失败 info.utils.table_point_count_tool_result = 已处理 %1$s：%2$s +error.write.aead_raw_page_copy_unsupported = Page AEAD 不支持原样复制页面；必须解密后使用目标文件上下文重新加密 +error.encrypt.file_context_already_used = 加密上下文已用于创建新文件；请生成新的 DEK 和文件身份 +error.encrypt.provider_discovery_failed = 加密 Provider 发现失败 +error.encrypt.page_aead_operations_missing = Page AEAD Provider 必须提供加密器和解密器 +error.encrypt.page_aead_requires_page_context = Page AEAD 必须使用带页面上下文的加解密接口 +error.encrypt.page_aead_decryption_failed = Page AEAD 解密失败 +error.encrypt.chunk_context_invalid = Chunk 位置认证需要文件身份、非负的 marker 偏移和 Chunk 序号 +error.encrypt.chunk_authentication_unsupported = Provider %1$s 不支持 Chunk 位置认证 +error.encrypt.invalid_chunk_authentication = 加密 Chunk 的位置认证凭据缺失或长度无效 +error.encrypt.invalid_chunk_authentication_size = 无效的 Chunk 认证长度：%1$d +error.encrypt.chunk_authentication_failed = Chunk 位置认证失败 +error.encrypt.invalid_page_body_lengths = 认证 Page 正文长度编码无效 diff --git a/java/tsfile/src/main/java/org/apache/tsfile/encrypt/ChunkCryptoContext.java b/java/tsfile/src/main/java/org/apache/tsfile/encrypt/ChunkCryptoContext.java new file mode 100644 index 0000000000..cced05c056 --- /dev/null +++ b/java/tsfile/src/main/java/org/apache/tsfile/encrypt/ChunkCryptoContext.java @@ -0,0 +1,66 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.tsfile.encrypt; + +import org.apache.tsfile.common.conf.TSFileConfig; +import org.apache.tsfile.exception.encrypt.EncryptException; +import org.apache.tsfile.i18n.Messages; + +import java.nio.ByteBuffer; +import java.util.Arrays; + +/** Binds a chunk's page identity to the absolute position of its on-disk marker. */ +public final class ChunkCryptoContext { + /** Format bound for opaque proofs; the immutable provider profile defines their exact size. */ + public static final int MAX_AUTHENTICATION_SIZE = 1024; + + private static final int ASSOCIATED_DATA_VERSION = 1; + private static final byte[] DOMAIN = "TsFile/ChunkPosition".getBytes(TSFileConfig.STRING_CHARSET); + private final byte[] associatedData; + + private ChunkCryptoContext(byte[] fileCryptoId, long chunkOffset, long chunkOrdinal) { + associatedData = + ByteBuffer.allocate( + Integer.BYTES * 2 + DOMAIN.length + fileCryptoId.length + Long.BYTES * 2) + .putInt(DOMAIN.length) + .put(DOMAIN) + .putInt(ASSOCIATED_DATA_VERSION) + .put(fileCryptoId) + .putLong(chunkOffset) + .putLong(chunkOrdinal) + .array(); + } + + public static ChunkCryptoContext forChunk( + EncryptParameter parameter, long chunkOffset, long chunkOrdinal) { + if (parameter == null || !parameter.isTdePageAead() || chunkOffset < 0 || chunkOrdinal < 0) { + throw new EncryptException(Messages.get("error.encrypt.chunk_context_invalid")); + } + byte[] fileId = parameter.getFileCryptoId(); + if (fileId == null || fileId.length != EncryptParameter.FILE_CRYPTO_ID_LENGTH) { + throw new EncryptException(Messages.get("error.encrypt.chunk_context_invalid")); + } + return new ChunkCryptoContext(fileId, chunkOffset, chunkOrdinal); + } + + public byte[] getAssociatedData() { + return Arrays.copyOf(associatedData, associatedData.length); + } +} diff --git a/java/tsfile/src/main/java/org/apache/tsfile/encrypt/EncryptParameter.java b/java/tsfile/src/main/java/org/apache/tsfile/encrypt/EncryptParameter.java index f7adb52217..815bec5534 100644 --- a/java/tsfile/src/main/java/org/apache/tsfile/encrypt/EncryptParameter.java +++ b/java/tsfile/src/main/java/org/apache/tsfile/encrypt/EncryptParameter.java @@ -18,14 +18,165 @@ */ package org.apache.tsfile.encrypt; -public class EncryptParameter { +import org.apache.tsfile.exception.encrypt.EncryptException; +import org.apache.tsfile.i18n.Messages; + +import javax.security.auth.Destroyable; + +import java.nio.ByteBuffer; +import java.util.Arrays; +import java.util.Objects; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicLong; + +/** + * Parameters required to encrypt or decrypt one TsFile. + * + *

The two-argument constructor preserves the legacy encryption API. New file-level encryption + * schemes should use {@link #pageAeadBuilder()} and identify an {@link IEncryptProvider} and one of + * its immutable encryption profiles by stable ids. Runtime implementation class names and algorithm + * implementation details are deliberately not part of the persisted parameters. + */ +public class EncryptParameter implements AutoCloseable, Destroyable { + + public static final int FILE_CRYPTO_ID_LENGTH = 16; private final String type; private final byte[] key; + private final boolean pageAead; + private final String providerId; + private final String profileId; + private final String keyId; + private final String keyVersion; + private final byte[] wrappedDataKey; + private final byte[] fileCryptoId; + private final AtomicLong nextChunkOrdinal; + private final AtomicBoolean newFileClaimed; + + private final FileEncryptContext fileEncryptContext; + private transient volatile boolean destroyed; public EncryptParameter(String type, byte[] key) { this.type = type; this.key = key; + this.pageAead = false; + this.providerId = null; + this.profileId = null; + this.keyId = null; + this.keyVersion = null; + this.wrappedDataKey = null; + this.fileCryptoId = null; + this.nextChunkOrdinal = new AtomicLong(); + this.newFileClaimed = new AtomicBoolean(); + this.fileEncryptContext = null; + } + + private EncryptParameter(Builder builder) { + this.type = null; + this.key = copy(builder.key); + this.pageAead = true; + this.providerId = builder.providerId; + this.profileId = builder.profileId; + this.keyId = builder.keyId; + this.keyVersion = builder.keyVersion; + this.wrappedDataKey = copy(builder.wrappedDataKey); + this.fileCryptoId = copy(builder.fileCryptoId); + this.nextChunkOrdinal = + builder.nextChunkOrdinal == null ? new AtomicLong() : builder.nextChunkOrdinal; + this.newFileClaimed = + builder.newFileClaimed == null ? new AtomicBoolean() : builder.newFileClaimed; + this.fileEncryptContext = new FileEncryptContext(); + } + + private EncryptParameter(EncryptParameter source) { + this.type = source.type; + this.key = copy(source.key); + this.pageAead = source.pageAead; + this.providerId = source.providerId; + this.profileId = source.profileId; + this.keyId = source.keyId; + this.keyVersion = source.keyVersion; + this.wrappedDataKey = copy(source.wrappedDataKey); + this.fileCryptoId = copy(source.fileCryptoId); + this.nextChunkOrdinal = source.nextChunkOrdinal; + this.newFileClaimed = source.newFileClaimed; + this.fileEncryptContext = source.fileEncryptContext; + this.fileEncryptContext.retain(); + } + + public static Builder pageAeadBuilder() { + return new Builder(); + } + + /** + * Copies this file's descriptor without initialized provider state. This does not generate a new + * data key or file identity and must not be used to create another new file. + */ + public synchronized EncryptParameter copy() { + if (destroyed) { + throw new EncryptException(Messages.get("error.encrypt.parameter_destroyed")); + } + if (!isTdePageAead()) { + return new EncryptParameter(type, key); + } + return pageAeadBuilder() + .key(key) + .providerId(providerId) + .profileId(profileId) + .keyId(keyId) + .keyVersion(keyVersion) + .wrappedDataKey(wrappedDataKey) + .fileCryptoId(fileCryptoId) + .nextChunkOrdinal(nextChunkOrdinal) + .newFileClaimed(newFileClaimed) + .build(); + } + + /** Reserves this context for one new file. Append/recovery use the existing on-disk header. */ + public synchronized void claimForNewFile() { + if (destroyed) { + throw new EncryptException(Messages.get("error.encrypt.parameter_destroyed")); + } + if (!pageAead || !newFileClaimed.compareAndSet(false, true)) { + throw new EncryptException(Messages.get("error.encrypt.file_context_already_used")); + } + } + + /** + * Returns an owned reference sharing this file's runtime provider. Each reference must be closed; + * the provider is destroyed only after the last reference closes. Legacy parameters are copied. + */ + public synchronized EncryptParameter retain() { + if (destroyed) { + throw new EncryptException(Messages.get("error.encrypt.parameter_destroyed")); + } + return pageAead ? new EncryptParameter(this) : copy(); + } + + /** Allocates a file-unique identity before the first page of a chunk is encrypted. */ + public long nextChunkOrdinal() { + if (!pageAead || destroyed) { + throw new EncryptException(Messages.get("error.encrypt.page_context_invalid_parameter")); + } + long ordinal = + nextChunkOrdinal.getAndUpdate(value -> value == Long.MAX_VALUE ? value : value + 1); + if (ordinal == Long.MAX_VALUE) { + throw new EncryptException(Messages.get("error.encrypt.chunk_ordinal_exhausted")); + } + return ordinal; + } + + /** Continues allocation after the highest retained chunk during append or crash recovery. */ + public void resumeAfterChunkOrdinal(long ordinal) { + if (!pageAead || destroyed) { + throw new EncryptException(Messages.get("error.encrypt.page_context_invalid_parameter")); + } + if (ordinal == Long.MAX_VALUE) { + throw new EncryptException(Messages.get("error.encrypt.chunk_ordinal_exhausted")); + } + if (ordinal >= 0) { + nextChunkOrdinal.accumulateAndGet(ordinal + 1, Math::max); + } } public byte[] getKey() { @@ -35,4 +186,245 @@ public byte[] getKey() { public String getType() { return type; } + + public String getProviderId() { + return providerId; + } + + public String getProfileId() { + return profileId; + } + + public String getKeyId() { + return keyId; + } + + public String getKeyVersion() { + return keyVersion; + } + + public byte[] getWrappedDataKey() { + return copy(wrappedDataKey); + } + + public byte[] getFileCryptoId() { + return copy(fileCryptoId); + } + + boolean hasValidFileCryptoId() { + return fileCryptoId != null && fileCryptoId.length == FILE_CRYPTO_ID_LENGTH; + } + + synchronized ByteBuffer newPageAssociatedDataBuffer(int pageFieldsSize) { + if (destroyed) { + throw new EncryptException(Messages.get("error.encrypt.parameter_destroyed")); + } + return fileEncryptContext.newPageAssociatedDataBuffer(this, pageFieldsSize); + } + + public int getPageBodyOverhead() { + int overhead = getOrCreateFileEncrypt().getPageBodyOverhead(); + if (overhead <= 0) { + throw new EncryptException( + Messages.format("error.encrypt.invalid_page_body_overhead", overhead, profileId)); + } + return overhead; + } + + public boolean isTdePageAead() { + return pageAead; + } + + public int getEncryptedPageBodySize(int plaintextSize) { + int size = getOrCreateFileEncrypt().getEncryptedPageBodySize(plaintextSize); + if (plaintextSize < 0 || size <= plaintextSize) { + throw new EncryptException(Messages.get("error.encrypt.invalid_page_body_lengths")); + } + return size; + } + + public int getPagePlaintextSize(ByteBuffer encryptedBody) { + int size = getOrCreateFileEncrypt().getPagePlaintextSize(encryptedBody.asReadOnlyBuffer()); + if (size < 0 + || size >= encryptedBody.remaining() + || getEncryptedPageBodySize(size) != encryptedBody.remaining()) { + throw new EncryptException(Messages.get("error.encrypt.invalid_page_body_lengths")); + } + return size; + } + + public int getChunkAuthenticationSize() { + return getOrCreateFileEncrypt().getChunkAuthenticationSize(); + } + + public byte[] authenticateChunk(long chunkOffset, long chunkOrdinal) { + return getOrCreateFileEncrypt() + .authenticateChunk(ChunkCryptoContext.forChunk(this, chunkOffset, chunkOrdinal)); + } + + public void verifyChunkAuthentication( + byte[] authentication, long chunkOffset, long chunkOrdinal) { + getOrCreateFileEncrypt() + .verifyChunkAuthentication( + authentication, ChunkCryptoContext.forChunk(this, chunkOffset, chunkOrdinal)); + } + + /** Copies of one file share the ordinal allocator; independently created files never do. */ + public boolean sharesPageAeadFileContext(EncryptParameter other) { + return other != null + && pageAead + && other.pageAead + && !destroyed + && !other.destroyed + && nextChunkOrdinal == other.nextChunkOrdinal + && Objects.equals(providerId, other.providerId) + && Objects.equals(profileId, other.profileId) + && Objects.equals(keyId, other.keyId) + && Objects.equals(keyVersion, other.keyVersion) + && Arrays.equals(fileCryptoId, other.fileCryptoId) + && Arrays.equals(wrappedDataKey, other.wrappedDataKey) + && Arrays.equals(key, other.key); + } + + synchronized IEncrypt getOrCreateFileEncrypt() { + if (destroyed) { + throw new EncryptException(Messages.get("error.encrypt.parameter_destroyed")); + } + return fileEncryptContext.getOrCreate(this); + } + + @Override + public boolean isDestroyed() { + return destroyed; + } + + @Override + public void destroy() { + synchronized (this) { + if (destroyed) { + return; + } + destroyed = true; + if (key != null && isTdePageAead()) { + Arrays.fill(key, (byte) 0); + } + } + if (fileEncryptContext != null) { + fileEncryptContext.release(); + } + } + + @Override + public void close() { + destroy(); + } + + private static byte[] copy(byte[] value) { + return value == null ? null : Arrays.copyOf(value, value.length); + } + + private static final class FileEncryptContext { + private int references = 1; + private IEncrypt fileEncrypt; + private byte[] pageAssociatedDataPrefix; + + private synchronized ByteBuffer newPageAssociatedDataBuffer( + EncryptParameter parameter, int pageFieldsSize) { + if (pageAssociatedDataPrefix == null) { + pageAssociatedDataPrefix = PageCryptoContext.buildFileAssociatedDataPrefix(parameter); + } + // Never expose the shared prefix to providers or individual page contexts. + return ByteBuffer.allocate(Math.addExact(pageAssociatedDataPrefix.length, pageFieldsSize)) + .put(pageAssociatedDataPrefix); + } + + private synchronized void retain() { + references++; + } + + private synchronized IEncrypt getOrCreate(EncryptParameter parameter) { + if (fileEncrypt == null) { + fileEncrypt = EncryptionProviderRegistry.create(parameter); + } + return fileEncrypt; + } + + private void release() { + IEncrypt toClose = null; + synchronized (this) { + if (--references == 0) { + toClose = fileEncrypt; + fileEncrypt = null; + pageAssociatedDataPrefix = null; + } + } + if (toClose != null) { + toClose.close(); + } + } + } + + public static final class Builder { + + private byte[] key; + private String providerId; + private String profileId; + private String keyId; + private String keyVersion; + private byte[] wrappedDataKey; + private byte[] fileCryptoId; + private AtomicLong nextChunkOrdinal; + private AtomicBoolean newFileClaimed; + + private Builder() {} + + public Builder key(byte[] key) { + this.key = copy(key); + return this; + } + + public Builder providerId(String providerId) { + this.providerId = providerId; + return this; + } + + public Builder profileId(String profileId) { + this.profileId = profileId; + return this; + } + + public Builder keyId(String keyId) { + this.keyId = keyId; + return this; + } + + public Builder keyVersion(String keyVersion) { + this.keyVersion = keyVersion; + return this; + } + + public Builder wrappedDataKey(byte[] wrappedDataKey) { + this.wrappedDataKey = copy(wrappedDataKey); + return this; + } + + public Builder fileCryptoId(byte[] fileCryptoId) { + this.fileCryptoId = copy(fileCryptoId); + return this; + } + + private Builder nextChunkOrdinal(AtomicLong nextChunkOrdinal) { + this.nextChunkOrdinal = nextChunkOrdinal; + return this; + } + + private Builder newFileClaimed(AtomicBoolean newFileClaimed) { + this.newFileClaimed = newFileClaimed; + return this; + } + + public EncryptParameter build() { + return new EncryptParameter(this); + } + } } diff --git a/java/tsfile/src/main/java/org/apache/tsfile/encrypt/EncryptUtils.java b/java/tsfile/src/main/java/org/apache/tsfile/encrypt/EncryptUtils.java index 7c8dbcd253..bd303d728a 100644 --- a/java/tsfile/src/main/java/org/apache/tsfile/encrypt/EncryptUtils.java +++ b/java/tsfile/src/main/java/org/apache/tsfile/encrypt/EncryptUtils.java @@ -258,8 +258,17 @@ public static EncryptParameter getEncryptParameter() { return encryptParam; } - /** Get the second EncryptParameter object according to the given type and first key. */ + /** + * Returns a page-AEAD file descriptor unchanged, or derives the legacy second key from the given + * type and first key. A null parameter selects unencrypted operation. + */ public static EncryptParameter getEncryptParameter(EncryptParameter param) { + if (param == null) { + return new EncryptParameter("org.apache.tsfile.encrypt.UNENCRYPTED", null); + } + if (param.isTdePageAead()) { + return param; + } return encryptParamCache.computeIfAbsent(param, EncryptUtils::generateEncryptParameter); } @@ -297,6 +306,16 @@ public static IEncrypt getEncrypt() { return getEncrypt(TSFileDescriptor.getInstance().getConfig()); } + public static IEncrypt getEncrypt(EncryptParameter encryptParameter) { + if (encryptParameter == null) { + return getEncrypt("org.apache.tsfile.encrypt.UNENCRYPTED", null); + } + if (encryptParameter.isTdePageAead()) { + return encryptParameter.getOrCreateFileEncrypt(); + } + return getEncrypt(encryptParameter.getType(), encryptParameter.getKey()); + } + public static IEncrypt getEncrypt(String encryptType, byte[] dataEncryptKey) { try { String className = getEncryptClass(encryptType); diff --git a/java/tsfile/src/main/java/org/apache/tsfile/encrypt/EncryptionProviderRegistry.java b/java/tsfile/src/main/java/org/apache/tsfile/encrypt/EncryptionProviderRegistry.java new file mode 100644 index 0000000000..3634c0827b --- /dev/null +++ b/java/tsfile/src/main/java/org/apache/tsfile/encrypt/EncryptionProviderRegistry.java @@ -0,0 +1,116 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ +package org.apache.tsfile.encrypt; + +import org.apache.tsfile.exception.encrypt.EncryptException; +import org.apache.tsfile.i18n.Messages; + +import java.util.Iterator; +import java.util.ServiceConfigurationError; +import java.util.ServiceLoader; +import java.util.concurrent.ConcurrentHashMap; + +/** Registry for file-level encryption providers identified by stable provider ids. */ +public final class EncryptionProviderRegistry { + + private static final ConcurrentHashMap PROVIDERS = + new ConcurrentHashMap<>(); + private static volatile boolean servicesLoaded; + + private EncryptionProviderRegistry() {} + + public static void registerProvider(IEncryptProvider provider) { + String providerId = provider == null ? null : provider.getProviderId(); + if (providerId == null || providerId.trim().isEmpty()) { + throw new EncryptException(Messages.get("error.encrypt.invalid_provider_id")); + } + providerId = providerId.trim(); + IEncryptProvider existing = PROVIDERS.putIfAbsent(providerId, provider); + if (existing != null && !existing.getClass().equals(provider.getClass())) { + throw new EncryptException(Messages.format("error.encrypt.duplicate_provider", providerId)); + } + } + + public static void unregisterProvider(String providerId) { + if (providerId != null) { + PROVIDERS.remove(providerId.trim()); + } + } + + static IEncrypt create(EncryptParameter encryptParameter) { + loadServices(); + String providerId = encryptParameter.getProviderId(); + IEncryptProvider provider = providerId == null ? null : PROVIDERS.get(providerId.trim()); + if (provider == null) { + throw new EncryptException(Messages.format("error.encrypt.provider_not_found", providerId)); + } + IEncrypt encrypt = provider.create(encryptParameter); + if (encrypt == null) { + throw new EncryptException( + Messages.format("error.encrypt.provider_returned_null", providerId)); + } + try { + return new PageAeadEncryption(encrypt); + } catch (RuntimeException | Error e) { + try { + encrypt.close(); + } catch (RuntimeException | Error closeFailure) { + e.addSuppressed(closeFailure); + } + throw e; + } + } + + private static synchronized void loadServices() { + if (servicesLoaded) { + return; + } + Iterator providers = ServiceLoader.load(IEncryptProvider.class).iterator(); + EncryptException failure = null; + while (true) { + try { + if (!providers.hasNext()) { + break; + } + } catch (ServiceConfigurationError e) { + if (failure == null) { + failure = + new EncryptException(Messages.get("error.encrypt.provider_discovery_failed"), e); + } else { + failure.addSuppressed(e); + } + break; + } + try { + registerProvider(providers.next()); + } catch (ServiceConfigurationError | RuntimeException e) { + if (failure == null) { + failure = + new EncryptException(Messages.get("error.encrypt.provider_discovery_failed"), e); + } else { + failure.addSuppressed(e); + } + } + } + if (failure != null) { + throw failure; + } + servicesLoaded = true; + } +} diff --git a/java/tsfile/src/main/java/org/apache/tsfile/encrypt/IDecryptor.java b/java/tsfile/src/main/java/org/apache/tsfile/encrypt/IDecryptor.java index de7735daee..1d31d3fc8e 100644 --- a/java/tsfile/src/main/java/org/apache/tsfile/encrypt/IDecryptor.java +++ b/java/tsfile/src/main/java/org/apache/tsfile/encrypt/IDecryptor.java @@ -57,6 +57,12 @@ static IDecryptor getDecryptor(String type, byte[] key) { } static IDecryptor getDecryptor(EncryptParameter encryptParam) { + if (encryptParam == null) { + return getDecryptor("org.apache.tsfile.encrypt.UNENCRYPTED", null); + } + if (encryptParam.isTdePageAead()) { + return EncryptUtils.getEncrypt(encryptParam).getDecryptor(); + } String type = encryptParam.getType(); byte[] key = encryptParam.getKey(); return getDecryptor(type, key); @@ -66,5 +72,11 @@ static IDecryptor getDecryptor(EncryptParameter encryptParam) { byte[] decrypt(byte[] data, int offset, int size); + default byte[] decryptPage( + byte[] data, int offset, int size, PageCryptoContext pageCryptoContext) { + throw new EncryptException( + Messages.format("error.encrypt.page_aead_decrypt_unsupported", getClass().getName())); + } + EncryptionType getEncryptionType(); } diff --git a/java/tsfile/src/main/java/org/apache/tsfile/encrypt/IEncrypt.java b/java/tsfile/src/main/java/org/apache/tsfile/encrypt/IEncrypt.java index 0ea6bbef7e..66d33d014d 100644 --- a/java/tsfile/src/main/java/org/apache/tsfile/encrypt/IEncrypt.java +++ b/java/tsfile/src/main/java/org/apache/tsfile/encrypt/IEncrypt.java @@ -18,9 +18,13 @@ */ package org.apache.tsfile.encrypt; +import org.apache.tsfile.exception.encrypt.EncryptException; +import org.apache.tsfile.i18n.Messages; + +import java.nio.ByteBuffer; import java.util.concurrent.ConcurrentHashMap; -public interface IEncrypt { +public interface IEncrypt extends AutoCloseable { static ConcurrentHashMap> encryptMap = new ConcurrentHashMap<>(); @@ -30,4 +34,64 @@ public interface IEncrypt { IDecryptor getDecryptor(); IEncryptor getEncryptor(); + + /** + * Returns the fixed number of bytes added to every encrypted page body by this profile. + * + *

Page AEAD providers must override this method. The value is part of the profile contract and + * must not change while files using that profile exist. + */ + default int getPageBodyOverhead() { + return -1; + } + + /** Returns the exact encoded size, including profile-specific framing and padding. */ + default int getEncryptedPageBodySize(int plaintextSize) { + int overhead = getPageBodyOverhead(); + if (plaintextSize < 0 || overhead <= 0) { + throw new EncryptException(Messages.get("error.encrypt.invalid_page_body_lengths")); + } + return Math.addExact(plaintextSize, overhead); + } + + /** + * Returns the candidate compressed plaintext size from a bounded, read-only encoded body. This is + * untrusted framing until decryptPage authenticates it together with the page context. + * Implementations must not decrypt, allocate based on this value, or modify the input. + */ + default int getPagePlaintextSize(ByteBuffer encryptedBody) { + int overhead = getPageBodyOverhead(); + if (overhead <= 0 || encryptedBody.remaining() < overhead) { + throw new EncryptException(Messages.get("error.encrypt.invalid_page_body_lengths")); + } + return encryptedBody.remaining() - overhead; + } + + /** + * Returns this immutable profile's chunk-position proof size. Page-AEAD providers must override + * this method; unsupported providers fail closed. The size must be positive and no larger than + * {@link ChunkCryptoContext#MAX_AUTHENTICATION_SIZE}. + */ + default int getChunkAuthenticationSize() { + return -1; + } + + /** + * Returns a fixed-size opaque proof authenticating a chunk's physical position. The proof is + * exactly {@link #getChunkAuthenticationSize()} bytes. Providers using AEAD must use a fresh + * nonce under the file DEK, independently of page encryption calls. + */ + default byte[] authenticateChunk(ChunkCryptoContext context) { + throw new EncryptException( + Messages.format("error.encrypt.chunk_authentication_unsupported", getClass().getName())); + } + + /** Verifies the proof before a chunk identity is used; unsupported providers fail closed. */ + default void verifyChunkAuthentication(byte[] authentication, ChunkCryptoContext context) { + throw new EncryptException( + Messages.format("error.encrypt.chunk_authentication_unsupported", getClass().getName())); + } + + @Override + default void close() {} } diff --git a/java/tsfile/src/main/java/org/apache/tsfile/encrypt/IEncryptProvider.java b/java/tsfile/src/main/java/org/apache/tsfile/encrypt/IEncryptProvider.java new file mode 100644 index 0000000000..a4a3311f13 --- /dev/null +++ b/java/tsfile/src/main/java/org/apache/tsfile/encrypt/IEncryptProvider.java @@ -0,0 +1,49 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ +package org.apache.tsfile.encrypt; + +/** + * Service provider interface for file-level encryption implementations. + * + *

Provider ids are persisted in a TsFile and therefore must remain stable across releases. A + * provider implementation may be registered explicitly through {@link EncryptionProviderRegistry} + * or discovered through {@link java.util.ServiceLoader}. Providers own the encryption profiles + * referenced by {@link EncryptParameter#getProfileId()}; a profile's algorithms, page layout, key + * envelope format, and page body overhead must remain immutable while files using it exist. + * + *

The caller must generate an independent DEK and fileCryptoId for every new file. An + * EncryptParameter, including its copies, describes one file, not a reusable encryption template. + * Providers must ensure that nonces are not reused under the same DEK, including after append, + * recovery, or recreation of a runtime context. PageCryptoContext is authenticated context, not a + * guaranteed-unique nonce source. Random-nonce profiles must define per-key usage limits and reject + * use beyond those limits, counting both page encryption and chunk position authentication. + * Version-1 encrypted chunk headers also require the fixed-size opaque position proof defined by + * {@link ChunkCryptoContext}; unsupported providers must reject these operations. + */ +public interface IEncryptProvider { + + String getProviderId(); + + /** + * Creates a file context that may be shared by concurrent chunk readers. Implementations must own + * their runtime key resources independently of the supplied parameter's lifetime. TsFile closes + * the context after its last retained reference is released. + */ + IEncrypt create(EncryptParameter encryptParameter); +} diff --git a/java/tsfile/src/main/java/org/apache/tsfile/encrypt/IEncryptor.java b/java/tsfile/src/main/java/org/apache/tsfile/encrypt/IEncryptor.java index 7d60ecfc45..5da7262116 100644 --- a/java/tsfile/src/main/java/org/apache/tsfile/encrypt/IEncryptor.java +++ b/java/tsfile/src/main/java/org/apache/tsfile/encrypt/IEncryptor.java @@ -57,6 +57,12 @@ static IEncryptor getEncryptor(String type, byte[] key) { } static IEncryptor getEncryptor(EncryptParameter encryptParam) { + if (encryptParam == null) { + return getEncryptor("org.apache.tsfile.encrypt.UNENCRYPTED", null); + } + if (encryptParam.isTdePageAead()) { + return EncryptUtils.getEncrypt(encryptParam).getEncryptor(); + } String type = encryptParam.getType(); byte[] key = encryptParam.getKey(); return getEncryptor(type, key); @@ -66,5 +72,11 @@ static IEncryptor getEncryptor(EncryptParameter encryptParam) { byte[] encrypt(byte[] data, int offset, int size); + default byte[] encryptPage( + byte[] data, int offset, int size, PageCryptoContext pageCryptoContext) { + throw new EncryptException( + Messages.format("error.encrypt.page_aead_encrypt_unsupported", getClass().getName())); + } + EncryptionType getEncryptionType(); } diff --git a/java/tsfile/src/main/java/org/apache/tsfile/encrypt/PageAeadEncryption.java b/java/tsfile/src/main/java/org/apache/tsfile/encrypt/PageAeadEncryption.java new file mode 100644 index 0000000000..c0d1934320 --- /dev/null +++ b/java/tsfile/src/main/java/org/apache/tsfile/encrypt/PageAeadEncryption.java @@ -0,0 +1,182 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ +package org.apache.tsfile.encrypt; + +import org.apache.tsfile.exception.encrypt.EncryptException; +import org.apache.tsfile.file.metadata.enums.EncryptionType; +import org.apache.tsfile.i18n.Messages; + +import java.nio.ByteBuffer; +import java.util.concurrent.atomic.AtomicBoolean; + +/** Enforces page-only, fail-closed semantics independently of a provider's legacy methods. */ +final class PageAeadEncryption implements IEncrypt { + private final IEncrypt delegate; + private final IEncryptor encryptor; + private final IDecryptor decryptor; + private final AtomicBoolean closed = new AtomicBoolean(); + + PageAeadEncryption(IEncrypt delegate) { + this.delegate = delegate; + IEncryptor providerEncryptor = delegate.getEncryptor(); + IDecryptor providerDecryptor = delegate.getDecryptor(); + if (providerEncryptor == null || providerDecryptor == null) { + throw new EncryptException(Messages.get("error.encrypt.page_aead_operations_missing")); + } + encryptor = + new IEncryptor() { + @Override + public byte[] encrypt(byte[] data) { + throw legacyOperation(); + } + + @Override + public byte[] encrypt(byte[] data, int offset, int size) { + throw legacyOperation(); + } + + @Override + public byte[] encryptPage(byte[] data, int offset, int size, PageCryptoContext context) { + ensureOpen(); + return providerEncryptor.encryptPage(data, offset, size, context); + } + + @Override + public EncryptionType getEncryptionType() { + return EncryptionType.NewWay; + } + }; + decryptor = + new IDecryptor() { + @Override + public byte[] decrypt(byte[] data) { + throw legacyOperation(); + } + + @Override + public byte[] decrypt(byte[] data, int offset, int size) { + throw legacyOperation(); + } + + @Override + public byte[] decryptPage(byte[] data, int offset, int size, PageCryptoContext context) { + ensureOpen(); + try { + return providerDecryptor.decryptPage(data, offset, size, context); + } catch (EncryptException e) { + throw e; + } catch (RuntimeException e) { + throw new EncryptException( + Messages.get("error.encrypt.page_aead_decryption_failed"), e); + } + } + + @Override + public EncryptionType getEncryptionType() { + return EncryptionType.NewWay; + } + }; + } + + private static EncryptException legacyOperation() { + return new EncryptException(Messages.get("error.encrypt.page_aead_requires_page_context")); + } + + private void ensureOpen() { + if (closed.get()) { + throw new EncryptException(Messages.get("error.encrypt.parameter_destroyed")); + } + } + + @Override + public IDecryptor getDecryptor() { + ensureOpen(); + return decryptor; + } + + @Override + public IEncryptor getEncryptor() { + ensureOpen(); + return encryptor; + } + + @Override + public int getPageBodyOverhead() { + ensureOpen(); + return delegate.getPageBodyOverhead(); + } + + @Override + public int getEncryptedPageBodySize(int plaintextSize) { + ensureOpen(); + return delegate.getEncryptedPageBodySize(plaintextSize); + } + + @Override + public int getPagePlaintextSize(ByteBuffer encryptedBody) { + ensureOpen(); + return delegate.getPagePlaintextSize(encryptedBody.asReadOnlyBuffer()); + } + + @Override + public int getChunkAuthenticationSize() { + ensureOpen(); + int size = delegate.getChunkAuthenticationSize(); + if (size <= 0 || size > ChunkCryptoContext.MAX_AUTHENTICATION_SIZE) { + throw new EncryptException( + Messages.format("error.encrypt.invalid_chunk_authentication_size", size)); + } + return size; + } + + @Override + public byte[] authenticateChunk(ChunkCryptoContext context) { + ensureOpen(); + int expectedSize = getChunkAuthenticationSize(); + byte[] authentication = delegate.authenticateChunk(context); + validateChunkAuthentication(authentication, expectedSize); + return authentication; + } + + @Override + public void verifyChunkAuthentication(byte[] authentication, ChunkCryptoContext context) { + ensureOpen(); + validateChunkAuthentication(authentication, getChunkAuthenticationSize()); + try { + delegate.verifyChunkAuthentication(authentication, context); + } catch (EncryptException e) { + throw e; + } catch (RuntimeException e) { + throw new EncryptException(Messages.get("error.encrypt.chunk_authentication_failed"), e); + } + } + + private static void validateChunkAuthentication(byte[] authentication, int expectedSize) { + if (authentication == null || authentication.length != expectedSize) { + throw new EncryptException(Messages.get("error.encrypt.invalid_chunk_authentication")); + } + } + + @Override + public void close() { + if (closed.compareAndSet(false, true)) { + delegate.close(); + } + } +} diff --git a/java/tsfile/src/main/java/org/apache/tsfile/encrypt/PageCryptoContext.java b/java/tsfile/src/main/java/org/apache/tsfile/encrypt/PageCryptoContext.java new file mode 100644 index 0000000000..3b3628cf30 --- /dev/null +++ b/java/tsfile/src/main/java/org/apache/tsfile/encrypt/PageCryptoContext.java @@ -0,0 +1,215 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ +package org.apache.tsfile.encrypt; + +import org.apache.tsfile.common.conf.TSFileConfig; +import org.apache.tsfile.exception.encrypt.EncryptException; +import org.apache.tsfile.i18n.Messages; + +import java.nio.ByteBuffer; +import java.util.Arrays; + +/** + * Canonical context for page-level authenticated encryption. + * + *

Encryption and decryption construct the same associated data from persisted file encryption + * metadata and stable page fields. Providers should pass {@link #getAssociatedData()} to their AEAD + * implementation unchanged. + */ +public final class PageCryptoContext { + + private static final int ASSOCIATED_DATA_VERSION = 1; + private static final byte[] EMPTY_BYTES = new byte[0]; + + private final int uncompressedSize; + private final int compressedPlaintextSize; + private final int encryptedPageBodySize; + private final int pageIndex; + private final byte[] associatedData; + + private PageCryptoContext( + EncryptParameter parameter, + int uncompressedSize, + int compressedPlaintextSize, + int encryptedPageBodySize, + int pageIndex, + long chunkOrdinal) { + validate(parameter, uncompressedSize, compressedPlaintextSize, pageIndex, chunkOrdinal); + this.uncompressedSize = uncompressedSize; + this.compressedPlaintextSize = compressedPlaintextSize; + this.encryptedPageBodySize = encryptedPageBodySize; + this.pageIndex = pageIndex; + this.associatedData = + buildAssociatedData( + parameter, uncompressedSize, compressedPlaintextSize, pageIndex, chunkOrdinal); + } + + public static PageCryptoContext forEncryption( + EncryptParameter parameter, + int uncompressedSize, + int compressedPlaintextSize, + int pageIndex, + long chunkOrdinal) { + validate(parameter, uncompressedSize, compressedPlaintextSize, pageIndex, chunkOrdinal); + int encryptedPageBodySize = parameter.getEncryptedPageBodySize(compressedPlaintextSize); + return new PageCryptoContext( + parameter, + uncompressedSize, + compressedPlaintextSize, + encryptedPageBodySize, + pageIndex, + chunkOrdinal); + } + + public static PageCryptoContext forDecryption( + EncryptParameter parameter, + int uncompressedSize, + int encryptedPageBodySize, + int pageIndex, + long chunkOrdinal) { + validate(parameter, uncompressedSize, 0, pageIndex, chunkOrdinal); + int overhead = parameter.getPageBodyOverhead(); + int compressedPlaintextSize = encryptedPageBodySize - overhead; + if (compressedPlaintextSize < 0) { + throw new EncryptException( + Messages.format("error.encrypt.page_body_too_short", encryptedPageBodySize, overhead)); + } + return new PageCryptoContext( + parameter, + uncompressedSize, + compressedPlaintextSize, + encryptedPageBodySize, + pageIndex, + chunkOrdinal); + } + + public int getUncompressedSize() { + return uncompressedSize; + } + + /** Builds AAD using profile framing; the candidate length is authenticated by decryptPage. */ + public static PageCryptoContext forDecryption( + EncryptParameter parameter, + int uncompressedSize, + ByteBuffer encryptedBody, + int pageIndex, + long chunkOrdinal) { + validate(parameter, uncompressedSize, 0, pageIndex, chunkOrdinal); + int plaintextSize = parameter.getPagePlaintextSize(encryptedBody); + return new PageCryptoContext( + parameter, + uncompressedSize, + plaintextSize, + encryptedBody.remaining(), + pageIndex, + chunkOrdinal); + } + + /** Returns the compressed plaintext size on both encryption and decryption paths. */ + public int getInputSize() { + return compressedPlaintextSize; + } + + public int getCompressedPlaintextSize() { + return compressedPlaintextSize; + } + + public int getEncryptedPageBodySize() { + return encryptedPageBodySize; + } + + public int getPageIndex() { + return pageIndex; + } + + public boolean isFirstPage() { + return pageIndex == 0; + } + + /** Returns a defensive copy so providers cannot mutate the authenticated page identity. */ + public byte[] getAssociatedData() { + return Arrays.copyOf(associatedData, associatedData.length); + } + + private static void validate( + EncryptParameter parameter, + int uncompressedSize, + int compressedPlaintextSize, + int pageIndex, + long chunkOrdinal) { + if (parameter == null || !parameter.isTdePageAead()) { + throw new EncryptException(Messages.get("error.encrypt.page_context_invalid_parameter")); + } + if (!parameter.hasValidFileCryptoId()) { + throw new EncryptException(Messages.get("error.encrypt.page_context_invalid_file_id")); + } + if (chunkOrdinal < 0) { + throw new EncryptException(Messages.get("error.encrypt.page_context_invalid_chunk_ordinal")); + } + if (uncompressedSize < 0 || compressedPlaintextSize < 0 || pageIndex < 0) { + throw new EncryptException( + Messages.format( + "error.encrypt.page_context_invalid_sizes", + uncompressedSize, + compressedPlaintextSize, + pageIndex)); + } + } + + private static byte[] buildAssociatedData( + EncryptParameter parameter, + int uncompressedSize, + int compressedPlaintextSize, + int pageIndex, + long chunkOrdinal) { + ByteBuffer buffer = parameter.newPageAssociatedDataBuffer(Long.BYTES + Integer.BYTES * 3); + buffer.putLong(chunkOrdinal); + buffer.putInt(pageIndex); + buffer.putInt(uncompressedSize); + buffer.putInt(compressedPlaintextSize); + return buffer.array(); + } + + static byte[] buildFileAssociatedDataPrefix(EncryptParameter parameter) { + byte[][] components = + new byte[][] { + bytes(parameter.getProviderId()), + bytes(parameter.getProfileId()), + bytes(parameter.getKeyId()), + bytes(parameter.getKeyVersion()), + parameter.getFileCryptoId() + }; + int size = Integer.BYTES * (1 + components.length); + for (byte[] component : components) { + size = Math.addExact(size, component.length); + } + + ByteBuffer buffer = ByteBuffer.allocate(size); + buffer.putInt(ASSOCIATED_DATA_VERSION); + for (byte[] component : components) { + buffer.putInt(component.length); + buffer.put(component); + } + return buffer.array(); + } + + private static byte[] bytes(String value) { + return value == null ? EMPTY_BYTES : value.getBytes(TSFileConfig.STRING_CHARSET); + } +} diff --git a/java/tsfile/src/main/java/org/apache/tsfile/exception/encrypt/EncryptException.java b/java/tsfile/src/main/java/org/apache/tsfile/exception/encrypt/EncryptException.java index 1c540ede01..e099a1c37b 100644 --- a/java/tsfile/src/main/java/org/apache/tsfile/exception/encrypt/EncryptException.java +++ b/java/tsfile/src/main/java/org/apache/tsfile/exception/encrypt/EncryptException.java @@ -20,6 +20,16 @@ package org.apache.tsfile.exception.encrypt; public class EncryptException extends RuntimeException { + /** Crypto failures are fatal even when a read API historically tolerates ordinary I/O errors. */ + public static boolean isEncryptionFailure(Throwable failure) { + for (Throwable cause = failure; cause != null; cause = cause.getCause()) { + if (cause instanceof EncryptException) { + return true; + } + } + return false; + } + public EncryptException(String message) { super("Encrypt problem: " + message); } diff --git a/java/tsfile/src/main/java/org/apache/tsfile/file/MetaMarker.java b/java/tsfile/src/main/java/org/apache/tsfile/file/MetaMarker.java index 5cca705945..ce8bf8b379 100644 --- a/java/tsfile/src/main/java/org/apache/tsfile/file/MetaMarker.java +++ b/java/tsfile/src/main/java/org/apache/tsfile/file/MetaMarker.java @@ -51,6 +51,9 @@ public class MetaMarker { /** Chunk header marker and this chunk has only one page. */ public static final byte ONLY_ONE_PAGE_CHUNK_HEADER = 5; + /** Optional file encryption header immediately following the TsFile version byte. */ + public static final byte ENCRYPTION_HEADER = 6; + /** Time Chunk header marker and this chunk has more than one page. */ public static final byte TIME_CHUNK_HEADER = (byte) (CHUNK_HEADER | TsFileConstant.TIME_COLUMN_MASK); diff --git a/java/tsfile/src/main/java/org/apache/tsfile/file/header/ChunkHeader.java b/java/tsfile/src/main/java/org/apache/tsfile/file/header/ChunkHeader.java index a03209fdca..a862a670ed 100644 --- a/java/tsfile/src/main/java/org/apache/tsfile/file/header/ChunkHeader.java +++ b/java/tsfile/src/main/java/org/apache/tsfile/file/header/ChunkHeader.java @@ -22,11 +22,14 @@ import org.apache.tsfile.common.conf.TSFileConfig; import org.apache.tsfile.encoding.decoder.Decoder; import org.apache.tsfile.encoding.decoder.DecoderWrapper; +import org.apache.tsfile.encrypt.ChunkCryptoContext; import org.apache.tsfile.enums.TSDataType; +import org.apache.tsfile.exception.encrypt.EncryptException; import org.apache.tsfile.file.MetaMarker; import org.apache.tsfile.file.metadata.TimeseriesMetadata; import org.apache.tsfile.file.metadata.enums.CompressionType; import org.apache.tsfile.file.metadata.enums.TSEncoding; +import org.apache.tsfile.i18n.Messages; import org.apache.tsfile.read.TsFileSequenceReader; import org.apache.tsfile.read.reader.TsFileInput; import org.apache.tsfile.utils.Pair; @@ -34,10 +37,13 @@ import org.apache.tsfile.utils.ReadWriteForEncodingUtils; import org.apache.tsfile.utils.ReadWriteIOUtils; +import java.io.EOFException; import java.io.IOException; import java.io.InputStream; import java.io.OutputStream; +import java.nio.BufferUnderflowException; import java.nio.ByteBuffer; +import java.util.Arrays; import java.util.function.Function; import java.util.function.LongConsumer; @@ -63,6 +69,8 @@ public class ChunkHeader { private final TSDataType dataType; private final CompressionType compressionType; private final TSEncoding encodingType; + private final long chunkOrdinal; + private byte[] chunkAuthentication; // the following fields do not need to be serialized. private int numOfPages; @@ -87,16 +95,29 @@ public ChunkHeader( TSEncoding encoding, int numOfPages, int mask) { + this(measurementID, dataSize, dataType, compressionType, encoding, numOfPages, mask, -1); + } + + public ChunkHeader( + String measurementID, + int dataSize, + TSDataType dataType, + CompressionType compressionType, + TSEncoding encoding, + int numOfPages, + int mask, + long chunkOrdinal) { this( (byte) ((numOfPages <= 1 ? MetaMarker.ONLY_ONE_PAGE_CHUNK_HEADER : MetaMarker.CHUNK_HEADER) | (byte) mask), measurementID, dataSize, - getSerializedSize(measurementID, dataSize), + getSerializedPrefixSize(measurementID, dataSize, chunkOrdinal >= 0), dataType, compressionType, - encoding); + encoding, + chunkOrdinal); this.numOfPages = numOfPages; } @@ -117,6 +138,25 @@ public ChunkHeader( encoding); } + public ChunkHeader( + byte chunkType, + String measurementID, + int dataSize, + TSDataType dataType, + CompressionType compressionType, + TSEncoding encoding, + long chunkOrdinal) { + this( + chunkType, + measurementID, + dataSize, + getSerializedPrefixSize(measurementID, dataSize, chunkOrdinal >= 0), + dataType, + compressionType, + encoding, + chunkOrdinal); + } + public ChunkHeader( byte chunkType, String measurementID, @@ -125,12 +165,25 @@ public ChunkHeader( TSDataType dataType, CompressionType compressionType, TSEncoding encoding) { + this(chunkType, measurementID, dataSize, headerSize, dataType, compressionType, encoding, -1); + } + + private ChunkHeader( + byte chunkType, + String measurementID, + int dataSize, + int headerSize, + TSDataType dataType, + CompressionType compressionType, + TSEncoding encoding, + long chunkOrdinal) { this.chunkType = chunkType; this.measurementID = measurementID; this.dataSize = dataSize; this.dataType = dataType; this.compressionType = compressionType; this.encodingType = encoding; + this.chunkOrdinal = chunkOrdinal; this.serializedSize = headerSize; } @@ -147,6 +200,19 @@ public static int getSerializedSize(String measurementID, int dataSize) { + TSEncoding.getSerializedSize(); // encodingType } + /** Exact size with a profile-defined proof; zero selects the plaintext header format. */ + public static int getSerializedSize(String measurementID, int dataSize, int authenticationSize) { + if (authenticationSize != 0) { + validateAuthenticationSize(authenticationSize, 0); + } + return getSerializedPrefixSize(measurementID, dataSize, authenticationSize != 0) + + authenticationSize; + } + + private static int getSerializedPrefixSize(String measurementID, int dataSize, boolean pageAead) { + return getSerializedSize(measurementID, dataSize) + (pageAead ? Long.BYTES + Integer.BYTES : 0); + } + /** * The estimated serialized size of chunk header. Only used when we don't know the actual dataSize * attribute @@ -165,7 +231,10 @@ public static int getSerializedSize(String measurementID) { } public int getSerializedSize() { - return serializedSize; + if (chunkOrdinal >= 0 && chunkAuthentication == null) { + throw new EncryptException(Messages.get("error.encrypt.invalid_chunk_authentication")); + } + return serializedSize + getChunkAuthenticationSize(); } /** @@ -176,13 +245,59 @@ public int getSerializedSize() { */ public static ChunkHeader deserializeFrom(InputStream inputStream, byte chunkType) throws IOException { + return deserializeFrom(inputStream, chunkType, false); + } + + public static ChunkHeader deserializeFrom( + InputStream inputStream, byte chunkType, boolean pageAead) throws IOException { + return deserializeFrom(inputStream, chunkType, pageAead, 0); + } + + /** Reads and bounds-checks the proof length against the file's provider profile. */ + public static ChunkHeader deserializeFrom( + InputStream inputStream, byte chunkType, int authenticationSize) throws IOException { + if (authenticationSize != 0) { + validateAuthenticationSize(authenticationSize, 0); + } + return deserializeFrom(inputStream, chunkType, authenticationSize != 0, authenticationSize); + } + + private static ChunkHeader deserializeFrom( + InputStream inputStream, byte chunkType, boolean pageAead, int expectedAuthenticationSize) + throws IOException { // read measurementID String measurementID = ReadWriteIOUtils.readVarIntString(inputStream); int dataSize = ReadWriteForEncodingUtils.readUnsignedVarInt(inputStream); TSDataType dataType = ReadWriteIOUtils.readDataType(inputStream); CompressionType type = ReadWriteIOUtils.readCompressionType(inputStream); TSEncoding encoding = ReadWriteIOUtils.readEncoding(inputStream); - return new ChunkHeader(chunkType, measurementID, dataSize, dataType, type, encoding); + long chunkOrdinal = -1; + byte[] authentication = null; + if (pageAead) { + chunkOrdinal = ReadWriteIOUtils.readLong(inputStream); + byte[] sizeBytes = inputStream.readNBytes(Integer.BYTES); + if (sizeBytes.length != Integer.BYTES) { + throw new EOFException(); + } + int size = ByteBuffer.wrap(sizeBytes).getInt(); + validateAuthenticationSize(size, expectedAuthenticationSize); + authentication = inputStream.readNBytes(size); + if (authentication.length != size) { + throw new EOFException(); + } + } + ChunkHeader header = + new ChunkHeader( + chunkType, + measurementID, + dataSize, + getSerializedPrefixSize(measurementID, dataSize, pageAead), + dataType, + type, + encoding, + chunkOrdinal); + header.chunkAuthentication = authentication; + return header; } /** @@ -194,7 +309,7 @@ public static ChunkHeader deserializeFrom(InputStream inputStream, byte chunkTyp * @throws IOException IOException */ public static ChunkHeader deserializeFrom(TsFileInput input, long offset) throws IOException { - return deserializeFrom(input, offset, null); + return deserializeFrom(input, offset, null, false); } /** @@ -208,43 +323,128 @@ public static ChunkHeader deserializeFrom(TsFileInput input, long offset) throws */ public static ChunkHeader deserializeFrom( TsFileInput input, long offset, LongConsumer ioSizeRecorder) throws IOException { + return deserializeFrom(input, offset, ioSizeRecorder, false); + } + + public static ChunkHeader deserializeFrom( + TsFileInput input, long offset, LongConsumer ioSizeRecorder, boolean pageAead) + throws IOException { + return deserializeFrom(input, offset, ioSizeRecorder, pageAead, 0); + } - // only 6 bytes, no need to call ioSizeRecorder.accept alone, combine into the remaining read - // operation - ByteBuffer buffer = ByteBuffer.allocate(Byte.BYTES + Integer.BYTES + 1); - input.read(buffer, offset); - buffer.flip(); - - // read chunk header from input to buffer - byte chunkType = buffer.get(); - int strLength = ReadWriteForEncodingUtils.readVarInt(buffer); - int alreadyReadLength = buffer.position(); - - int remainingBytes = - strLength - + Integer.BYTES - + 1 // uVarInt dataSize - + TSDataType.getSerializedSize() // dataType - + CompressionType.getSerializedSize() // compressionType - + TSEncoding.getSerializedSize(); - buffer = ByteBuffer.allocate(remainingBytes); - - if (ioSizeRecorder != null) { - ioSizeRecorder.accept((long) alreadyReadLength + remainingBytes); + /** Positional counterpart with profile validation before reading the proof payload. */ + public static ChunkHeader deserializeFrom( + TsFileInput input, long offset, LongConsumer ioSizeRecorder, int authenticationSize) + throws IOException { + if (authenticationSize != 0) { + validateAuthenticationSize(authenticationSize, 0); } + return deserializeFrom( + input, offset, ioSizeRecorder, authenticationSize != 0, authenticationSize); + } - input.read(buffer, offset + alreadyReadLength); - buffer.flip(); + private static ChunkHeader deserializeFrom( + TsFileInput input, + long offset, + LongConsumer ioSizeRecorder, + boolean pageAead, + int expectedAuthenticationSize) + throws IOException { + try { + + // only 6 bytes, no need to call ioSizeRecorder.accept alone, combine into the remaining read + // operation + ByteBuffer buffer = ByteBuffer.allocate(Byte.BYTES + Integer.BYTES + 1); + readAtMost(input, offset, buffer); + buffer.flip(); + + // read chunk header from input to buffer + byte chunkType = buffer.get(); + int strLength = ReadWriteForEncodingUtils.readVarInt(buffer); + int alreadyReadLength = buffer.position(); + + int remainingBytes = + strLength + + Integer.BYTES + + 1 // uVarInt dataSize + + TSDataType.getSerializedSize() // dataType + + CompressionType.getSerializedSize() // compressionType + + TSEncoding.getSerializedSize() + + (pageAead ? Long.BYTES + Integer.BYTES : 0); + buffer = ByteBuffer.allocate(remainingBytes); + + if (ioSizeRecorder != null) { + ioSizeRecorder.accept((long) alreadyReadLength + remainingBytes); + } + + readAtMost(input, offset + alreadyReadLength, buffer); + buffer.flip(); + + // read measurementID + String measurementID = ReadWriteIOUtils.readStringWithLength(buffer, strLength); + int dataSize = ReadWriteForEncodingUtils.readUnsignedVarInt(buffer); + TSDataType dataType = ReadWriteIOUtils.readDataType(buffer); + CompressionType type = ReadWriteIOUtils.readCompressionType(buffer); + TSEncoding encoding = ReadWriteIOUtils.readEncoding(buffer); + long chunkOrdinal = -1; + byte[] authentication = null; + if (pageAead) { + if (buffer.remaining() < Long.BYTES + Integer.BYTES) { + throw new EOFException(); + } + chunkOrdinal = buffer.getLong(); + int size = buffer.getInt(); + validateAuthenticationSize(size, expectedAuthenticationSize); + authentication = new byte[size]; + } + int chunkHeaderSize = alreadyReadLength + buffer.position(); + if (authentication != null) { + ByteBuffer proof = ByteBuffer.wrap(authentication); + int bufferedBytes = Math.min(buffer.remaining(), proof.remaining()); + proof.put(buffer.array(), buffer.position(), bufferedBytes); + if (ioSizeRecorder != null) { + ioSizeRecorder.accept(proof.remaining()); + } + readAtMost(input, offset + chunkHeaderSize + bufferedBytes, proof); + if (proof.hasRemaining()) { + throw new EOFException(); + } + } + ChunkHeader header = + new ChunkHeader( + chunkType, + measurementID, + dataSize, + chunkHeaderSize, + dataType, + type, + encoding, + chunkOrdinal); + header.chunkAuthentication = authentication; + return header; + } catch (BufferUnderflowException e) { + throw new EOFException(); + } + } - // read measurementID - String measurementID = ReadWriteIOUtils.readStringWithLength(buffer, strLength); - int dataSize = ReadWriteForEncodingUtils.readUnsignedVarInt(buffer); - TSDataType dataType = ReadWriteIOUtils.readDataType(buffer); - CompressionType type = ReadWriteIOUtils.readCompressionType(buffer); - TSEncoding encoding = ReadWriteIOUtils.readEncoding(buffer); - int chunkHeaderSize = alreadyReadLength + buffer.position(); - return new ChunkHeader( - chunkType, measurementID, dataSize, chunkHeaderSize, dataType, type, encoding); + private static void readAtMost(TsFileInput input, long offset, ByteBuffer target) + throws IOException { + int initialPosition = target.position(); + while (target.hasRemaining()) { + int count = input.read(target, offset + target.position() - initialPosition); + if (count <= 0) { + break; + } + } + } + + private static void validateAuthenticationSize(int size, int expectedSize) { + if (size <= 0 + || size > ChunkCryptoContext.MAX_AUTHENTICATION_SIZE + || (expectedSize > 0 && size != expectedSize)) { + throw new EncryptException( + Messages.format("error.encrypt.invalid_chunk_authentication_size", size)); + } } /** @@ -269,6 +469,20 @@ public String getMeasurementID() { return measurementID; } + public long getChunkOrdinal() { + return chunkOrdinal; + } + + public byte[] getChunkAuthentication() { + return chunkAuthentication == null + ? null + : Arrays.copyOf(chunkAuthentication, chunkAuthentication.length); + } + + public int getChunkAuthenticationSize() { + return chunkAuthentication == null ? 0 : chunkAuthentication.length; + } + public void setMeasurementID(String measurementID) { this.measurementID = measurementID; } @@ -305,6 +519,12 @@ public void setReplaceDecoder(Function replaceDecoder) * @throws IOException IOException */ public int serializeTo(OutputStream outputStream) throws IOException { + return serializeTo(outputStream, chunkAuthentication); + } + + /** Writes a proof generated for the destination position without mutating the source header. */ + public int serializeTo(OutputStream outputStream, byte[] authentication) throws IOException { + validateChunkAuthentication(authentication); int length = 0; length += ReadWriteIOUtils.write(chunkType, outputStream); length += ReadWriteIOUtils.writeVar(measurementID, outputStream); @@ -312,6 +532,12 @@ public int serializeTo(OutputStream outputStream) throws IOException { length += ReadWriteIOUtils.write(dataType, outputStream); length += ReadWriteIOUtils.write(compressionType, outputStream); length += ReadWriteIOUtils.write(encodingType, outputStream); + if (chunkOrdinal >= 0) { + length += ReadWriteIOUtils.write(chunkOrdinal, outputStream); + length += ReadWriteIOUtils.write(authentication.length, outputStream); + outputStream.write(authentication); + length += authentication.length; + } return length; } @@ -322,6 +548,7 @@ public int serializeTo(OutputStream outputStream) throws IOException { * @return length */ public int serializeTo(ByteBuffer buffer) { + validateChunkAuthentication(chunkAuthentication); int length = 0; length += ReadWriteIOUtils.write(chunkType, buffer); length += ReadWriteIOUtils.writeVar(measurementID, buffer); @@ -329,9 +556,25 @@ public int serializeTo(ByteBuffer buffer) { length += ReadWriteIOUtils.write(dataType, buffer); length += ReadWriteIOUtils.write(compressionType, buffer); length += ReadWriteIOUtils.write(encodingType, buffer); + if (chunkOrdinal >= 0) { + length += ReadWriteIOUtils.write(chunkOrdinal, buffer); + length += ReadWriteIOUtils.write(chunkAuthentication.length, buffer); + buffer.put(chunkAuthentication); + length += chunkAuthentication.length; + } return length; } + private void validateChunkAuthentication(byte[] authentication) { + if ((chunkOrdinal >= 0 + && (authentication == null + || authentication.length <= 0 + || authentication.length > ChunkCryptoContext.MAX_AUTHENTICATION_SIZE)) + || (chunkOrdinal < 0 && authentication != null)) { + throw new EncryptException(Messages.get("error.encrypt.invalid_chunk_authentication")); + } + } + public int getNumOfPages() { return numOfPages; } @@ -362,6 +605,8 @@ public String toString() { + numOfPages + ", serializedSize=" + serializedSize + + ", chunkOrdinal=" + + chunkOrdinal + '}'; } diff --git a/java/tsfile/src/main/java/org/apache/tsfile/file/header/FileEncryptionHeader.java b/java/tsfile/src/main/java/org/apache/tsfile/file/header/FileEncryptionHeader.java new file mode 100644 index 0000000000..bd95056a83 --- /dev/null +++ b/java/tsfile/src/main/java/org/apache/tsfile/file/header/FileEncryptionHeader.java @@ -0,0 +1,194 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ +package org.apache.tsfile.file.header; + +import org.apache.tsfile.encrypt.EncryptParameter; +import org.apache.tsfile.exception.encrypt.EncryptException; +import org.apache.tsfile.file.MetaMarker; +import org.apache.tsfile.i18n.Messages; +import org.apache.tsfile.utils.PublicBAOS; +import org.apache.tsfile.utils.ReadWriteForEncodingUtils; +import org.apache.tsfile.utils.ReadWriteIOUtils; + +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import java.nio.ByteBuffer; + +/** + * Versioned file encryption metadata stored before the first TsFile data marker. + * + *

The header contains only public identifiers and a wrapped data key. Plaintext key material and + * runtime implementation class names are never serialized. + */ +public final class FileEncryptionHeader { + + public static final byte VERSION = 1; + public static final int MAX_HEADER_SIZE = 1024 * 1024; + + private FileEncryptionHeader() {} + + public static int serialize(EncryptParameter parameter, OutputStream outputStream) + throws IOException { + validate(parameter); + + byte[] fileCryptoId = parameter.getFileCryptoId(); + byte[] wrappedDataKey = parameter.getWrappedDataKey(); + long expectedSize = + (long) Byte.BYTES + + ReadWriteForEncodingUtils.varIntStringSize(parameter.getProviderId()) + + ReadWriteForEncodingUtils.varIntStringSize(parameter.getProfileId()) + + ReadWriteForEncodingUtils.varIntStringSize(parameter.getKeyId()) + + ReadWriteForEncodingUtils.varIntStringSize(parameter.getKeyVersion()) + + ReadWriteForEncodingUtils.uVarIntSize(fileCryptoId.length) + + fileCryptoId.length + + ReadWriteForEncodingUtils.uVarIntSize(wrappedDataKey.length) + + wrappedDataKey.length; + if (expectedSize > MAX_HEADER_SIZE) { + throw new EncryptException( + Messages.format("error.file.encryption_header_invalid_size", expectedSize)); + } + PublicBAOS payload = new PublicBAOS((int) expectedSize); + int payloadSize = 0; + payloadSize += ReadWriteIOUtils.write(VERSION, payload); + payloadSize += ReadWriteIOUtils.writeVar(parameter.getProviderId(), payload); + payloadSize += ReadWriteIOUtils.writeVar(parameter.getProfileId(), payload); + payloadSize += ReadWriteIOUtils.writeVar(parameter.getKeyId(), payload); + payloadSize += ReadWriteIOUtils.writeVar(parameter.getKeyVersion(), payload); + payloadSize += writeBytes(fileCryptoId, payload); + payloadSize += writeBytes(wrappedDataKey, payload); + + if (payloadSize != payload.size() || payloadSize > MAX_HEADER_SIZE) { + throw new EncryptException( + Messages.format("error.file.encryption_header_invalid_size", payloadSize)); + } + + int size = ReadWriteIOUtils.write(MetaMarker.ENCRYPTION_HEADER, outputStream); + size += ReadWriteForEncodingUtils.writeUnsignedVarInt(payloadSize, outputStream); + payload.writeTo(outputStream); + return size + payloadSize; + } + + /** Deserializes a header after its {@link MetaMarker#ENCRYPTION_HEADER} marker was consumed. */ + public static EncryptParameter deserialize(InputStream inputStream) throws IOException { + int payloadSize = ReadWriteForEncodingUtils.readUnsignedVarInt(inputStream); + if (payloadSize <= 0 || payloadSize > MAX_HEADER_SIZE) { + throw new EncryptException( + Messages.format("error.file.encryption_header_invalid_size", payloadSize)); + } + + byte[] payload = readExactly(inputStream, payloadSize); + ByteBuffer buffer = ByteBuffer.wrap(payload); + byte version = buffer.get(); + if (version != VERSION) { + throw new EncryptException( + Messages.format("error.file.encryption_header_unsupported_version", version)); + } + + EncryptParameter parameter; + try { + parameter = + EncryptParameter.pageAeadBuilder() + .providerId(ReadWriteIOUtils.readVarIntString(buffer)) + .profileId(ReadWriteIOUtils.readVarIntString(buffer)) + .keyId(ReadWriteIOUtils.readVarIntString(buffer)) + .keyVersion(ReadWriteIOUtils.readVarIntString(buffer)) + .fileCryptoId(readBytes(buffer)) + .wrappedDataKey(readBytes(buffer)) + .build(); + } catch (RuntimeException e) { + throw new EncryptException(Messages.get("error.file.encryption_header_malformed"), e); + } + if (buffer.hasRemaining()) { + throw new EncryptException(Messages.get("error.file.encryption_header_trailing_bytes")); + } + validate(parameter); + return parameter; + } + + private static int writeBytes(byte[] bytes, OutputStream outputStream) throws IOException { + int size = ReadWriteForEncodingUtils.writeUnsignedVarInt(bytes.length, outputStream); + outputStream.write(bytes); + return size + bytes.length; + } + + private static byte[] readBytes(ByteBuffer buffer) { + int length = ReadWriteForEncodingUtils.readUnsignedVarInt(buffer); + if (length < 0 || length > MAX_HEADER_SIZE || length > buffer.remaining()) { + throw new EncryptException( + Messages.format("error.file.encryption_header_invalid_component_size", length)); + } + byte[] bytes = new byte[length]; + buffer.get(bytes); + return bytes; + } + + private static byte[] readExactly(InputStream inputStream, int size) throws IOException { + byte[] bytes = new byte[size]; + int offset = 0; + while (offset < size) { + int read = inputStream.read(bytes, offset, size - offset); + if (read < 0) { + throw new EncryptException( + Messages.format("error.file.encryption_header_truncated", size, offset)); + } + if (read == 0) { + int value = inputStream.read(); + if (value < 0) { + throw new EncryptException( + Messages.format("error.file.encryption_header_truncated", size, offset)); + } + bytes[offset++] = (byte) value; + } else { + offset += read; + } + } + return bytes; + } + + private static void validate(EncryptParameter parameter) { + if (parameter == null || !parameter.isTdePageAead()) { + invalidField("pageAead"); + } + requireText(parameter.getProviderId(), "providerId"); + requireText(parameter.getProfileId(), "profileId"); + requireText(parameter.getKeyId(), "keyId"); + requireText(parameter.getKeyVersion(), "keyVersion"); + + byte[] fileCryptoId = parameter.getFileCryptoId(); + if (fileCryptoId == null || fileCryptoId.length != EncryptParameter.FILE_CRYPTO_ID_LENGTH) { + invalidField("fileCryptoId"); + } + byte[] wrappedDataKey = parameter.getWrappedDataKey(); + if (wrappedDataKey == null || wrappedDataKey.length == 0) { + invalidField("wrappedDataKey"); + } + } + + private static void requireText(String value, String field) { + if (value == null || value.trim().isEmpty()) { + invalidField(field); + } + } + + private static void invalidField(String field) { + throw new EncryptException( + Messages.format("error.file.encryption_header_invalid_field", field)); + } +} diff --git a/java/tsfile/src/main/java/org/apache/tsfile/read/TsFileAlignedSeriesReaderIterator.java b/java/tsfile/src/main/java/org/apache/tsfile/read/TsFileAlignedSeriesReaderIterator.java index ed0254a691..c04fe77cbe 100644 --- a/java/tsfile/src/main/java/org/apache/tsfile/read/TsFileAlignedSeriesReaderIterator.java +++ b/java/tsfile/src/main/java/org/apache/tsfile/read/TsFileAlignedSeriesReaderIterator.java @@ -65,25 +65,34 @@ public NextAlignedChunkInfo nextReader() throws IOException { long totalSize = 0; long totalPointNum = 0; int notNullChunkNum = 0; - for (IChunkMetadata valueChunkMetadata : valueChunkMetadataList) { - if (valueChunkMetadata == null) { - continue; + try { + for (IChunkMetadata valueChunkMetadata : valueChunkMetadataList) { + if (valueChunkMetadata == null) { + continue; + } + while (!valueChunkMetadata + .getMeasurementUid() + .equals(schemaList.get(schemaIdx).getMeasurementName())) { + schemaIdx++; + } + Chunk chunk = reader.readMemChunk((ChunkMetadata) valueChunkMetadata); + valueChunks[schemaIdx++] = chunk; + notNullChunkNum++; + totalPointNum += ((ChunkMetadata) valueChunkMetadata).getNumOfPoints(); + totalSize += chunk.getHeader().getSerializedSize() + chunk.getHeader().getDataSize(); } - while (!valueChunkMetadata - .getMeasurementUid() - .equals(schemaList.get(schemaIdx).getMeasurementName())) { - schemaIdx++; + AlignedChunkReader chunkReader = + new AlignedChunkReader(timeChunk, Arrays.asList(valueChunks)); + return new NextAlignedChunkInfo(chunkReader, totalSize, notNullChunkNum, totalPointNum); + } catch (IOException | RuntimeException | Error e) { + timeChunk.close(); + for (Chunk valueChunk : valueChunks) { + if (valueChunk != null) { + valueChunk.close(); + } } - Chunk chunk = reader.readMemChunk((ChunkMetadata) valueChunkMetadata); - valueChunks[schemaIdx++] = chunk; - notNullChunkNum++; - totalPointNum += ((ChunkMetadata) valueChunkMetadata).getNumOfPoints(); - totalSize += chunk.getHeader().getSerializedSize() + chunk.getHeader().getDataSize(); + throw e; } - - AlignedChunkReader chunkReader = new AlignedChunkReader(timeChunk, Arrays.asList(valueChunks)); - - return new NextAlignedChunkInfo(chunkReader, totalSize, notNullChunkNum, totalPointNum); } public class NextAlignedChunkInfo { diff --git a/java/tsfile/src/main/java/org/apache/tsfile/read/TsFileReader.java b/java/tsfile/src/main/java/org/apache/tsfile/read/TsFileReader.java index f333afed6b..5e216b7bc4 100644 --- a/java/tsfile/src/main/java/org/apache/tsfile/read/TsFileReader.java +++ b/java/tsfile/src/main/java/org/apache/tsfile/read/TsFileReader.java @@ -34,8 +34,11 @@ import java.io.File; import java.io.IOException; import java.util.ArrayList; +import java.util.Collections; +import java.util.HashSet; import java.util.List; import java.util.Map; +import java.util.Set; import java.util.stream.Collectors; public class TsFileReader implements AutoCloseable { @@ -44,6 +47,7 @@ public class TsFileReader implements AutoCloseable { private IMetadataQuerier metadataQuerier; private IChunkLoader chunkLoader; private TsFileExecutor tsfileExecutor; + private final Set queryDataSets = Collections.synchronizedSet(new HashSet<>()); @TsFileApi public TsFileReader(File file) throws IOException { @@ -59,13 +63,24 @@ public TsFileReader(TsFileSequenceReader fileReader) throws IOException { } public QueryDataSet query(QueryExpression queryExpression) throws IOException { - return tsfileExecutor.execute(queryExpression); + return retainQuery(tsfileExecutor.execute(queryExpression)); } public QueryDataSet query( QueryExpression queryExpression, long partitionStartOffset, long partitionEndOffset) throws IOException { - return tsfileExecutor.execute(queryExpression, partitionStartOffset, partitionEndOffset); + return retainQuery( + tsfileExecutor.execute(queryExpression, partitionStartOffset, partitionEndOffset)); + } + + private QueryDataSet retainQuery(QueryDataSet dataSet) throws IOException { + if (fileReader.hasFileEncryptionHeader()) { + synchronized (queryDataSets) { + queryDataSets.removeIf(QueryDataSet::isClosed); + queryDataSets.add(dataSet); + } + } + return dataSet; } public List getAllDeviceIds() throws IOException { @@ -87,6 +102,35 @@ public List getMeasurement(IDeviceID deviceID) throws IOExcep @Override @TsFileApi public void close() throws IOException { - fileReader.close(); + IOException failure = null; + try { + synchronized (queryDataSets) { + for (QueryDataSet dataSet : queryDataSets) { + try { + dataSet.close(); + } catch (IOException e) { + if (failure == null) { + failure = e; + } else { + failure.addSuppressed(e); + } + } + } + } + } finally { + queryDataSets.clear(); + try { + fileReader.close(); + } catch (IOException e) { + if (failure == null) { + failure = e; + } else { + failure.addSuppressed(e); + } + } + } + if (failure != null) { + throw failure; + } } } diff --git a/java/tsfile/src/main/java/org/apache/tsfile/read/TsFileSequenceReader.java b/java/tsfile/src/main/java/org/apache/tsfile/read/TsFileSequenceReader.java index 2eea4089fd..a32f783211 100644 --- a/java/tsfile/src/main/java/org/apache/tsfile/read/TsFileSequenceReader.java +++ b/java/tsfile/src/main/java/org/apache/tsfile/read/TsFileSequenceReader.java @@ -30,16 +30,19 @@ import org.apache.tsfile.encrypt.EncryptParameter; import org.apache.tsfile.encrypt.EncryptUtils; import org.apache.tsfile.encrypt.IDecryptor; +import org.apache.tsfile.encrypt.PageCryptoContext; import org.apache.tsfile.enums.TSDataType; import org.apache.tsfile.exception.NotCompatibleTsFileException; import org.apache.tsfile.exception.StopReadTsFileByInterruptException; import org.apache.tsfile.exception.TsFileRuntimeException; import org.apache.tsfile.exception.TsFileStatisticsMistakesException; +import org.apache.tsfile.exception.encrypt.EncryptException; import org.apache.tsfile.exception.read.FileVersionTooOldException; import org.apache.tsfile.file.IMetadataIndexEntry; import org.apache.tsfile.file.MetaMarker; import org.apache.tsfile.file.header.ChunkGroupHeader; import org.apache.tsfile.file.header.ChunkHeader; +import org.apache.tsfile.file.header.FileEncryptionHeader; import org.apache.tsfile.file.header.PageHeader; import org.apache.tsfile.file.metadata.AbstractAlignedChunkMetadata; import org.apache.tsfile.file.metadata.AbstractAlignedTimeSeriesMetadata; @@ -85,6 +88,7 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import java.io.EOFException; import java.io.File; import java.io.IOException; import java.io.Serializable; @@ -114,6 +118,9 @@ public class TsFileSequenceReader implements AutoCloseable { + private static final int BASE_FILE_HEADER_SIZE = + TSFileConfig.MAGIC_STRING.getBytes(TSFileConfig.STRING_CHARSET).length + Byte.BYTES; + private static final Logger logger = LoggerFactory.getLogger(TsFileSequenceReader.class); private static final Logger resourceLogger = LoggerFactory.getLogger("FileMonitor"); protected static final TSFileConfig config = TSFileDescriptor.getInstance().getConfig(); @@ -147,6 +154,13 @@ public class TsFileSequenceReader implements AutoCloseable { private EncryptParameter dataEncryptParam = null; + private EncryptParameter fileEncryptionParam; + private long maxChunkOrdinal = -1; + + private long dataStartOffset = BASE_FILE_HEADER_SIZE; + private boolean fileEncryptionHeaderLoaded; + private int firstDataMarker = -1; + /** * Create a file reader of the given file. The reader will read the tail of the file to get the * file metadata size.Then the reader will skip the first @@ -168,7 +182,7 @@ public Map countChunksPerChunkGroup() throws IOException { return result; } - int headerLength = TSFileConfig.MAGIC_STRING.getBytes().length + Byte.BYTES; + long headerLength = getDataStartOffset(); if (checkFile.length() < headerLength) { return result; } @@ -315,6 +329,7 @@ public TsFileSequenceReader(String file, boolean loadMetadataSize, LongConsumer try { loadFileVersion(ioSizeRecorder); + ensureFileEncryptionHeaderLoaded(ioSizeRecorder); if (loadMetadataSize) { loadMetadataSize(ioSizeRecorder); } @@ -379,6 +394,7 @@ public TsFileSequenceReader(TsFileInput input, boolean loadMetadataSize) throws this.tsFileInput = input; this.file = input.getFilePath(); try { + ensureFileEncryptionHeaderLoaded(null); if (loadMetadataSize) { // NOTE no autoRepair here loadMetadataSize(); } @@ -403,6 +419,7 @@ public TsFileSequenceReader( try { loadFileVersion(ioSizeRecorder); + ensureFileEncryptionHeaderLoaded(ioSizeRecorder); if (loadMetadataSize) { loadMetadataSize(ioSizeRecorder); } @@ -430,6 +447,7 @@ public TsFileSequenceReader( */ public TsFileSequenceReader(TsFileInput input, long fileMetadataPos, int fileMetadataSize) { this.tsFileInput = input; + this.file = input.getFilePath(); this.fileMetadataPos = fileMetadataPos; this.fileMetadataSize = fileMetadataSize; } @@ -447,13 +465,14 @@ public TsFileSequenceReader( private void loadFileVersion(LongConsumer ioSizeRecorder) throws IOException { try { tsFileInput.position(TSFileConfig.MAGIC_STRING.getBytes(TSFileConfig.STRING_CHARSET).length); - final ByteBuffer buffer = ByteBuffer.allocate(Byte.BYTES); + final ByteBuffer buffer = ByteBuffer.allocate(2 * Byte.BYTES); + int bytesRead = ReadWriteIOUtils.readAsPossible(tsFileInput, buffer); if (ioSizeRecorder != null) { - ioSizeRecorder.accept(Byte.BYTES); + ioSizeRecorder.accept(bytesRead); } - tsFileInput.read(buffer); buffer.flip(); fileVersion = buffer.get(); + firstDataMarker = buffer.hasRemaining() ? Byte.toUnsignedInt(buffer.get()) : -1; checkFileVersion(); configDeserializer(); @@ -465,6 +484,71 @@ private void loadFileVersion(LongConsumer ioSizeRecorder) throws IOException { } } + private void loadFileEncryptionHeader(LongConsumer ioSizeRecorder) throws IOException { + long savedPosition = tsFileInput.position(); + dataStartOffset = BASE_FILE_HEADER_SIZE; + fileEncryptionParam = null; + try { + if (tsFileInput.size() <= BASE_FILE_HEADER_SIZE) { + return; + } + boolean markerPrefetched = firstDataMarker >= 0; + tsFileInput.position(BASE_FILE_HEADER_SIZE + (markerPrefetched ? Byte.BYTES : 0)); + int marker = markerPrefetched ? firstDataMarker : tsFileInput.wrapAsInputStream().read(); + if (marker < 0) { + throw new EOFException(); + } + if (marker != MetaMarker.ENCRYPTION_HEADER) { + switch ((byte) marker) { + case MetaMarker.CHUNK_GROUP_HEADER, + MetaMarker.CHUNK_HEADER, + MetaMarker.SEPARATOR, + MetaMarker.VERSION, + MetaMarker.OPERATION_INDEX_RANGE, + MetaMarker.ONLY_ONE_PAGE_CHUNK_HEADER, + MetaMarker.TIME_CHUNK_HEADER, + MetaMarker.VALUE_CHUNK_HEADER, + MetaMarker.ONLY_ONE_PAGE_TIME_CHUNK_HEADER, + MetaMarker.ONLY_ONE_PAGE_VALUE_CHUNK_HEADER -> {} + default -> MetaMarker.handleUnexpectedMarker((byte) marker); + } + if (ioSizeRecorder != null && !markerPrefetched) { + ioSizeRecorder.accept(Byte.BYTES); + } + return; + } + + fileEncryptionParam = FileEncryptionHeader.deserialize(tsFileInput.wrapAsInputStream()); + dataStartOffset = tsFileInput.position(); + if (ioSizeRecorder != null) { + ioSizeRecorder.accept( + dataStartOffset - BASE_FILE_HEADER_SIZE - (markerPrefetched ? Byte.BYTES : 0)); + } + } catch (EncryptException e) { + throw new IOException(e.getMessage(), e); + } finally { + tsFileInput.position(savedPosition); + } + } + + public long getDataStartOffset() throws IOException { + ensureFileEncryptionHeaderLoaded(null); + return dataStartOffset; + } + + public boolean hasFileEncryptionHeader() throws IOException { + ensureFileEncryptionHeaderLoaded(null); + return fileEncryptionParam != null; + } + + private synchronized void ensureFileEncryptionHeaderLoaded(LongConsumer ioSizeRecorder) + throws IOException { + if (!fileEncryptionHeaderLoaded) { + loadFileEncryptionHeader(ioSizeRecorder); + fileEncryptionHeaderLoaded = true; + } + } + private void configDeserializer() { if (fileVersion == TSFileConfig.VERSION_NUMBER_V3) { deserializeConfig = CompatibilityUtils.v3DeserializeConfig; @@ -670,8 +754,12 @@ public EncryptParameter getEncryptParam() throws IOException { * @param ioSizeRecorder can be null */ public EncryptParameter getEncryptParam(LongConsumer ioSizeRecorder) throws IOException { + ensureFileEncryptionHeaderLoaded(ioSizeRecorder); if (dataEncryptParam != null) { return dataEncryptParam; + } else if (fileEncryptionParam != null) { + dataEncryptParam = fileEncryptionParam; + return dataEncryptParam; } else { if (fileMetadataSize != 0) { readFileMetadata(ioSizeRecorder); @@ -692,7 +780,10 @@ public EncryptParameter getEncryptParam(LongConsumer ioSizeRecorder) throws IOEx return dataEncryptParam; } } - return EncryptUtils.getEncryptParameter(firstEncryptParam); + return EncryptUtils.getEncryptParameter( + firstEncryptParam != null && firstEncryptParam.isTdePageAead() + ? new EncryptParameter("org.apache.tsfile.encrypt.UNENCRYPTED", null) + : firstEncryptParam); } } @@ -2035,7 +2126,15 @@ public void readPlanIndex() throws IOException { */ public ChunkHeader readChunkHeader(byte chunkType) throws IOException { try { - return ChunkHeader.deserializeFrom(tsFileInput.wrapAsInputStream(), chunkType); + ensureFileEncryptionHeaderLoaded(null); + long chunkOffset = tsFileInput.position() - Byte.BYTES; + ChunkHeader header = + ChunkHeader.deserializeFrom( + tsFileInput.wrapAsInputStream(), + chunkType, + fileEncryptionParam == null ? 0 : fileEncryptionParam.getChunkAuthenticationSize()); + authenticateChunkHeader(header, chunkOffset); + return header; } catch (StopReadTsFileByInterruptException e) { throw e; } catch (Throwable t) { @@ -2054,7 +2153,15 @@ public ChunkHeader readChunkHeader(byte chunkType) throws IOException { private ChunkHeader readChunkHeader(long position, LongConsumer ioSizeRecorder) throws IOException { try { - return ChunkHeader.deserializeFrom(tsFileInput, position, ioSizeRecorder); + ensureFileEncryptionHeaderLoaded(ioSizeRecorder); + ChunkHeader header = + ChunkHeader.deserializeFrom( + tsFileInput, + position, + ioSizeRecorder, + fileEncryptionParam == null ? 0 : fileEncryptionParam.getChunkAuthenticationSize()); + authenticateChunkHeader(header, position); + return header; } catch (StopReadTsFileByInterruptException e) { throw e; } catch (Throwable t) { @@ -2064,6 +2171,27 @@ private ChunkHeader readChunkHeader(long position, LongConsumer ioSizeRecorder) } } + /** Reads and authenticates a chunk header at its absolute marker offset. */ + public ChunkHeader readChunkHeader(long position) throws IOException { + return readChunkHeader(position, null); + } + + private void authenticateChunkHeader(ChunkHeader header, long chunkOffset) throws IOException { + if (fileEncryptionParam != null && header.getChunkOrdinal() < 0) { + throw new org.apache.tsfile.exception.encrypt.EncryptException( + Messages.get("error.write.invalid_chunk_ordinal")); + } + if (fileEncryptionParam != null) { + fileEncryptionParam.verifyChunkAuthentication( + header.getChunkAuthentication(), chunkOffset, header.getChunkOrdinal()); + } + maxChunkOrdinal = Math.max(maxChunkOrdinal, header.getChunkOrdinal()); + } + + public long getMaxChunkOrdinal() { + return maxChunkOrdinal; + } + /** * notice, this function will modify channel's position. * @@ -2235,17 +2363,63 @@ public ByteBuffer readCompressedPage(PageHeader header) throws IOException { } public ByteBuffer readPage(PageHeader header, CompressionType type) throws IOException { + return readPage(header, type, 0); + } + + public ByteBuffer readPage(PageHeader header, CompressionType type, int pageIndex) + throws IOException { + if (fileEncryptionParam != null && fileEncryptionParam.isTdePageAead()) { + throw new IOException(Messages.get("error.encrypt.page_read_requires_chunk_ordinal")); + } + return readPage(header, type, pageIndex, -1); + } + + public ByteBuffer readPage( + PageHeader header, CompressionType type, int pageIndex, long chunkOrdinal) + throws IOException { ByteBuffer buffer = readData(-1, header.getCompressedSize()); - IDecryptor decryptor = IDecryptor.getDecryptor(getEncryptParam()); + EncryptParameter encryptParameter = getEncryptParam(); + IDecryptor decryptor = IDecryptor.getDecryptor(encryptParameter); if (header.getUncompressedSize() == 0) { return buffer; } - ByteBuffer finalBuffer = decrypt(decryptor, buffer); + ByteBuffer finalBuffer = + decrypt(decryptor, encryptParameter, buffer, header, pageIndex, chunkOrdinal); finalBuffer = uncompress(type, finalBuffer, header.getUncompressedSize()); return finalBuffer; } - private static ByteBuffer decrypt(IDecryptor decryptor, ByteBuffer buffer) { + private static ByteBuffer decrypt( + IDecryptor decryptor, + EncryptParameter encryptParameter, + ByteBuffer buffer, + PageHeader header, + int pageIndex, + long chunkOrdinal) + throws IOException { + if (encryptParameter != null && encryptParameter.isTdePageAead()) { + PageCryptoContext pageCryptoContext = + PageCryptoContext.forDecryption( + encryptParameter, + header.getUncompressedSize(), + buffer.asReadOnlyBuffer(), + pageIndex, + chunkOrdinal); + byte[] plaintext = + decryptor.decryptPage( + buffer.array(), + buffer.arrayOffset() + buffer.position(), + buffer.remaining(), + pageCryptoContext); + if (plaintext.length != pageCryptoContext.getCompressedPlaintextSize()) { + throw new IOException( + Messages.format( + "error.encrypt.page_plaintext_size_mismatch", + pageCryptoContext.getCompressedPlaintextSize(), + plaintext.length)); + } + return ByteBuffer.wrap(plaintext); + } if (decryptor == null || decryptor.getEncryptionType() == EncryptionType.UNENCRYPTED) { return buffer; } @@ -2296,7 +2470,16 @@ public void close() throws IOException { if (resourceLogger.isDebugEnabled()) { resourceLogger.debug("{} reader is closed.", file); } - this.tsFileInput.close(); + try { + this.tsFileInput.close(); + } finally { + if (dataEncryptParam != null && dataEncryptParam.isTdePageAead()) { + dataEncryptParam.close(); + } + if (fileEncryptionParam != null && fileEncryptionParam != dataEncryptParam) { + fileEncryptionParam.close(); + } + } } public String getFileName() { @@ -2437,7 +2620,7 @@ public long selfCheck( // ChunkMetadata of current ChunkGroup List chunkMetadataList = new ArrayList<>(); - int headerLength = TSFileConfig.MAGIC_STRING.getBytes().length + Byte.BYTES; + long headerLength = getDataStartOffset(); if (fileSize < headerLength) { return TsFileCheckStatus.INCOMPATIBLE_FILE; } @@ -2524,7 +2707,12 @@ public long selfCheck( Decoder valueDecoder = Decoder.getDecoderByType( chunkHeader.getEncodingType(), chunkHeader.getDataType()); - ByteBuffer pageData = readPage(pageHeader, chunkHeader.getCompressionType()); + ByteBuffer pageData = + readPage( + pageHeader, + chunkHeader.getCompressionType(), + 0, + chunkHeader.getChunkOrdinal()); TSEncoding configuredTimeEncoding = TSEncoding.valueOf(TSFileDescriptor.getInstance().getConfig().getTimeEncoder()); boolean isTimeColumn = @@ -2663,6 +2851,9 @@ public long selfCheck( truncatedSize = this.position() - 1; } } catch (Exception e) { + if (org.apache.tsfile.exception.encrypt.EncryptException.isEncryptionFailure(e)) { + throw new IOException(e.getMessage(), e); + } logger.warn( Messages.get("log.read.sequence_reader_self_check_warn"), file, @@ -2699,7 +2890,7 @@ public long selfCheckWithInfo( long fileSize = checkFile.length(); logger.info(Messages.get("log.read.sequence_reader_file_length"), fileSize); - int headerLength = TSFileConfig.MAGIC_STRING.getBytes().length + Byte.BYTES; + long headerLength = getDataStartOffset(); if (fileSize < headerLength) { return TsFileCheckStatus.INCOMPATIBLE_FILE; } @@ -2776,7 +2967,8 @@ public long checkChunkAndPagesStatistics(IChunkMetadata chunkMetadata) throws IO PageHeader pageHeader = this.readPageHeader(chunkHeader.getDataType(), false); Decoder valueDecoder = Decoder.getDecoderByType(chunkHeader.getEncodingType(), chunkHeader.getDataType()); - ByteBuffer pageData = readPage(pageHeader, chunkHeader.getCompressionType()); + ByteBuffer pageData = + readPage(pageHeader, chunkHeader.getCompressionType(), 0, chunkHeader.getChunkOrdinal()); Decoder timeDecoder = Decoder.getDecoderByType( TSEncoding.valueOf(TSFileDescriptor.getInstance().getConfig().getTimeEncoder()), diff --git a/java/tsfile/src/main/java/org/apache/tsfile/read/UnClosedTsFileReader.java b/java/tsfile/src/main/java/org/apache/tsfile/read/UnClosedTsFileReader.java index 230a2e1643..fcf11c9d6e 100644 --- a/java/tsfile/src/main/java/org/apache/tsfile/read/UnClosedTsFileReader.java +++ b/java/tsfile/src/main/java/org/apache/tsfile/read/UnClosedTsFileReader.java @@ -21,7 +21,6 @@ import org.apache.tsfile.common.conf.TSFileDescriptor; import org.apache.tsfile.encrypt.EncryptParameter; -import org.apache.tsfile.encrypt.EncryptUtils; import org.apache.tsfile.exception.NotImplementedException; import org.apache.tsfile.file.metadata.TsFileMetadata; @@ -46,7 +45,7 @@ public UnClosedTsFileReader(String file, LongConsumer ioSizeRecorder) throws IOE public UnClosedTsFileReader( String file, EncryptParameter encryptParam, LongConsumer ioSizeRecorder) throws IOException { super(file, false, ioSizeRecorder, encryptParam); - this.dataEncryptParam = EncryptUtils.getEncryptParameter(encryptParam); + this.dataEncryptParam = super.getEncryptParam(); } /** unclosed file has no tail magic data. */ diff --git a/java/tsfile/src/main/java/org/apache/tsfile/read/common/Chunk.java b/java/tsfile/src/main/java/org/apache/tsfile/read/common/Chunk.java index d41ad14a32..91e89ec913 100644 --- a/java/tsfile/src/main/java/org/apache/tsfile/read/common/Chunk.java +++ b/java/tsfile/src/main/java/org/apache/tsfile/read/common/Chunk.java @@ -49,8 +49,11 @@ import static org.apache.tsfile.utils.RamUsageEstimator.sizeOfByteArray; -/** used in query. */ -public class Chunk { +/** + * A materialized chunk owning a page-AEAD reference that survives file reader closure. Close it + * after use, or transfer it to a chunk reader that releases the reference when closed. + */ +public class Chunk implements AutoCloseable { private static final long INSTANCE_SIZE = RamUsageEstimator.shallowSizeOfInstance(Chunk.class) @@ -88,7 +91,7 @@ public Chunk( this.chunkData = buffer; this.deleteIntervalList = deleteIntervalList; this.chunkStatistic = chunkStatistic; - this.encryptParam = encryptParam; + this.encryptParam = retainPageAeadParameter(encryptParam); } public Chunk(ChunkHeader header, ByteBuffer buffer) { @@ -100,7 +103,18 @@ public Chunk(ChunkHeader header, ByteBuffer buffer) { public Chunk(ChunkHeader header, ByteBuffer buffer, EncryptParameter encryptParam) { this.chunkHeader = header; this.chunkData = buffer; - this.encryptParam = encryptParam; + this.encryptParam = retainPageAeadParameter(encryptParam); + } + + private static EncryptParameter retainPageAeadParameter(EncryptParameter parameter) { + return parameter != null && parameter.isTdePageAead() ? parameter.retain() : parameter; + } + + @Override + public void close() { + if (encryptParam != null && encryptParam.isTdePageAead()) { + encryptParam.close(); + } } public EncryptParameter getEncryptParam() { @@ -124,6 +138,10 @@ public void setDeleteIntervalList(List list) { } public void mergeChunkByAppendPage(Chunk chunk) throws IOException { + if ((encryptParam != null && encryptParam.isTdePageAead()) + || (chunk.encryptParam != null && chunk.encryptParam.isTdePageAead())) { + throw new IOException(Messages.get("error.chunk.merge_page_aead_unsupported")); + } int dataSize = 0; // from where the page data of the merged chunk starts, if -1, it means the merged chunk has // more than one page @@ -213,13 +231,20 @@ public Statistics getChunkStatistic() { * and deleteIntervalList are all null in cache */ public long getRetainedSizeInBytes() { - return INSTANCE_SIZE + sizeOfByteArray(chunkData.capacity()); + int authenticationSize = chunkHeader.getChunkAuthenticationSize(); + return INSTANCE_SIZE + + sizeOfByteArray(chunkData.capacity()) + + (authenticationSize == 0 ? 0 : sizeOfByteArray(authenticationSize)); } public Chunk rewrite(TSDataType newType, Chunk timeChunk) throws IOException { if (newType == null || newType == chunkHeader.getDataType()) { return this; } + if (encryptParam != null && encryptParam.isTdePageAead()) { + encryptParam.resumeAfterChunkOrdinal(chunkHeader.getChunkOrdinal()); + encryptParam.resumeAfterChunkOrdinal(timeChunk.getHeader().getChunkOrdinal()); + } TSEncoding encoding = TSFileDescriptor.getInstance().getConfig().getValueEncoder(newType); IMeasurementSchema schema = new MeasurementSchema( @@ -233,34 +258,35 @@ public Chunk rewrite(TSDataType newType, Chunk timeChunk) throws IOException { encoding, schema.getValueEncoder(), encryptParam); - List valueChunks = new ArrayList<>(); - valueChunks.add(this); - TableChunkReader chunkReader = - new TableChunkReader( + try (Chunk readTimeChunk = new Chunk( timeChunk.getHeader(), timeChunk.getData(), null, timeChunk.getChunkStatistic(), - timeChunk.getEncryptParam()), - valueChunks, - null); - List pages = chunkReader.loadPageReaderList(); - for (IPageReader page : pages) { - IPointReader pointReader = page.getAllSatisfiedPageData().getBatchDataIterator(); - while (pointReader.hasNextTimeValuePair()) { - TimeValuePair point = pointReader.nextTimeValuePair(); - Object convertedValue = null; - if (point.getValue().getVector()[0] != null) { - convertedValue = - newType.castFromSingleValue( - chunkHeader.getDataType(), point.getValue().getVector()[0].getValue()); + timeChunk.getEncryptParam()); + Chunk readValueChunk = + new Chunk(chunkHeader, chunkData, deleteIntervalList, chunkStatistic, encryptParam)) { + List valueChunks = new ArrayList<>(); + valueChunks.add(readValueChunk); + TableChunkReader chunkReader = new TableChunkReader(readTimeChunk, valueChunks, null); + List pages = chunkReader.loadPageReaderList(); + for (IPageReader page : pages) { + IPointReader pointReader = page.getAllSatisfiedPageData().getBatchDataIterator(); + while (pointReader.hasNextTimeValuePair()) { + TimeValuePair point = pointReader.nextTimeValuePair(); + Object convertedValue = null; + if (point.getValue().getVector()[0] != null) { + convertedValue = + newType.castFromSingleValue( + chunkHeader.getDataType(), point.getValue().getVector()[0].getValue()); + } + long timestamp = point.getTimestamp(); + Type.fromTsDataType(newType) + .write(chunkWriter, timestamp, convertedValue, convertedValue == null); } - long timestamp = point.getTimestamp(); - Type.fromTsDataType(newType) - .write(chunkWriter, timestamp, convertedValue, convertedValue == null); + chunkWriter.sealCurrentPage(); } - chunkWriter.sealCurrentPage(); } ByteBuffer newChunkData = chunkWriter.getByteBuffer(); ChunkHeader newChunkHeader = @@ -270,7 +296,8 @@ public Chunk rewrite(TSDataType newType, Chunk timeChunk) throws IOException { newChunkData.capacity(), newType, chunkHeader.getCompressionType(), - encoding); + encoding, + chunkWriter.getChunkOrdinal()); chunkData.flip(); timeChunk.chunkData.flip(); return new Chunk( @@ -285,27 +312,33 @@ public Chunk rewrite(TSDataType newType) throws IOException { if (newType == null || newType == chunkHeader.getDataType()) { return this; } + if (encryptParam != null && encryptParam.isTdePageAead()) { + encryptParam.resumeAfterChunkOrdinal(chunkHeader.getChunkOrdinal()); + } TSEncoding encoding = TSFileDescriptor.getInstance().getConfig().getValueEncoder(newType); IMeasurementSchema schema = new MeasurementSchema( chunkHeader.getMeasurementID(), newType, encoding, chunkHeader.getCompressionType()); ChunkWriterImpl chunkWriter = new ChunkWriterImpl(schema, encryptParam); - ChunkReader chunkReader = new ChunkReader(this); - List pages = chunkReader.loadPageReaderList(); - for (IPageReader page : pages) { - BatchData batchData = page.getAllSatisfiedPageData(); - IPointReader pointReader = batchData.getBatchDataIterator(); - while (pointReader.hasNextTimeValuePair()) { - TimeValuePair point = pointReader.nextTimeValuePair(); - Object convertedValue = - newType.castFromSingleValue(chunkHeader.getDataType(), point.getValue().getValue()); - long timestamp = point.getTimestamp(); - if (convertedValue == null) { - throw new IOException(Messages.format("error.read.chunk_non_aligned_null", timestamp)); + try (Chunk readChunk = + new Chunk(chunkHeader, chunkData, deleteIntervalList, chunkStatistic, encryptParam)) { + ChunkReader chunkReader = new ChunkReader(readChunk); + List pages = chunkReader.loadPageReaderList(); + for (IPageReader page : pages) { + BatchData batchData = page.getAllSatisfiedPageData(); + IPointReader pointReader = batchData.getBatchDataIterator(); + while (pointReader.hasNextTimeValuePair()) { + TimeValuePair point = pointReader.nextTimeValuePair(); + Object convertedValue = + newType.castFromSingleValue(chunkHeader.getDataType(), point.getValue().getValue()); + long timestamp = point.getTimestamp(); + if (convertedValue == null) { + throw new IOException(Messages.format("error.read.chunk_non_aligned_null", timestamp)); + } + Type.fromTsDataType(newType).write(chunkWriter, timestamp, convertedValue); } - Type.fromTsDataType(newType).write(chunkWriter, timestamp, convertedValue); + chunkWriter.sealCurrentPage(); } - chunkWriter.sealCurrentPage(); } ByteBuffer newChunkData = chunkWriter.getByteBuffer(); ChunkHeader newChunkHeader = @@ -317,7 +350,8 @@ public Chunk rewrite(TSDataType newType) throws IOException { newChunkData.capacity(), newType, chunkHeader.getCompressionType(), - encoding); + encoding, + chunkWriter.getChunkOrdinal()); chunkData.flip(); return new Chunk( newChunkHeader, diff --git a/java/tsfile/src/main/java/org/apache/tsfile/read/controller/CachedChunkLoaderImpl.java b/java/tsfile/src/main/java/org/apache/tsfile/read/controller/CachedChunkLoaderImpl.java index 92925ea82f..43ef003c8c 100644 --- a/java/tsfile/src/main/java/org/apache/tsfile/read/controller/CachedChunkLoaderImpl.java +++ b/java/tsfile/src/main/java/org/apache/tsfile/read/controller/CachedChunkLoaderImpl.java @@ -20,6 +20,7 @@ package org.apache.tsfile.read.controller; import org.apache.tsfile.common.cache.LRUCache; +import org.apache.tsfile.encrypt.EncryptParameter; import org.apache.tsfile.file.metadata.ChunkMetadata; import org.apache.tsfile.file.metadata.IChunkMetadata; import org.apache.tsfile.file.metadata.statistics.Statistics; @@ -61,7 +62,9 @@ public CachedChunkLoaderImpl(TsFileSequenceReader fileSequenceReader, int cacheS @Override protected Chunk loadObjectByKey(ChunkCacheKey chunkCacheKey) throws IOException { - return reader.readMemChunk(chunkCacheKey); + try (Chunk chunk = reader.readMemChunk(chunkCacheKey)) { + return new Chunk(chunk.getHeader(), chunk.getData(), (EncryptParameter) null); + } } }; } diff --git a/java/tsfile/src/main/java/org/apache/tsfile/read/query/dataset/DataSetWithTimeGenerator.java b/java/tsfile/src/main/java/org/apache/tsfile/read/query/dataset/DataSetWithTimeGenerator.java index 63999a91ad..a5ca445e41 100644 --- a/java/tsfile/src/main/java/org/apache/tsfile/read/query/dataset/DataSetWithTimeGenerator.java +++ b/java/tsfile/src/main/java/org/apache/tsfile/read/query/dataset/DataSetWithTimeGenerator.java @@ -65,6 +65,41 @@ public boolean hasNextWithoutConstraint() throws IOException { return timeGenerator.hasNext(); } + @Override + public void close() throws IOException { + if (isClosed()) { + return; + } + super.close(); + IOException failure = null; + for (FileSeriesReaderByTimestamp reader : readers) { + if (reader == null) { + continue; + } + try { + reader.close(); + } catch (IOException e) { + if (failure == null) { + failure = e; + } else { + failure.addSuppressed(e); + } + } + } + try { + timeGenerator.close(); + } catch (IOException e) { + if (failure == null) { + failure = e; + } else { + failure.addSuppressed(e); + } + } + if (failure != null) { + throw failure; + } + } + @Override public RowRecord nextWithoutConstraint() throws IOException { long timestamp = timeGenerator.next(); diff --git a/java/tsfile/src/main/java/org/apache/tsfile/read/query/dataset/DataSetWithoutTimeGenerator.java b/java/tsfile/src/main/java/org/apache/tsfile/read/query/dataset/DataSetWithoutTimeGenerator.java index 00faed8fdb..d588421fb4 100644 --- a/java/tsfile/src/main/java/org/apache/tsfile/read/query/dataset/DataSetWithoutTimeGenerator.java +++ b/java/tsfile/src/main/java/org/apache/tsfile/read/query/dataset/DataSetWithoutTimeGenerator.java @@ -60,7 +60,16 @@ public DataSetWithoutTimeGenerator( throws IOException { super(paths, dataTypes); this.readers = readers; - initHeap(); + try { + initHeap(); + } catch (IOException | RuntimeException | Error e) { + try { + close(); + } catch (IOException closeError) { + e.addSuppressed(closeError); + } + throw e; + } } private void initHeap() throws IOException { @@ -93,6 +102,29 @@ public boolean hasNextWithoutConstraint() { return !timeHeap.isEmpty(); } + @Override + public void close() throws IOException { + if (isClosed()) { + return; + } + super.close(); + IOException failure = null; + for (AbstractFileSeriesReader reader : readers) { + try { + reader.closeChunkReader(); + } catch (IOException e) { + if (failure == null) { + failure = e; + } else { + failure.addSuppressed(e); + } + } + } + if (failure != null) { + throw failure; + } + } + @SuppressWarnings("squid:S3776") // Suppress high Cognitive Complexity warning @Override public RowRecord nextWithoutConstraint() throws IOException { diff --git a/java/tsfile/src/main/java/org/apache/tsfile/read/query/dataset/QueryDataSet.java b/java/tsfile/src/main/java/org/apache/tsfile/read/query/dataset/QueryDataSet.java index 0f50d59ec9..70c0a9c25d 100644 --- a/java/tsfile/src/main/java/org/apache/tsfile/read/query/dataset/QueryDataSet.java +++ b/java/tsfile/src/main/java/org/apache/tsfile/read/query/dataset/QueryDataSet.java @@ -28,7 +28,8 @@ import java.util.List; import java.util.Set; -public abstract class QueryDataSet { +public abstract class QueryDataSet implements AutoCloseable { + private boolean closed; protected List paths; protected List dataTypes; @@ -122,6 +123,37 @@ public void setWithoutNullColumnsIndex(Set withoutNullColumnsIndex) { } public boolean hasNext() throws IOException { + if (closed) { + return false; + } + boolean hasNext; + try { + hasNext = hasNextRow(); + } catch (IOException | RuntimeException | Error e) { + try { + close(); + } catch (IOException closeError) { + e.addSuppressed(closeError); + } + throw e; + } + if (!hasNext) { + close(); + } + return hasNext; + } + + /** Releases query-owned resources; the underlying file remains owned by its file reader. */ + @Override + public void close() throws IOException { + closed = true; + } + + public final boolean isClosed() { + return closed; + } + + private boolean hasNextRow() throws IOException { // proceed to the OFFSET row by skipping rows while (rowOffset > 0) { if (hasNextWithoutConstraint()) { diff --git a/java/tsfile/src/main/java/org/apache/tsfile/read/query/dataset/TreeResultSet.java b/java/tsfile/src/main/java/org/apache/tsfile/read/query/dataset/TreeResultSet.java index 6321e06226..618424bc02 100644 --- a/java/tsfile/src/main/java/org/apache/tsfile/read/query/dataset/TreeResultSet.java +++ b/java/tsfile/src/main/java/org/apache/tsfile/read/query/dataset/TreeResultSet.java @@ -21,11 +21,15 @@ import org.apache.tsfile.annotations.TsFileApi; import org.apache.tsfile.file.metadata.StringArrayDeviceID; +import org.apache.tsfile.i18n.Messages; import org.apache.tsfile.read.common.Field; import org.apache.tsfile.read.common.Path; import org.apache.tsfile.read.common.type.Type; import org.apache.tsfile.write.record.TSRecord; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + import java.io.IOException; import java.util.HashMap; import java.util.Iterator; @@ -37,6 +41,7 @@ import java.util.stream.IntStream; public class TreeResultSet extends AbstractResultSet { + private static final Logger LOGGER = LoggerFactory.getLogger(TreeResultSet.class); private QueryDataSet queryDataSet; private List deviceList; private List measurementList; @@ -79,7 +84,11 @@ public boolean next() throws IOException { @TsFileApi public void close() { - // nothing to be done + try { + queryDataSet.close(); + } catch (IOException e) { + LOGGER.error(Messages.get("log.read.tree_result_set_close_error"), e); + } } @Override diff --git a/java/tsfile/src/main/java/org/apache/tsfile/read/query/timegenerator/TimeGenerator.java b/java/tsfile/src/main/java/org/apache/tsfile/read/query/timegenerator/TimeGenerator.java index c073a45b71..75269efc8e 100644 --- a/java/tsfile/src/main/java/org/apache/tsfile/read/query/timegenerator/TimeGenerator.java +++ b/java/tsfile/src/main/java/org/apache/tsfile/read/query/timegenerator/TimeGenerator.java @@ -44,7 +44,10 @@ * that satisfies the filter condition. Then we use this timestamp to get values in other series * that are not included in IExpression */ -public abstract class TimeGenerator { +public abstract class TimeGenerator implements AutoCloseable { + + @Override + public void close() throws IOException {} private HashMap> leafNodeCache = new HashMap<>(); private HashMap> leafValuesCache; diff --git a/java/tsfile/src/main/java/org/apache/tsfile/read/query/timegenerator/TsFileTimeGenerator.java b/java/tsfile/src/main/java/org/apache/tsfile/read/query/timegenerator/TsFileTimeGenerator.java index c4e021e96e..e623b60cb3 100644 --- a/java/tsfile/src/main/java/org/apache/tsfile/read/query/timegenerator/TsFileTimeGenerator.java +++ b/java/tsfile/src/main/java/org/apache/tsfile/read/query/timegenerator/TsFileTimeGenerator.java @@ -29,12 +29,14 @@ import org.apache.tsfile.read.reader.series.FileSeriesReader; import java.io.IOException; +import java.util.ArrayList; import java.util.List; public class TsFileTimeGenerator extends TimeGenerator { private IChunkLoader chunkLoader; private IMetadataQuerier metadataQuerier; + private final List readers = new ArrayList<>(); public TsFileTimeGenerator( IExpression iexpression, IChunkLoader chunkLoader, IMetadataQuerier metadataQuerier) @@ -42,7 +44,16 @@ public TsFileTimeGenerator( this.chunkLoader = chunkLoader; this.metadataQuerier = metadataQuerier; - super.constructNode(iexpression); + try { + super.constructNode(iexpression); + } catch (IOException | RuntimeException | Error e) { + try { + close(); + } catch (IOException closeError) { + e.addSuppressed(closeError); + } + throw e; + } } @Override @@ -50,7 +61,29 @@ protected IBatchReader generateNewBatchReader(SingleSeriesExpression expression) throws IOException { List chunkMetadataList = metadataQuerier.getChunkMetaDataList(expression.getSeriesPath()); - return new FileSeriesReader(chunkLoader, chunkMetadataList, expression.getFilter()); + FileSeriesReader reader = + new FileSeriesReader(chunkLoader, chunkMetadataList, expression.getFilter()); + readers.add(reader); + return reader; + } + + @Override + public void close() throws IOException { + IOException failure = null; + for (FileSeriesReader reader : readers) { + try { + reader.closeChunkReader(); + } catch (IOException e) { + if (failure == null) { + failure = e; + } else { + failure.addSuppressed(e); + } + } + } + if (failure != null) { + throw failure; + } } @Override diff --git a/java/tsfile/src/main/java/org/apache/tsfile/read/reader/IChunkReader.java b/java/tsfile/src/main/java/org/apache/tsfile/read/reader/IChunkReader.java index e13d32d974..cec2c361b9 100644 --- a/java/tsfile/src/main/java/org/apache/tsfile/read/reader/IChunkReader.java +++ b/java/tsfile/src/main/java/org/apache/tsfile/read/reader/IChunkReader.java @@ -24,7 +24,7 @@ import java.io.IOException; import java.util.List; -public interface IChunkReader { +public interface IChunkReader extends AutoCloseable { boolean hasNextSatisfiedPage() throws IOException; diff --git a/java/tsfile/src/main/java/org/apache/tsfile/read/reader/TsFileLastReader.java b/java/tsfile/src/main/java/org/apache/tsfile/read/reader/TsFileLastReader.java index 8700a850e4..0b30ea6341 100644 --- a/java/tsfile/src/main/java/org/apache/tsfile/read/reader/TsFileLastReader.java +++ b/java/tsfile/src/main/java/org/apache/tsfile/read/reader/TsFileLastReader.java @@ -19,14 +19,13 @@ package org.apache.tsfile.read.reader; -import org.apache.tsfile.compress.IUnCompressor; import org.apache.tsfile.encoding.decoder.Decoder; +import org.apache.tsfile.encrypt.IDecryptor; import org.apache.tsfile.enums.TSDataType; import org.apache.tsfile.file.header.PageHeader; import org.apache.tsfile.file.metadata.ChunkMetadata; import org.apache.tsfile.file.metadata.IDeviceID; import org.apache.tsfile.file.metadata.TimeseriesMetadata; -import org.apache.tsfile.file.metadata.enums.CompressionType; import org.apache.tsfile.i18n.Messages; import org.apache.tsfile.read.TimeValuePair; import org.apache.tsfile.read.TsFileSequenceReader; @@ -181,6 +180,8 @@ private TimeValuePair readAlignedLastPoint(Chunk chunk, ChunkMetadata chunkMetad ByteBuffer chunkData = chunk.getData(); PageHeader lastPageHeader = null; ByteBuffer lastPageData = null; + int pageIndex = 0; + int lastPageIndex = -1; while (chunkData.hasRemaining()) { PageHeader pageHeader; if (chunk.isSinglePageChunk()) { @@ -196,17 +197,20 @@ private TimeValuePair readAlignedLastPoint(Chunk chunk, ChunkMetadata chunkMetad || (pageHeader.getStatistics() != null && pageHeader.getStatistics().getCount() > 0)) { lastPageHeader = pageHeader; lastPageData = pageData; + lastPageIndex = pageIndex; } + pageIndex++; } if (lastPageHeader != null) { - CompressionType compressionType = chunk.getHeader().getCompressionType(); - if (compressionType != CompressionType.UNCOMPRESSED) { - ByteBuffer uncompressedPage = ByteBuffer.allocate(lastPageHeader.getUncompressedSize()); - IUnCompressor.getUnCompressor(compressionType).uncompress(lastPageData, uncompressedPage); - lastPageData = uncompressedPage; - lastPageData.flip(); - } + lastPageData = + ChunkReader.deserializePageData( + lastPageHeader, + lastPageData, + chunk.getHeader(), + IDecryptor.getDecryptor(chunk.getEncryptParam()), + chunk.getEncryptParam(), + lastPageIndex); ValuePageReader valuePageReader = new ValuePageReader( @@ -265,15 +269,15 @@ private Pair readLastPoint( return new Pair<>(seriesMeta.getMeasurementId(), null); } - Chunk chunk = sequenceReader.readMemChunk(lastNonEmptyChunkMetadata); - - if (!isAligned) { - return new Pair<>(seriesMeta.getMeasurementId(), readNonAlignedLastPoint(chunk)); - } else { - return new Pair<>( - seriesMeta.getMeasurementId(), - readAlignedLastPoint( - chunk, lastNonEmptyChunkMetadata, seriesMeta.getStatistics().getEndTime())); + try (Chunk chunk = sequenceReader.readMemChunk(lastNonEmptyChunkMetadata)) { + if (!isAligned) { + return new Pair<>(seriesMeta.getMeasurementId(), readNonAlignedLastPoint(chunk)); + } else { + return new Pair<>( + seriesMeta.getMeasurementId(), + readAlignedLastPoint( + chunk, lastNonEmptyChunkMetadata, seriesMeta.getStatistics().getEndTime())); + } } } diff --git a/java/tsfile/src/main/java/org/apache/tsfile/read/reader/block/DeviceOrderedTsBlockReader.java b/java/tsfile/src/main/java/org/apache/tsfile/read/reader/block/DeviceOrderedTsBlockReader.java index 3bf2943a81..4495df4d61 100644 --- a/java/tsfile/src/main/java/org/apache/tsfile/read/reader/block/DeviceOrderedTsBlockReader.java +++ b/java/tsfile/src/main/java/org/apache/tsfile/read/reader/block/DeviceOrderedTsBlockReader.java @@ -19,6 +19,7 @@ package org.apache.tsfile.read.reader.block; +import org.apache.tsfile.exception.encrypt.EncryptException; import org.apache.tsfile.i18n.Messages; import org.apache.tsfile.read.common.block.TsBlock; import org.apache.tsfile.read.controller.IChunkLoader; @@ -31,6 +32,7 @@ import org.slf4j.LoggerFactory; import java.io.IOException; +import java.io.UncheckedIOException; import java.util.NoSuchElementException; public class DeviceOrderedTsBlockReader implements TsBlockReader { @@ -43,6 +45,7 @@ public class DeviceOrderedTsBlockReader implements TsBlockReader { private SingleDeviceTsBlockReader currentReader; private ExpressionTree timeFilter; private ExpressionTree measurementFilter; + private boolean closed; public DeviceOrderedTsBlockReader( DeviceTaskIterator taskIterator, @@ -61,9 +64,13 @@ public DeviceOrderedTsBlockReader( @Override public boolean hasNext() { + if (closed) { + return false; + } if (currentReader != null && currentReader.hasNext()) { return true; } + closeCurrentReaderForIteration(); while (taskIterator.hasNext()) { final DeviceQueryTask nextTask = taskIterator.next(); try { @@ -71,11 +78,16 @@ public boolean hasNext() { new SingleDeviceTsBlockReader( nextTask, metadataQuerier, chunkLoader, blockSize, timeFilter, measurementFilter); } catch (IOException e) { + if (EncryptException.isEncryptionFailure(e)) { + closed = true; + throw new UncheckedIOException(e); + } LOGGER.error(Messages.get("log.read.block_reader_construct_error"), nextTask, e); } if (currentReader != null && currentReader.hasNext()) { return true; } + closeCurrentReaderForIteration(); } return false; } @@ -90,8 +102,25 @@ public TsBlock next() throws IOException { @Override public void close() throws Exception { + closed = true; + closeCurrentReader(); + } + + private void closeCurrentReaderForIteration() { + try { + closeCurrentReader(); + } catch (Exception e) { + throw new IllegalStateException(e); + } + } + + private void closeCurrentReader() throws Exception { if (currentReader != null) { - currentReader.close(); + try { + currentReader.close(); + } finally { + currentReader = null; + } } } } diff --git a/java/tsfile/src/main/java/org/apache/tsfile/read/reader/block/SingleDeviceTsBlockReader.java b/java/tsfile/src/main/java/org/apache/tsfile/read/reader/block/SingleDeviceTsBlockReader.java index da4f05e116..4fd474a90f 100644 --- a/java/tsfile/src/main/java/org/apache/tsfile/read/reader/block/SingleDeviceTsBlockReader.java +++ b/java/tsfile/src/main/java/org/apache/tsfile/read/reader/block/SingleDeviceTsBlockReader.java @@ -39,6 +39,7 @@ import org.slf4j.LoggerFactory; import java.io.IOException; +import java.io.UncheckedIOException; import java.util.ArrayList; import java.util.HashMap; import java.util.List; @@ -59,6 +60,7 @@ public class SingleDeviceTsBlockReader implements TsBlockReader { private final Map idColumnContextMap; private Long nextTime; + private boolean closed; public SingleDeviceTsBlockReader( DeviceQueryTask task, @@ -77,22 +79,27 @@ public SingleDeviceTsBlockReader( this.measureColumnContextMap = new HashMap<>(); this.idColumnContextMap = new HashMap<>(); - final List> chunkMetadataLists = - metadataQuerier.getChunkMetadataLists( - task.getDeviceID(), - task.getColumnMapping().getMeasurementColumns(), - task.getIndexRoot()); + try { + final List> chunkMetadataLists = + metadataQuerier.getChunkMetadataLists( + task.getDeviceID(), + task.getColumnMapping().getMeasurementColumns(), + task.getIndexRoot()); - Filter timeFilter = timeExpression == null ? null : timeExpression.toFilter(); - for (List chunkMetadataList : chunkMetadataLists) { - constructColumnContext(chunkMetadataList, chunkLoader, timeFilter); - } + Filter timeFilter = timeExpression == null ? null : timeExpression.toFilter(); + for (List chunkMetadataList : chunkMetadataLists) { + constructColumnContext(chunkMetadataList, chunkLoader, timeFilter); + } - for (String idColumn : task.getColumnMapping().getIdColumns()) { - final List columnPosInResult = task.getColumnMapping().getColumnPos(idColumn); - // the first segment in DeviceId is the table name - final int columnPosInId = task.getTableSchema().findIdColumnOrder(idColumn) + 1; - idColumnContextMap.put(idColumn, new IdColumnContext(columnPosInResult, columnPosInId)); + for (String idColumn : task.getColumnMapping().getIdColumns()) { + final List columnPosInResult = task.getColumnMapping().getColumnPos(idColumn); + // the first segment in DeviceId is the table name + final int columnPosInId = task.getTableSchema().findIdColumnOrder(idColumn) + 1; + idColumnContextMap.put(idColumn, new IdColumnContext(columnPosInResult, columnPosInId)); + } + } catch (IOException | RuntimeException | Error e) { + closeAfterFailure(e); + throw e; } } @@ -105,33 +112,63 @@ private void constructColumnContext( final IChunkMetadata chunkMetadata = chunkMetadataList.get(0); AbstractFileSeriesReader seriesReader = new FileSeriesReader(chunkLoader, chunkMetadataList, timeFilter, false); - if (seriesReader.hasNextBatch()) { - if (chunkMetadata instanceof AbstractAlignedChunkMetadata) { - final List currentChunkMeasurementNames = - seriesReader.getCurrentChunkMeasurementNames(); - List> posInResult = new ArrayList<>(); - for (String currentChunkMeasurementName : currentChunkMeasurementNames) { - posInResult.add(task.getColumnMapping().getColumnPos(currentChunkMeasurementName)); + try { + if (seriesReader.hasNextBatch()) { + if (chunkMetadata instanceof AbstractAlignedChunkMetadata) { + final List currentChunkMeasurementNames = + seriesReader.getCurrentChunkMeasurementNames(); + List> posInResult = new ArrayList<>(); + for (String currentChunkMeasurementName : currentChunkMeasurementNames) { + posInResult.add(task.getColumnMapping().getColumnPos(currentChunkMeasurementName)); + } + measureColumnContextMap.put( + "", + new VectorMeasurementColumnContext( + posInResult, seriesReader.nextBatch(), seriesReader)); + } else { + final String measurementUid = chunkMetadata.getMeasurementUid(); + measureColumnContextMap.put( + measurementUid, + new SingleMeasurementColumnContext( + measurementUid, + task.getColumnMapping().getColumnPos(measurementUid), + seriesReader.nextBatch(), + seriesReader)); } - measureColumnContextMap.put( - "", - new VectorMeasurementColumnContext( - posInResult, seriesReader.nextBatch(), seriesReader)); } else { - final String measurementUid = chunkMetadata.getMeasurementUid(); - measureColumnContextMap.put( - measurementUid, - new SingleMeasurementColumnContext( - measurementUid, - task.getColumnMapping().getColumnPos(measurementUid), - seriesReader.nextBatch(), - seriesReader)); + seriesReader.closeChunkReader(); } + } catch (IOException | RuntimeException | Error e) { + try { + seriesReader.closeChunkReader(); + } catch (IOException | RuntimeException | Error closeError) { + e.addSuppressed(closeError); + } + throw e; } } @Override public boolean hasNext() { + if (closed) { + return false; + } + try { + return prepareNextBlock(); + } catch (IOException e) { + closeAfterFailure(e); + if (org.apache.tsfile.exception.encrypt.EncryptException.isEncryptionFailure(e)) { + throw new UncheckedIOException(e); + } + LOGGER.error(Messages.get("log.read.block_reader_fill_measurements_error"), e); + return false; + } catch (RuntimeException | Error e) { + closeAfterFailure(e); + throw e; + } + } + + private boolean prepareNextBlock() throws IOException { if (!lastBlockReturned) { return true; } @@ -158,13 +195,8 @@ public boolean hasNext() { } } - try { - fillMeasurements(minTimeColumns); - nextTime = null; - } catch (IOException e) { - LOGGER.error(Messages.get("log.read.block_reader_fill_measurements_error"), e); - return false; - } + fillMeasurements(minTimeColumns); + nextTime = null; // all columns have exhausted if (measureColumnContextMap.isEmpty()) { @@ -251,7 +283,46 @@ public TsBlock next() throws IOException { @Override public void close() throws Exception { - // nothing to be done + closeColumnReaders(); + } + + private void closeAfterFailure(Throwable failure) { + try { + closeColumnReaders(); + } catch (IOException | RuntimeException | Error closeError) { + failure.addSuppressed(closeError); + } + } + + private void closeColumnReaders() throws IOException { + if (closed) { + return; + } + closed = true; + Throwable failure = null; + for (MeasurementColumnContext context : measureColumnContextMap.values()) { + try { + context.seriesReader.closeChunkReader(); + } catch (IOException | RuntimeException | Error e) { + if (failure == null) { + failure = e; + } else { + failure.addSuppressed(e); + } + } + } + measureColumnContextMap.clear(); + idColumnContextMap.clear(); + nextTime = null; + if (failure instanceof IOException) { + throw (IOException) failure; + } + if (failure instanceof RuntimeException) { + throw (RuntimeException) failure; + } + if (failure instanceof Error) { + throw (Error) failure; + } } public abstract static class MeasurementColumnContext { diff --git a/java/tsfile/src/main/java/org/apache/tsfile/read/reader/chunk/AbstractAlignedChunkReader.java b/java/tsfile/src/main/java/org/apache/tsfile/read/reader/chunk/AbstractAlignedChunkReader.java index acc9789e40..f42fae00c9 100644 --- a/java/tsfile/src/main/java/org/apache/tsfile/read/reader/chunk/AbstractAlignedChunkReader.java +++ b/java/tsfile/src/main/java/org/apache/tsfile/read/reader/chunk/AbstractAlignedChunkReader.java @@ -41,7 +41,10 @@ import java.util.List; import java.util.function.LongConsumer; +/** Owns the supplied time and value chunks' encryption references until {@link #close()}. */ public abstract class AbstractAlignedChunkReader extends AbstractChunkReader { + private final Chunk timeChunk; + private final List valueChunks; // chunk header of the time column private final ChunkHeader timeChunkHeader; // chunk data of the time column @@ -67,6 +70,8 @@ public abstract class AbstractAlignedChunkReader extends AbstractChunkReader { LongConsumer filterRowsRecorder) throws IOException { super(readStopTime, queryFilter, filterRowsRecorder); + this.timeChunk = timeChunk; + this.valueChunks = new ArrayList<>(valueChunkList); this.timeChunkHeader = timeChunk.getHeader(); this.timeChunkDataBuffer = timeChunk.getData(); this.timeDeleteIntervalList = timeChunk.getDeleteIntervalList(); @@ -81,7 +86,22 @@ public abstract class AbstractAlignedChunkReader extends AbstractChunkReader { valueChunkStatisticsList.add(chunk == null ? null : chunk.getChunkStatistic()); }); this.encryptParam = timeChunk.getEncryptParam(); - initAllPageReaders(timeChunk.getChunkStatistic(), valueChunkStatisticsList); + try { + initAllPageReaders(timeChunk.getChunkStatistic(), valueChunkStatisticsList); + } catch (IOException | RuntimeException | Error e) { + close(); + throw e; + } + } + + @Override + public void close() { + timeChunk.close(); + for (Chunk valueChunk : valueChunks) { + if (valueChunk != null) { + valueChunk.close(); + } + } } /** construct all the page readers in this chunk */ @@ -90,12 +110,15 @@ private void initAllPageReaders( List> valueChunkStatisticsList) throws IOException { // construct next satisfied page header + int pageIndex = 0; while (timeChunkDataBuffer.remaining() > 0) { // deserialize PageHeader from chunkDataBuffer AbstractAlignedPageReader alignedPageReader = isSinglePageChunk() - ? deserializeFromSinglePageChunk(timeChunkStatistics, valueChunkStatisticsList) - : deserializeFromMultiPageChunk(); + ? deserializeFromSinglePageChunk( + timeChunkStatistics, valueChunkStatisticsList, pageIndex) + : deserializeFromMultiPageChunk(pageIndex); + pageIndex++; if (alignedPageReader != null) { pageReaderList.add(alignedPageReader); } @@ -108,7 +131,8 @@ private boolean isSinglePageChunk() { private AbstractAlignedPageReader deserializeFromSinglePageChunk( Statistics timeChunkStatistics, - List> valueChunkStatisticsList) + List> valueChunkStatisticsList, + int pageIndex) throws IOException { PageHeader timePageHeader = PageHeader.deserializeFrom(timeChunkDataBuffer, timeChunkStatistics); @@ -132,12 +156,13 @@ private AbstractAlignedPageReader deserializeFromSinglePageChunk( skipCurrentPage(timePageHeader, valuePageHeaderList); return null; } - return constructAlignedPageReader(timePageHeader, valuePageHeaderList); + return constructAlignedPageReader(timePageHeader, valuePageHeaderList, pageIndex); } abstract boolean needSkipForSinglePageChunk(boolean isAllNull, PageHeader timePageHeader); - private AbstractAlignedPageReader deserializeFromMultiPageChunk() throws IOException { + private AbstractAlignedPageReader deserializeFromMultiPageChunk(int pageIndex) + throws IOException { PageHeader timePageHeader = PageHeader.deserializeFrom(timeChunkDataBuffer, timeChunkHeader.getDataType()); List valuePageHeaderList = new ArrayList<>(); @@ -158,7 +183,7 @@ private AbstractAlignedPageReader deserializeFromMultiPageChunk() throws IOExcep skipCurrentPage(timePageHeader, valuePageHeaderList); return null; } - return constructAlignedPageReader(timePageHeader, valuePageHeaderList); + return constructAlignedPageReader(timePageHeader, valuePageHeaderList, pageIndex); } abstract boolean needSkipForMultiPageChunk(boolean isAllNull, PageHeader timePageHeader); @@ -193,11 +218,17 @@ private void skipCurrentPage(PageHeader timePageHeader, List valuePa } private AbstractAlignedPageReader constructAlignedPageReader( - PageHeader timePageHeader, List rawValuePageHeaderList) throws IOException { + PageHeader timePageHeader, List rawValuePageHeaderList, int pageIndex) + throws IOException { IDecryptor decrytor = IDecryptor.getDecryptor(encryptParam); ByteBuffer timePageData = ChunkReader.deserializePageData( - timePageHeader, timeChunkDataBuffer, timeChunkHeader, decrytor); + timePageHeader, + timeChunkDataBuffer, + timeChunkHeader, + decrytor, + encryptParam, + pageIndex); List valuePageHeaderList = new ArrayList<>(); LazyLoadPageData[] lazyLoadPageDataArray = new LazyLoadPageData[rawValuePageHeaderList.size()]; @@ -238,7 +269,9 @@ private AbstractAlignedPageReader constructAlignedPageReader( valueChunkDataBufferList.get(i).array(), currentPagePosition, IUnCompressor.getUnCompressor(valueChunkHeader.getCompressionType()), - encryptParam); + encryptParam, + pageIndex, + valueChunkHeader.getChunkOrdinal()); valueDataTypeList.add(valueChunkHeader.getDataType()); valueDecoderList.add(valueChunkHeader.calculateDecoderForNonTimeChunk()); isAllNull = false; diff --git a/java/tsfile/src/main/java/org/apache/tsfile/read/reader/chunk/ChunkReader.java b/java/tsfile/src/main/java/org/apache/tsfile/read/reader/chunk/ChunkReader.java index b555a25e1b..ff4c68d7ce 100644 --- a/java/tsfile/src/main/java/org/apache/tsfile/read/reader/chunk/ChunkReader.java +++ b/java/tsfile/src/main/java/org/apache/tsfile/read/reader/chunk/ChunkReader.java @@ -22,6 +22,7 @@ import org.apache.tsfile.compress.IUnCompressor; import org.apache.tsfile.encrypt.EncryptParameter; import org.apache.tsfile.encrypt.IDecryptor; +import org.apache.tsfile.encrypt.PageCryptoContext; import org.apache.tsfile.file.MetaMarker; import org.apache.tsfile.file.header.ChunkHeader; import org.apache.tsfile.file.header.PageHeader; @@ -40,8 +41,10 @@ import java.util.List; import java.util.function.LongConsumer; +/** Owns the supplied chunk's encryption reference until {@link #close()}. */ public class ChunkReader extends AbstractChunkReader { + private final Chunk chunk; private final ChunkHeader chunkHeader; private final ByteBuffer chunkDataBuffer; private final List deleteIntervalList; @@ -52,11 +55,22 @@ public class ChunkReader extends AbstractChunkReader { public ChunkReader( Chunk chunk, long readStopTime, Filter queryFilter, LongConsumer filterRowsRecorder) { super(readStopTime, queryFilter, filterRowsRecorder); + this.chunk = chunk; this.chunkHeader = chunk.getHeader(); this.chunkDataBuffer = chunk.getData(); this.deleteIntervalList = chunk.getDeleteIntervalList(); this.encryptParam = chunk.getEncryptParam(); - initAllPageReaders(chunk.getChunkStatistic()); + try { + initAllPageReaders(chunk.getChunkStatistic()); + } catch (RuntimeException | Error e) { + chunk.close(); + throw e; + } + } + + @Override + public void close() { + chunk.close(); } public ChunkReader(Chunk chunk) throws IOException { @@ -81,7 +95,9 @@ public ChunkReader(Chunk chunk, long readStopTime) { private void initAllPageReaders(Statistics chunkStatistic) { // construct next satisfied page header + int pageIndex = 0; while (chunkDataBuffer.remaining() > 0) { + int currentPageIndex = pageIndex++; // deserialize a PageHeader from chunkDataBuffer PageHeader pageHeader; if (((byte) (chunkHeader.getChunkType() & 0x3F)) == MetaMarker.ONLY_ONE_PAGE_CHUNK_HEADER) { @@ -100,7 +116,7 @@ private void initAllPageReaders(Statistics chunkStatisti if (pageDeleted(pageHeader)) { skipCurrentPage(pageHeader); } else { - pageReaderList.add(constructPageReader(pageHeader)); + pageReaderList.add(constructPageReader(pageHeader, currentPageIndex)); } } } @@ -141,7 +157,7 @@ private void skipCurrentPage(PageHeader pageHeader) { chunkDataBuffer.position(chunkDataBuffer.position() + pageHeader.getCompressedSize()); } - private PageReader constructPageReader(PageHeader pageHeader) { + private PageReader constructPageReader(PageHeader pageHeader, int pageIndex) { IUnCompressor unCompressor = IUnCompressor.getUnCompressor(chunkHeader.getCompressionType()); // record the current position of chunkDataBuffer, use this to get the page data in PageReader // through directly accessing the buffer array @@ -151,7 +167,12 @@ private PageReader constructPageReader(PageHeader pageHeader) { new PageReader( pageHeader, new LazyLoadPageData( - chunkDataBuffer.array(), currentPagePosition, unCompressor, encryptParam), + chunkDataBuffer.array(), + currentPagePosition, + unCompressor, + encryptParam, + pageIndex, + chunkHeader.getChunkOrdinal()), chunkHeader.getDataType(), chunkHeader.calculateDecoderForNonTimeChunk(), getConfiguredTimeDecoder(), @@ -221,7 +242,55 @@ public static ByteBuffer decryptAndUncompressPageData( compressedPageData.arrayOffset() + compressedPageData.position(), compressedPageBodyLength); unCompressor.uncompress( - decryptedPageData, 0, compressedPageBodyLength, uncompressedPageData, 0); + decryptedPageData, 0, decryptedPageData.length, uncompressedPageData, 0); + } catch (Exception e) { + throw new IOException( + Messages.format( + "error.read.uncompress_error_with_header", + pageHeader.getUncompressedSize(), + pageHeader.getCompressedSize(), + pageHeader, + e.getMessage()), + e); + } + compressedPageData.position(compressedPageData.position() + compressedPageBodyLength); + return ByteBuffer.wrap(uncompressedPageData); + } + + public static ByteBuffer decryptAndUncompressPageData( + PageHeader pageHeader, + IUnCompressor unCompressor, + ByteBuffer compressedPageData, + IDecryptor decryptor, + EncryptParameter encryptParameter, + int pageIndex, + long chunkOrdinal) + throws IOException { + int compressedPageBodyLength = pageHeader.getCompressedSize(); + byte[] uncompressedPageData = new byte[pageHeader.getUncompressedSize()]; + try { + PageCryptoContext pageCryptoContext = + PageCryptoContext.forDecryption( + encryptParameter, + pageHeader.getUncompressedSize(), + compressedPageData.slice().limit(compressedPageBodyLength), + pageIndex, + chunkOrdinal); + byte[] decryptedPageData = + decryptor.decryptPage( + compressedPageData.array(), + compressedPageData.arrayOffset() + compressedPageData.position(), + compressedPageBodyLength, + pageCryptoContext); + if (decryptedPageData.length != pageCryptoContext.getCompressedPlaintextSize()) { + throw new IOException( + Messages.format( + "error.encrypt.page_plaintext_size_mismatch", + pageCryptoContext.getCompressedPlaintextSize(), + decryptedPageData.length)); + } + unCompressor.uncompress( + decryptedPageData, 0, decryptedPageData.length, uncompressedPageData, 0); } catch (Exception e) { throw new IOException( Messages.format( @@ -243,6 +312,31 @@ public static ByteBuffer deserializePageData( ByteBuffer compressedPageBody = readCompressedPageData(pageHeader, chunkBuffer); if (decryptor == null || decryptor.getEncryptionType() == EncryptionType.UNENCRYPTED) { return uncompressPageData(pageHeader, unCompressor, compressedPageBody); + } + return decryptAndUncompressPageData(pageHeader, unCompressor, compressedPageBody, decryptor); + } + + public static ByteBuffer deserializePageData( + PageHeader pageHeader, + ByteBuffer chunkBuffer, + ChunkHeader chunkHeader, + IDecryptor decryptor, + EncryptParameter encryptParameter, + int pageIndex) + throws IOException { + IUnCompressor unCompressor = IUnCompressor.getUnCompressor(chunkHeader.getCompressionType()); + ByteBuffer compressedPageBody = readCompressedPageData(pageHeader, chunkBuffer); + if (encryptParameter != null && encryptParameter.isTdePageAead()) { + return decryptAndUncompressPageData( + pageHeader, + unCompressor, + compressedPageBody, + decryptor, + encryptParameter, + pageIndex, + chunkHeader.getChunkOrdinal()); + } else if (decryptor == null || decryptor.getEncryptionType() == EncryptionType.UNENCRYPTED) { + return uncompressPageData(pageHeader, unCompressor, compressedPageBody); } else { return decryptAndUncompressPageData(pageHeader, unCompressor, compressedPageBody, decryptor); } diff --git a/java/tsfile/src/main/java/org/apache/tsfile/read/reader/page/LazyLoadPageData.java b/java/tsfile/src/main/java/org/apache/tsfile/read/reader/page/LazyLoadPageData.java index 7b01a67f54..3eea62b90b 100644 --- a/java/tsfile/src/main/java/org/apache/tsfile/read/reader/page/LazyLoadPageData.java +++ b/java/tsfile/src/main/java/org/apache/tsfile/read/reader/page/LazyLoadPageData.java @@ -23,11 +23,14 @@ import org.apache.tsfile.encrypt.EncryptParameter; import org.apache.tsfile.encrypt.EncryptUtils; import org.apache.tsfile.encrypt.IDecryptor; +import org.apache.tsfile.encrypt.PageCryptoContext; import org.apache.tsfile.file.header.PageHeader; +import org.apache.tsfile.file.metadata.enums.EncryptionType; import org.apache.tsfile.i18n.Messages; import java.io.IOException; import java.nio.ByteBuffer; +import java.util.Arrays; public class LazyLoadPageData { /** Reference to the data of original chunkDataBuffer. * */ @@ -39,11 +42,16 @@ public class LazyLoadPageData { private final EncryptParameter encryptParam; + private final int pageIndex; + private final long chunkOrdinal; + public LazyLoadPageData(byte[] data, int offset, IUnCompressor unCompressor) { this.chunkData = data; this.pageDataOffset = offset; this.unCompressor = unCompressor; this.encryptParam = EncryptUtils.getEncryptParameter(); + this.pageIndex = -1; + this.chunkOrdinal = -1; } public LazyLoadPageData( @@ -52,17 +60,57 @@ public LazyLoadPageData( this.pageDataOffset = offset; this.unCompressor = unCompressor; this.encryptParam = encryptParam; + this.pageIndex = -1; + this.chunkOrdinal = -1; + } + + public LazyLoadPageData( + byte[] data, + int offset, + IUnCompressor unCompressor, + EncryptParameter encryptParam, + int pageIndex, + long chunkOrdinal) { + this.chunkData = data; + this.pageDataOffset = offset; + this.unCompressor = unCompressor; + this.encryptParam = encryptParam; + this.pageIndex = pageIndex; + this.chunkOrdinal = chunkOrdinal; } public ByteBuffer uncompressPageData(PageHeader pageHeader) throws IOException { int compressedPageBodyLength = pageHeader.getCompressedSize(); byte[] uncompressedPageData = new byte[pageHeader.getUncompressedSize()]; IDecryptor decryptor = IDecryptor.getDecryptor(encryptParam); - byte[] decryptedPageData = - decryptor.decrypt(chunkData, pageDataOffset, compressedPageBodyLength); + byte[] decryptedPageData; + if (encryptParam != null && encryptParam.isTdePageAead()) { + PageCryptoContext pageCryptoContext = + PageCryptoContext.forDecryption( + encryptParam, + pageHeader.getUncompressedSize(), + ByteBuffer.wrap(chunkData, pageDataOffset, compressedPageBodyLength).slice(), + pageIndex, + chunkOrdinal); + decryptedPageData = + decryptor.decryptPage( + chunkData, pageDataOffset, compressedPageBodyLength, pageCryptoContext); + if (decryptedPageData.length != pageCryptoContext.getCompressedPlaintextSize()) { + throw new IOException( + Messages.format( + "error.encrypt.page_plaintext_size_mismatch", + pageCryptoContext.getCompressedPlaintextSize(), + decryptedPageData.length)); + } + } else if (decryptor.getEncryptionType() == EncryptionType.UNENCRYPTED) { + decryptedPageData = + Arrays.copyOfRange(chunkData, pageDataOffset, pageDataOffset + compressedPageBodyLength); + } else { + decryptedPageData = decryptor.decrypt(chunkData, pageDataOffset, compressedPageBodyLength); + } try { unCompressor.uncompress( - decryptedPageData, 0, compressedPageBodyLength, uncompressedPageData, 0); + decryptedPageData, 0, decryptedPageData.length, uncompressedPageData, 0); } catch (Exception e) { throw new IOException( Messages.format( diff --git a/java/tsfile/src/main/java/org/apache/tsfile/read/reader/series/AbstractFileSeriesReader.java b/java/tsfile/src/main/java/org/apache/tsfile/read/reader/series/AbstractFileSeriesReader.java index 3616150150..545071b8be 100644 --- a/java/tsfile/src/main/java/org/apache/tsfile/read/reader/series/AbstractFileSeriesReader.java +++ b/java/tsfile/src/main/java/org/apache/tsfile/read/reader/series/AbstractFileSeriesReader.java @@ -67,6 +67,7 @@ public boolean hasNextBatch() throws IOException { if (chunkReader != null && chunkReader.hasNextSatisfiedPage()) { return true; } + closeChunkReader(); // current chunk does not have additional batch, init new chunk reader while (chunkToRead < chunkMetadataList.size()) { @@ -82,6 +83,7 @@ public boolean hasNextBatch() throws IOException { if (chunkReader.hasNextSatisfiedPage()) { return true; } + closeChunkReader(); } return false; } @@ -97,7 +99,22 @@ public BatchData nextBatch() throws IOException { @Override public void close() throws IOException { - chunkLoader.close(); + try { + closeChunkReader(); + } finally { + chunkLoader.close(); + } + } + + /** Releases the current chunk without closing the shared file loader. */ + public void closeChunkReader() throws IOException { + if (chunkReader != null) { + try { + chunkReader.close(); + } finally { + chunkReader = null; + } + } } private IChunkMetadata nextChunkMeta() { diff --git a/java/tsfile/src/main/java/org/apache/tsfile/read/reader/series/FileSeriesReader.java b/java/tsfile/src/main/java/org/apache/tsfile/read/reader/series/FileSeriesReader.java index 9fa26c8298..13b2297a37 100644 --- a/java/tsfile/src/main/java/org/apache/tsfile/read/reader/series/FileSeriesReader.java +++ b/java/tsfile/src/main/java/org/apache/tsfile/read/reader/series/FileSeriesReader.java @@ -65,19 +65,29 @@ protected void initChunkReader(IChunkMetadata chunkMetaData) throws IOException Chunk timeChunk = chunkLoader.loadChunk((ChunkMetadata) (alignedChunkMetadata.getTimeChunkMetadata())); List valueChunkList = new ArrayList<>(); - for (IChunkMetadata metadata : alignedChunkMetadata.getValueChunkMetadataList()) { - if (metadata != null) { - valueChunkList.add(chunkLoader.loadChunk((ChunkMetadata) metadata)); - currentChunkMeasurementNames.add(metadata.getMeasurementUid()); - continue; + try { + for (IChunkMetadata metadata : alignedChunkMetadata.getValueChunkMetadataList()) { + if (metadata != null) { + valueChunkList.add(chunkLoader.loadChunk((ChunkMetadata) metadata)); + currentChunkMeasurementNames.add(metadata.getMeasurementUid()); + continue; + } + valueChunkList.add(null); + currentChunkMeasurementNames.add(null); } - valueChunkList.add(null); - currentChunkMeasurementNames.add(null); - } - if (ignoreAllNullRows) { - this.chunkReader = new AlignedChunkReader(timeChunk, valueChunkList, filter); - } else { - this.chunkReader = new TableChunkReader(timeChunk, valueChunkList, filter); + if (ignoreAllNullRows) { + this.chunkReader = new AlignedChunkReader(timeChunk, valueChunkList, filter); + } else { + this.chunkReader = new TableChunkReader(timeChunk, valueChunkList, filter); + } + } catch (IOException | RuntimeException | Error e) { + timeChunk.close(); + for (Chunk valueChunk : valueChunkList) { + if (valueChunk != null) { + valueChunk.close(); + } + } + throw e; } } } diff --git a/java/tsfile/src/main/java/org/apache/tsfile/read/reader/series/FileSeriesReaderByTimestamp.java b/java/tsfile/src/main/java/org/apache/tsfile/read/reader/series/FileSeriesReaderByTimestamp.java index ae6c394ca7..06c38a10a0 100644 --- a/java/tsfile/src/main/java/org/apache/tsfile/read/reader/series/FileSeriesReaderByTimestamp.java +++ b/java/tsfile/src/main/java/org/apache/tsfile/read/reader/series/FileSeriesReaderByTimestamp.java @@ -38,7 +38,7 @@ * Series reader is used to query one series of one tsfile, using this reader to query the value of * a series with given timestamps. */ -public class FileSeriesReaderByTimestamp { +public class FileSeriesReaderByTimestamp implements AutoCloseable { protected IChunkLoader chunkLoader; protected List chunkMetadataList; @@ -99,6 +99,8 @@ public Object getValueInTimestamp(long timestamp) throws IOException { data = chunkReader.nextPageData(); } else if (!constructNextSatisfiedChunkReader()) { return null; + } else { + data = chunkReader.nextPageData(); } } } @@ -137,11 +139,15 @@ public boolean hasNext() throws IOException { } private boolean constructNextSatisfiedChunkReader() throws IOException { + closeChunkReader(); while (currentChunkIndex < chunkMetadataList.size()) { IChunkMetadata chunkMetaData = chunkMetadataList.get(currentChunkIndex++); if (chunkSatisfied(chunkMetaData)) { initChunkReader(chunkMetaData); - return true; + if (chunkReader.hasNextSatisfiedPage()) { + return true; + } + closeChunkReader(); } } return false; @@ -156,14 +162,51 @@ private void initChunkReader(IChunkMetadata chunkMetaData) throws IOException { Chunk timeChunk = chunkLoader.loadChunk((ChunkMetadata) (alignedChunkMetadata.getTimeChunkMetadata())); List valueChunkList = new ArrayList<>(); - for (IChunkMetadata metadata : alignedChunkMetadata.getValueChunkMetadataList()) { - valueChunkList.add(chunkLoader.loadChunk((ChunkMetadata) metadata)); + try { + for (IChunkMetadata metadata : alignedChunkMetadata.getValueChunkMetadataList()) { + valueChunkList.add(chunkLoader.loadChunk((ChunkMetadata) metadata)); + } + this.chunkReader = new AlignedChunkReader(timeChunk, valueChunkList, currentTimestamp); + } catch (IOException | RuntimeException | Error e) { + closeChunkOnFailure(timeChunk, e); + for (Chunk valueChunk : valueChunkList) { + closeChunkOnFailure(valueChunk, e); + } + throw e; } - this.chunkReader = new AlignedChunkReader(timeChunk, valueChunkList, currentTimestamp); } } private boolean chunkSatisfied(IChunkMetadata chunkMetaData) { return chunkMetaData.getEndTime() >= currentTimestamp; } + + private static void closeChunkOnFailure(Chunk chunk, Throwable failure) { + if (chunk == null) { + return; + } + try { + chunk.close(); + } catch (RuntimeException | Error closeFailure) { + if (closeFailure != failure) { + failure.addSuppressed(closeFailure); + } + } + } + + private void closeChunkReader() throws IOException { + if (chunkReader != null) { + try { + chunkReader.close(); + } finally { + chunkReader = null; + data = null; + } + } + } + + @Override + public void close() throws IOException { + closeChunkReader(); + } } diff --git a/java/tsfile/src/main/java/org/apache/tsfile/utils/ReadWriteIOUtils.java b/java/tsfile/src/main/java/org/apache/tsfile/utils/ReadWriteIOUtils.java index c98be63456..f27d355537 100644 --- a/java/tsfile/src/main/java/org/apache/tsfile/utils/ReadWriteIOUtils.java +++ b/java/tsfile/src/main/java/org/apache/tsfile/utils/ReadWriteIOUtils.java @@ -27,6 +27,7 @@ import org.apache.tsfile.read.reader.TsFileInput; import java.io.DataOutputStream; +import java.io.EOFException; import java.io.IOException; import java.io.InputStream; import java.io.OutputStream; @@ -602,9 +603,21 @@ public static int read(ByteBuffer buffer) { /** Read a long var from inputStream. */ public static long readLong(InputStream inputStream) throws IOException { byte[] bytes = new byte[LONG_LEN]; - int readLen = inputStream.read(bytes); - if (readLen != LONG_LEN) { - throw new IOException(Messages.format(RETURN_ERROR_KEY, LONG_LEN, readLen)); + int readLen = 0; + while (readLen < LONG_LEN) { + int count = inputStream.read(bytes, readLen, LONG_LEN - readLen); + if (count < 0) { + throw new EOFException(Messages.format(RETURN_ERROR_KEY, LONG_LEN, readLen)); + } + if (count == 0) { + int value = inputStream.read(); + if (value < 0) { + throw new EOFException(Messages.format(RETURN_ERROR_KEY, LONG_LEN, readLen)); + } + bytes[readLen++] = (byte) value; + } else { + readLen += count; + } } return BytesUtils.bytesToLong(bytes); } diff --git a/java/tsfile/src/main/java/org/apache/tsfile/utils/TsFileSketchTool.java b/java/tsfile/src/main/java/org/apache/tsfile/utils/TsFileSketchTool.java index a20c0e8883..a04c13aa5e 100644 --- a/java/tsfile/src/main/java/org/apache/tsfile/utils/TsFileSketchTool.java +++ b/java/tsfile/src/main/java/org/apache/tsfile/utils/TsFileSketchTool.java @@ -311,8 +311,7 @@ private void printFileInfo() { private void printChunk(List allChunkGroupMetadata) { try { - long nextChunkGroupHeaderPos = - (long) TSFileConfig.MAGIC_STRING.getBytes().length + Byte.BYTES; + long nextChunkGroupHeaderPos = reader.getDataStartOffset(); // ChunkGroup begins for (ChunkGroupMetadata chunkGroupMetadata : allChunkGroupMetadata) { printlnBoth( diff --git a/java/tsfile/src/main/java/org/apache/tsfile/write/TsFileWriter.java b/java/tsfile/src/main/java/org/apache/tsfile/write/TsFileWriter.java index 5521a4f723..b2c7563de4 100644 --- a/java/tsfile/src/main/java/org/apache/tsfile/write/TsFileWriter.java +++ b/java/tsfile/src/main/java/org/apache/tsfile/write/TsFileWriter.java @@ -118,7 +118,7 @@ public TsFileWriter(File file) throws IOException { public TsFileWriter(File file, EncryptParameter firstEncryptParam) throws IOException { this( - new TsFileIOWriter(file), + new TsFileIOWriter(file, firstEncryptParam), new Schema(), TSFileDescriptor.getInstance().getConfig(), firstEncryptParam); @@ -151,7 +151,7 @@ public TsFileWriter(File file, Schema schema) throws IOException { public TsFileWriter(File file, Schema schema, EncryptParameter firstEncryptParam) throws IOException { this( - new TsFileIOWriter(file), + new TsFileIOWriter(file, firstEncryptParam), schema, TSFileDescriptor.getInstance().getConfig(), firstEncryptParam); @@ -170,7 +170,7 @@ public TsFileWriter(TsFileOutput output, Schema schema) throws IOException { public TsFileWriter(TsFileOutput output, Schema schema, EncryptParameter firstEncryptParam) throws IOException { this( - new TsFileIOWriter(output), + new TsFileIOWriter(output, firstEncryptParam), schema, TSFileDescriptor.getInstance().getConfig(), firstEncryptParam); @@ -190,7 +190,7 @@ public TsFileWriter(File file, Schema schema, TSFileConfig conf) throws IOExcept public TsFileWriter( File file, Schema schema, TSFileConfig conf, EncryptParameter firstEncryptParam) throws IOException { - this(new TsFileIOWriter(file), schema, conf, firstEncryptParam); + this(new TsFileIOWriter(file, conf, firstEncryptParam), schema, conf, firstEncryptParam); } /** @@ -231,20 +231,41 @@ protected TsFileWriter( if (this.pageSize >= chunkGroupSizeThreshold) { LOG.warn(Messages.get("log.write.page_size_warn"), pageSize, chunkGroupSizeThreshold); } - this.secondEncryptParam = EncryptUtils.getEncryptParameter(firstEncryptParam); - String encryptLevel; - if (firstEncryptParam != null - && !Objects.equals(firstEncryptParam.getType(), "UNENCRYPTED") - && !Objects.equals(firstEncryptParam.getType(), "org.apache.tsfile.encrypt.UNENCRYPTED")) { - encryptLevel = "2"; - String str = - EncryptUtils.getKeyStr( - IEncryptor.getEncryptor(firstEncryptParam.getType(), firstEncryptParam.getKey()) - .encrypt(secondEncryptParam.getKey())); - fileWriter.setEncryptParam(encryptLevel, secondEncryptParam.getType(), str); + EncryptParameter fileEncryptParameter = fileWriter.getEncryptParameter(); + if (fileWriter instanceof RestorableTsFileIOWriter + && firstEncryptParam != null + && firstEncryptParam.isTdePageAead() + && (fileEncryptParameter == null || !fileEncryptParameter.isTdePageAead())) { + firstEncryptParam = fileEncryptParameter; + } + EncryptParameter effectiveEncryptParameter = + fileEncryptParameter != null && fileEncryptParameter.isTdePageAead() + ? fileEncryptParameter + : firstEncryptParam; + if (effectiveEncryptParameter != null && effectiveEncryptParameter.isTdePageAead()) { + this.secondEncryptParam = + effectiveEncryptParameter == fileEncryptParameter + ? effectiveEncryptParameter + : effectiveEncryptParameter.copy(); + fileWriter.setEncryptParam(this.secondEncryptParam); + fileWriter.writeEncryptionHeaderIfNecessary(); } else { - encryptLevel = "0"; - fileWriter.setEncryptParam(encryptLevel, "org.apache.tsfile.encrypt.UNENCRYPTED", ""); + this.secondEncryptParam = EncryptUtils.getEncryptParameter(firstEncryptParam); + String encryptLevel; + if (firstEncryptParam != null + && !Objects.equals(firstEncryptParam.getType(), "UNENCRYPTED") + && !Objects.equals( + firstEncryptParam.getType(), "org.apache.tsfile.encrypt.UNENCRYPTED")) { + encryptLevel = "2"; + String str = + EncryptUtils.getKeyStr( + IEncryptor.getEncryptor(firstEncryptParam.getType(), firstEncryptParam.getKey()) + .encrypt(secondEncryptParam.getKey())); + fileWriter.setEncryptParam(encryptLevel, secondEncryptParam.getType(), str); + } else { + encryptLevel = "0"; + fileWriter.setEncryptParam(encryptLevel, "org.apache.tsfile.encrypt.UNENCRYPTED", ""); + } } } @@ -729,8 +750,14 @@ private void reset() { @TsFileApi public void close() throws IOException { LOG.info(Messages.get("log.write.close_file")); - flush(); - fileWriter.endFile(); + try { + flush(); + fileWriter.endFile(); + } finally { + if (secondEncryptParam != null && secondEncryptParam.isTdePageAead()) { + secondEncryptParam.close(); + } + } } /** diff --git a/java/tsfile/src/main/java/org/apache/tsfile/write/chunk/ChunkWriterImpl.java b/java/tsfile/src/main/java/org/apache/tsfile/write/chunk/ChunkWriterImpl.java index 7d7912b145..b4683382b0 100644 --- a/java/tsfile/src/main/java/org/apache/tsfile/write/chunk/ChunkWriterImpl.java +++ b/java/tsfile/src/main/java/org/apache/tsfile/write/chunk/ChunkWriterImpl.java @@ -55,6 +55,7 @@ public class ChunkWriterImpl implements IChunkWriter { private final ICompressor compressor; private final EncryptParameter encryptParam; + private long chunkOrdinal = -1; /** all pages of this chunk. */ private final PublicBAOS pageBuffer; @@ -106,6 +107,7 @@ public ChunkWriterImpl(IMeasurementSchema schema) { this.measurementSchema = schema; this.compressor = ICompressor.getCompressor(schema.getCompressor()); this.encryptParam = EncryptUtils.getEncryptParameter(); + this.chunkOrdinal = nextChunkOrdinal(); this.pageBuffer = new PublicBAOS(); this.pageSizeThreshold = TSFileDescriptor.getInstance().getConfig().getPageSizeInByte(); @@ -118,6 +120,7 @@ public ChunkWriterImpl(IMeasurementSchema schema) { this.statistics = Statistics.getStatsByType(measurementSchema.getType()); this.pageWriter = new PageWriter(measurementSchema, encryptParam); + this.pageWriter.setChunkOrdinal(chunkOrdinal); this.pageWriter.setTimeEncoder(measurementSchema.getTimeEncoder()); this.pageWriter.setValueEncoder(measurementSchema.getValueEncoder()); @@ -130,6 +133,7 @@ public ChunkWriterImpl(IMeasurementSchema schema, EncryptParameter encryptParam) this.measurementSchema = schema; this.compressor = ICompressor.getCompressor(schema.getCompressor()); this.encryptParam = encryptParam; + this.chunkOrdinal = nextChunkOrdinal(); this.pageBuffer = new PublicBAOS(); this.pageSizeThreshold = TSFileDescriptor.getInstance().getConfig().getPageSizeInByte(); @@ -142,6 +146,7 @@ public ChunkWriterImpl(IMeasurementSchema schema, EncryptParameter encryptParam) this.statistics = Statistics.getStatsByType(measurementSchema.getType()); this.pageWriter = new PageWriter(measurementSchema, this.encryptParam); + this.pageWriter.setChunkOrdinal(chunkOrdinal); this.pageWriter.setTimeEncoder(measurementSchema.getTimeEncoder()); this.pageWriter.setValueEncoder(measurementSchema.getValueEncoder()); @@ -161,6 +166,12 @@ public ChunkWriterImpl( this.isMerging = isMerging; } + private long nextChunkOrdinal() { + return encryptParam != null && encryptParam.isTdePageAead() + ? encryptParam.nextChunkOrdinal() + : -1; + } + private void checkSdtEncoding() { if (measurementSchema.getProps() != null && !isMerging) { if (measurementSchema.getProps().getOrDefault(LOSS, "").equals(SDT)) { @@ -322,17 +333,18 @@ private void writePageToPageBuffer() { try { if (numOfPages == 0) { // record the firstPageStatistics this.firstPageStatistics = pageWriter.getStatistics(); - this.sizeWithoutStatistic = pageWriter.writePageHeaderAndDataIntoBuff(pageBuffer, true); + this.sizeWithoutStatistic = + pageWriter.writePageHeaderAndDataIntoBuff(pageBuffer, true, numOfPages); } else if (numOfPages == 1) { // put the firstPageStatistics into pageBuffer byte[] b = pageBuffer.toByteArray(); pageBuffer.reset(); pageBuffer.write(b, 0, this.sizeWithoutStatistic); firstPageStatistics.serialize(pageBuffer); pageBuffer.write(b, this.sizeWithoutStatistic, b.length - this.sizeWithoutStatistic); - pageWriter.writePageHeaderAndDataIntoBuff(pageBuffer, false); + pageWriter.writePageHeaderAndDataIntoBuff(pageBuffer, false, numOfPages); firstPageStatistics = null; } else { - pageWriter.writePageHeaderAndDataIntoBuff(pageBuffer, false); + pageWriter.writePageHeaderAndDataIntoBuff(pageBuffer, false, numOfPages); } // update statistics of this chunk @@ -364,6 +376,10 @@ public void writeToFileWriter( sizeWithoutStatistic = 0; firstPageStatistics = null; this.statistics = Statistics.getStatsByType(measurementSchema.getType()); + this.chunkOrdinal = nextChunkOrdinal(); + if (pageWriter != null) { + pageWriter.setChunkOrdinal(chunkOrdinal); + } } @Override @@ -380,7 +396,10 @@ public long getSerializedChunkSize() { return 0; } // return the serialized size of the chunk header + all pages - return ChunkHeader.getSerializedSize(measurementSchema.getMeasurementName(), pageBuffer.size()) + return ChunkHeader.getSerializedSize( + measurementSchema.getMeasurementName(), + pageBuffer.size(), + chunkOrdinal >= 0 ? encryptParam.getChunkAuthenticationSize() : 0) + (long) pageBuffer.size(); } @@ -432,6 +451,10 @@ public TSDataType getDataType() { */ public void writePageHeaderAndDataIntoBuff(ByteBuffer data, PageHeader header) throws PageException { + if (chunkOrdinal >= 0) { + throw new org.apache.tsfile.exception.encrypt.EncryptException( + Messages.get("error.write.aead_raw_page_copy_unsupported")); + } // write the page header to pageBuffer try { logger.debug(Messages.get("log.write.page_header_flush_start"), pageBuffer.size()); @@ -505,7 +528,9 @@ private void writeAllPagesOfChunkToTsFile( statistics, pageBuffer.size(), numOfPages, - 0); + 0, + chunkOrdinal, + encryptParam); long dataOffset = writer.getPos(); @@ -542,6 +567,10 @@ public int getNumOfPages() { return numOfPages; } + public long getChunkOrdinal() { + return chunkOrdinal; + } + public ByteBuffer getByteBuffer() { return ByteBuffer.wrap(pageBuffer.toByteArray()); } diff --git a/java/tsfile/src/main/java/org/apache/tsfile/write/chunk/TimeChunkWriter.java b/java/tsfile/src/main/java/org/apache/tsfile/write/chunk/TimeChunkWriter.java index 392a0cf458..8f4a01d00a 100644 --- a/java/tsfile/src/main/java/org/apache/tsfile/write/chunk/TimeChunkWriter.java +++ b/java/tsfile/src/main/java/org/apache/tsfile/write/chunk/TimeChunkWriter.java @@ -57,6 +57,7 @@ public class TimeChunkWriter { private CompressionType compressionType; private EncryptParameter encryptParam; + private long chunkOrdinal = -1; /** all pages of this chunk. */ private PublicBAOS pageBuffer; @@ -96,6 +97,7 @@ public TimeChunkWriter( this.encodingType = encodingType; this.compressionType = compressionType; this.encryptParam = EncryptUtils.getEncryptParameter(); + this.chunkOrdinal = nextChunkOrdinal(); this.pageBuffer = new PublicBAOS(); this.pageSizeThreshold = TSFileDescriptor.getInstance().getConfig().getPageSizeInByte(); @@ -110,6 +112,7 @@ public TimeChunkWriter( this.pageWriter = new TimePageWriter( timeEncoder, ICompressor.getCompressor(compressionType), this.encryptParam); + this.pageWriter.setChunkOrdinal(chunkOrdinal); } public TimeChunkWriter( @@ -122,6 +125,7 @@ public TimeChunkWriter( this.encodingType = encodingType; this.compressionType = compressionType; this.encryptParam = encryptParam; + this.chunkOrdinal = nextChunkOrdinal(); this.pageBuffer = new PublicBAOS(); this.pageSizeThreshold = TSFileDescriptor.getInstance().getConfig().getPageSizeInByte(); @@ -135,12 +139,19 @@ public TimeChunkWriter( this.pageWriter = new TimePageWriter(timeEncoder, ICompressor.getCompressor(compressionType), encryptParam); + this.pageWriter.setChunkOrdinal(chunkOrdinal); } public void write(long time) { pageWriter.write(time); } + private long nextChunkOrdinal() { + return encryptParam != null && encryptParam.isTdePageAead() + ? encryptParam.nextChunkOrdinal() + : -1; + } + public void write(long[] timestamps, int batchSize, int arrayOffset) { pageWriter.write(timestamps, batchSize, arrayOffset); } @@ -184,17 +195,18 @@ public void writePageToPageBuffer() { try { if (numOfPages == 0) { // record the firstPageStatistics this.firstPageStatistics = pageWriter.getStatistics(); - this.sizeWithoutStatistic = pageWriter.writePageHeaderAndDataIntoBuff(pageBuffer, true); + this.sizeWithoutStatistic = + pageWriter.writePageHeaderAndDataIntoBuff(pageBuffer, true, numOfPages); } else if (numOfPages == 1) { // put the firstPageStatistics into pageBuffer byte[] b = pageBuffer.toByteArray(); pageBuffer.reset(); pageBuffer.write(b, 0, this.sizeWithoutStatistic); firstPageStatistics.serialize(pageBuffer); pageBuffer.write(b, this.sizeWithoutStatistic, b.length - this.sizeWithoutStatistic); - pageWriter.writePageHeaderAndDataIntoBuff(pageBuffer, false); + pageWriter.writePageHeaderAndDataIntoBuff(pageBuffer, false, numOfPages); firstPageStatistics = null; } else { - pageWriter.writePageHeaderAndDataIntoBuff(pageBuffer, false); + pageWriter.writePageHeaderAndDataIntoBuff(pageBuffer, false, numOfPages); } // update statistics of this chunk @@ -210,6 +222,10 @@ public void writePageToPageBuffer() { public void writePageHeaderAndDataIntoBuff(ByteBuffer data, PageHeader header) throws PageException { + if (chunkOrdinal >= 0) { + throw new org.apache.tsfile.exception.encrypt.EncryptException( + Messages.get("error.write.aead_raw_page_copy_unsupported")); + } // write the page header to pageBuffer try { logger.debug(Messages.get("log.write.page_header_flush_start"), pageBuffer.size()); @@ -262,6 +278,10 @@ public void writeToFileWriter(TsFileIOWriter tsfileWriter) throws IOException { sizeWithoutStatistic = 0; firstPageStatistics = null; this.statistics = new TimeStatistics(); + this.chunkOrdinal = nextChunkOrdinal(); + if (pageWriter != null) { + pageWriter.setChunkOrdinal(chunkOrdinal); + } } public long estimateMaxSeriesMemSize() { @@ -276,7 +296,10 @@ public long getCurrentChunkSize() { return 0; } // return the serialized size of the chunk header + all pages - return ChunkHeader.getSerializedSize(measurementId, pageBuffer.size()) + return ChunkHeader.getSerializedSize( + measurementId, + pageBuffer.size(), + chunkOrdinal >= 0 ? encryptParam.getChunkAuthenticationSize() : 0) + (long) pageBuffer.size(); } @@ -322,7 +345,9 @@ public void writeAllPagesOfChunkToTsFile(TsFileIOWriter writer) throws IOExcepti statistics, pageBuffer.size(), numOfPages, - TsFileConstant.TIME_COLUMN_MASK); + TsFileConstant.TIME_COLUMN_MASK, + chunkOrdinal, + encryptParam); long dataOffset = writer.getPos(); diff --git a/java/tsfile/src/main/java/org/apache/tsfile/write/chunk/ValueChunkWriter.java b/java/tsfile/src/main/java/org/apache/tsfile/write/chunk/ValueChunkWriter.java index 382df6e808..12eebc59f6 100644 --- a/java/tsfile/src/main/java/org/apache/tsfile/write/chunk/ValueChunkWriter.java +++ b/java/tsfile/src/main/java/org/apache/tsfile/write/chunk/ValueChunkWriter.java @@ -61,6 +61,7 @@ public class ValueChunkWriter { private final CompressionType compressionType; private final EncryptParameter encryptParam; + private long chunkOrdinal = -1; /** all pages of this chunk. */ private final PublicBAOS pageBuffer; @@ -100,6 +101,7 @@ public ValueChunkWriter( this.dataType = dataType; this.compressionType = compressionType; this.encryptParam = EncryptUtils.getEncryptParameter(); + this.chunkOrdinal = nextChunkOrdinal(); this.pageBuffer = new PublicBAOS(); this.pageSizeThreshold = TSFileDescriptor.getInstance().getConfig().getPageSizeInByte(); this.maxNumberOfPointsInPage = @@ -112,6 +114,7 @@ public ValueChunkWriter( this.pageWriter = new ValuePageWriter( valueEncoder, ICompressor.getCompressor(compressionType), dataType, this.encryptParam); + this.pageWriter.setChunkOrdinal(chunkOrdinal); } public ValueChunkWriter( @@ -126,6 +129,7 @@ public ValueChunkWriter( this.dataType = dataType; this.compressionType = compressionType; this.encryptParam = encryptParam; + this.chunkOrdinal = nextChunkOrdinal(); this.pageBuffer = new PublicBAOS(); this.pageSizeThreshold = TSFileDescriptor.getInstance().getConfig().getPageSizeInByte(); this.maxNumberOfPointsInPage = @@ -138,6 +142,13 @@ public ValueChunkWriter( this.pageWriter = new ValuePageWriter( valueEncoder, ICompressor.getCompressor(compressionType), dataType, this.encryptParam); + this.pageWriter.setChunkOrdinal(chunkOrdinal); + } + + private long nextChunkOrdinal() { + return encryptParam != null && encryptParam.isTdePageAead() + ? encryptParam.nextChunkOrdinal() + : -1; } public void write(long time, long value, boolean isNull) { @@ -214,7 +225,8 @@ public void writePageToPageBuffer() { // record the firstPageStatistics if it is not empty page this.firstPageStatistics = pageWriter.getStatistics(); } - this.sizeWithoutStatistic = pageWriter.writePageHeaderAndDataIntoBuff(pageBuffer, true); + this.sizeWithoutStatistic = + pageWriter.writePageHeaderAndDataIntoBuff(pageBuffer, true, numOfPages); } else if (numOfPages == 1) { // put the firstPageStatistics into pageBuffer if (firstPageStatistics != null) { // Consider previous page is an empty page byte[] b = pageBuffer.toByteArray(); @@ -223,10 +235,10 @@ public void writePageToPageBuffer() { firstPageStatistics.serialize(pageBuffer); pageBuffer.write(b, this.sizeWithoutStatistic, b.length - this.sizeWithoutStatistic); } - pageWriter.writePageHeaderAndDataIntoBuff(pageBuffer, false); + pageWriter.writePageHeaderAndDataIntoBuff(pageBuffer, false, numOfPages); firstPageStatistics = null; } else { - pageWriter.writePageHeaderAndDataIntoBuff(pageBuffer, false); + pageWriter.writePageHeaderAndDataIntoBuff(pageBuffer, false, numOfPages); } // update statistics of this chunk @@ -242,6 +254,10 @@ public void writePageToPageBuffer() { public void writePageHeaderAndDataIntoBuff(ByteBuffer data, PageHeader header) throws PageException { + if (chunkOrdinal >= 0) { + throw new org.apache.tsfile.exception.encrypt.EncryptException( + Messages.get("error.write.aead_raw_page_copy_unsupported")); + } // write the page header to pageBuffer try { logger.debug(Messages.get("log.write.page_header_flush_start"), pageBuffer.size()); @@ -314,6 +330,10 @@ public void writeToFileWriter( sizeWithoutStatistic = 0; firstPageStatistics = null; this.statistics = Statistics.getStatsByType(dataType); + this.chunkOrdinal = nextChunkOrdinal(); + if (pageWriter != null) { + pageWriter.setChunkOrdinal(chunkOrdinal); + } } public long estimateMaxSeriesMemSize() { @@ -335,11 +355,15 @@ public long getCurrentChunkSize() { // Empty chunk, it may happen if pageBuffer stores empty bits and only chunk header will be // flushed. if (statistics.getCount() == 0) { - return ChunkHeader.getSerializedSize(measurementId, 0); + return ChunkHeader.getSerializedSize( + measurementId, 0, chunkOrdinal >= 0 ? encryptParam.getChunkAuthenticationSize() : 0); } // return the serialized size of the chunk header + all pages - return ChunkHeader.getSerializedSize(measurementId, pageBuffer.size()) + return ChunkHeader.getSerializedSize( + measurementId, + pageBuffer.size(), + chunkOrdinal >= 0 ? encryptParam.getChunkAuthenticationSize() : 0) + (long) pageBuffer.size(); } @@ -385,6 +409,10 @@ public int getNumOfPages() { return numOfPages; } + public long getChunkOrdinal() { + return chunkOrdinal; + } + public TSDataType getDataType() { return dataType; } @@ -421,7 +449,9 @@ public void writeAllPagesOfChunkToTsFile( statistics, 0, 0, - TsFileConstant.VALUE_COLUMN_MASK); + TsFileConstant.VALUE_COLUMN_MASK, + chunkOrdinal, + encryptParam); writer.endCurrentChunk(); return; } @@ -435,7 +465,9 @@ public void writeAllPagesOfChunkToTsFile( statistics, pageBuffer.size(), numOfPages, - TsFileConstant.VALUE_COLUMN_MASK); + TsFileConstant.VALUE_COLUMN_MASK, + chunkOrdinal, + encryptParam); long dataOffset = writer.getPos(); diff --git a/java/tsfile/src/main/java/org/apache/tsfile/write/page/PageBodyEncoder.java b/java/tsfile/src/main/java/org/apache/tsfile/write/page/PageBodyEncoder.java new file mode 100644 index 0000000000..e86946b5e1 --- /dev/null +++ b/java/tsfile/src/main/java/org/apache/tsfile/write/page/PageBodyEncoder.java @@ -0,0 +1,86 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ +package org.apache.tsfile.write.page; + +import org.apache.tsfile.compress.ICompressor; +import org.apache.tsfile.encrypt.EncryptParameter; +import org.apache.tsfile.encrypt.IEncryptor; +import org.apache.tsfile.encrypt.PageCryptoContext; +import org.apache.tsfile.exception.encrypt.EncryptException; +import org.apache.tsfile.file.metadata.enums.CompressionType; +import org.apache.tsfile.file.metadata.enums.EncryptionType; +import org.apache.tsfile.i18n.Messages; + +import java.io.IOException; +import java.nio.ByteBuffer; + +class PageBodyEncoder { + + private PageBodyEncoder() {} + + /** Returns an array-backed buffer whose remaining bytes are the encoded page body. */ + static ByteBuffer encode( + ByteBuffer pageData, + int uncompressedSize, + ICompressor compressor, + EncryptParameter encryptParameter, + int pageIndex, + long chunkOrdinal) + throws IOException { + byte[] plaintext; + int plaintextOffset = 0; + int compressedSize; + int inputOffset = pageData.arrayOffset() + pageData.position(); + + if (compressor.getType().equals(CompressionType.UNCOMPRESSED)) { + plaintext = pageData.array(); + plaintextOffset = inputOffset; + compressedSize = uncompressedSize; + } else if (compressor.getType().equals(CompressionType.GZIP)) { + plaintext = compressor.compress(pageData.array(), inputOffset, uncompressedSize); + compressedSize = plaintext.length; + } else { + plaintext = new byte[compressor.getMaxBytesForCompression(uncompressedSize)]; + compressedSize = + compressor.compress(pageData.array(), inputOffset, uncompressedSize, plaintext); + } + + IEncryptor encryptor = IEncryptor.getEncryptor(encryptParameter); + byte[] encryptedPageBody; + if (encryptParameter != null && encryptParameter.isTdePageAead()) { + PageCryptoContext pageCryptoContext = + PageCryptoContext.forEncryption( + encryptParameter, uncompressedSize, compressedSize, pageIndex, chunkOrdinal); + encryptedPageBody = + encryptor.encryptPage(plaintext, plaintextOffset, compressedSize, pageCryptoContext); + if (encryptedPageBody.length != pageCryptoContext.getEncryptedPageBodySize()) { + throw new EncryptException( + Messages.format( + "error.encrypt.page_output_size_mismatch", + pageCryptoContext.getEncryptedPageBodySize(), + encryptedPageBody.length)); + } + } else if (encryptor.getEncryptionType() == EncryptionType.UNENCRYPTED) { + return ByteBuffer.wrap(plaintext, plaintextOffset, compressedSize); + } else { + encryptedPageBody = encryptor.encrypt(plaintext, plaintextOffset, compressedSize); + } + return ByteBuffer.wrap(encryptedPageBody); + } +} diff --git a/java/tsfile/src/main/java/org/apache/tsfile/write/page/PageWriter.java b/java/tsfile/src/main/java/org/apache/tsfile/write/page/PageWriter.java index d479ad8d43..ce840dfff0 100644 --- a/java/tsfile/src/main/java/org/apache/tsfile/write/page/PageWriter.java +++ b/java/tsfile/src/main/java/org/apache/tsfile/write/page/PageWriter.java @@ -22,10 +22,7 @@ import org.apache.tsfile.encoding.encoder.Encoder; import org.apache.tsfile.encrypt.EncryptParameter; import org.apache.tsfile.encrypt.EncryptUtils; -import org.apache.tsfile.encrypt.IEncryptor; import org.apache.tsfile.enums.TSDataType; -import org.apache.tsfile.file.metadata.enums.CompressionType; -import org.apache.tsfile.file.metadata.enums.EncryptionType; import org.apache.tsfile.file.metadata.statistics.Statistics; import org.apache.tsfile.utils.Binary; import org.apache.tsfile.utils.PublicBAOS; @@ -38,8 +35,6 @@ import java.io.IOException; import java.io.Serializable; import java.nio.ByteBuffer; -import java.nio.channels.Channels; -import java.nio.channels.WritableByteChannel; /** * This writer is used to write time-value into a page. It consists of a time encoder, a value @@ -52,6 +47,7 @@ public class PageWriter { private ICompressor compressor; private EncryptParameter encryptParam; + private long chunkOrdinal = -1; // time private Encoder timeEncoder; @@ -231,28 +227,22 @@ public ByteBuffer getUncompressedBytes() throws IOException { /** write the page header and data into the PageWriter's output stream. */ public int writePageHeaderAndDataIntoBuff(PublicBAOS pageBuffer, boolean first) throws IOException { + return writePageHeaderAndDataIntoBuff(pageBuffer, first, first ? 0 : -1); + } + + /** write the page header and data into the PageWriter's output stream. */ + public int writePageHeaderAndDataIntoBuff(PublicBAOS pageBuffer, boolean first, int pageIndex) + throws IOException { if (statistics.getCount() == 0) { return 0; } ByteBuffer pageData = getUncompressedBytes(); int uncompressedSize = pageData.remaining(); - int compressedSize; - byte[] compressedBytes = null; - - if (compressor.getType().equals(CompressionType.UNCOMPRESSED)) { - compressedSize = uncompressedSize; - } else if (compressor.getType().equals(CompressionType.GZIP)) { - compressedBytes = - compressor.compress(pageData.array(), pageData.position(), uncompressedSize); - compressedSize = compressedBytes.length; - } else { - compressedBytes = new byte[compressor.getMaxBytesForCompression(uncompressedSize)]; - // data is never a directByteBuffer now, so we can use data.array() - compressedSize = - compressor.compress( - pageData.array(), pageData.position(), uncompressedSize, compressedBytes); - } + ByteBuffer pageBody = + PageBodyEncoder.encode( + pageData, uncompressedSize, compressor, encryptParam, pageIndex, chunkOrdinal); + int pageBodySize = pageBody.remaining(); // write the page header to IOWriter int sizeWithoutStatistic = 0; @@ -260,41 +250,16 @@ public int writePageHeaderAndDataIntoBuff(PublicBAOS pageBuffer, boolean first) sizeWithoutStatistic += ReadWriteForEncodingUtils.writeUnsignedVarInt(uncompressedSize, pageBuffer); sizeWithoutStatistic += - ReadWriteForEncodingUtils.writeUnsignedVarInt(compressedSize, pageBuffer); + ReadWriteForEncodingUtils.writeUnsignedVarInt(pageBodySize, pageBuffer); } else { ReadWriteForEncodingUtils.writeUnsignedVarInt(uncompressedSize, pageBuffer); - ReadWriteForEncodingUtils.writeUnsignedVarInt(compressedSize, pageBuffer); + ReadWriteForEncodingUtils.writeUnsignedVarInt(pageBodySize, pageBuffer); statistics.serialize(pageBuffer); } - IEncryptor encryptor = IEncryptor.getEncryptor(encryptParam); - // write page content to temp PBAOS logger.trace("start to flush a page data into buffer, buffer position {} ", pageBuffer.size()); - if (compressor.getType().equals(CompressionType.UNCOMPRESSED)) { - if (encryptor.getEncryptionType().equals(EncryptionType.UNENCRYPTED)) { - try (WritableByteChannel channel = Channels.newChannel(pageBuffer)) { - channel.write(pageData); - } - } else { - byte[] encryptedBytes = null; - encryptedBytes = encryptor.encrypt(pageData.array(), pageData.position(), uncompressedSize); - // data is never a directByteBuffer now, so we can use data.array() - int encryptedSize = encryptedBytes.length; - pageBuffer.write(encryptedBytes, 0, encryptedSize); - } - - } else { - if (encryptor.getEncryptionType().equals(EncryptionType.UNENCRYPTED)) { - pageBuffer.write(compressedBytes, 0, compressedSize); - } else { - byte[] encryptedBytes = null; - encryptedBytes = encryptor.encrypt(compressedBytes, 0, compressedSize); - // data is never a directByteBuffer now, so we can use data.array() - int encryptedSize = encryptedBytes.length; - pageBuffer.write(encryptedBytes, 0, encryptedSize); - } - } + pageBuffer.write(pageBody.array(), pageBody.arrayOffset() + pageBody.position(), pageBodySize); logger.trace("start to flush a page data into buffer, buffer position {} ", pageBuffer.size()); return sizeWithoutStatistic; } @@ -319,6 +284,10 @@ public void reset(IMeasurementSchema measurementSchema) { statistics = Statistics.getStatsByType(measurementSchema.getType()); } + public void setChunkOrdinal(long chunkOrdinal) { + this.chunkOrdinal = chunkOrdinal; + } + public void setTimeEncoder(Encoder encoder) { this.timeEncoder = encoder; } diff --git a/java/tsfile/src/main/java/org/apache/tsfile/write/page/TimePageWriter.java b/java/tsfile/src/main/java/org/apache/tsfile/write/page/TimePageWriter.java index e09e66cdf4..6463ef51be 100644 --- a/java/tsfile/src/main/java/org/apache/tsfile/write/page/TimePageWriter.java +++ b/java/tsfile/src/main/java/org/apache/tsfile/write/page/TimePageWriter.java @@ -22,9 +22,6 @@ import org.apache.tsfile.encoding.encoder.Encoder; import org.apache.tsfile.encrypt.EncryptParameter; import org.apache.tsfile.encrypt.EncryptUtils; -import org.apache.tsfile.encrypt.IEncryptor; -import org.apache.tsfile.file.metadata.enums.CompressionType; -import org.apache.tsfile.file.metadata.enums.EncryptionType; import org.apache.tsfile.file.metadata.statistics.TimeStatistics; import org.apache.tsfile.utils.PublicBAOS; import org.apache.tsfile.utils.ReadWriteForEncodingUtils; @@ -34,8 +31,6 @@ import java.io.IOException; import java.nio.ByteBuffer; -import java.nio.channels.Channels; -import java.nio.channels.WritableByteChannel; /** * This writer is used to write time into a page. It consists of a time encoder and respective @@ -48,6 +43,7 @@ public class TimePageWriter { private final ICompressor compressor; private final EncryptParameter encryptParam; + private long chunkOrdinal = -1; // time private Encoder timeEncoder; @@ -114,28 +110,22 @@ public ByteBuffer getUncompressedBytes() throws IOException { /** write the page header and data into the PageWriter's output stream. */ public int writePageHeaderAndDataIntoBuff(PublicBAOS pageBuffer, boolean first) throws IOException { + return writePageHeaderAndDataIntoBuff(pageBuffer, first, first ? 0 : -1); + } + + /** write the page header and data into the PageWriter's output stream. */ + public int writePageHeaderAndDataIntoBuff(PublicBAOS pageBuffer, boolean first, int pageIndex) + throws IOException { if (statistics.getCount() == 0) { return 0; } ByteBuffer pageData = getUncompressedBytes(); int uncompressedSize = pageData.remaining(); - int compressedSize; - byte[] compressedBytes = null; - - if (compressor.getType().equals(CompressionType.UNCOMPRESSED)) { - compressedSize = uncompressedSize; - } else if (compressor.getType().equals(CompressionType.GZIP)) { - compressedBytes = - compressor.compress(pageData.array(), pageData.position(), uncompressedSize); - compressedSize = compressedBytes.length; - } else { - compressedBytes = new byte[compressor.getMaxBytesForCompression(uncompressedSize)]; - // data is never a directByteBuffer now, so we can use data.array() - compressedSize = - compressor.compress( - pageData.array(), pageData.position(), uncompressedSize, compressedBytes); - } + ByteBuffer pageBody = + PageBodyEncoder.encode( + pageData, uncompressedSize, compressor, encryptParam, pageIndex, chunkOrdinal); + int pageBodySize = pageBody.remaining(); // write the page header to IOWriter int sizeWithoutStatistic = 0; @@ -143,41 +133,17 @@ public int writePageHeaderAndDataIntoBuff(PublicBAOS pageBuffer, boolean first) sizeWithoutStatistic += ReadWriteForEncodingUtils.writeUnsignedVarInt(uncompressedSize, pageBuffer); sizeWithoutStatistic += - ReadWriteForEncodingUtils.writeUnsignedVarInt(compressedSize, pageBuffer); + ReadWriteForEncodingUtils.writeUnsignedVarInt(pageBodySize, pageBuffer); } else { ReadWriteForEncodingUtils.writeUnsignedVarInt(uncompressedSize, pageBuffer); - ReadWriteForEncodingUtils.writeUnsignedVarInt(compressedSize, pageBuffer); + ReadWriteForEncodingUtils.writeUnsignedVarInt(pageBodySize, pageBuffer); statistics.serialize(pageBuffer); } // write page content to temp PBAOS logger.trace( "start to flush a time page data into buffer, buffer position {} ", pageBuffer.size()); - IEncryptor encryptor = IEncryptor.getEncryptor(encryptParam); - if (compressor.getType().equals(CompressionType.UNCOMPRESSED)) { - if (encryptor.getEncryptionType().equals(EncryptionType.UNENCRYPTED)) { - try (WritableByteChannel channel = Channels.newChannel(pageBuffer)) { - channel.write(pageData); - } - } else { - byte[] encryptedBytes = null; - encryptedBytes = encryptor.encrypt(pageData.array(), pageData.position(), uncompressedSize); - // data is never a directByteBuffer now, so we can use data.array() - int encryptedSize = encryptedBytes.length; - pageBuffer.write(encryptedBytes, 0, encryptedSize); - } - - } else { - if (encryptor.getEncryptionType().equals(EncryptionType.UNENCRYPTED)) { - pageBuffer.write(compressedBytes, 0, compressedSize); - } else { - byte[] encryptedBytes = null; - encryptedBytes = encryptor.encrypt(compressedBytes, 0, compressedSize); - // data is never a directByteBuffer now, so we can use data.array() - int encryptedSize = encryptedBytes.length; - pageBuffer.write(encryptedBytes, 0, encryptedSize); - } - } + pageBuffer.write(pageBody.array(), pageBody.arrayOffset() + pageBody.position(), pageBodySize); logger.trace( "finish flushing a time page data into buffer, buffer position {} ", pageBuffer.size()); return sizeWithoutStatistic; @@ -199,6 +165,10 @@ public void reset() { statistics = new TimeStatistics(); } + public void setChunkOrdinal(long chunkOrdinal) { + this.chunkOrdinal = chunkOrdinal; + } + public void setTimeEncoder(Encoder encoder) { this.timeEncoder = encoder; } diff --git a/java/tsfile/src/main/java/org/apache/tsfile/write/page/ValuePageWriter.java b/java/tsfile/src/main/java/org/apache/tsfile/write/page/ValuePageWriter.java index f7f375759a..16c101ee67 100644 --- a/java/tsfile/src/main/java/org/apache/tsfile/write/page/ValuePageWriter.java +++ b/java/tsfile/src/main/java/org/apache/tsfile/write/page/ValuePageWriter.java @@ -22,10 +22,7 @@ import org.apache.tsfile.encoding.encoder.Encoder; import org.apache.tsfile.encrypt.EncryptParameter; import org.apache.tsfile.encrypt.EncryptUtils; -import org.apache.tsfile.encrypt.IEncryptor; import org.apache.tsfile.enums.TSDataType; -import org.apache.tsfile.file.metadata.enums.CompressionType; -import org.apache.tsfile.file.metadata.enums.EncryptionType; import org.apache.tsfile.file.metadata.statistics.Statistics; import org.apache.tsfile.utils.Binary; import org.apache.tsfile.utils.BitMap; @@ -38,8 +35,6 @@ import java.io.IOException; import java.io.Serializable; import java.nio.ByteBuffer; -import java.nio.channels.Channels; -import java.nio.channels.WritableByteChannel; /** * This writer is used to write value into a page. It consists of a value encoder and respective @@ -51,6 +46,7 @@ public class ValuePageWriter { private final ICompressor compressor; private final EncryptParameter encryptParam; + private long chunkOrdinal = -1; // value private Encoder valueEncoder; @@ -405,6 +401,12 @@ public int writeEmptyPageIntoBuff(PublicBAOS pageBuffer) { /** write the page header and data into the PageWriter's output stream. */ public int writePageHeaderAndDataIntoBuff(PublicBAOS pageBuffer, boolean first) throws IOException { + return writePageHeaderAndDataIntoBuff(pageBuffer, first, first ? 0 : -1); + } + + /** write the page header and data into the PageWriter's output stream. */ + public int writePageHeaderAndDataIntoBuff(PublicBAOS pageBuffer, boolean first, int pageIndex) + throws IOException { if (size == 0) { return 0; } else if (statistics.getCount() == 0) { @@ -414,22 +416,10 @@ public int writePageHeaderAndDataIntoBuff(PublicBAOS pageBuffer, boolean first) ByteBuffer pageData = getUncompressedBytes(); int uncompressedSize = pageData.remaining(); - int compressedSize; - byte[] compressedBytes = null; - - if (compressor.getType().equals(CompressionType.UNCOMPRESSED)) { - compressedSize = uncompressedSize; - } else if (compressor.getType().equals(CompressionType.GZIP)) { - compressedBytes = - compressor.compress(pageData.array(), pageData.position(), uncompressedSize); - compressedSize = compressedBytes.length; - } else { - compressedBytes = new byte[compressor.getMaxBytesForCompression(uncompressedSize)]; - // data is never a directByteBuffer now, so we can use data.array() - compressedSize = - compressor.compress( - pageData.array(), pageData.position(), uncompressedSize, compressedBytes); - } + ByteBuffer pageBody = + PageBodyEncoder.encode( + pageData, uncompressedSize, compressor, encryptParam, pageIndex, chunkOrdinal); + int pageBodySize = pageBody.remaining(); // write the page header to IOWriter int sizeWithoutStatistic = 0; @@ -437,39 +427,15 @@ public int writePageHeaderAndDataIntoBuff(PublicBAOS pageBuffer, boolean first) sizeWithoutStatistic += ReadWriteForEncodingUtils.writeUnsignedVarInt(uncompressedSize, pageBuffer); sizeWithoutStatistic += - ReadWriteForEncodingUtils.writeUnsignedVarInt(compressedSize, pageBuffer); + ReadWriteForEncodingUtils.writeUnsignedVarInt(pageBodySize, pageBuffer); } else { ReadWriteForEncodingUtils.writeUnsignedVarInt(uncompressedSize, pageBuffer); - ReadWriteForEncodingUtils.writeUnsignedVarInt(compressedSize, pageBuffer); + ReadWriteForEncodingUtils.writeUnsignedVarInt(pageBodySize, pageBuffer); statistics.serialize(pageBuffer); } - IEncryptor encryptor = IEncryptor.getEncryptor(encryptParam); // write page content to temp PBAOS logger.trace("start to flush a page data into buffer, buffer position {} ", pageBuffer.size()); - if (compressor.getType().equals(CompressionType.UNCOMPRESSED)) { - if (encryptor.getEncryptionType().equals(EncryptionType.UNENCRYPTED)) { - try (WritableByteChannel channel = Channels.newChannel(pageBuffer)) { - channel.write(pageData); - } - } else { - byte[] encryptedBytes = null; - encryptedBytes = encryptor.encrypt(pageData.array(), pageData.position(), uncompressedSize); - // data is never a directByteBuffer now, so we can use data.array() - int encryptedSize = encryptedBytes.length; - pageBuffer.write(encryptedBytes, 0, encryptedSize); - } - - } else { - if (encryptor.getEncryptionType().equals(EncryptionType.UNENCRYPTED)) { - pageBuffer.write(compressedBytes, 0, compressedSize); - } else { - byte[] encryptedBytes = null; - encryptedBytes = encryptor.encrypt(compressedBytes, 0, compressedSize); - // data is never a directByteBuffer now, so we can use data.array() - int encryptedSize = encryptedBytes.length; - pageBuffer.write(encryptedBytes, 0, encryptedSize); - } - } + pageBuffer.write(pageBody.array(), pageBody.arrayOffset() + pageBody.position(), pageBodySize); logger.trace("start to flush a page data into buffer, buffer position {} ", pageBuffer.size()); return sizeWithoutStatistic; } @@ -493,6 +459,10 @@ public void reset(TSDataType dataType) { statistics = Statistics.getStatsByType(dataType); } + public void setChunkOrdinal(long chunkOrdinal) { + this.chunkOrdinal = chunkOrdinal; + } + public void setValueEncoder(Encoder encoder) { this.valueEncoder = encoder; } diff --git a/java/tsfile/src/main/java/org/apache/tsfile/write/v4/AbstractTableModelTsFileWriter.java b/java/tsfile/src/main/java/org/apache/tsfile/write/v4/AbstractTableModelTsFileWriter.java index cf7e3b6513..886d6e07e5 100644 --- a/java/tsfile/src/main/java/org/apache/tsfile/write/v4/AbstractTableModelTsFileWriter.java +++ b/java/tsfile/src/main/java/org/apache/tsfile/write/v4/AbstractTableModelTsFileWriter.java @@ -94,7 +94,7 @@ protected AbstractTableModelTsFileWriter(File file, long chunkGroupSizeThreshold protected AbstractTableModelTsFileWriter( File file, long chunkGroupSizeThreshold, EncryptParameter firstEncryptParam) throws IOException { - this(new TsFileIOWriter(file), chunkGroupSizeThreshold, firstEncryptParam); + this(new TsFileIOWriter(file, firstEncryptParam), chunkGroupSizeThreshold, firstEncryptParam); } @TsFileApi @@ -123,20 +123,35 @@ protected AbstractTableModelTsFileWriter( LOG.warn(Messages.get("log.write.page_size_warn"), pageSize, chunkGroupSizeThreshold); } - this.secondEncryptParam = EncryptUtils.getEncryptParameter(firstEncryptParam); - String encryptLevel; - if (firstEncryptParam != null - && !Objects.equals(firstEncryptParam.getType(), "UNENCRYPTED") - && !Objects.equals(firstEncryptParam.getType(), "org.apache.tsfile.encrypt.UNENCRYPTED")) { - encryptLevel = "2"; - String str = - EncryptUtils.getKeyStr( - IEncryptor.getEncryptor(firstEncryptParam.getType(), firstEncryptParam.getKey()) - .encrypt(secondEncryptParam.getKey())); - fileWriter.setEncryptParam(encryptLevel, secondEncryptParam.getType(), str); + EncryptParameter fileEncryptParameter = fileWriter.getEncryptParameter(); + EncryptParameter effectiveEncryptParameter = + fileEncryptParameter != null && fileEncryptParameter.isTdePageAead() + ? fileEncryptParameter + : firstEncryptParam; + if (effectiveEncryptParameter != null && effectiveEncryptParameter.isTdePageAead()) { + this.secondEncryptParam = + effectiveEncryptParameter == fileEncryptParameter + ? effectiveEncryptParameter + : effectiveEncryptParameter.copy(); + fileWriter.setEncryptParam(this.secondEncryptParam); + fileWriter.writeEncryptionHeaderIfNecessary(); } else { - encryptLevel = "0"; - fileWriter.setEncryptParam(encryptLevel, "org.apache.tsfile.encrypt.UNENCRYPTED", ""); + this.secondEncryptParam = EncryptUtils.getEncryptParameter(firstEncryptParam); + String encryptLevel; + if (firstEncryptParam != null + && !Objects.equals(firstEncryptParam.getType(), "UNENCRYPTED") + && !Objects.equals( + firstEncryptParam.getType(), "org.apache.tsfile.encrypt.UNENCRYPTED")) { + encryptLevel = "2"; + String str = + EncryptUtils.getKeyStr( + IEncryptor.getEncryptor(firstEncryptParam.getType(), firstEncryptParam.getKey()) + .encrypt(secondEncryptParam.getKey())); + fileWriter.setEncryptParam(encryptLevel, secondEncryptParam.getType(), str); + } else { + encryptLevel = "0"; + fileWriter.setEncryptParam(encryptLevel, "org.apache.tsfile.encrypt.UNENCRYPTED", ""); + } } } @@ -273,7 +288,13 @@ public void addTsFileProperty(String key, byte[] value) { @TsFileApi public void close() throws IOException { LOG.info(Messages.get("log.write.close_file")); - flush(); - fileWriter.endFile(); + try { + flush(); + fileWriter.endFile(); + } finally { + if (secondEncryptParam != null && secondEncryptParam.isTdePageAead()) { + secondEncryptParam.close(); + } + } } } diff --git a/java/tsfile/src/main/java/org/apache/tsfile/write/writer/ForceAppendTsFileWriter.java b/java/tsfile/src/main/java/org/apache/tsfile/write/writer/ForceAppendTsFileWriter.java index 1f627d693c..cc94ea8574 100644 --- a/java/tsfile/src/main/java/org/apache/tsfile/write/writer/ForceAppendTsFileWriter.java +++ b/java/tsfile/src/main/java/org/apache/tsfile/write/writer/ForceAppendTsFileWriter.java @@ -45,6 +45,7 @@ public class ForceAppendTsFileWriter extends TsFileIOWriter { private long truncatePosition; private static Logger logger = LoggerFactory.getLogger(ForceAppendTsFileWriter.class); + private EncryptParameter ownedEncryptParameter; public ForceAppendTsFileWriter(File file) throws IOException { this( @@ -58,18 +59,27 @@ public ForceAppendTsFileWriter(File file, EncryptParameter param) throws IOExcep if (logger.isDebugEnabled()) { logger.debug(Messages.get("log.write.writer_opened"), file.getName()); } - this.out = FSFactoryProducer.getFileOutputFactory().getTsFileOutput(file.getPath(), true); this.file = file; - setEncryptParam(param); // file doesn't exist if (file.length() == 0 || !file.exists()) { throw new TsFileNotCompleteException( Messages.format("error.write.force_append_not_complete", file.getPath())); } + this.out = FSFactoryProducer.getFileOutputFactory().getTsFileOutput(file.getPath(), true); + setEncryptParam(param); try (TsFileSequenceReader reader = new TsFileSequenceReader(file.getAbsolutePath(), param, true)) { + if (!reader.hasFileEncryptionHeader() && param != null && param.isTdePageAead()) { + throw new IOException(Messages.get("error.write.aead_append_requires_file_header")); + } + EncryptParameter recoveredParameter = reader.getEncryptParam(); + if (recoveredParameter != null && recoveredParameter.isTdePageAead()) { + ownedEncryptParameter = recoveredParameter.copy(); + setEncryptParam(ownedEncryptParameter); + } + markExistingFileStarted(reader.hasFileEncryptionHeader()); // this tsfile is not complete if (!reader.isComplete()) { @@ -85,9 +95,21 @@ public ForceAppendTsFileWriter(File file, EncryptParameter param) throws IOExcep for (IDeviceID device : devices) { List chunkMetadataList = new ArrayList<>(); reader.readChunkMetadataInDevice(device).values().forEach(chunkMetadataList::addAll); + if (ownedEncryptParameter != null) { + for (ChunkMetadata chunkMetadata : chunkMetadataList) { + reader.position(chunkMetadata.getOffsetOfChunkHeader()); + reader.readChunkHeader(reader.readMarker()); + } + } ChunkGroupMetadata chunkGroupMetadata = new ChunkGroupMetadata(device, chunkMetadataList); chunkGroupMetadataList.add(chunkGroupMetadata); } + if (ownedEncryptParameter != null) { + ownedEncryptParameter.resumeAfterChunkOrdinal(reader.getMaxChunkOrdinal()); + } + } catch (IOException | RuntimeException e) { + closeAfterFailedInitialization(e); + throw e; } } @@ -98,4 +120,38 @@ public void doTruncate() throws IOException { public long getTruncatePosition() { return truncatePosition; } + + @Override + public void endFile() throws IOException { + try { + super.endFile(); + } finally { + closeOwnedEncryptParameter(); + } + } + + @Override + public void close() throws IOException { + try { + super.close(); + } finally { + closeOwnedEncryptParameter(); + } + } + + private void closeOwnedEncryptParameter() { + if (ownedEncryptParameter != null) { + ownedEncryptParameter.close(); + } + } + + private void closeAfterFailedInitialization(Exception exception) { + try { + out.close(); + } catch (IOException e) { + exception.addSuppressed(e); + } finally { + closeOwnedEncryptParameter(); + } + } } diff --git a/java/tsfile/src/main/java/org/apache/tsfile/write/writer/RestorableTsFileIOWriter.java b/java/tsfile/src/main/java/org/apache/tsfile/write/writer/RestorableTsFileIOWriter.java index 09866b3eb3..13dd66d471 100644 --- a/java/tsfile/src/main/java/org/apache/tsfile/write/writer/RestorableTsFileIOWriter.java +++ b/java/tsfile/src/main/java/org/apache/tsfile/write/writer/RestorableTsFileIOWriter.java @@ -126,23 +126,41 @@ public RestorableTsFileIOWriter(File file, boolean truncate, EncryptParameter pa } this.file = file; this.out = FSFactoryProducer.getFileOutputFactory().getTsFileOutput(file.getPath(), true); - this.param = param; - setEncryptParam(param); + this.param = param != null && param.isTdePageAead() ? param.copy() : param; + setEncryptParam(this.param); // file doesn't exist if (file.length() == 0) { - startFile(); - crashed = true; - canWrite = true; - return; + try { + startFile(); + crashed = true; + canWrite = true; + return; + } catch (IOException | RuntimeException e) { + closeAfterFailedInitialization(e); + throw e; + } } try { if (file.exists()) { try (TsFileSequenceReader reader = new TsFileSequenceReader(file.getAbsolutePath(), param, false)) { + EncryptParameter recoveredParameter = reader.getEncryptParam(); + if ((recoveredParameter != null && recoveredParameter.isTdePageAead()) + || (this.param != null && this.param.isTdePageAead())) { + if (this.param != null && this.param.isTdePageAead()) { + this.param.close(); + } + this.param = recoveredParameter == null ? null : recoveredParameter.copy(); + setEncryptParam(this.param); + } + markExistingFileStarted(reader.hasFileEncryptionHeader()); schema.setEnabledUpdateSchema(false); truncatedSize = reader.selfCheck(schema, chunkGroupMetadataList, true); + if (this.param != null && this.param.isTdePageAead()) { + this.param.resumeAfterChunkOrdinal(reader.getMaxChunkOrdinal()); + } minPlanIndex = reader.getMinPlanIndex(); maxPlanIndex = reader.getMaxPlanIndex(); if (truncatedSize == TsFileCheckStatus.COMPLETE_FILE) { @@ -165,7 +183,7 @@ public RestorableTsFileIOWriter(File file, boolean truncate, EncryptParameter pa } } } catch (Exception e) { - out.close(); + closeAfterFailedInitialization(e); throw e; } } @@ -209,6 +227,45 @@ public Schema getKnownSchema() { return schema; } + @Override + public EncryptParameter getEncryptParameter() { + return param; + } + + @Override + public void endFile() throws IOException { + try { + super.endFile(); + } finally { + closeOwnedEncryptParameter(); + } + } + + @Override + public void close() throws IOException { + try { + super.close(); + } finally { + closeOwnedEncryptParameter(); + } + } + + private void closeOwnedEncryptParameter() { + if (param != null && param.isTdePageAead()) { + param.close(); + } + } + + private void closeAfterFailedInitialization(Exception exception) { + try { + out.close(); + } catch (IOException e) { + exception.addSuppressed(e); + } finally { + closeOwnedEncryptParameter(); + } + } + /** * For query. * diff --git a/java/tsfile/src/main/java/org/apache/tsfile/write/writer/TsFileIOWriter.java b/java/tsfile/src/main/java/org/apache/tsfile/write/writer/TsFileIOWriter.java index 5e1d33c21b..abf45b0766 100644 --- a/java/tsfile/src/main/java/org/apache/tsfile/write/writer/TsFileIOWriter.java +++ b/java/tsfile/src/main/java/org/apache/tsfile/write/writer/TsFileIOWriter.java @@ -25,10 +25,12 @@ import org.apache.tsfile.encrypt.EncryptUtils; import org.apache.tsfile.encrypt.IEncryptor; import org.apache.tsfile.enums.TSDataType; +import org.apache.tsfile.exception.encrypt.EncryptException; import org.apache.tsfile.external.commons.io.FileUtils; import org.apache.tsfile.file.MetaMarker; import org.apache.tsfile.file.header.ChunkGroupHeader; import org.apache.tsfile.file.header.ChunkHeader; +import org.apache.tsfile.file.header.FileEncryptionHeader; import org.apache.tsfile.file.metadata.ChunkGroupMetadata; import org.apache.tsfile.file.metadata.ChunkMetadata; import org.apache.tsfile.file.metadata.IChunkMetadata; @@ -136,6 +138,14 @@ public class TsFileIOWriter implements AutoCloseable { protected String encryptKey; + protected EncryptParameter encryptParameter; + private EncryptParameter ownedEncryptParameter; + private boolean newFileContextClaimed; + + private boolean fileStarted; + + private boolean encryptionHeaderWritten; + private final List flushListeners = new ArrayList<>(); protected String currentTable; @@ -187,13 +197,18 @@ public TsFileIOWriter(File file, TSFileConfig conf) throws IOException { } public TsFileIOWriter(File file, TSFileConfig conf, EncryptParameter param) throws IOException { - this.out = FSFactoryProducer.getFileOutputFactory().getTsFileOutput(file.getPath(), false); + prepareNewFileEncryption(param); this.file = file; - if (resourceLogger.isDebugEnabled()) { - resourceLogger.debug("{} writer is opened.", file.getName()); + try { + this.out = FSFactoryProducer.getFileOutputFactory().getTsFileOutput(file.getPath(), false); + if (resourceLogger.isDebugEnabled()) { + resourceLogger.debug("{} writer is opened.", file.getName()); + } + startFile(); + } catch (IOException | RuntimeException | Error e) { + closeAfterNewFileFailure(e); + throw e; } - setEncryptParam(param); - startFile(); } /** @@ -210,8 +225,47 @@ public TsFileIOWriter(TsFileOutput output) throws IOException { public TsFileIOWriter(TsFileOutput output, EncryptParameter param) throws IOException { this.out = output; - setEncryptParam(param); - startFile(); + try { + prepareNewFileEncryption(param); + startFile(); + } catch (IOException | RuntimeException | Error e) { + closeAfterNewFileFailure(e); + throw e; + } + } + + private void prepareNewFileEncryption(EncryptParameter parameter) { + if (parameter != null && parameter.isTdePageAead()) { + ownedEncryptParameter = parameter.copy(); + try { + ownedEncryptParameter.claimForNewFile(); + } catch (RuntimeException | Error e) { + ownedEncryptParameter.close(); + throw e; + } + newFileContextClaimed = true; + setEncryptParam(ownedEncryptParameter); + } else { + setEncryptParam(parameter); + } + } + + private void closeAfterNewFileFailure(Throwable failure) { + try { + if (out != null) { + out.close(); + } + } catch (IOException | RuntimeException e) { + failure.addSuppressed(e); + } finally { + closeOwnedEncryptParameter(); + } + } + + private void closeOwnedEncryptParameter() { + if (ownedEncryptParameter != null) { + ownedEncryptParameter.close(); + } } /** for test only */ @@ -236,14 +290,28 @@ public TsFileIOWriter(File file, long maxMetadataSize, EncryptParameter param) } public void setEncryptParam(String encryptLevel, String encryptType, String encryptKey) { + if (encryptionHeaderWritten) { + throw new EncryptException(Messages.get("error.write.encryption_context_immutable")); + } this.encryptLevel = encryptLevel; this.encryptType = encryptType; this.encryptKey = encryptKey; + this.encryptParameter = new EncryptParameter(encryptType, null); } public void setEncryptParam(EncryptParameter param) { + if (encryptionHeaderWritten && encryptParameter != param) { + throw new EncryptException(Messages.get("error.write.encryption_context_immutable")); + } + this.encryptParameter = param; if (param == null) { setEncryptParam("0", "org.apache.tsfile.encrypt.UNENCRYPTED", null); + } else if (param.isTdePageAead()) { + // File-level encryption metadata is persisted in FileEncryptionHeader. Legacy footer fields + // stay explicitly unencrypted so wrapped key material is not duplicated. + encryptLevel = "0"; + encryptType = "org.apache.tsfile.encrypt.UNENCRYPTED"; + encryptKey = ""; } else { if (!Objects.equals(param.getType(), "UNENCRYPTED") && !Objects.equals(param.getType(), "org.apache.tsfile.encrypt.UNENCRYPTED")) { @@ -258,6 +326,34 @@ public void setEncryptParam(EncryptParameter param) { setEncryptParam("0", "org.apache.tsfile.encrypt.UNENCRYPTED", null); } } + this.encryptParameter = param; + } + + public EncryptParameter getEncryptParameter() { + return encryptParameter; + } + + protected void markExistingFileStarted(boolean hasEncryptionHeader) { + fileStarted = true; + encryptionHeaderWritten = hasEncryptionHeader; + } + + public void writeEncryptionHeaderIfNecessary() throws IOException { + if (encryptParameter == null || !encryptParameter.isTdePageAead() || encryptionHeaderWritten) { + return; + } + long baseHeaderSize = MAGIC_STRING_BYTES.length + Byte.BYTES; + if (!fileStarted || out.getPosition() != baseHeaderSize) { + throw new IOException(Messages.get("error.write.encryption_header_must_precede_data")); + } + if (!newFileContextClaimed) { + encryptParameter.claimForNewFile(); + newFileContextClaimed = true; + } + FileEncryptionHeader.serialize(encryptParameter, out.wrapAsStream()); + // Make the header visible to readers without a separate durability synchronization. + out.flush(); + encryptionHeaderWritten = true; } /** Add a custom property to the TsFile metadata. */ @@ -277,15 +373,19 @@ public void addFlushListener(FlushChunkMetadataListener listener) { * @throws IOException if an I/O error occurs. */ public void writeBytesToStream(PublicBAOS bytes) throws IOException { + writeEncryptionHeaderIfNecessary(); bytes.writeTo(out.wrapAsStream()); } protected void startFile() throws IOException { out.write(MAGIC_STRING_BYTES); out.write(VERSION_NUMBER_BYTE); + fileStarted = true; + writeEncryptionHeaderIfNecessary(); } public int startChunkGroup(IDeviceID deviceId) throws IOException { + writeEncryptionHeaderIfNecessary(); updateTableSize(deviceId); this.currentChunkGroupDeviceId = deviceId; if (logger.isDebugEnabled()) { @@ -346,6 +446,41 @@ public void startFlushChunk( int numOfPages, int mask) throws IOException { + startFlushChunk( + measurementId, + compressionCodecName, + tsDataType, + encodingType, + statistics, + dataSize, + numOfPages, + mask, + -1, + null); + } + + public void startFlushChunk( + String measurementId, + CompressionType compressionCodecName, + TSDataType tsDataType, + TSEncoding encodingType, + Statistics statistics, + int dataSize, + int numOfPages, + int mask, + long chunkOrdinal, + EncryptParameter chunkEncryptParameter) + throws IOException { + + boolean pageAead = encryptParameter != null && encryptParameter.isTdePageAead(); + if (pageAead != (chunkOrdinal >= 0)) { + throw new IOException(Messages.get("error.write.invalid_chunk_ordinal")); + } + if (pageAead && !encryptParameter.sharesPageAeadFileContext(chunkEncryptParameter)) { + throw new IOException(Messages.get("error.write.encrypted_chunk_context_mismatch")); + } + + writeEncryptionHeaderIfNecessary(); currentChunkMetadata = new ChunkMetadata( @@ -365,12 +500,15 @@ public void startFlushChunk( compressionCodecName, encodingType, numOfPages, - mask); - header.serializeTo(out.wrapAsStream()); + mask, + chunkOrdinal); + writeChunkHeader(header); } /** Write a whole chunk in another file into this file. Providing fast merge for IoTDB. */ public void writeChunk(Chunk chunk, ChunkMetadata chunkMetadata) throws IOException { + writeEncryptionHeaderIfNecessary(); + validateChunkEncryptionContext(chunk); ChunkHeader chunkHeader = chunk.getHeader(); currentChunkMetadata = new ChunkMetadata( @@ -381,7 +519,7 @@ public void writeChunk(Chunk chunk, ChunkMetadata chunkMetadata) throws IOExcept out.getPosition(), chunkMetadata.getStatistics()); currentChunkMetadata.setMask(chunkMetadata.getMask()); - chunkHeader.serializeTo(out.wrapAsStream()); + writeChunkHeader(chunkHeader); out.write(chunk.getData()); endCurrentChunk(); if (logger.isDebugEnabled()) { @@ -400,6 +538,7 @@ public void writeEmptyValueChunk( TSEncoding encodingType, Statistics statistics) throws IOException { + writeEncryptionHeaderIfNecessary(); currentChunkMetadata = new ChunkMetadata( measurementId, @@ -417,12 +556,17 @@ public void writeEmptyValueChunk( compressionType, encodingType, 0, - TsFileConstant.VALUE_COLUMN_MASK); - emptyChunkHeader.serializeTo(out.wrapAsStream()); + TsFileConstant.VALUE_COLUMN_MASK, + encryptParameter != null && encryptParameter.isTdePageAead() + ? encryptParameter.nextChunkOrdinal() + : -1); + writeChunkHeader(emptyChunkHeader); endCurrentChunk(); } public void writeChunk(Chunk chunk) throws IOException { + writeEncryptionHeaderIfNecessary(); + validateChunkEncryptionContext(chunk); ChunkHeader chunkHeader = chunk.getHeader(); currentChunkMetadata = new ChunkMetadata( @@ -436,11 +580,19 @@ public void writeChunk(Chunk chunk) throws IOException { (byte) (chunkHeader.getChunkType() & (TsFileConstant.TIME_COLUMN_MASK | TsFileConstant.VALUE_COLUMN_MASK))); - chunkHeader.serializeTo(out.wrapAsStream()); + writeChunkHeader(chunkHeader); out.write(chunk.getData()); endCurrentChunk(); } + private void writeChunkHeader(ChunkHeader header) throws IOException { + byte[] authentication = + encryptParameter != null && encryptParameter.isTdePageAead() + ? encryptParameter.authenticateChunk(out.getPosition(), header.getChunkOrdinal()) + : null; + header.serializeTo(out.wrapAsStream(), authentication); + } + /** end chunk and write some log. */ public void endCurrentChunk() { this.currentChunkMetadataSize += currentChunkMetadata.getRetainedSizeInBytes(); @@ -468,6 +620,7 @@ public void endFile() throws IOException { if (!canWrite) { return; } + writeEncryptionHeaderIfNecessary(); updateTableSize(null); checkInMemoryPathCount(); @@ -484,7 +637,11 @@ public void endFile() throws IOException { // flush page cache data to disk out.force(); // close file - out.close(); + try { + out.close(); + } finally { + closeOwnedEncryptParameter(); + } if (resourceLogger.isDebugEnabled() && file != null) { resourceLogger.debug("{} writer is closed.", file.getName()); @@ -689,9 +846,13 @@ public void reset() throws IOException { */ public void close() throws IOException { canWrite = false; - out.close(); - if (tempOutput != null) { - this.tempOutput.close(); + try { + out.close(); + if (tempOutput != null) { + this.tempOutput.close(); + } + } finally { + closeOwnedEncryptParameter(); } } @@ -712,12 +873,26 @@ public void setFile(File file) { } public void writePlanIndices() throws IOException { + writeEncryptionHeaderIfNecessary(); ReadWriteIOUtils.write(MetaMarker.OPERATION_INDEX_RANGE, out.wrapAsStream()); ReadWriteIOUtils.write(minPlanIndex, out.wrapAsStream()); ReadWriteIOUtils.write(maxPlanIndex, out.wrapAsStream()); out.flush(); } + private void validateChunkEncryptionContext(Chunk chunk) throws IOException { + EncryptParameter sourceParameter = chunk.getEncryptParam(); + boolean targetUsesPageAead = encryptParameter != null && encryptParameter.isTdePageAead(); + boolean sourceUsesPageAead = sourceParameter != null && sourceParameter.isTdePageAead(); + if (!targetUsesPageAead && !sourceUsesPageAead) { + if (chunk.getHeader().getChunkOrdinal() >= 0) { + throw new IOException(Messages.get("error.write.unencrypted_chunk_has_ordinal")); + } + return; + } + throw new IOException(Messages.get("error.write.encrypted_chunk_copy_unsupported")); + } + public void truncate(long offset) throws IOException { out.truncate(offset); } diff --git a/java/tsfile/src/test/java/org/apache/tsfile/encrypt/EncryptionProviderRegistryTest.java b/java/tsfile/src/test/java/org/apache/tsfile/encrypt/EncryptionProviderRegistryTest.java new file mode 100644 index 0000000000..96f46e571e --- /dev/null +++ b/java/tsfile/src/test/java/org/apache/tsfile/encrypt/EncryptionProviderRegistryTest.java @@ -0,0 +1,228 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ +package org.apache.tsfile.encrypt; + +import org.apache.tsfile.exception.encrypt.EncryptException; + +import org.junit.After; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.TemporaryFolder; + +import java.lang.reflect.Field; +import java.net.URLClassLoader; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; + +public class EncryptionProviderRegistryTest { + private static final String SLOW_ID = "registry-slow-provider"; + private static final String GOOD_ID = "registry-good-provider"; + private static final String BAD_ID = "registry-retry-provider"; + private static volatile CountDownLatch entered; + private static volatile CountDownLatch release; + private static volatile boolean failConstruction; + + @Rule public TemporaryFolder temporaryFolder = new TemporaryFolder(); + + @Before + public void setUp() throws Exception { + resetDiscovery(); + entered = new CountDownLatch(1); + release = new CountDownLatch(1); + failConstruction = false; + } + + @After + public void tearDown() throws Exception { + release.countDown(); + for (String id : new String[] {SLOW_ID, GOOD_ID, BAD_ID}) { + EncryptionProviderRegistry.unregisterProvider(id); + } + resetDiscovery(); + } + + private static void resetDiscovery() throws Exception { + Field field = EncryptionProviderRegistry.class.getDeclaredField("servicesLoaded"); + field.setAccessible(true); + field.setBoolean(null, false); + } + + @Test + public void testConcurrentFirstUseWaitsForServiceRegistration() throws Exception { + try (URLClassLoader loader = serviceLoader(SlowProvider.class)) { + ExecutorService pool = Executors.newFixedThreadPool(2); + CountDownLatch secondStarted = new CountDownLatch(1); + try { + Future first = pool.submit(() -> createThrough(loader, SLOW_ID)); + assertTrue(entered.await(10, TimeUnit.SECONDS)); + Future second = + pool.submit( + () -> { + secondStarted.countDown(); + createThrough(loader, SLOW_ID); + }); + assertTrue(secondStarted.await(10, TimeUnit.SECONDS)); + Thread.sleep(100); + assertFalse(second.isDone()); + release.countDown(); + first.get(10, TimeUnit.SECONDS); + second.get(10, TimeUnit.SECONDS); + } finally { + release.countDown(); + pool.shutdownNow(); + assertTrue(pool.awaitTermination(10, TimeUnit.SECONDS)); + } + } + } + + @Test + public void testFailedDiscoveryIsRetryableAndContinuesRegistration() throws Exception { + failConstruction = true; + try (URLClassLoader loader = serviceLoader(RetryProvider.class, GoodProvider.class)) { + assertThrows(EncryptException.class, () -> createThrough(loader, GOOD_ID)); + Field field = EncryptionProviderRegistry.class.getDeclaredField("PROVIDERS"); + field.setAccessible(true); + assertTrue(((java.util.Map) field.get(null)).containsKey(GOOD_ID)); + failConstruction = false; + createThrough(loader, GOOD_ID); + createThrough(loader, BAD_ID); + } + } + + @Test + public void testUnsupportedChunkAuthenticationNamesImplementation() { + IEncrypt legacy = new UNENCRYPTED(null); + try (legacy) { + assertTrue( + assertThrows(EncryptException.class, () -> legacy.authenticateChunk(null)) + .getMessage() + .contains(legacy.getClass().getName())); + assertTrue( + assertThrows( + EncryptException.class, () -> legacy.verifyChunkAuthentication(new byte[0], null)) + .getMessage() + .contains(legacy.getClass().getName())); + } + } + + @Test + public void testLegacyOperationsAreRejectedEvenForPermissiveProvider() { + IEncrypt permissive = new UNENCRYPTED(null); + try (IEncrypt guarded = new PageAeadEncryption(permissive)) { + assertThrows(EncryptException.class, () -> guarded.getEncryptor().encrypt(new byte[1])); + assertThrows(EncryptException.class, () -> guarded.getEncryptor().encrypt(new byte[1], 0, 1)); + assertThrows(EncryptException.class, () -> guarded.getDecryptor().decrypt(new byte[1])); + assertThrows(EncryptException.class, () -> guarded.getDecryptor().decrypt(new byte[1], 0, 1)); + } + } + + private URLClassLoader serviceLoader(Class... providers) throws Exception { + Path root = temporaryFolder.newFolder().toPath(); + Path descriptor = root.resolve("META-INF/services/" + IEncryptProvider.class.getName()); + Files.createDirectories(descriptor.getParent()); + StringBuilder names = new StringBuilder(); + for (Class provider : providers) { + names.append(provider.getName()).append('\n'); + } + Files.write(descriptor, names.toString().getBytes(StandardCharsets.UTF_8)); + return new URLClassLoader( + new java.net.URL[] {root.toUri().toURL()}, getClass().getClassLoader()); + } + + private static void createThrough(ClassLoader loader, String providerId) { + ClassLoader original = Thread.currentThread().getContextClassLoader(); + Thread.currentThread().setContextClassLoader(loader); + try (EncryptParameter parameter = + EncryptParameter.pageAeadBuilder() + .key(new byte[16]) + .providerId(providerId) + .profileId(TestAeadEncryptionProvider.PROFILE_ID) + .keyId("test-key") + .keyVersion("1") + .wrappedDataKey(new byte[24]) + .fileCryptoId(new byte[16]) + .build(); + IEncrypt encrypt = EncryptionProviderRegistry.create(parameter)) { + assertEquals(28, encrypt.getPageBodyOverhead()); + } finally { + Thread.currentThread().setContextClassLoader(original); + } + } + + public static class SlowProvider implements IEncryptProvider { + public SlowProvider() { + entered.countDown(); + try { + if (!release.await(10, TimeUnit.SECONDS)) { + throw new IllegalStateException("Service registration probe timed out"); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IllegalStateException(e); + } + } + + @Override + public String getProviderId() { + return SLOW_ID; + } + + @Override + public IEncrypt create(EncryptParameter parameter) { + return TestAeadEncryptionProvider.INSTANCE.create(parameter); + } + } + + public static class GoodProvider implements IEncryptProvider { + @Override + public String getProviderId() { + return GOOD_ID; + } + + @Override + public IEncrypt create(EncryptParameter parameter) { + return TestAeadEncryptionProvider.INSTANCE.create(parameter); + } + } + + public static class RetryProvider extends GoodProvider { + public RetryProvider() { + if (failConstruction) { + throw new IllegalStateException("Injected service discovery failure"); + } + } + + @Override + public String getProviderId() { + return BAD_ID; + } + } +} diff --git a/java/tsfile/src/test/java/org/apache/tsfile/encrypt/PageCryptoContextTest.java b/java/tsfile/src/test/java/org/apache/tsfile/encrypt/PageCryptoContextTest.java new file mode 100644 index 0000000000..efd4f3a46f --- /dev/null +++ b/java/tsfile/src/test/java/org/apache/tsfile/encrypt/PageCryptoContextTest.java @@ -0,0 +1,426 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ +package org.apache.tsfile.encrypt; + +import org.apache.tsfile.common.conf.TSFileConfig; +import org.apache.tsfile.exception.encrypt.EncryptException; +import org.apache.tsfile.file.metadata.enums.EncryptionType; + +import org.junit.AfterClass; +import org.junit.Before; +import org.junit.BeforeClass; +import org.junit.Test; + +import java.io.ByteArrayOutputStream; +import java.io.DataOutputStream; +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; + +import static org.junit.Assert.assertArrayEquals; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertSame; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; + +public class PageCryptoContextTest { + + @BeforeClass + public static void setUpProvider() { + EncryptionProviderRegistry.registerProvider(TestAeadEncryptionProvider.INSTANCE); + } + + @AfterClass + public static void tearDownProvider() { + EncryptionProviderRegistry.unregisterProvider(TestAeadEncryptionProvider.PROVIDER_ID); + } + + @Before + public void resetProviderCount() { + TestAeadEncryptionProvider.resetCreateCount(); + } + + @Test + public void testAeadRejectsTamperingAndPageSwap() { + byte[] dataKey = new byte[16]; + Arrays.fill(dataKey, (byte) 0x5A); + byte[] fileCryptoId = new byte[EncryptParameter.FILE_CRYPTO_ID_LENGTH]; + Arrays.fill(fileCryptoId, (byte) 0x2C); + EncryptParameter parameter = TestAeadEncryptionProvider.createParameter(dataKey, fileCryptoId); + byte[] plaintext = new byte[] {1, 3, 5, 7, 9, 11, 13, 15}; + + try { + IEncryptor encryptor = IEncryptor.getEncryptor(parameter); + IDecryptor decryptor = IDecryptor.getDecryptor(parameter); + PageCryptoContext page0 = + PageCryptoContext.forEncryption(parameter, plaintext.length, plaintext.length, 0, 0); + PageCryptoContext page1 = + PageCryptoContext.forEncryption(parameter, plaintext.length, plaintext.length, 1, 0); + PageCryptoContext anotherChunkPage0 = + PageCryptoContext.forEncryption(parameter, plaintext.length, plaintext.length, 0, 1); + + byte[] encryptedPage0 = encryptor.encryptPage(plaintext, 0, plaintext.length, page0); + byte[] encryptedPage1 = encryptor.encryptPage(plaintext, 0, plaintext.length, page1); + assertArrayEquals( + plaintext, decryptor.decryptPage(encryptedPage0, 0, encryptedPage0.length, page0)); + assertArrayEquals( + plaintext, decryptor.decryptPage(encryptedPage1, 0, encryptedPage1.length, page1)); + + assertThrows( + EncryptException.class, + () -> decryptor.decryptPage(encryptedPage0, 0, encryptedPage0.length, page1)); + assertThrows( + EncryptException.class, + () -> decryptor.decryptPage(encryptedPage0, 0, encryptedPage0.length, anotherChunkPage0)); + + encryptedPage0[encryptedPage0.length - 1] ^= 1; + assertThrows( + EncryptException.class, + () -> decryptor.decryptPage(encryptedPage0, 0, encryptedPage0.length, page0)); + + assertEquals(1, TestAeadEncryptionProvider.getCreateCount()); + } finally { + parameter.close(); + } + } + + @Test + public void testPaddedLengthFramingAndAuthenticatedCandidate() { + try (EncryptParameter parameter = + TestAeadEncryptionProvider.createPaddedParameter(new byte[16], new byte[16])) { + for (int size : new int[] {0, 1, 15, 16, 17, 31, 32, 65536}) { + byte[] plaintext = new byte[size]; + PageCryptoContext write = PageCryptoContext.forEncryption(parameter, size, size, 0, 2); + byte[] body = IEncryptor.getEncryptor(parameter).encryptPage(plaintext, 0, size, write); + assertEquals(write.getEncryptedPageBodySize(), body.length); + ByteBuffer slice = ByteBuffer.allocateDirect(body.length + 9); + slice.position(5); + slice.put(body).flip().position(5); + PageCryptoContext read = PageCryptoContext.forDecryption(parameter, size, slice, 0, 2); + assertEquals(5, slice.position()); + assertEquals(size, read.getCompressedPlaintextSize()); + assertArrayEquals(write.getAssociatedData(), read.getAssociatedData()); + assertArrayEquals( + plaintext, IDecryptor.getDecryptor(parameter).decryptPage(body, 0, body.length, read)); + if (size > 1 && size % 16 != 0) { + ByteBuffer.wrap(body).putInt(size - 1); + PageCryptoContext modified = + PageCryptoContext.forDecryption(parameter, size, ByteBuffer.wrap(body), 0, 2); + assertThrows( + EncryptException.class, + () -> IDecryptor.getDecryptor(parameter).decryptPage(body, 0, body.length, modified)); + } + } + assertThrows( + EncryptException.class, + () -> PageCryptoContext.forDecryption(parameter, 1, ByteBuffer.allocate(3), 0, 2)); + assertThrows( + EncryptException.class, + () -> + PageCryptoContext.forDecryption( + parameter, 1, ByteBuffer.allocate(48).putInt(-1).flip(), 0, 2)); + } + } + + @Test + public void testChunkOrdinalAllocationSurvivesCopyAndRecovery() { + EncryptParameter parameter = + TestAeadEncryptionProvider.createParameter( + new byte[16], new byte[EncryptParameter.FILE_CRYPTO_ID_LENGTH]); + try { + assertEquals(0, parameter.nextChunkOrdinal()); + try (EncryptParameter copy = parameter.copy()) { + assertEquals(1, copy.nextChunkOrdinal()); + parameter.resumeAfterChunkOrdinal(10); + assertEquals(11, copy.nextChunkOrdinal()); + } + assertEquals(12, parameter.nextChunkOrdinal()); + } finally { + parameter.close(); + } + } + + @Test + public void testRetainedReferencesShareProviderUntilLastClose() { + EncryptParameter parameter = + TestAeadEncryptionProvider.createParameter( + new byte[16], new byte[EncryptParameter.FILE_CRYPTO_ID_LENGTH]); + int previousCloseCount = TestAeadEncryptionProvider.getCloseCount(); + try (EncryptParameter first = parameter.retain(); + EncryptParameter second = parameter.retain()) { + assertSame(first.getOrCreateFileEncrypt(), second.getOrCreateFileEncrypt()); + assertEquals(1, TestAeadEncryptionProvider.getCreateCount()); + parameter.close(); + first.close(); + assertEquals(previousCloseCount, TestAeadEncryptionProvider.getCloseCount()); + assertEquals(28, second.getPageBodyOverhead()); + assertThrows(EncryptException.class, parameter::retain); + assertThrows(EncryptException.class, first::getOrCreateFileEncrypt); + } finally { + parameter.close(); + } + assertEquals(previousCloseCount + 1, TestAeadEncryptionProvider.getCloseCount()); + } + + @Test + public void testRetainedReferenceCanInitializeAfterOriginalCloses() { + EncryptParameter parameter = + TestAeadEncryptionProvider.createParameter( + new byte[16], new byte[EncryptParameter.FILE_CRYPTO_ID_LENGTH]); + try (EncryptParameter retained = parameter.retain()) { + parameter.close(); + byte[] plaintext = {1, 2, 3}; + PageCryptoContext context = PageCryptoContext.forEncryption(retained, 3, 3, 0, 0); + byte[] encrypted = IEncryptor.getEncryptor(retained).encryptPage(plaintext, 0, 3, context); + assertArrayEquals( + plaintext, + IDecryptor.getDecryptor(retained).decryptPage(encrypted, 0, encrypted.length, context)); + assertEquals(1, TestAeadEncryptionProvider.getCreateCount()); + } finally { + parameter.close(); + } + } + + @Test + public void testConcurrentRetainedOwnersCreateAndCloseProviderOnce() throws Exception { + EncryptParameter parameter = + TestAeadEncryptionProvider.createParameter( + new byte[16], new byte[EncryptParameter.FILE_CRYPTO_ID_LENGTH]); + List owners = new ArrayList<>(); + ExecutorService executor = Executors.newFixedThreadPool(4); + CountDownLatch start = new CountDownLatch(1); + int previousCloseCount = TestAeadEncryptionProvider.getCloseCount(); + try { + for (int i = 0; i < 4; i++) { + owners.add(parameter.retain()); + } + parameter.close(); + List> results = new ArrayList<>(); + for (EncryptParameter owner : owners) { + results.add( + executor.submit( + () -> { + try (EncryptParameter reference = owner) { + start.await(); + return reference.getOrCreateFileEncrypt(); + } + })); + } + start.countDown(); + IEncrypt provider = results.get(0).get(10, TimeUnit.SECONDS); + for (Future result : results) { + assertSame(provider, result.get(10, TimeUnit.SECONDS)); + } + assertEquals(1, TestAeadEncryptionProvider.getCreateCount()); + assertEquals(previousCloseCount + 1, TestAeadEncryptionProvider.getCloseCount()); + } finally { + start.countDown(); + executor.shutdownNow(); + for (EncryptParameter owner : owners) { + owner.close(); + } + parameter.close(); + } + } + + @Test + public void testCopyStillOwnsAnIndependentProvider() { + try (EncryptParameter parameter = + TestAeadEncryptionProvider.createParameter( + new byte[16], new byte[EncryptParameter.FILE_CRYPTO_ID_LENGTH]); + EncryptParameter copy = parameter.copy()) { + parameter.getPageBodyOverhead(); + copy.getPageBodyOverhead(); + assertEquals(2, TestAeadEncryptionProvider.getCreateCount()); + parameter.close(); + assertEquals(28, copy.getPageBodyOverhead()); + } + } + + @Test + public void testProviderLookupNormalizesIdLikeRegistration() { + try (EncryptParameter parameter = + TestAeadEncryptionProvider.createParameter( + new byte[16], new byte[EncryptParameter.FILE_CRYPTO_ID_LENGTH]); + EncryptParameter nonCanonical = + EncryptParameter.pageAeadBuilder() + .key(parameter.getKey()) + .providerId(" " + TestAeadEncryptionProvider.PROVIDER_ID + " ") + .profileId(parameter.getProfileId()) + .keyId(parameter.getKeyId()) + .keyVersion(parameter.getKeyVersion()) + .wrappedDataKey(parameter.getWrappedDataKey()) + .fileCryptoId(parameter.getFileCryptoId()) + .build()) { + assertEquals(28, nonCanonical.getPageBodyOverhead()); + } + } + + @Test + public void testPageAeadFileContextRequiresSharedAllocatorAndKey() { + byte[] fileId = new byte[EncryptParameter.FILE_CRYPTO_ID_LENGTH]; + EncryptParameter parameter = TestAeadEncryptionProvider.createParameter(new byte[16], fileId); + EncryptParameter independent = TestAeadEncryptionProvider.createParameter(new byte[16], fileId); + try (EncryptParameter copy = parameter.copy()) { + assertTrue(parameter.sharesPageAeadFileContext(copy)); + assertFalse(parameter.sharesPageAeadFileContext(independent)); + copy.getKey()[0] ^= 1; + assertFalse(parameter.sharesPageAeadFileContext(copy)); + } finally { + parameter.close(); + independent.close(); + } + } + + @Test + public void testAssociatedDataCannotBeModifiedByProvider() { + EncryptParameter parameter = + TestAeadEncryptionProvider.createParameter( + new byte[16], new byte[EncryptParameter.FILE_CRYPTO_ID_LENGTH]); + try { + PageCryptoContext context = PageCryptoContext.forEncryption(parameter, 8, 8, 0, 0); + byte[] expected = context.getAssociatedData(); + byte[] providerCopy = context.getAssociatedData(); + providerCopy[0] ^= 1; + assertArrayEquals(expected, context.getAssociatedData()); + } finally { + parameter.close(); + } + } + + @Test + public void testCachedAssociatedDataPreservesEncodingAndReferenceLifetime() throws Exception { + try (EncryptParameter parameter = + TestAeadEncryptionProvider.createParameter(new byte[16], new byte[16])) { + PageCryptoContext.forEncryption(parameter, 100, 80, 0, 0); + try (EncryptParameter retained = parameter.retain()) { + parameter.close(); + for (int page = 0; page < 100; page++) { + ByteArrayOutputStream expected = new ByteArrayOutputStream(); + DataOutputStream output = new DataOutputStream(expected); + output.writeInt(1); + for (String field : + new String[] { + retained.getProviderId(), retained.getProfileId(), + retained.getKeyId(), retained.getKeyVersion() + }) { + byte[] bytes = field.getBytes(TSFileConfig.STRING_CHARSET); + output.writeInt(bytes.length); + output.write(bytes); + } + output.writeInt(EncryptParameter.FILE_CRYPTO_ID_LENGTH); + output.write(retained.getFileCryptoId()); + output.writeLong(7); + output.writeInt(page); + output.writeInt(100); + output.writeInt(80); + PageCryptoContext context = PageCryptoContext.forEncryption(retained, 100, 80, page, 7); + assertArrayEquals(expected.toByteArray(), context.getAssociatedData()); + context.getAssociatedData()[0] ^= 1; + assertArrayEquals(expected.toByteArray(), context.getAssociatedData()); + assertArrayEquals( + expected.toByteArray(), + PageCryptoContext.forDecryption( + retained, 100, context.getEncryptedPageBodySize(), page, 7) + .getAssociatedData()); + } + } + } + } + + @Test + public void testLegacyCipherCannotSilentlyHandleAuthenticatedPages() { + EncryptParameter parameter = + TestAeadEncryptionProvider.createParameter( + new byte[16], new byte[EncryptParameter.FILE_CRYPTO_ID_LENGTH]); + try { + PageCryptoContext context = PageCryptoContext.forEncryption(parameter, 1, 1, 0, 0); + IEncryptor legacyEncryptor = + new IEncryptor() { + @Override + public byte[] encrypt(byte[] data) { + return data; + } + + @Override + public byte[] encrypt(byte[] data, int offset, int size) { + throw new AssertionError("Legacy encryption must not be called for AEAD pages"); + } + + @Override + public EncryptionType getEncryptionType() { + return EncryptionType.NewWay; + } + }; + IDecryptor legacyDecryptor = + new IDecryptor() { + @Override + public byte[] decrypt(byte[] data) { + return data; + } + + @Override + public byte[] decrypt(byte[] data, int offset, int size) { + throw new AssertionError("Legacy decryption must not be called for AEAD pages"); + } + + @Override + public EncryptionType getEncryptionType() { + return EncryptionType.NewWay; + } + }; + byte[] page = {1}; + assertTrue( + assertThrows( + EncryptException.class, () -> legacyEncryptor.encryptPage(page, 0, 1, context)) + .getMessage() + .contains(legacyEncryptor.getClass().getName())); + assertTrue( + assertThrows( + EncryptException.class, () -> legacyDecryptor.decryptPage(page, 0, 1, context)) + .getMessage() + .contains(legacyDecryptor.getClass().getName())); + } finally { + parameter.close(); + } + } + + @Test + public void testUnregisterProviderTrimsId() { + EncryptParameter parameter = + TestAeadEncryptionProvider.createParameter( + new byte[16], new byte[EncryptParameter.FILE_CRYPTO_ID_LENGTH]); + try { + EncryptionProviderRegistry.create(parameter); + EncryptionProviderRegistry.unregisterProvider( + " " + TestAeadEncryptionProvider.PROVIDER_ID + " "); + assertThrows(EncryptException.class, () -> EncryptionProviderRegistry.create(parameter)); + } finally { + EncryptionProviderRegistry.registerProvider(TestAeadEncryptionProvider.INSTANCE); + parameter.close(); + } + } +} diff --git a/java/tsfile/src/test/java/org/apache/tsfile/encrypt/TestAeadEncryptionProvider.java b/java/tsfile/src/test/java/org/apache/tsfile/encrypt/TestAeadEncryptionProvider.java new file mode 100644 index 0000000000..9ec770392a --- /dev/null +++ b/java/tsfile/src/test/java/org/apache/tsfile/encrypt/TestAeadEncryptionProvider.java @@ -0,0 +1,350 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ +package org.apache.tsfile.encrypt; + +import org.apache.tsfile.exception.encrypt.EncryptException; +import org.apache.tsfile.file.metadata.enums.EncryptionType; + +import javax.crypto.Cipher; +import javax.crypto.Mac; +import javax.crypto.SecretKey; +import javax.crypto.spec.GCMParameterSpec; +import javax.crypto.spec.SecretKeySpec; + +import java.security.GeneralSecurityException; +import java.security.SecureRandom; +import java.util.Arrays; +import java.util.concurrent.atomic.AtomicInteger; + +/** JDK-only AEAD provider used to exercise the generic TsFile encryption SPI. */ +public final class TestAeadEncryptionProvider implements IEncryptProvider { + + public static final TestAeadEncryptionProvider INSTANCE = new TestAeadEncryptionProvider(); + public static final String PROVIDER_ID = "test-aead-provider"; + public static final String PROFILE_ID = "AES_GCM_128_AES_WRAP_128_V1"; + public static final String MAC_PROFILE_ID = "TEST_AES_GCM_HMAC_SHA256_V1"; + public static final String PADDED_PROFILE_ID = "TEST_PADDED_AES_GCM_V1"; + + private static final int IV_LENGTH = 12; + private static final int TAG_LENGTH = 16; + + private static final byte[] TEST_KEY_ENCRYPTION_KEY = + new byte[] { + 0x10, 0x32, 0x54, 0x76, 0x01, 0x23, 0x45, 0x67, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, + 0x00 + }; + private static final SecureRandom SECURE_RANDOM = new SecureRandom(); + private static final AtomicInteger CREATE_COUNT = new AtomicInteger(); + private static final AtomicInteger CLOSE_COUNT = new AtomicInteger(); + + private TestAeadEncryptionProvider() {} + + @Override + public String getProviderId() { + return PROVIDER_ID; + } + + @Override + public IEncrypt create(EncryptParameter encryptParameter) { + if (!PROFILE_ID.equals(encryptParameter.getProfileId()) + && !MAC_PROFILE_ID.equals(encryptParameter.getProfileId()) + && !PADDED_PROFILE_ID.equals(encryptParameter.getProfileId())) { + throw new EncryptException("unsupported test encryption profile"); + } + CREATE_COUNT.incrementAndGet(); + return new TestAeadEncrypt( + resolveDataKey(encryptParameter), + MAC_PROFILE_ID.equals(encryptParameter.getProfileId()), + PADDED_PROFILE_ID.equals(encryptParameter.getProfileId())); + } + + public static EncryptParameter createParameter(byte[] dataKey, byte[] fileCryptoId) { + return createParameter(dataKey, fileCryptoId, PROFILE_ID); + } + + /** Test-only 32-byte proof profile; page bodies still use the existing AES-GCM test cipher. */ + public static EncryptParameter createMacParameter(byte[] dataKey, byte[] fileCryptoId) { + return createParameter(dataKey, fileCryptoId, MAC_PROFILE_ID); + } + + public static EncryptParameter createPaddedParameter(byte[] dataKey, byte[] fileCryptoId) { + return createParameter(dataKey, fileCryptoId, PADDED_PROFILE_ID); + } + + private static EncryptParameter createParameter( + byte[] dataKey, byte[] fileCryptoId, String profile) { + return EncryptParameter.pageAeadBuilder() + .key(dataKey) + .providerId(PROVIDER_ID) + .profileId(profile) + .keyId("test-key") + .keyVersion("1") + .wrappedDataKey(wrapDataKey(dataKey)) + .fileCryptoId(fileCryptoId) + .build(); + } + + public static int getCreateCount() { + return CREATE_COUNT.get(); + } + + public static int getCloseCount() { + return CLOSE_COUNT.get(); + } + + public static void resetCreateCount() { + CREATE_COUNT.set(0); + } + + private static byte[] resolveDataKey(EncryptParameter parameter) { + byte[] dataKey = parameter.getKey(); + if (dataKey != null) { + return Arrays.copyOf(dataKey, dataKey.length); + } + try { + Cipher cipher = Cipher.getInstance("AESWrap"); + cipher.init(Cipher.UNWRAP_MODE, new SecretKeySpec(TEST_KEY_ENCRYPTION_KEY, "AES")); + SecretKey unwrapped = + (SecretKey) cipher.unwrap(parameter.getWrappedDataKey(), "AES", Cipher.SECRET_KEY); + return unwrapped.getEncoded(); + } catch (GeneralSecurityException e) { + throw new EncryptException("test data key unwrap failed", e); + } + } + + private static byte[] wrapDataKey(byte[] dataKey) { + try { + Cipher cipher = Cipher.getInstance("AESWrap"); + cipher.init(Cipher.WRAP_MODE, new SecretKeySpec(TEST_KEY_ENCRYPTION_KEY, "AES")); + return cipher.wrap(new SecretKeySpec(dataKey, "AES")); + } catch (GeneralSecurityException e) { + throw new EncryptException("test data key wrap failed", e); + } + } + + private static final class TestAeadEncrypt implements IEncrypt { + + private final byte[] dataKey; + private final boolean macProof; + private final boolean padded; + + private TestAeadEncrypt(byte[] dataKey, boolean macProof, boolean padded) { + this.dataKey = dataKey; + this.macProof = macProof; + this.padded = padded; + } + + @Override + public IDecryptor getDecryptor() { + return new AeadDecryptor(dataKey, padded); + } + + @Override + public IEncryptor getEncryptor() { + return new AeadEncryptor(dataKey, padded); + } + + @Override + public int getPageBodyOverhead() { + return padded ? -1 : IV_LENGTH + TAG_LENGTH; + } + + @Override + public int getEncryptedPageBodySize(int plaintextSize) { + return padded + ? 32 + (plaintextSize / 16 + 1) * 16 + : IEncrypt.super.getEncryptedPageBodySize(plaintextSize); + } + + @Override + public int getPagePlaintextSize(java.nio.ByteBuffer body) { + if (!padded) { + return IEncrypt.super.getPagePlaintextSize(body); + } + if (body.remaining() < 48) { + throw new EncryptException("Truncated test padded page"); + } + return body.getInt(); + } + + @Override + public int getChunkAuthenticationSize() { + return macProof ? 32 : IV_LENGTH + TAG_LENGTH; + } + + @Override + public byte[] authenticateChunk(ChunkCryptoContext context) { + if (macProof) { + return chunkMac(context); + } + return new AeadEncryptor(dataKey) + .encryptWithAad(new byte[0], 0, 0, context.getAssociatedData()); + } + + @Override + public void verifyChunkAuthentication(byte[] authentication, ChunkCryptoContext context) { + if (macProof) { + if (!java.security.MessageDigest.isEqual(authentication, chunkMac(context))) { + throw new EncryptException("Test chunk MAC mismatch"); + } + return; + } + byte[] plaintext = + new AeadDecryptor(dataKey) + .decryptWithAad( + authentication, 0, authentication.length, context.getAssociatedData()); + if (plaintext.length != 0) { + throw new EncryptException("Unexpected chunk authentication plaintext"); + } + } + + private byte[] chunkMac(ChunkCryptoContext context) { + try { + Mac mac = Mac.getInstance("HmacSHA256"); + mac.init(new SecretKeySpec(dataKey, "HmacSHA256")); + return mac.doFinal(context.getAssociatedData()); + } catch (GeneralSecurityException e) { + throw new EncryptException("Test chunk MAC failed", e); + } + } + + @Override + public void close() { + Arrays.fill(dataKey, (byte) 0); + CLOSE_COUNT.incrementAndGet(); + } + } + + private static final class AeadEncryptor implements IEncryptor { + + private final byte[] dataKey; + private final boolean padded; + + private AeadEncryptor(byte[] dataKey) { + this(dataKey, false); + } + + private AeadEncryptor(byte[] dataKey, boolean padded) { + this.dataKey = dataKey; + this.padded = padded; + } + + @Override + public byte[] encrypt(byte[] data) { + return encrypt(data, 0, data.length); + } + + @Override + public byte[] encrypt(byte[] data, int offset, int size) { + throw new UnsupportedOperationException(); + } + + @Override + public byte[] encryptPage( + byte[] data, int offset, int size, PageCryptoContext pageCryptoContext) { + if (!padded) { + return encryptWithAad(data, offset, size, pageCryptoContext.getAssociatedData()); + } + byte[] input = new byte[(size / 16 + 1) * 16]; + System.arraycopy(data, offset, input, 0, size); + byte[] encrypted = + encryptWithAad(input, 0, input.length, pageCryptoContext.getAssociatedData()); + return java.nio.ByteBuffer.allocate(4 + encrypted.length).putInt(size).put(encrypted).array(); + } + + private byte[] encryptWithAad(byte[] data, int offset, int size, byte[] associatedData) { + byte[] iv = new byte[IV_LENGTH]; + SECURE_RANDOM.nextBytes(iv); + try { + Cipher cipher = Cipher.getInstance("AES/GCM/NoPadding"); + cipher.init( + Cipher.ENCRYPT_MODE, + new SecretKeySpec(dataKey, "AES"), + new GCMParameterSpec(TAG_LENGTH * Byte.SIZE, iv)); + cipher.updateAAD(associatedData); + byte[] ciphertext = cipher.doFinal(data, offset, size); + byte[] pageBody = new byte[iv.length + ciphertext.length]; + System.arraycopy(iv, 0, pageBody, 0, iv.length); + System.arraycopy(ciphertext, 0, pageBody, iv.length, ciphertext.length); + return pageBody; + } catch (GeneralSecurityException e) { + throw new EncryptException("test page encryption failed", e); + } + } + + @Override + public EncryptionType getEncryptionType() { + return EncryptionType.NewWay; + } + } + + private static final class AeadDecryptor implements IDecryptor { + + private final byte[] dataKey; + private final boolean padded; + + private AeadDecryptor(byte[] dataKey) { + this(dataKey, false); + } + + private AeadDecryptor(byte[] dataKey, boolean padded) { + this.dataKey = dataKey; + this.padded = padded; + } + + @Override + public byte[] decrypt(byte[] data) { + return decrypt(data, 0, data.length); + } + + @Override + public byte[] decrypt(byte[] data, int offset, int size) { + throw new UnsupportedOperationException(); + } + + @Override + public byte[] decryptPage( + byte[] data, int offset, int size, PageCryptoContext pageCryptoContext) { + if (!padded) { + return decryptWithAad(data, offset, size, pageCryptoContext.getAssociatedData()); + } + byte[] decoded = + decryptWithAad(data, offset + 4, size - 4, pageCryptoContext.getAssociatedData()); + return Arrays.copyOf(decoded, pageCryptoContext.getCompressedPlaintextSize()); + } + + private byte[] decryptWithAad(byte[] data, int offset, int size, byte[] associatedData) { + try { + GCMParameterSpec spec = + new GCMParameterSpec(TAG_LENGTH * Byte.SIZE, data, offset, IV_LENGTH); + Cipher cipher = Cipher.getInstance("AES/GCM/NoPadding"); + cipher.init(Cipher.DECRYPT_MODE, new SecretKeySpec(dataKey, "AES"), spec); + cipher.updateAAD(associatedData); + return cipher.doFinal(data, offset + IV_LENGTH, size - IV_LENGTH); + } catch (GeneralSecurityException e) { + throw new EncryptException("test page decryption failed", e); + } + } + + @Override + public EncryptionType getEncryptionType() { + return EncryptionType.NewWay; + } + } +} diff --git a/java/tsfile/src/test/java/org/apache/tsfile/file/header/FileEncryptionHeaderTest.java b/java/tsfile/src/test/java/org/apache/tsfile/file/header/FileEncryptionHeaderTest.java new file mode 100644 index 0000000000..38ab61650e --- /dev/null +++ b/java/tsfile/src/test/java/org/apache/tsfile/file/header/FileEncryptionHeaderTest.java @@ -0,0 +1,140 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ +package org.apache.tsfile.file.header; + +import org.apache.tsfile.encrypt.EncryptParameter; +import org.apache.tsfile.encrypt.TestAeadEncryptionProvider; +import org.apache.tsfile.exception.encrypt.EncryptException; +import org.apache.tsfile.file.MetaMarker; +import org.apache.tsfile.utils.ReadWriteForEncodingUtils; + +import org.junit.Test; + +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.IOException; + +import static org.junit.Assert.assertArrayEquals; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertThrows; + +public class FileEncryptionHeaderTest { + + @Test + public void testRoundTrip() throws IOException { + byte[] dataKey = new byte[16]; + byte[] fileCryptoId = new byte[EncryptParameter.FILE_CRYPTO_ID_LENGTH]; + for (int i = 0; i < dataKey.length; i++) { + dataKey[i] = (byte) (i + 1); + fileCryptoId[i] = (byte) (0x40 + i); + } + EncryptParameter parameter = TestAeadEncryptionProvider.createParameter(dataKey, fileCryptoId); + + ByteArrayOutputStream output = new ByteArrayOutputStream(); + FileEncryptionHeader.serialize(parameter, output); + byte[] serialized = output.toByteArray(); + assertEquals(MetaMarker.ENCRYPTION_HEADER, serialized[0]); + assertFalse(contains(serialized, dataKey)); + + ByteArrayInputStream input = new ByteArrayInputStream(serialized); + assertEquals(MetaMarker.ENCRYPTION_HEADER, input.read()); + EncryptParameter restored = FileEncryptionHeader.deserialize(input); + + assertNull(restored.getType()); + assertNull(restored.getKey()); + assertEquals(parameter.getProviderId(), restored.getProviderId()); + assertEquals(parameter.getProfileId(), restored.getProfileId()); + assertEquals(parameter.getKeyId(), restored.getKeyId()); + assertEquals(parameter.getKeyVersion(), restored.getKeyVersion()); + assertArrayEquals(parameter.getFileCryptoId(), restored.getFileCryptoId()); + assertArrayEquals(parameter.getWrappedDataKey(), restored.getWrappedDataKey()); + } + + @Test + public void testLargeHeaderUsesExactSerializedSize() throws IOException { + try (EncryptParameter parameter = + EncryptParameter.pageAeadBuilder() + .providerId("provider") + .profileId("profile") + .keyId("k".repeat(300)) + .keyVersion("1") + .fileCryptoId(new byte[16]) + .wrappedDataKey(new byte[20000]) + .build()) { + ByteArrayOutputStream output = new ByteArrayOutputStream(); + int serializedSize = FileEncryptionHeader.serialize(parameter, output); + assertEquals(serializedSize, output.size()); + ByteArrayInputStream input = new ByteArrayInputStream(output.toByteArray()); + assertEquals(MetaMarker.ENCRYPTION_HEADER, input.read()); + int payloadSize = ReadWriteForEncodingUtils.readUnsignedVarInt(input); + assertEquals(payloadSize, input.available()); + input = new ByteArrayInputStream(output.toByteArray()); + input.read(); + try (EncryptParameter restored = FileEncryptionHeader.deserialize(input)) { + assertEquals(parameter.getKeyId(), restored.getKeyId()); + assertArrayEquals(parameter.getWrappedDataKey(), restored.getWrappedDataKey()); + assertEquals(0, input.available()); + } + } + } + + @Test + public void testOversizedHeaderDoesNotWriteOutput() { + try (EncryptParameter parameter = + EncryptParameter.pageAeadBuilder() + .providerId("provider") + .profileId("profile") + .keyId("key") + .keyVersion("1") + .fileCryptoId(new byte[16]) + .wrappedDataKey(new byte[FileEncryptionHeader.MAX_HEADER_SIZE]) + .build()) { + ByteArrayOutputStream output = new ByteArrayOutputStream(); + assertThrows(EncryptException.class, () -> FileEncryptionHeader.serialize(parameter, output)); + assertEquals(0, output.size()); + } + } + + @Test + public void testRejectInvalidFileCryptoId() { + EncryptParameter parameter = + TestAeadEncryptionProvider.createParameter(new byte[16], new byte[8]); + assertThrows( + EncryptException.class, + () -> FileEncryptionHeader.serialize(parameter, new ByteArrayOutputStream())); + } + + private static boolean contains(byte[] data, byte[] target) { + for (int i = 0; i <= data.length - target.length; i++) { + boolean matches = true; + for (int j = 0; j < target.length; j++) { + if (data[i + j] != target[j]) { + matches = false; + break; + } + } + if (matches) { + return true; + } + } + return false; + } +} diff --git a/java/tsfile/src/test/java/org/apache/tsfile/read/PageAeadReaderLifecycleTest.java b/java/tsfile/src/test/java/org/apache/tsfile/read/PageAeadReaderLifecycleTest.java new file mode 100644 index 0000000000..1d2a497a30 --- /dev/null +++ b/java/tsfile/src/test/java/org/apache/tsfile/read/PageAeadReaderLifecycleTest.java @@ -0,0 +1,477 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ +package org.apache.tsfile.read; + +import org.apache.tsfile.encrypt.EncryptParameter; +import org.apache.tsfile.encrypt.EncryptionProviderRegistry; +import org.apache.tsfile.encrypt.TestAeadEncryptionProvider; +import org.apache.tsfile.enums.TSDataType; +import org.apache.tsfile.file.header.ChunkHeader; +import org.apache.tsfile.file.metadata.AlignedChunkMetadata; +import org.apache.tsfile.file.metadata.ChunkMetadata; +import org.apache.tsfile.file.metadata.IDeviceID; +import org.apache.tsfile.file.metadata.enums.TSEncoding; +import org.apache.tsfile.read.common.Chunk; +import org.apache.tsfile.read.common.Path; +import org.apache.tsfile.read.controller.CachedChunkLoaderImpl; +import org.apache.tsfile.read.controller.IChunkLoader; +import org.apache.tsfile.read.expression.QueryExpression; +import org.apache.tsfile.read.query.dataset.QueryDataSet; +import org.apache.tsfile.read.query.dataset.TreeResultSet; +import org.apache.tsfile.read.reader.BufferedTsFileInput; +import org.apache.tsfile.read.reader.IChunkReader; +import org.apache.tsfile.read.reader.chunk.AlignedChunkReader; +import org.apache.tsfile.read.reader.chunk.ChunkReader; +import org.apache.tsfile.read.reader.series.FileSeriesReader; +import org.apache.tsfile.read.reader.series.FileSeriesReaderByTimestamp; +import org.apache.tsfile.utils.TsPrimitiveType; +import org.apache.tsfile.write.TsFileWriter; +import org.apache.tsfile.write.record.TSRecord; +import org.apache.tsfile.write.record.datapoint.LongDataPoint; +import org.apache.tsfile.write.schema.MeasurementSchema; + +import org.junit.AfterClass; +import org.junit.BeforeClass; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.TemporaryFolder; + +import java.io.File; +import java.io.IOException; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.Map; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertSame; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +public class PageAeadReaderLifecycleTest { + @Rule public final TemporaryFolder temporaryFolder = new TemporaryFolder(); + + @BeforeClass + public static void registerProvider() { + EncryptionProviderRegistry.registerProvider(TestAeadEncryptionProvider.INSTANCE); + } + + @AfterClass + public static void unregisterProvider() { + EncryptionProviderRegistry.unregisterProvider(TestAeadEncryptionProvider.PROVIDER_ID); + } + + @Test + public void testPlaintextQueriesAreNotStronglyRetained() throws Exception { + File file = new File(temporaryFolder.getRoot(), "plain.tsfile"); + try (TsFileWriter writer = new TsFileWriter(file)) { + writer.registerTimeseries(new Path("d1"), new MeasurementSchema("s1", TSDataType.INT64)); + for (int i = 1; i <= 3; i++) { + writer.writeRecord(new TSRecord("d1", i).addTuple(new LongDataPoint("s1", i))); + } + } + try (TsFileReader reader = new TsFileReader(file)) { + for (int i = 0; i < 100; i++) { + QueryDataSet query = reader.query(queryExpression()); + assertTrue(query.hasNext()); + query.next(); + } + java.lang.reflect.Field field = TsFileReader.class.getDeclaredField("queryDataSets"); + field.setAccessible(true); + assertTrue(((java.util.Set) field.get(reader)).isEmpty()); + } + } + + @Test + public void testPaddedBodiesInAlignedAndNonAlignedQueries() throws Exception { + for (boolean aligned : new boolean[] {false, true}) { + File file = writeFile(aligned, true); + try (TsFileReader reader = new TsFileReader(file)) { + QueryDataSet query = reader.query(queryExpression()); + int count = 0; + while (query.hasNext()) { + assertEquals(++count * 10L, query.next().getFields().get(0).getLongV()); + } + assertEquals(3, count); + } + } + } + + @Test + public void testSeriesScanReleasesEachChunkAndCreatesOneProvider() throws Exception { + File file = writeFile(false); + int createCount = TestAeadEncryptionProvider.getCreateCount(); + int closeCount = TestAeadEncryptionProvider.getCloseCount(); + try (TsFileSequenceReader sequence = new TsFileSequenceReader(file.getPath())) { + TrackingChunkLoader loader = new TrackingChunkLoader(sequence); + FileSeriesReader reader = + new FileSeriesReader( + loader, + new ArrayList<>(sequence.getChunkMetadataList(new Path("d1", "s1", true))), + null); + try { + int count = 0; + while (reader.hasNextBatch()) { + assertEquals(++count * 10L, reader.nextBatch().currentValue()); + } + assertEquals(3, count); + assertEquals(3, loader.loaded.size()); + assertAllDestroyed(loader.loaded); + assertEquals(createCount + 1, TestAeadEncryptionProvider.getCreateCount()); + } finally { + reader.close(); + } + } + assertEquals(closeCount + 1, TestAeadEncryptionProvider.getCloseCount()); + } + + @Test + public void testTimestampReaderReleasesChunksOnSwitchAndEarlyClose() throws Exception { + File file = writeFile(false); + try (TsFileSequenceReader sequence = new TsFileSequenceReader(file.getPath())) { + TrackingChunkLoader loader = new TrackingChunkLoader(sequence); + try (FileSeriesReaderByTimestamp reader = + new FileSeriesReaderByTimestamp( + loader, new ArrayList<>(sequence.getChunkMetadataList(new Path("d1", "s1", true))))) { + assertEquals(10L, reader.getValueInTimestamp(1)); + assertFalse(loader.loaded.get(0).getEncryptParam().isDestroyed()); + assertEquals(20L, reader.getValueInTimestamp(2)); + assertTrue(loader.loaded.get(0).getEncryptParam().isDestroyed()); + } + assertAllDestroyed(loader.loaded); + } + } + + @Test + public void testTimestampReaderPreservesLoadFailureAndClosesRemainingChunks() throws Exception { + IChunkLoader loader = mock(IChunkLoader.class); + AlignedChunkMetadata aligned = mock(AlignedChunkMetadata.class); + ChunkMetadata timeMetadata = mock(ChunkMetadata.class); + ChunkMetadata firstMetadata = mock(ChunkMetadata.class); + ChunkMetadata secondMetadata = mock(ChunkMetadata.class); + Chunk timeChunk = mock(Chunk.class); + Chunk firstChunk = mock(Chunk.class); + IOException loadFailure = new IOException("load failed"); + IllegalStateException closeFailure = new IllegalStateException("close failed"); + when(aligned.getEndTime()).thenReturn(Long.MAX_VALUE); + when(aligned.getTimeChunkMetadata()).thenReturn(timeMetadata); + when(aligned.getValueChunkMetadataList()) + .thenReturn(Arrays.asList(firstMetadata, secondMetadata)); + when(loader.loadChunk(timeMetadata)).thenReturn(timeChunk); + when(loader.loadChunk(firstMetadata)).thenReturn(firstChunk); + when(loader.loadChunk(secondMetadata)).thenThrow(loadFailure); + doThrow(closeFailure).when(timeChunk).close(); + try (FileSeriesReaderByTimestamp reader = + new FileSeriesReaderByTimestamp(loader, Arrays.asList(aligned))) { + IOException actual = assertThrows(IOException.class, () -> reader.getValueInTimestamp(1)); + assertSame(loadFailure, actual); + assertEquals(1, actual.getSuppressed().length); + assertSame(closeFailure, actual.getSuppressed()[0]); + verify(timeChunk).close(); + verify(firstChunk).close(); + } + } + + @Test + public void testCachedChunkReaderSurvivesFileCloseAndReleasesProvider() throws Exception { + File file = writeFile(false); + int closeCount = TestAeadEncryptionProvider.getCloseCount(); + IChunkReader chunkReader; + try (TsFileSequenceReader sequence = new TsFileSequenceReader(file.getPath())) { + CachedChunkLoaderImpl loader = new CachedChunkLoaderImpl(sequence); + chunkReader = + loader.getChunkReader( + sequence.getChunkMetadataList(new Path("d1", "s1", true)).get(0), null); + } + try { + assertEquals(10L, chunkReader.nextPageData().currentValue()); + assertEquals(closeCount, TestAeadEncryptionProvider.getCloseCount()); + } finally { + chunkReader.close(); + } + assertEquals(closeCount + 1, TestAeadEncryptionProvider.getCloseCount()); + } + + @Test + public void testAlignedReaderClosesTimeAndAllValueContexts() throws Exception { + File file = writeFile(true); + int closeCount = TestAeadEncryptionProvider.getCloseCount(); + Chunk time; + List values = new ArrayList<>(); + AlignedChunkReader reader; + try (TsFileSequenceReader sequence = new TsFileSequenceReader(file.getPath())) { + Map> metadata = + sequence.readChunkMetadataInDevice(IDeviceID.Factory.DEFAULT_FACTORY.create("d1")); + time = sequence.readMemChunk(metadata.get("").get(0)); + values.add(sequence.readMemChunk(metadata.get("s1").get(0))); + values.add(sequence.readMemChunk(metadata.get("s2").get(0))); + values.add(null); + reader = new AlignedChunkReader(time, values); + } + try { + assertTrue(reader.hasNextSatisfiedPage()); + reader.nextPageData(); + } finally { + reader.close(); + reader.close(); + } + assertTrue(time.getEncryptParam().isDestroyed()); + assertAllDestroyed(values); + assertEquals(closeCount + 1, TestAeadEncryptionProvider.getCloseCount()); + } + + @Test + public void testMalformedChunkReaderConstructionReleasesContext() throws Exception { + File file = writeFile(false); + try (TsFileSequenceReader sequence = new TsFileSequenceReader(file.getPath()); + Chunk chunk = + sequence.readMemChunk( + sequence.getChunkMetadataList(new Path("d1", "s1", true)).get(0))) { + chunk.getData().limit(2); + assertThrows(RuntimeException.class, () -> new ChunkReader(chunk)); + assertTrue(chunk.getEncryptParam().isDestroyed()); + } + } + + @Test + public void testMalformedAlignedReaderConstructionReleasesAllContexts() throws Exception { + File file = writeFile(true); + try (TsFileSequenceReader sequence = new TsFileSequenceReader(file.getPath())) { + Map> metadata = + sequence.readChunkMetadataInDevice(IDeviceID.Factory.DEFAULT_FACTORY.create("d1")); + try (Chunk time = sequence.readMemChunk(metadata.get("").get(0)); + Chunk first = sequence.readMemChunk(metadata.get("s1").get(0)); + Chunk second = sequence.readMemChunk(metadata.get("s2").get(0))) { + time.getData().limit(2); + assertThrows( + IOException.class, + () -> new AlignedChunkReader(time, Arrays.asList(first, second, null))); + assertAllDestroyed(Arrays.asList(time, first, second)); + } + } + } + + @Test + public void testTreeResultSetEarlyCloseReleasesQueryWithoutClosingFile() throws Exception { + File file = writeFile(false); + int closeCount = TestAeadEncryptionProvider.getCloseCount(); + try (TsFileReader reader = new TsFileReader(file)) { + QueryDataSet query = reader.query(queryExpression()); + TreeResultSet result = new TreeResultSet(query, Arrays.asList("d1"), Arrays.asList("s1")); + assertTrue(result.next()); + result.close(); + result.close(); + assertTrue(query.isClosed()); + assertFalse(query.hasNext()); + try (QueryDataSet another = reader.query(queryExpression())) { + assertTrue(another.hasNext()); + assertEquals(10L, another.next().getFields().get(0).getLongV()); + } + } + assertEquals(closeCount + 1, TestAeadEncryptionProvider.getCloseCount()); + } + + @Test + public void testFileReaderClosesAnAbandonedQuery() throws Exception { + File file = writeFile(false); + int closeCount = TestAeadEncryptionProvider.getCloseCount(); + QueryDataSet query; + try (TsFileReader reader = new TsFileReader(file)) { + query = reader.query(queryExpression()); + assertTrue(query.hasNext()); + } + assertTrue(query.isClosed()); + assertEquals(closeCount + 1, TestAeadEncryptionProvider.getCloseCount()); + } + + @Test + public void testQueryLimitClosesRemainingChunkReferences() throws Exception { + File file = writeFile(false); + int closeCount = TestAeadEncryptionProvider.getCloseCount(); + try (TsFileReader reader = new TsFileReader(file)) { + QueryDataSet query = reader.query(queryExpression()); + query.setRowLimit(1); + assertTrue(query.hasNext()); + assertEquals(10L, query.next().getFields().get(0).getLongV()); + assertFalse(query.hasNext()); + assertTrue(query.isClosed()); + } + assertEquals(closeCount + 1, TestAeadEncryptionProvider.getCloseCount()); + } + + @Test + public void testAlignedIteratorFailureClosesPreviouslyLoadedChunks() throws Exception { + File file = writeFile(true); + int closeCount = TestAeadEncryptionProvider.getCloseCount(); + try (FailingSequenceReader sequence = new FailingSequenceReader(file)) { + sequence.getEncryptParam().getPageBodyOverhead(); + Map> metadata = + sequence.readChunkMetadataInDevice(IDeviceID.Factory.DEFAULT_FACTORY.create("d1")); + AlignedChunkMetadata aligned = + new AlignedChunkMetadata( + metadata.get("").get(0), + Arrays.asList(metadata.get("s1").get(0), metadata.get("s2").get(0))); + TsFileAlignedSeriesReaderIterator iterator = + new TsFileAlignedSeriesReaderIterator( + sequence, + Arrays.asList(aligned), + Arrays.asList( + new MeasurementSchema("s1", TSDataType.INT64, TSEncoding.PLAIN), + new MeasurementSchema("s2", TSDataType.INT64, TSEncoding.PLAIN))); + assertThrows(IOException.class, iterator::nextReader); + assertEquals(2, sequence.loaded.size()); + assertAllDestroyed(sequence.loaded); + } + assertEquals(closeCount + 1, TestAeadEncryptionProvider.getCloseCount()); + } + + @Test + public void testAlignedRewriteReleasesTemporaryReferencesAndKeepsSources() throws Exception { + File file = writeFile(true); + int closeCount = TestAeadEncryptionProvider.getCloseCount(); + try (TsFileSequenceReader sequence = new TsFileSequenceReader(file.getPath())) { + Map> metadata = + sequence.readChunkMetadataInDevice(IDeviceID.Factory.DEFAULT_FACTORY.create("d1")); + try (Chunk time = sequence.readMemChunk(metadata.get("").get(0)); + Chunk value = sequence.readMemChunk(metadata.get("s1").get(0)); + Chunk rewritten = value.rewrite(TSDataType.DOUBLE, time)) { + assertFalse(time.getEncryptParam().isDestroyed()); + assertFalse(value.getEncryptParam().isDestroyed()); + AlignedChunkReader chunkReader = new AlignedChunkReader(time, Arrays.asList(rewritten)); + TsPrimitiveType[] values = (TsPrimitiveType[]) chunkReader.nextPageData().currentValue(); + assertEquals(10.0, values[0].getDouble(), 0.0); + } + } + assertEquals(closeCount + 1, TestAeadEncryptionProvider.getCloseCount()); + } + + @Test + public void testMetadataOffsetConstructorsRemainCompatibleAndLoadLazily() throws Exception { + File file = writeFile(false); + try (TsFileSequenceReader original = new TsFileSequenceReader(file.getPath())) { + ChunkMetadata metadata = original.getChunkMetadataList(new Path("d1", "s1", true)).get(0); + for (boolean withParameter : new boolean[] {false, true}) { + BufferedTsFileInput input = new BufferedTsFileInput(file.toPath()); + try (TsFileSequenceReader reader = + metadataOffsetReader( + input, + original.getFileMetadataPos(), + original.getTsFileMetadataSize(), + withParameter)) { + reader.position(metadata.getOffsetOfChunkHeader()); + ChunkHeader header = reader.readChunkHeader(reader.readMarker()); + assertTrue(header.getChunkOrdinal() >= 0); + try (Chunk chunk = reader.readMemChunk(metadata)) { + assertEquals(10L, new ChunkReader(chunk).nextPageData().currentValue()); + } + } + } + } + } + + // Deliberately has no throws clause: callers of the old public constructors must still compile. + private static TsFileSequenceReader metadataOffsetReader( + BufferedTsFileInput input, long position, int size, boolean withParameter) { + return withParameter + ? new TsFileSequenceReader(input, position, size, null) + : new TsFileSequenceReader(input, position, size); + } + + private File writeFile(boolean aligned) throws Exception { + return writeFile(aligned, false); + } + + private File writeFile(boolean aligned, boolean padded) throws Exception { + File file = + new File(temporaryFolder.getRoot(), "lifecycle-" + aligned + "-" + padded + ".tsfile"); + try (EncryptParameter parameter = + padded + ? TestAeadEncryptionProvider.createPaddedParameter(new byte[16], new byte[16]) + : TestAeadEncryptionProvider.createParameter(new byte[16], new byte[16]); + TsFileWriter writer = new TsFileWriter(file, parameter)) { + MeasurementSchema first = new MeasurementSchema("s1", TSDataType.INT64, TSEncoding.PLAIN); + if (aligned) { + writer.registerAlignedTimeseries( + new Path("d1"), + Arrays.asList(first, new MeasurementSchema("s2", TSDataType.INT64, TSEncoding.PLAIN))); + } else { + writer.registerTimeseries(new Path("d1"), first); + } + for (int i = 1; i <= 3; i++) { + TSRecord record = new TSRecord("d1", i).addTuple(new LongDataPoint("s1", i * 10L)); + if (aligned) { + record.addTuple(new LongDataPoint("s2", i * 100L)); + } + writer.writeRecord(record); + writer.flush(); + } + } + return file; + } + + private static QueryExpression queryExpression() { + return QueryExpression.create(Arrays.asList(new Path("d1", "s1", true)), null); + } + + private static void assertAllDestroyed(List chunks) { + for (Chunk chunk : chunks) { + if (chunk != null) { + assertTrue(chunk.getEncryptParam().isDestroyed()); + } + } + } + + private static class TrackingChunkLoader extends CachedChunkLoaderImpl { + private final List loaded = new ArrayList<>(); + + private TrackingChunkLoader(TsFileSequenceReader reader) { + super(reader); + } + + @Override + public Chunk loadChunk(ChunkMetadata metadata) throws IOException { + Chunk chunk = super.loadChunk(metadata); + loaded.add(chunk); + return chunk; + } + } + + private static class FailingSequenceReader extends TsFileSequenceReader { + private final List loaded = new ArrayList<>(); + + private FailingSequenceReader(File file) throws IOException { + super(file.getPath()); + } + + @Override + public Chunk readMemChunk(ChunkMetadata metadata) throws IOException { + if (loaded.size() == 2) { + throw new IOException("Injected chunk read failure"); + } + Chunk chunk = super.readMemChunk(metadata); + loaded.add(chunk); + return chunk; + } + } +} diff --git a/java/tsfile/src/test/java/org/apache/tsfile/read/TsFileSequenceReaderHeaderTest.java b/java/tsfile/src/test/java/org/apache/tsfile/read/TsFileSequenceReaderHeaderTest.java new file mode 100644 index 0000000000..0530c316a1 --- /dev/null +++ b/java/tsfile/src/test/java/org/apache/tsfile/read/TsFileSequenceReaderHeaderTest.java @@ -0,0 +1,164 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ +package org.apache.tsfile.read; + +import org.apache.tsfile.common.conf.TSFileConfig; +import org.apache.tsfile.encrypt.EncryptParameter; +import org.apache.tsfile.encrypt.TestAeadEncryptionProvider; +import org.apache.tsfile.enums.TSDataType; +import org.apache.tsfile.file.MetaMarker; +import org.apache.tsfile.file.header.ChunkHeader; +import org.apache.tsfile.file.header.FileEncryptionHeader; +import org.apache.tsfile.file.metadata.enums.CompressionType; +import org.apache.tsfile.file.metadata.enums.TSEncoding; +import org.apache.tsfile.read.reader.BufferedTsFileInput; + +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.TemporaryFolder; + +import java.io.ByteArrayOutputStream; +import java.io.File; +import java.nio.ByteBuffer; +import java.nio.file.Files; +import java.util.concurrent.atomic.AtomicLong; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; + +public class TsFileSequenceReaderHeaderTest { + + @Rule public TemporaryFolder folder = new TemporaryFolder(); + + @Test + public void testPlainMarkerIsReadWithVersion() throws Exception { + ByteArrayOutputStream bytes = baseHeader(); + bytes.write(MetaMarker.SEPARATOR); + BufferedTsFileInput input = input(bytes); + AtomicLong ioBytes = new AtomicLong(); + try (TsFileSequenceReader reader = + new TsFileSequenceReader(input, false, false, ioBytes::addAndGet)) { + assertFalse(reader.hasFileEncryptionHeader()); + assertEquals(7, reader.getDataStartOffset()); + assertEquals(0, reader.position()); + assertEquals(2, ioBytes.get()); + verify(input, times(1)).read(any(ByteBuffer.class)); + verify(input, never()).position(7); + verify(input, never()).wrapAsInputStream(); + } + } + + @Test + public void testEncryptedMarkerIsReadWithVersion() throws Exception { + ByteArrayOutputStream bytes = baseHeader(); + try (EncryptParameter parameter = + TestAeadEncryptionProvider.createParameter(new byte[16], new byte[16])) { + FileEncryptionHeader.serialize(parameter, bytes); + } + BufferedTsFileInput input = input(bytes); + AtomicLong ioBytes = new AtomicLong(); + try (TsFileSequenceReader reader = + new TsFileSequenceReader(input, false, false, ioBytes::addAndGet)) { + assertTrue(reader.hasFileEncryptionHeader()); + assertEquals(bytes.size(), reader.getDataStartOffset()); + assertEquals(0, reader.position()); + assertEquals(bytes.size() - 6, ioBytes.get()); + verify(input, never()).position(7); + } + } + + @Test + public void testBaseHeaderWithoutMarker() throws Exception { + AtomicLong ioBytes = new AtomicLong(); + try (TsFileSequenceReader reader = + new TsFileSequenceReader(input(baseHeader()), false, false, ioBytes::addAndGet)) { + assertFalse(reader.hasFileEncryptionHeader()); + assertEquals(7, reader.getDataStartOffset()); + assertEquals(1, ioBytes.get()); + } + } + + @Test + public void testShortReadsStillPrefetchMarker() throws Exception { + ByteArrayOutputStream bytes = baseHeader(); + bytes.write(MetaMarker.SEPARATOR); + BufferedTsFileInput input = input(bytes); + doAnswer( + invocation -> { + ByteBuffer buffer = invocation.getArgument(0); + int limit = buffer.limit(); + buffer.limit(buffer.position() + 1); + try { + return invocation.callRealMethod(); + } finally { + buffer.limit(limit); + } + }) + .when(input) + .read(any(ByteBuffer.class)); + try (TsFileSequenceReader reader = new TsFileSequenceReader(input, false, false, null)) { + assertFalse(reader.hasFileEncryptionHeader()); + verify(input, times(2)).read(any(ByteBuffer.class)); + verify(input, never()).wrapAsInputStream(); + } + } + + @Test + public void testLazyHeaderLoadingPreservesPosition() throws Exception { + ByteArrayOutputStream bytes = baseHeader(); + bytes.write(MetaMarker.SEPARATOR); + BufferedTsFileInput input = input(bytes); + input.position(8); + try (TsFileSequenceReader reader = new TsFileSequenceReader(input, 8L, 0)) { + assertFalse(reader.hasFileEncryptionHeader()); + assertEquals(8, reader.position()); + } + } + + @Test + public void testChunkHeaderIncludesOrdinalInDescription() { + ChunkHeader header = + new ChunkHeader( + "s1", 10, TSDataType.INT64, CompressionType.LZ4, TSEncoding.PLAIN, 1, 0, 42); + assertTrue(header.toString().contains("chunkOrdinal=42")); + ChunkHeader plainHeader = + new ChunkHeader("s1", 10, TSDataType.INT64, CompressionType.LZ4, TSEncoding.PLAIN, 1); + assertTrue(plainHeader.toString().contains("chunkOrdinal=-1")); + } + + private ByteArrayOutputStream baseHeader() throws Exception { + ByteArrayOutputStream bytes = new ByteArrayOutputStream(); + bytes.write(TSFileConfig.MAGIC_STRING.getBytes(TSFileConfig.STRING_CHARSET)); + bytes.write(TSFileConfig.VERSION_NUMBER); + return bytes; + } + + private BufferedTsFileInput input(ByteArrayOutputStream bytes) throws Exception { + File file = folder.newFile(); + Files.write(file.toPath(), bytes.toByteArray()); + return spy(new BufferedTsFileInput(file.toPath())); + } +} diff --git a/java/tsfile/src/test/java/org/apache/tsfile/read/UnClosedTsFileReaderTest.java b/java/tsfile/src/test/java/org/apache/tsfile/read/UnClosedTsFileReaderTest.java index 8e865413a1..5fadd06dca 100644 --- a/java/tsfile/src/test/java/org/apache/tsfile/read/UnClosedTsFileReaderTest.java +++ b/java/tsfile/src/test/java/org/apache/tsfile/read/UnClosedTsFileReaderTest.java @@ -20,6 +20,9 @@ package org.apache.tsfile.read; import org.apache.tsfile.constant.TestConstant; +import org.apache.tsfile.encrypt.EncryptParameter; +import org.apache.tsfile.encrypt.EncryptionProviderRegistry; +import org.apache.tsfile.encrypt.TestAeadEncryptionProvider; import org.apache.tsfile.enums.TSDataType; import org.apache.tsfile.file.metadata.ChunkMetadata; import org.apache.tsfile.file.metadata.IDeviceID.Factory; @@ -30,6 +33,8 @@ import org.apache.tsfile.write.schema.MeasurementSchema; import org.apache.tsfile.write.writer.TsFileIOWriter; +import org.junit.AfterClass; +import org.junit.BeforeClass; import org.junit.Test; import java.io.File; @@ -41,6 +46,16 @@ public class UnClosedTsFileReaderTest { + @BeforeClass + public static void setUpEncryptionProvider() { + EncryptionProviderRegistry.registerProvider(TestAeadEncryptionProvider.INSTANCE); + } + + @AfterClass + public static void tearDownEncryptionProvider() { + EncryptionProviderRegistry.unregisterProvider(TestAeadEncryptionProvider.PROVIDER_ID); + } + @Test public void testRead() throws IOException { File file = new File(TestConstant.BASE_OUTPUT_PATH + File.separator + "test.tsfile"); @@ -71,4 +86,47 @@ public void testRead() throws IOException { file.delete(); } + + @Test + public void testReadEncryptedPagesUsingFileHeader() throws IOException { + File file = new File(TestConstant.BASE_OUTPUT_PATH + File.separator + "test-encrypted.tsfile"); + byte[] fileCryptoId = new byte[EncryptParameter.FILE_CRYPTO_ID_LENGTH]; + fileCryptoId[0] = 1; + EncryptParameter encryptParameter = + TestAeadEncryptionProvider.createParameter(new byte[16], fileCryptoId); + + try { + TsFileIOWriter writer = new TsFileIOWriter(file, encryptParameter); + writer.startChunkGroup(Factory.DEFAULT_FACTORY.create("root.sg1.d1")); + ChunkWriterImpl chunkWriter = + new ChunkWriterImpl(new MeasurementSchema("s1", TSDataType.INT64), encryptParameter); + chunkWriter.write(1, 1L); + chunkWriter.sealCurrentPage(); + chunkWriter.write(2, 2L); + chunkWriter.sealCurrentPage(); + chunkWriter.writeToFileWriter(writer); + writer.endChunkGroup(); + writer.close(); + ChunkMetadata chunkMetadata = + writer.getChunkGroupMetadataList().get(0).getChunkMetadataList().get(0); + + try (UnClosedTsFileReader reader = new UnClosedTsFileReader(file.getAbsolutePath(), null)) { + assertEquals("test-key", reader.getEncryptParam().getKeyId()); + ChunkReader chunkReader = new ChunkReader(reader.readMemChunk(chunkMetadata)); + + BatchData firstPage = chunkReader.nextPageData(); + assertTrue(firstPage.hasCurrent()); + assertEquals(1, firstPage.currentTime()); + assertEquals(1L, firstPage.currentValue()); + + BatchData secondPage = chunkReader.nextPageData(); + assertTrue(secondPage.hasCurrent()); + assertEquals(2, secondPage.currentTime()); + assertEquals(2L, secondPage.currentValue()); + } + } finally { + encryptParameter.close(); + file.delete(); + } + } } diff --git a/java/tsfile/src/test/java/org/apache/tsfile/read/reader/block/PageAeadBlockReaderLifecycleTest.java b/java/tsfile/src/test/java/org/apache/tsfile/read/reader/block/PageAeadBlockReaderLifecycleTest.java new file mode 100644 index 0000000000..d02fef0c60 --- /dev/null +++ b/java/tsfile/src/test/java/org/apache/tsfile/read/reader/block/PageAeadBlockReaderLifecycleTest.java @@ -0,0 +1,387 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ +package org.apache.tsfile.read.reader.block; + +import org.apache.tsfile.encrypt.EncryptParameter; +import org.apache.tsfile.encrypt.EncryptionProviderRegistry; +import org.apache.tsfile.encrypt.TestAeadEncryptionProvider; +import org.apache.tsfile.enums.ColumnCategory; +import org.apache.tsfile.enums.TSDataType; +import org.apache.tsfile.exception.encrypt.EncryptException; +import org.apache.tsfile.file.metadata.ChunkMetadata; +import org.apache.tsfile.file.metadata.IChunkMetadata; +import org.apache.tsfile.file.metadata.IDeviceID; +import org.apache.tsfile.file.metadata.TableSchema; +import org.apache.tsfile.file.metadata.enums.TSEncoding; +import org.apache.tsfile.read.TsFileSequenceReader; +import org.apache.tsfile.read.common.Chunk; +import org.apache.tsfile.read.common.Path; +import org.apache.tsfile.read.controller.CachedChunkLoaderImpl; +import org.apache.tsfile.read.controller.IMetadataQuerier; +import org.apache.tsfile.read.controller.MetadataQuerierByFileImpl; +import org.apache.tsfile.read.query.dataset.TableResultSet; +import org.apache.tsfile.read.query.executor.TableQueryExecutor; +import org.apache.tsfile.read.query.executor.TableQueryExecutor.ColumnMapping; +import org.apache.tsfile.read.query.executor.TableQueryExecutor.TableQueryOrdering; +import org.apache.tsfile.read.query.executor.task.DeviceQueryTask; +import org.apache.tsfile.write.TsFileWriter; +import org.apache.tsfile.write.record.TSRecord; +import org.apache.tsfile.write.record.Tablet; +import org.apache.tsfile.write.record.datapoint.LongDataPoint; +import org.apache.tsfile.write.schema.MeasurementSchema; + +import org.junit.AfterClass; +import org.junit.BeforeClass; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.TemporaryFolder; + +import java.io.File; +import java.io.IOException; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertSame; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.anySet; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.ArgumentMatchers.isNull; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.when; + +public class PageAeadBlockReaderLifecycleTest { + private static final List TABLE_COLUMNS = Arrays.asList("device_id", "s1", "s2"); + private static final List TABLE_TYPES = + Arrays.asList(TSDataType.STRING, TSDataType.INT64, TSDataType.INT64); + private static final List TABLE_CATEGORIES = + Arrays.asList(ColumnCategory.TAG, ColumnCategory.FIELD, ColumnCategory.FIELD); + private static final IDeviceID TREE_DEVICE = IDeviceID.Factory.DEFAULT_FACTORY.create("d1"); + + @Rule public final TemporaryFolder temporaryFolder = new TemporaryFolder(); + + @BeforeClass + public static void registerProvider() { + EncryptionProviderRegistry.registerProvider(TestAeadEncryptionProvider.INSTANCE); + } + + @AfterClass + public static void unregisterProvider() { + EncryptionProviderRegistry.unregisterProvider(TestAeadEncryptionProvider.PROVIDER_ID); + } + + @Test + public void testPublicTableQueryPropagatesAuthenticationFailure() throws Exception { + File file = writeTableFile(1); + try (TsFileSequenceReader sequence = new TsFileSequenceReader(file.getPath())) { + TrackingChunkLoader loader = new TrackingChunkLoader(sequence); + loader.tamperTimeColumn = true; + TableQueryExecutor executor = + new TableQueryExecutor( + new MetadataQuerierByFileImpl(sequence), loader, TableQueryOrdering.DEVICE); + try (TableResultSet result = tableResult(executor)) { + Exception failure = assertThrows(Exception.class, result::next); + assertTrue(EncryptException.isEncryptionFailure(failure)); + } + assertAllDestroyed(loader.loaded); + } + } + + @Test + public void testPublicTableQueryPropagatesLaterValuePageFailure() throws Exception { + File file = writeTableFile(1); + try (TsFileSequenceReader sequence = new TsFileSequenceReader(file.getPath())) { + TrackingChunkLoader loader = new TrackingChunkLoader(sequence); + loader.tamperSecondColumn = true; + TableQueryExecutor executor = + new TableQueryExecutor( + new MetadataQuerierByFileImpl(sequence), loader, TableQueryOrdering.DEVICE); + try (TableResultSet result = tableResult(executor)) { + Exception failure = + assertThrows( + Exception.class, + () -> { + while (result.next()) { + result.getLong("s2"); + } + }); + assertTrue(EncryptException.isEncryptionFailure(failure)); + } + assertAllDestroyed(loader.loaded); + } + } + + @Test + public void testTableResultEarlyCloseReleasesContextsAndKeepsSharedLoaderOpen() throws Exception { + File file = writeTableFile(1); + int closeCount = TestAeadEncryptionProvider.getCloseCount(); + try (TsFileSequenceReader sequence = new TsFileSequenceReader(file.getPath())) { + TrackingChunkLoader loader = new TrackingChunkLoader(sequence); + TableQueryExecutor executor = + new TableQueryExecutor( + new MetadataQuerierByFileImpl(sequence), loader, TableQueryOrdering.DEVICE); + try (TableResultSet result = tableResult(executor)) { + assertTrue(result.next()); + assertEquals(1L, result.getLong("s1")); + assertTrue(loader.loaded.stream().anyMatch(c -> !c.getEncryptParam().isDestroyed())); + result.close(); + result.close(); + assertAllDestroyed(loader.loaded); + assertEquals(closeCount, TestAeadEncryptionProvider.getCloseCount()); + } + // A result only owns its chunks, so another query can still use the same file loader. + try (TableResultSet another = tableResult(executor)) { + assertTrue(another.next()); + assertEquals(1L, another.getLong("s1")); + } + assertAllDestroyed(loader.loaded); + } + assertEquals(closeCount + 1, TestAeadEncryptionProvider.getCloseCount()); + } + + @Test + public void testDeviceSwitchAndExhaustionReleaseAllContexts() throws Exception { + File file = writeTableFile(2); + int closeCount = TestAeadEncryptionProvider.getCloseCount(); + try (TsFileSequenceReader sequence = new TsFileSequenceReader(file.getPath())) { + TrackingChunkLoader loader = new TrackingChunkLoader(sequence); + TableQueryExecutor executor = + new TableQueryExecutor( + new MetadataQuerierByFileImpl(sequence), loader, TableQueryOrdering.DEVICE); + try (TableResultSet result = tableResult(executor)) { + int rows = 0; + while (result.next()) { + rows++; + } + assertEquals(4096, rows); + assertAllDestroyed(loader.loaded); + } + } + assertEquals(closeCount + 1, TestAeadEncryptionProvider.getCloseCount()); + } + + @Test + public void testSingleDeviceEarlyCloseReleasesEveryScalarColumn() throws Exception { + File file = writeScalarFile(); + int closeCount = TestAeadEncryptionProvider.getCloseCount(); + try (TsFileSequenceReader sequence = new TsFileSequenceReader(file.getPath())) { + TrackingChunkLoader loader = new TrackingChunkLoader(sequence); + IMetadataQuerier metadata = scalarMetadata(sequence); + DeviceQueryTask task = scalarTask(metadata, loader); + try (SingleDeviceTsBlockReader reader = + new SingleDeviceTsBlockReader(task, metadata, loader, 1, null, null)) { + assertTrue(reader.hasNext()); + assertEquals(1, reader.next().getPositionCount()); + reader.close(); + reader.close(); + assertFalse(reader.hasNext()); + assertEquals(2, loader.loaded.size()); + assertAllDestroyed(loader.loaded); + assertTrue(sequence.fileSize() > 0); + } + } + assertEquals(closeCount + 1, TestAeadEncryptionProvider.getCloseCount()); + } + + @Test + public void testLaterColumnLoadFailureReleasesPreviouslyInitializedColumn() throws Exception { + File file = writeScalarFile(); + int closeCount = TestAeadEncryptionProvider.getCloseCount(); + try (TsFileSequenceReader sequence = new TsFileSequenceReader(file.getPath())) { + TrackingChunkLoader loader = new TrackingChunkLoader(sequence); + IOException failure = new IOException("Injected column load failure"); + loader.loadFailure = failure; + IMetadataQuerier metadata = scalarMetadata(sequence); + DeviceQueryTask task = scalarTask(metadata, loader); + assertSame( + failure, + assertThrows( + IOException.class, + () -> new SingleDeviceTsBlockReader(task, metadata, loader, 1, null, null))); + assertEquals(1, loader.loaded.size()); + assertAllDestroyed(loader.loaded); + assertTrue(sequence.fileSize() > 0); + } + assertEquals(closeCount + 1, TestAeadEncryptionProvider.getCloseCount()); + } + + @Test + public void testLaterColumnDecodeFailureAlsoReleasesTheUnregisteredReader() throws Exception { + File file = writeScalarFile(); + int closeCount = TestAeadEncryptionProvider.getCloseCount(); + try (TsFileSequenceReader sequence = new TsFileSequenceReader(file.getPath())) { + TrackingChunkLoader loader = new TrackingChunkLoader(sequence); + loader.tamperSecondColumn = true; + IMetadataQuerier metadata = scalarMetadata(sequence); + DeviceQueryTask task = scalarTask(metadata, loader); + assertThrows( + EncryptException.class, + () -> new SingleDeviceTsBlockReader(task, metadata, loader, 1, null, null)); + assertEquals(2, loader.loaded.size()); + assertAllDestroyed(loader.loaded); + assertTrue(sequence.fileSize() > 0); + } + assertEquals(closeCount + 1, TestAeadEncryptionProvider.getCloseCount()); + } + + @Test + public void testRuntimeColumnInitializationFailurePreservesCauseAndReleasesReaders() + throws Exception { + File file = writeScalarFile(); + int closeCount = TestAeadEncryptionProvider.getCloseCount(); + try (TsFileSequenceReader sequence = new TsFileSequenceReader(file.getPath())) { + TrackingChunkLoader loader = new TrackingChunkLoader(sequence); + IMetadataQuerier metadata = scalarMetadata(sequence); + DeviceQueryTask original = scalarTask(metadata, loader); + ColumnMapping mapping = spy(original.getColumnMapping()); + IllegalStateException failure = new IllegalStateException("Injected column mapping failure"); + doThrow(failure).when(mapping).getColumnPos("s2"); + DeviceQueryTask task = + new DeviceQueryTask( + TREE_DEVICE, original.getColumnNames(), mapping, null, original.getTableSchema()); + assertSame( + failure, + assertThrows( + IllegalStateException.class, + () -> new SingleDeviceTsBlockReader(task, metadata, loader, 1, null, null))); + assertEquals(2, loader.loaded.size()); + assertAllDestroyed(loader.loaded); + } + assertEquals(closeCount + 1, TestAeadEncryptionProvider.getCloseCount()); + } + + private TableResultSet tableResult(TableQueryExecutor executor) throws Exception { + return new TableResultSet( + executor.query("t1", TABLE_COLUMNS, null, null, null), TABLE_COLUMNS, TABLE_TYPES, "t1"); + } + + private File writeTableFile(int devices) throws Exception { + File file = new File(temporaryFolder.getRoot(), "table-lifecycle.tsfile"); + try (EncryptParameter parameter = + TestAeadEncryptionProvider.createParameter( + new byte[16], new byte[EncryptParameter.FILE_CRYPTO_ID_LENGTH]); + TsFileWriter writer = new TsFileWriter(file, parameter)) { + writer.setGenerateTableSchema(true); + writer.registerTableSchema( + new TableSchema( + "t1", + Arrays.asList( + new MeasurementSchema("device_id", TSDataType.STRING), + new MeasurementSchema("s1", TSDataType.INT64), + new MeasurementSchema("s2", TSDataType.INT64)), + TABLE_CATEGORIES)); + for (int device = 0; device < devices; device++) { + Tablet tablet = new Tablet("t1", TABLE_COLUMNS, TABLE_TYPES, TABLE_CATEGORIES, 2048); + for (int row = 0; row < 2048; row++) { + tablet.addTimestamp(row, row + 1); + tablet.addValue("device_id", row, "device-" + device); + tablet.addValue("s1", row, (long) row + 1); + tablet.addValue("s2", row, (long) row + 101); + } + writer.writeTable(tablet); + } + } + return file; + } + + private File writeScalarFile() throws Exception { + File file = new File(temporaryFolder.getRoot(), "scalar-columns.tsfile"); + try (EncryptParameter parameter = + TestAeadEncryptionProvider.createParameter( + new byte[16], new byte[EncryptParameter.FILE_CRYPTO_ID_LENGTH]); + TsFileWriter writer = new TsFileWriter(file, parameter)) { + writer.registerTimeseries( + new Path("d1"), new MeasurementSchema("s1", TSDataType.INT64, TSEncoding.PLAIN)); + writer.registerTimeseries( + new Path("d1"), new MeasurementSchema("s2", TSDataType.INT64, TSEncoding.PLAIN)); + for (int row = 1; row <= 3; row++) { + writer.writeRecord( + new TSRecord("d1", row) + .addTuple(new LongDataPoint("s1", row)) + .addTuple(new LongDataPoint("s2", row + 100L))); + } + } + return file; + } + + private IMetadataQuerier scalarMetadata(TsFileSequenceReader sequence) throws IOException { + IMetadataQuerier metadata = mock(IMetadataQuerier.class); + List first = + new ArrayList<>(sequence.getChunkMetadataList(new Path("d1", "s1", true))); + List second = + new ArrayList<>(sequence.getChunkMetadataList(new Path("d1", "s2", true))); + when(metadata.getChunkMetadataLists(eq(TREE_DEVICE), anySet(), isNull())) + .thenReturn(Arrays.asList(first, second)); + return metadata; + } + + private DeviceQueryTask scalarTask(IMetadataQuerier metadata, TrackingChunkLoader loader) + throws Exception { + List columns = Arrays.asList("s1", "s2"); + TableSchema schema = + new TableSchema( + "t1", + Arrays.asList( + new MeasurementSchema("s1", TSDataType.INT64), + new MeasurementSchema("s2", TSDataType.INT64)), + Arrays.asList(ColumnCategory.FIELD, ColumnCategory.FIELD)); + ColumnMapping mapping = + new TableQueryExecutor(metadata, loader, TableQueryOrdering.DEVICE).new ColumnMapping(); + mapping.add("s1", 0, schema); + mapping.add("s2", 1, schema); + return new DeviceQueryTask(TREE_DEVICE, columns, mapping, null, schema); + } + + private static void assertAllDestroyed(List chunks) { + for (Chunk chunk : chunks) { + assertTrue(chunk.getEncryptParam().isDestroyed()); + } + } + + private static class TrackingChunkLoader extends CachedChunkLoaderImpl { + private final List loaded = new ArrayList<>(); + private IOException loadFailure; + private boolean tamperSecondColumn; + private boolean tamperTimeColumn; + + private TrackingChunkLoader(TsFileSequenceReader sequence) { + super(sequence); + } + + @Override + public Chunk loadChunk(ChunkMetadata metadata) throws IOException { + if (loadFailure != null && "s2".equals(metadata.getMeasurementUid())) { + throw loadFailure; + } + Chunk chunk = super.loadChunk(metadata); + if ((tamperSecondColumn && "s2".equals(metadata.getMeasurementUid())) + || (tamperTimeColumn && "".equals(metadata.getMeasurementUid()))) { + byte[] data = chunk.getData().array(); + data[data.length - 1] ^= 1; + } + loaded.add(chunk); + return chunk; + } + } +} diff --git a/java/tsfile/src/test/java/org/apache/tsfile/utils/ReadWriteIOUtilsTest.java b/java/tsfile/src/test/java/org/apache/tsfile/utils/ReadWriteIOUtilsTest.java index 4322903ce9..0d1bd36674 100644 --- a/java/tsfile/src/test/java/org/apache/tsfile/utils/ReadWriteIOUtilsTest.java +++ b/java/tsfile/src/test/java/org/apache/tsfile/utils/ReadWriteIOUtilsTest.java @@ -27,7 +27,9 @@ import java.io.ByteArrayInputStream; import java.io.ByteArrayOutputStream; import java.io.DataOutputStream; +import java.io.EOFException; import java.io.IOException; +import java.io.InputStream; import java.nio.BufferUnderflowException; import java.nio.ByteBuffer; import java.util.ArrayList; @@ -43,6 +45,35 @@ public class ReadWriteIOUtilsTest { protected static final int DEFAULT_BUFFER_SIZE = 4096; + @Test + public void testReadLongFromFragmentedStream() throws IOException { + long expected = 0x8123456789ABCDEFL; + byte[] encoded = ByteBuffer.allocate(Long.BYTES).putLong(expected).array(); + InputStream input = + new ByteArrayInputStream(encoded) { + private boolean zeroRead = true; + + @Override + public synchronized int read(byte[] bytes, int offset, int length) { + if (zeroRead) { + zeroRead = false; + return 0; + } + return super.read(bytes, offset, Math.min(length, 2)); + } + }; + Assert.assertEquals(expected, ReadWriteIOUtils.readLong(input)); + Assert.assertEquals(-1, input.read()); + } + + @Test + public void testReadLongRejectsEveryTruncatedLength() { + for (int length = 0; length < Long.BYTES; length++) { + InputStream input = new ByteArrayInputStream(new byte[length]); + Assert.assertThrows(EOFException.class, () -> ReadWriteIOUtils.readLong(input)); + } + } + @Test public void testReadStringRejectsOversizedLength() { for (int length : new int[] {2, Integer.MAX_VALUE}) { diff --git a/java/tsfile/src/test/java/org/apache/tsfile/write/ChunkPositionAuthenticationTest.java b/java/tsfile/src/test/java/org/apache/tsfile/write/ChunkPositionAuthenticationTest.java new file mode 100644 index 0000000000..80f4c6cfa9 --- /dev/null +++ b/java/tsfile/src/test/java/org/apache/tsfile/write/ChunkPositionAuthenticationTest.java @@ -0,0 +1,526 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.tsfile.write; + +import org.apache.tsfile.encrypt.ChunkCryptoContext; +import org.apache.tsfile.encrypt.EncryptParameter; +import org.apache.tsfile.encrypt.EncryptionProviderRegistry; +import org.apache.tsfile.encrypt.IEncrypt; +import org.apache.tsfile.encrypt.IEncryptProvider; +import org.apache.tsfile.encrypt.TestAeadEncryptionProvider; +import org.apache.tsfile.encrypt.UNENCRYPTED; +import org.apache.tsfile.enums.TSDataType; +import org.apache.tsfile.exception.encrypt.EncryptException; +import org.apache.tsfile.file.header.ChunkHeader; +import org.apache.tsfile.file.metadata.ChunkMetadata; +import org.apache.tsfile.file.metadata.IDeviceID; +import org.apache.tsfile.file.metadata.enums.CompressionType; +import org.apache.tsfile.file.metadata.enums.TSEncoding; +import org.apache.tsfile.file.metadata.statistics.Statistics; +import org.apache.tsfile.read.TsFileReader; +import org.apache.tsfile.read.TsFileSequenceReader; +import org.apache.tsfile.read.common.Chunk; +import org.apache.tsfile.read.common.Path; +import org.apache.tsfile.read.expression.QueryExpression; +import org.apache.tsfile.read.query.dataset.QueryDataSet; +import org.apache.tsfile.read.reader.LocalTsFileInput; +import org.apache.tsfile.utils.RamUsageEstimator; +import org.apache.tsfile.write.record.TSRecord; +import org.apache.tsfile.write.record.datapoint.LongDataPoint; +import org.apache.tsfile.write.schema.MeasurementSchema; +import org.apache.tsfile.write.writer.RestorableTsFileIOWriter; +import org.apache.tsfile.write.writer.TsFileIOWriter; + +import org.junit.AfterClass; +import org.junit.BeforeClass; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.TemporaryFolder; + +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.EOFException; +import java.io.File; +import java.io.IOException; +import java.nio.ByteBuffer; +import java.nio.file.Files; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; + +import static org.junit.Assert.assertArrayEquals; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; + +public class ChunkPositionAuthenticationTest { + @Rule public TemporaryFolder folder = new TemporaryFolder(); + + @BeforeClass + public static void registerProvider() { + EncryptionProviderRegistry.registerProvider(TestAeadEncryptionProvider.INSTANCE); + } + + @AfterClass + public static void unregisterProvider() { + EncryptionProviderRegistry.unregisterProvider(TestAeadEncryptionProvider.PROVIDER_ID); + } + + @Test + public void testProofBindsOffsetOrdinalAndFileIdentity() { + try (EncryptParameter parameter = parameter(); + EncryptParameter other = + TestAeadEncryptionProvider.createParameter( + new byte[16], new byte[] {1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0})) { + byte[] proof = parameter.authenticateChunk(100, 3); + assertEquals(parameter.getChunkAuthenticationSize(), proof.length); + parameter.verifyChunkAuthentication(proof, 100, 3); + assertThrows( + EncryptException.class, () -> parameter.verifyChunkAuthentication(proof, 101, 3)); + assertThrows( + EncryptException.class, () -> parameter.verifyChunkAuthentication(proof, 100, 4)); + assertThrows(EncryptException.class, () -> other.verifyChunkAuthentication(proof, 100, 3)); + byte[] corrupted = proof.clone(); + corrupted[0] ^= 1; + assertThrows( + EncryptException.class, () -> parameter.verifyChunkAuthentication(corrupted, 100, 3)); + corrupted[0] ^= 1; + corrupted[corrupted.length - 1] ^= 1; + assertThrows( + EncryptException.class, () -> parameter.verifyChunkAuthentication(corrupted, 100, 3)); + assertThrows( + EncryptException.class, () -> parameter.verifyChunkAuthentication(new byte[0], 100, 3)); + assertFalse(Arrays.equals(proof, parameter.authenticateChunk(100, 3))); + ChunkCryptoContext context = ChunkCryptoContext.forChunk(parameter, 100, 3); + byte[] aad = context.getAssociatedData(); + aad[0] ^= 1; + assertFalse(Arrays.equals(aad, context.getAssociatedData())); + assertThrows(EncryptException.class, () -> ChunkCryptoContext.forChunk(parameter, -1, 3)); + assertThrows(EncryptException.class, () -> ChunkCryptoContext.forChunk(parameter, 100, -1)); + parameter.close(); + assertThrows(EncryptException.class, () -> parameter.authenticateChunk(100, 3)); + } + } + + @Test + public void testUnsupportedProviderAndMissingHeaderProofFailClosed() throws Exception { + try (EncryptParameter parameter = parameter(); + IEncrypt unsupported = new UNENCRYPTED(null)) { + ChunkCryptoContext context = ChunkCryptoContext.forChunk(parameter, 100, 3); + assertThrows(EncryptException.class, () -> unsupported.authenticateChunk(context)); + assertThrows( + EncryptException.class, + () -> unsupported.verifyChunkAuthentication(new byte[28], context)); + ChunkHeader header = + new ChunkHeader( + "s1", 0, TSDataType.INT64, CompressionType.UNCOMPRESSED, TSEncoding.PLAIN, 0, 0, 3); + ByteArrayOutputStream output = new ByteArrayOutputStream(); + assertThrows(EncryptException.class, () -> header.serializeTo(output)); + assertEquals(0, output.size()); + assertThrows(EncryptException.class, header::getSerializedSize); + ByteBuffer buffer = ByteBuffer.allocate(128); + assertThrows(EncryptException.class, () -> header.serializeTo(buffer)); + assertEquals(0, buffer.position()); + } + } + + @Test + public void testWholeChunkSwapRejectedByBothReadPathsAndQuery() throws Exception { + File file = writeFile(); + List metadata = metadata(file); + swapChunks(file, metadata); + try (TsFileSequenceReader reader = new TsFileSequenceReader(file.getPath())) { + assertThrows( + EncryptException.class, + () -> reader.readChunkHeader(metadata.get(0).getOffsetOfChunkHeader())); + reader.position(metadata.get(1).getOffsetOfChunkHeader()); + byte marker = reader.readMarker(); + assertThrows(EncryptException.class, () -> reader.readChunkHeader(marker)); + assertThrows(EncryptException.class, () -> reader.readMemChunk(metadata.get(0))); + } + Exception failure = + assertThrows( + Exception.class, + () -> { + try (TsFileReader reader = + new TsFileReader(new TsFileSequenceReader(file.getPath()))) { + QueryDataSet data = + reader.query( + QueryExpression.create( + Collections.singletonList(new Path("d1", "s1", true)), null)); + while (data.hasNext()) { + data.next(); + } + } + }); + assertTrue(EncryptException.isEncryptionFailure(failure)); + } + + @Test + public void testRecoveryDoesNotTruncateAwayAuthenticationFailure() throws Exception { + File file = writeFile(); + List metadata = metadata(file); + int end; + try (TsFileSequenceReader reader = new TsFileSequenceReader(file.getPath())) { + ChunkHeader last = reader.readChunkHeader(metadata.get(1).getOffsetOfChunkHeader()); + end = + Math.toIntExact( + metadata.get(1).getOffsetOfChunkHeader() + + last.getSerializedSize() + + last.getDataSize()); + } + swapChunks(file, metadata); + byte[] bytes = Arrays.copyOf(Files.readAllBytes(file.toPath()), end); + Files.write(file.toPath(), bytes); + IOException failure = assertThrows(IOException.class, () -> new RestorableTsFileIOWriter(file)); + assertTrue(EncryptException.isEncryptionFailure(failure)); + assertArrayEquals(bytes, Files.readAllBytes(file.toPath())); + } + + @Test + public void testIncompleteTailProofCanBeRecovered() throws Exception { + File file = writeFile(); + List metadata = metadata(file); + int end; + long dataStart; + try (TsFileSequenceReader reader = new TsFileSequenceReader(file.getPath())) { + ChunkHeader first = reader.readChunkHeader(metadata.get(0).getOffsetOfChunkHeader()); + end = + Math.toIntExact( + metadata.get(0).getOffsetOfChunkHeader() + + first.getSerializedSize() + - first.getChunkAuthenticationSize() + + 10); + dataStart = reader.getDataStartOffset(); + } + Files.write(file.toPath(), Arrays.copyOf(Files.readAllBytes(file.toPath()), end)); + try (RestorableTsFileIOWriter recovered = new RestorableTsFileIOWriter(file)) { + assertTrue(recovered.canWrite()); + assertEquals(dataStart, file.length()); + } + } + + @Test + public void testEmptyValueChunkHasPositionProof() throws Exception { + File file = new File(folder.getRoot(), "empty.tsfile"); + long offset; + try (EncryptParameter parameter = parameter(); + TsFileIOWriter writer = new TsFileIOWriter(file, parameter)) { + writer.startChunkGroup(IDeviceID.Factory.DEFAULT_FACTORY.create("d1")); + offset = writer.getPos(); + writer.writeEmptyValueChunk( + "s1", + CompressionType.UNCOMPRESSED, + TSDataType.INT64, + TSEncoding.PLAIN, + Statistics.getStatsByType(TSDataType.INT64)); + writer.endChunkGroup(); + writer.endFile(); + } + try (TsFileSequenceReader reader = new TsFileSequenceReader(file.getPath())) { + ChunkHeader header = reader.readChunkHeader(offset); + assertEquals(0, header.getDataSize()); + assertEquals(28, header.getChunkAuthentication().length); + ChunkHeader plain = + new ChunkHeader( + "s1", 0, TSDataType.INT64, CompressionType.UNCOMPRESSED, TSEncoding.PLAIN, 0); + try (Chunk encrypted = new Chunk(header, ByteBuffer.allocate(0)); + Chunk plaintext = new Chunk(plain, ByteBuffer.allocate(0))) { + assertEquals( + RamUsageEstimator.sizeOfByteArray(28), + encrypted.getRetainedSizeInBytes() - plaintext.getRetainedSizeInBytes()); + } + ByteBuffer buffer = ByteBuffer.allocate(header.getSerializedSize()); + assertEquals(header.getSerializedSize(), header.serializeTo(buffer)); + } + } + + @Test + public void testVariableProofHeaderRoundTripAndFragmentedReads() throws Exception { + for (int size : new int[] {28, 32, 128, ChunkCryptoContext.MAX_AUTHENTICATION_SIZE}) { + ChunkHeader source = + new ChunkHeader( + "s1", 0, TSDataType.INT64, CompressionType.UNCOMPRESSED, TSEncoding.PLAIN, 0, 0, 3); + byte[] proof = new byte[size]; + Arrays.fill(proof, (byte) 0x5A); + ByteArrayOutputStream output = new ByteArrayOutputStream(); + int length = source.serializeTo(output, proof); + byte[] bytes = output.toByteArray(); + assertEquals(bytes.length, length); + assertEquals(ChunkHeader.getSerializedSize("s1", 0, size), length); + ByteArrayInputStream stream = + new ByteArrayInputStream(bytes, 1, bytes.length - 1) { + @Override + public synchronized int read(byte[] target, int offset, int count) { + return super.read(target, offset, Math.min(count, 2)); + } + }; + ChunkHeader sequential = ChunkHeader.deserializeFrom(stream, bytes[0], size); + assertArrayEquals(proof, sequential.getChunkAuthentication()); + assertEquals(length, sequential.getSerializedSize()); + assertEquals(0, stream.available()); + ByteBuffer buffer = ByteBuffer.allocate(length); + assertEquals(length, sequential.serializeTo(buffer)); + assertArrayEquals(bytes, buffer.array()); + File file = folder.newFile(); + Files.write(file.toPath(), bytes); + try (CloseableInput input = + new CloseableInput(file) { + @Override + public int read(ByteBuffer target, long position) throws IOException { + int limit = target.limit(); + target.limit(Math.min(limit, target.position() + 2)); + try { + return super.read(target, position); + } finally { + target.limit(limit); + } + } + }) { + input.position(2); + ChunkHeader positional = ChunkHeader.deserializeFrom(input, 0, null, size); + assertEquals(length, positional.getSerializedSize()); + assertArrayEquals(proof, positional.getChunkAuthentication()); + assertEquals(2, input.position()); + } + for (int truncatedLength : new int[] {length - size - 1, length - 1}) { + byte[] truncated = Arrays.copyOf(bytes, truncatedLength); + assertThrows( + EOFException.class, + () -> + ChunkHeader.deserializeFrom( + new ByteArrayInputStream(truncated, 1, truncated.length - 1), bytes[0], size)); + Files.write(file.toPath(), truncated); + try (CloseableInput input = new CloseableInput(file)) { + assertThrows(EOFException.class, () -> ChunkHeader.deserializeFrom(input, 0, null, size)); + } + } + } + } + + @Test + public void testInvalidOrMismatchedLengthIsRejectedBeforeProofRead() throws Exception { + ChunkHeader source = + new ChunkHeader( + "s1", 0, TSDataType.INT64, CompressionType.UNCOMPRESSED, TSEncoding.PLAIN, 0, 0, 3); + ByteArrayOutputStream output = new ByteArrayOutputStream(); + source.serializeTo(output, new byte[32]); + byte[] prefix = Arrays.copyOf(output.toByteArray(), output.size() - 32); + for (int size : new int[] {-1, 0, 28, 33, ChunkCryptoContext.MAX_AUTHENTICATION_SIZE + 1}) { + ByteBuffer.wrap(prefix).putInt(prefix.length - Integer.BYTES, size); + assertThrows( + EncryptException.class, + () -> + ChunkHeader.deserializeFrom( + new ByteArrayInputStream(prefix, 1, prefix.length - 1), prefix[0], 32)); + File file = folder.newFile(); + Files.write(file.toPath(), prefix); + try (CloseableInput input = new CloseableInput(file)) { + assertThrows(EncryptException.class, () -> ChunkHeader.deserializeFrom(input, 0, null, 32)); + } + } + } + + @Test + public void testProviderMustDeclareBoundedProofSizeAndReturnMatchingProof() { + for (int declared : new int[] {-1, 0, ChunkCryptoContext.MAX_AUTHENTICATION_SIZE + 1, 32}) { + String id = "test-proof-size-" + declared; + EncryptionProviderRegistry.registerProvider( + new IEncryptProvider() { + @Override + public String getProviderId() { + return id; + } + + @Override + public IEncrypt create(EncryptParameter parameter) { + return new UNENCRYPTED(null) { + @Override + public int getChunkAuthenticationSize() { + return declared; + } + + @Override + public byte[] authenticateChunk(ChunkCryptoContext context) { + return new byte[28]; + } + }; + } + }); + try (EncryptParameter parameter = + EncryptParameter.pageAeadBuilder() + .providerId(id) + .profileId("test") + .fileCryptoId(new byte[16]) + .build()) { + if (declared != 32) { + assertThrows(EncryptException.class, parameter::getChunkAuthenticationSize); + } + assertThrows(EncryptException.class, () -> parameter.authenticateChunk(100, 0)); + assertThrows( + EncryptException.class, + () -> parameter.verifyChunkAuthentication(new byte[28], 100, 0)); + } finally { + EncryptionProviderRegistry.unregisterProvider(id); + } + } + } + + @Test + public void test32ByteProofSupportsQueryBothReadersAndRecovery() throws Exception { + File file = writeFile(true); + List metadata = metadata(file); + long dataEnd; + try (TsFileSequenceReader reader = new TsFileSequenceReader(file.getPath())) { + for (ChunkMetadata item : metadata) { + ChunkHeader header = reader.readChunkHeader(item.getOffsetOfChunkHeader()); + assertEquals(32, header.getChunkAuthenticationSize()); + reader.position(item.getOffsetOfChunkHeader()); + assertArrayEquals( + header.getChunkAuthentication(), + reader.readChunkHeader(reader.readMarker()).getChunkAuthentication()); + } + ChunkMetadata last = metadata.get(1); + ChunkHeader header = reader.readChunkHeader(last.getOffsetOfChunkHeader()); + dataEnd = last.getOffsetOfChunkHeader() + header.getSerializedSize() + header.getDataSize(); + } + try (TsFileReader reader = new TsFileReader(file)) { + QueryDataSet data = + reader.query( + QueryExpression.create(Collections.singletonList(new Path("d1", "s1", true)), null)); + int points = 0; + while (data.hasNext()) { + assertEquals(++points, data.next().getTimestamp()); + } + assertEquals(2, points); + } + // Keep the separator so recovery knows the final chunk group is complete. + Files.write(file.toPath(), Arrays.copyOf(Files.readAllBytes(file.toPath()), (int) dataEnd + 1)); + try (RestorableTsFileIOWriter recovered = new RestorableTsFileIOWriter(file)) { + assertTrue(recovered.canWrite()); + assertEquals(dataEnd, file.length()); + } + } + + @Test + public void test32ByteWholeChunkSwapFailsAuthentication() throws Exception { + File file = writeFile(true); + List metadata = metadata(file); + swapChunks(file, metadata); + try (TsFileSequenceReader reader = new TsFileSequenceReader(file.getPath())) { + assertThrows( + EncryptException.class, + () -> reader.readChunkHeader(metadata.get(0).getOffsetOfChunkHeader())); + reader.position(metadata.get(1).getOffsetOfChunkHeader()); + byte marker = reader.readMarker(); + assertThrows(EncryptException.class, () -> reader.readChunkHeader(marker)); + } + } + + @Test + public void testAlteredProofLengthCannotBeRecoveredAsTruncatedTail() throws Exception { + File file = writeFile(true); + List metadata = metadata(file); + int sizeOffset; + try (TsFileSequenceReader reader = new TsFileSequenceReader(file.getPath())) { + long offset = metadata.get(0).getOffsetOfChunkHeader(); + ChunkHeader header = reader.readChunkHeader(offset); + sizeOffset = (int) offset + header.getSerializedSize() - 32 - Integer.BYTES; + } + byte[] bytes = Arrays.copyOf(Files.readAllBytes(file.toPath()), sizeOffset + Integer.BYTES); + ByteBuffer.wrap(bytes).putInt(sizeOffset, 33); + Files.write(file.toPath(), bytes); + IOException failure = assertThrows(IOException.class, () -> new RestorableTsFileIOWriter(file)); + assertTrue(EncryptException.isEncryptionFailure(failure)); + assertArrayEquals(bytes, Files.readAllBytes(file.toPath())); + } + + @Test + public void testFileRenameKeepsPositionsValid() throws Exception { + File original = writeFile(); + File renamed = new File(folder.getRoot(), "renamed.tsfile"); + Files.move(original.toPath(), renamed.toPath()); + try (TsFileSequenceReader reader = new TsFileSequenceReader(renamed.getPath())) { + for (ChunkMetadata chunk : reader.getChunkMetadataList(new Path("d1", "s1", true))) { + reader.readChunkHeader(chunk.getOffsetOfChunkHeader()); + } + } + } + + private File writeFile() throws Exception { + return writeFile(false); + } + + private File writeFile(boolean macProof) throws Exception { + File file = new File(folder.getRoot(), "positions" + macProof + ".tsfile"); + try (EncryptParameter parameter = + macProof + ? TestAeadEncryptionProvider.createMacParameter(new byte[16], new byte[16]) + : parameter(); + TsFileWriter writer = new TsFileWriter(file, parameter)) { + writer.registerTimeseries( + new Path("d1"), + new MeasurementSchema( + "s1", TSDataType.INT64, TSEncoding.PLAIN, CompressionType.UNCOMPRESSED)); + writer.writeRecord(new TSRecord("d1", 1).addTuple(new LongDataPoint("s1", 11))); + writer.flush(); + writer.getIOWriter().writePlanIndices(); + writer.writeRecord(new TSRecord("d1", 2).addTuple(new LongDataPoint("s1", 22))); + writer.flush(); + } + return file; + } + + private List metadata(File file) throws IOException { + try (TsFileSequenceReader reader = new TsFileSequenceReader(file.getPath())) { + List result = reader.getChunkMetadataList(new Path("d1", "s1", true)); + assertEquals(2, result.size()); + return result; + } + } + + private void swapChunks(File file, List metadata) throws IOException { + int length; + int first = Math.toIntExact(metadata.get(0).getOffsetOfChunkHeader()); + int second = Math.toIntExact(metadata.get(1).getOffsetOfChunkHeader()); + try (TsFileSequenceReader reader = new TsFileSequenceReader(file.getPath())) { + ChunkHeader a = reader.readChunkHeader((long) first); + ChunkHeader b = reader.readChunkHeader((long) second); + length = a.getSerializedSize() + a.getDataSize(); + assertEquals(length, b.getSerializedSize() + b.getDataSize()); + } + byte[] bytes = Files.readAllBytes(file.toPath()); + byte[] saved = Arrays.copyOfRange(bytes, first, first + length); + System.arraycopy(bytes, second, bytes, first, length); + System.arraycopy(saved, 0, bytes, second, length); + Files.write(file.toPath(), bytes); + } + + private EncryptParameter parameter() { + return TestAeadEncryptionProvider.createParameter(new byte[16], new byte[16]); + } + + private static class CloseableInput extends LocalTsFileInput implements AutoCloseable { + private CloseableInput(File file) throws IOException { + super(file.toPath()); + } + } +} diff --git a/java/tsfile/src/test/java/org/apache/tsfile/write/TDEPageAeadTsFileTest.java b/java/tsfile/src/test/java/org/apache/tsfile/write/TDEPageAeadTsFileTest.java new file mode 100644 index 0000000000..6baa9eafc1 --- /dev/null +++ b/java/tsfile/src/test/java/org/apache/tsfile/write/TDEPageAeadTsFileTest.java @@ -0,0 +1,780 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ +package org.apache.tsfile.write; + +import org.apache.tsfile.common.conf.TSFileConfig; +import org.apache.tsfile.common.conf.TSFileDescriptor; +import org.apache.tsfile.constant.TestConstant; +import org.apache.tsfile.encoding.encoder.PlainEncoder; +import org.apache.tsfile.encrypt.EncryptParameter; +import org.apache.tsfile.encrypt.EncryptionProviderRegistry; +import org.apache.tsfile.encrypt.TestAeadEncryptionProvider; +import org.apache.tsfile.enums.TSDataType; +import org.apache.tsfile.exception.encrypt.EncryptException; +import org.apache.tsfile.file.MetaMarker; +import org.apache.tsfile.file.header.ChunkHeader; +import org.apache.tsfile.file.header.PageHeader; +import org.apache.tsfile.file.metadata.ChunkMetadata; +import org.apache.tsfile.file.metadata.enums.CompressionType; +import org.apache.tsfile.file.metadata.enums.TSEncoding; +import org.apache.tsfile.read.TimeValuePair; +import org.apache.tsfile.read.TsFileReader; +import org.apache.tsfile.read.TsFileSequenceReader; +import org.apache.tsfile.read.common.BatchData; +import org.apache.tsfile.read.common.Chunk; +import org.apache.tsfile.read.common.Path; +import org.apache.tsfile.read.common.RowRecord; +import org.apache.tsfile.read.expression.QueryExpression; +import org.apache.tsfile.read.query.dataset.QueryDataSet; +import org.apache.tsfile.read.reader.BufferedTsFileInput; +import org.apache.tsfile.read.reader.IPageReader; +import org.apache.tsfile.read.reader.IPointReader; +import org.apache.tsfile.read.reader.TsFileLastReader; +import org.apache.tsfile.read.reader.chunk.ChunkReader; +import org.apache.tsfile.utils.Binary; +import org.apache.tsfile.utils.Pair; +import org.apache.tsfile.write.chunk.AlignedChunkWriterImpl; +import org.apache.tsfile.write.chunk.ChunkWriterImpl; +import org.apache.tsfile.write.chunk.TimeChunkWriter; +import org.apache.tsfile.write.chunk.ValueChunkWriter; +import org.apache.tsfile.write.record.TSRecord; +import org.apache.tsfile.write.record.Tablet; +import org.apache.tsfile.write.record.datapoint.LongDataPoint; +import org.apache.tsfile.write.schema.IMeasurementSchema; +import org.apache.tsfile.write.schema.MeasurementSchema; +import org.apache.tsfile.write.writer.LocalTsFileOutput; +import org.apache.tsfile.write.writer.TsFileIOWriter; + +import org.junit.After; +import org.junit.AfterClass; +import org.junit.Before; +import org.junit.BeforeClass; +import org.junit.Test; + +import java.io.File; +import java.io.FileOutputStream; +import java.io.IOException; +import java.io.RandomAccessFile; +import java.nio.ByteBuffer; +import java.nio.charset.StandardCharsets; +import java.util.Arrays; +import java.util.List; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotEquals; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.verify; + +public class TDEPageAeadTsFileTest { + + private final File file = + new File(TestConstant.BASE_OUTPUT_PATH + File.separator + "tde-page-aead.tsfile"); + + @BeforeClass + public static void setUpEncryptionProvider() { + EncryptionProviderRegistry.registerProvider(TestAeadEncryptionProvider.INSTANCE); + } + + @AfterClass + public static void tearDownEncryptionProvider() { + EncryptionProviderRegistry.unregisterProvider(TestAeadEncryptionProvider.PROVIDER_ID); + } + + @Before + public void setUp() { + if (!file.getParentFile().exists()) { + assertTrue(file.getParentFile().mkdirs()); + } + } + + @After + public void tearDown() { + if (file.exists()) { + assertTrue(file.delete()); + } + } + + @Test + public void testRawPageCopyIsRejectedBeforeChangingBuffers() throws Exception { + try (EncryptParameter parameter = + TestAeadEncryptionProvider.createParameter( + new byte[16], new byte[EncryptParameter.FILE_CRYPTO_ID_LENGTH])) { + MeasurementSchema schema = new MeasurementSchema("s1", TSDataType.INT64, TSEncoding.PLAIN); + ChunkWriterImpl scalar = new ChunkWriterImpl(schema, parameter); + TimeChunkWriter time = + new TimeChunkWriter( + "", + CompressionType.UNCOMPRESSED, + TSEncoding.PLAIN, + new PlainEncoder(TSDataType.INT64, 0), + parameter); + ValueChunkWriter value = + new ValueChunkWriter( + "s1", + CompressionType.UNCOMPRESSED, + TSDataType.INT64, + TSEncoding.PLAIN, + new PlainEncoder(TSDataType.INT64, 0), + parameter); + AlignedChunkWriterImpl aligned = new AlignedChunkWriterImpl(Arrays.asList(schema), parameter); + ByteBuffer data = ByteBuffer.wrap(new byte[] {1, 2, 3}); + PageHeader header = new PageHeader(3, 3, null); + assertThrows( + EncryptException.class, () -> scalar.writePageHeaderAndDataIntoBuff(data, header)); + assertThrows(EncryptException.class, () -> time.writePageHeaderAndDataIntoBuff(data, header)); + assertThrows( + EncryptException.class, () -> value.writePageHeaderAndDataIntoBuff(data, header)); + assertThrows( + EncryptException.class, () -> aligned.writePageHeaderAndDataIntoTimeBuff(data, header)); + assertThrows( + EncryptException.class, + () -> aligned.writePageHeaderAndDataIntoValueBuff(data, header, 0)); + assertEquals(0, data.position()); + } + } + + @Test + public void testIoWriterConstructorDoesNotDestroyCallerParameter() throws Exception { + try (EncryptParameter parameter = + TestAeadEncryptionProvider.createParameter( + new byte[16], new byte[EncryptParameter.FILE_CRYPTO_ID_LENGTH])) { + try (TsFileWriter writer = new TsFileWriter(new TsFileIOWriter(file), parameter)) { + writer.registerTimeseries(new Path("d1"), new MeasurementSchema("s1", TSDataType.INT64)); + writer.writeRecord(new TSRecord("d1", 1).addTuple(new LongDataPoint("s1", 11))); + } + assertFalse(parameter.isDestroyed()); + } + } + + @Test + public void testNewFileCannotReuseAeadParameterOrItsCopy() throws Exception { + File second = new File(file.getParentFile(), "tde-reused-context.tsfile"); + try (EncryptParameter parameter = + TestAeadEncryptionProvider.createParameter( + new byte[16], new byte[EncryptParameter.FILE_CRYPTO_ID_LENGTH])) { + try (TsFileWriter writer = new TsFileWriter(file, parameter)) { + writer.registerTimeseries(new Path("d1"), new MeasurementSchema("s1", TSDataType.INT64)); + writer.writeRecord(new TSRecord("d1", 1).addTuple(new LongDataPoint("s1", 11))); + } + long originalSize = file.length(); + assertThrows(EncryptException.class, () -> new TsFileWriter(file, parameter)); + assertEquals(originalSize, file.length()); + try (EncryptParameter copy = parameter.copy()) { + assertThrows(EncryptException.class, () -> new TsFileWriter(second, copy)); + } + assertFalse(second.exists()); + assertFalse(parameter.isDestroyed()); + } + } + + @Test + public void testTruncatedEncryptionHeaderReportsIOExceptionWithoutModifyingFile() + throws Exception { + try (EncryptParameter parameter = + TestAeadEncryptionProvider.createParameter( + new byte[16], new byte[EncryptParameter.FILE_CRYPTO_ID_LENGTH]); + TsFileIOWriter writer = new TsFileIOWriter(file, parameter)) { + writer.getIOWriterOut().flush(); + } + try (RandomAccessFile output = new RandomAccessFile(file, "rw")) { + output.setLength(output.length() - 5); + } + long truncatedSize = file.length(); + IOException failure = + assertThrows(IOException.class, () -> new TsFileSequenceReader(file.getPath(), false)); + assertTrue(EncryptException.isEncryptionFailure(failure)); + assertThrows( + IOException.class, () -> new org.apache.tsfile.write.writer.RestorableTsFileIOWriter(file)); + assertEquals(truncatedSize, file.length()); + } + + @Test + public void testUnknownFirstMarkerIsRejectedWithoutMovingLazyReader() throws Exception { + try (TsFileIOWriter writer = new TsFileIOWriter(file)) { + writer.endFile(); + } + try (RandomAccessFile output = new RandomAccessFile(file, "rw")) { + output.seek(TSFileConfig.MAGIC_STRING.getBytes(TSFileConfig.STRING_CHARSET).length + 1); + output.writeByte(0x7F); + } + assertThrows(IOException.class, () -> new TsFileSequenceReader(file.getPath(), false)); + BufferedTsFileInput input = new BufferedTsFileInput(file.toPath()); + try (TsFileSequenceReader reader = new TsFileSequenceReader(input, 0, 0)) { + input.position(3); + assertThrows(IOException.class, reader::hasFileEncryptionHeader); + assertEquals(3, input.position()); + } + } + + @Test + public void testEncryptionHeaderDoesNotForceOutput() throws Exception { + LocalTsFileOutput output = spy(new LocalTsFileOutput(new FileOutputStream(file))); + try (EncryptParameter parameter = + TestAeadEncryptionProvider.createParameter( + new byte[16], new byte[EncryptParameter.FILE_CRYPTO_ID_LENGTH]); + TsFileIOWriter writer = new TsFileIOWriter(output, parameter)) { + long dataStartOffset = output.getPosition(); + assertTrue(dataStartOffset > TSFileConfig.MAGIC_STRING.length() + Byte.BYTES); + writer.writeEncryptionHeaderIfNecessary(); + assertEquals(dataStartOffset, output.getPosition()); + verify(output, never()).force(); + verify(output).flush(); + + writer.getIOWriterOut().force(); + verify(output).force(); + } + } + + @Test + public void testAlignedMultiPageReadWrite() throws Exception { + TSFileConfig config = TSFileDescriptor.getInstance().getConfig(); + int previousMaxPointsInPage = config.getMaxNumberOfPointsInPage(); + config.setMaxNumberOfPointsInPage(1); + byte[] fileCryptoId = new byte[EncryptParameter.FILE_CRYPTO_ID_LENGTH]; + fileCryptoId[0] = 1; + EncryptParameter encryptParameter = + TestAeadEncryptionProvider.createParameter(new byte[16], fileCryptoId); + + try { + List schemas = + Arrays.asList( + new MeasurementSchema("s1", TSDataType.INT64, TSEncoding.RLE), + new MeasurementSchema("s2", TSDataType.INT64, TSEncoding.RLE)); + try (TsFileWriter writer = new TsFileWriter(file, encryptParameter)) { + writer.registerAlignedTimeseries(new Path("d1"), schemas); + for (int i = 1; i <= 3; i++) { + writer.writeRecord( + new TSRecord("d1", i) + .addTuple(new LongDataPoint("s1", i * 10L)) + .addTuple(new LongDataPoint("s2", i * 100L))); + } + } + + try (TsFileReader reader = new TsFileReader(new TsFileSequenceReader(file.getPath()))) { + QueryDataSet dataSet = + reader.query( + QueryExpression.create( + Arrays.asList(new Path("d1", "s1", true), new Path("d1", "s2", true)), null)); + for (int i = 1; i <= 3; i++) { + RowRecord record = dataSet.next(); + assertEquals(i, record.getTimestamp()); + assertEquals(i * 10L, record.getFields().get(0).getLongV()); + assertEquals(i * 100L, record.getFields().get(1).getLongV()); + } + assertFalse(dataSet.hasNext()); + } + } finally { + encryptParameter.close(); + config.setMaxNumberOfPointsInPage(previousMaxPointsInPage); + } + } + + @Test + public void testBufferedInputConstructorLoadsEncryptionHeader() throws Exception { + byte[] fileCryptoId = new byte[EncryptParameter.FILE_CRYPTO_ID_LENGTH]; + fileCryptoId[0] = 1; + EncryptParameter encryptParameter = + TestAeadEncryptionProvider.createParameter(new byte[16], fileCryptoId); + + try { + try (TsFileWriter writer = new TsFileWriter(file, encryptParameter)) { + writer.registerTimeseries( + new Path("d1"), new MeasurementSchema("s1", TSDataType.INT64, TSEncoding.RLE)); + writer.writeRecord(new TSRecord("d1", 1).addTuple(new LongDataPoint("s1", 1L))); + } + + try (TsFileSequenceReader reader = + new TsFileSequenceReader(new BufferedTsFileInput(file.toPath()), false, false, null)) { + reader.position(reader.getDataStartOffset()); + assertEquals(MetaMarker.CHUNK_GROUP_HEADER, reader.readMarker()); + } + } finally { + encryptParameter.close(); + } + } + + @Test + public void testMetadataOffsetConstructorLoadsEncryptionHeader() throws Exception { + try (EncryptParameter parameter = + TestAeadEncryptionProvider.createParameter( + new byte[16], new byte[EncryptParameter.FILE_CRYPTO_ID_LENGTH])) { + try (TsFileWriter writer = new TsFileWriter(file, parameter)) { + writer.registerTimeseries( + new Path("d1"), new MeasurementSchema("s1", TSDataType.INT64, TSEncoding.RLE)); + writer.writeRecord(new TSRecord("d1", 1).addTuple(new LongDataPoint("s1", 11L))); + } + try (TsFileSequenceReader original = new TsFileSequenceReader(file.getPath())) { + ChunkMetadata metadata = original.getChunkMetadataList(new Path("d1", "s1", true)).get(0); + try (TsFileSequenceReader reader = + new TsFileSequenceReader( + new BufferedTsFileInput(file.toPath()), + original.getFileMetadataPos(), + original.getTsFileMetadataSize())) { + assertTrue(reader.hasFileEncryptionHeader()); + original.position(metadata.getOffsetOfChunkHeader()); + reader.position(metadata.getOffsetOfChunkHeader()); + assertEquals( + original.readChunkHeader(original.readMarker()).getChunkOrdinal(), + reader.readChunkHeader(reader.readMarker()).getChunkOrdinal()); + } + } + } + } + + @Test + public void testMaterializedChunkSurvivesReaderClose() throws Exception { + try (EncryptParameter parameter = + TestAeadEncryptionProvider.createParameter( + new byte[16], new byte[EncryptParameter.FILE_CRYPTO_ID_LENGTH])) { + try (TsFileWriter writer = new TsFileWriter(file, parameter)) { + writer.registerTimeseries( + new Path("d1"), new MeasurementSchema("s1", TSDataType.INT64, TSEncoding.PLAIN)); + writer.writeRecord(new TSRecord("d1", 1).addTuple(new LongDataPoint("s1", 11L))); + } + Chunk chunk; + try (TsFileSequenceReader reader = new TsFileSequenceReader(file.getPath())) { + chunk = reader.readMemChunk(reader.getChunkMetadataList(new Path("d1", "s1", true)).get(0)); + } + try (Chunk ownedChunk = chunk) { + BatchData data = new ChunkReader(ownedChunk).nextPageData(); + IPointReader points = data.getBatchDataIterator(); + assertTrue(points.hasNextTimeValuePair()); + assertEquals(11L, points.nextTimeValuePair().getValue().getLong()); + } + assertTrue(chunk.getEncryptParam().isDestroyed()); + chunk.getData().rewind(); + assertThrows(EncryptException.class, () -> new ChunkReader(chunk).nextPageData()); + } + } + + @Test + public void testAlignedBlobLastPointDecryptsPage() throws Exception { + TSFileConfig config = TSFileDescriptor.getInstance().getConfig(); + int previousMaxPointsInPage = config.getMaxNumberOfPointsInPage(); + config.setMaxNumberOfPointsInPage(1); + try (EncryptParameter parameter = + TestAeadEncryptionProvider.createParameter( + new byte[16], new byte[EncryptParameter.FILE_CRYPTO_ID_LENGTH])) { + List schemas = + Arrays.asList(new MeasurementSchema("s1", TSDataType.BLOB, TSEncoding.PLAIN)); + try (TsFileWriter writer = new TsFileWriter(file, parameter)) { + writer.registerAlignedTimeseries(new Path("d1"), schemas); + Tablet tablet = new Tablet("d1", schemas, 2); + for (int i = 0; i < 2; i++) { + tablet.addTimestamp(i, i + 1); + tablet.addValue(i, 0, ("value-" + i).getBytes(StandardCharsets.UTF_8)); + } + writer.writeTree(tablet); + } + try (TsFileLastReader reader = new TsFileLastReader(file.getPath(), true, false)) { + Pair>> last = reader.next(); + assertEquals("s1", last.right.get(1).left); + assertEquals(2L, last.right.get(1).right.getTimestamp()); + assertEquals( + new Binary("value-1", StandardCharsets.UTF_8), + last.right.get(1).right.getValue().getBinary()); + } + } finally { + config.setMaxNumberOfPointsInPage(previousMaxPointsInPage); + } + } + + @Test + public void testUnencryptedWriterRejectsChunkWithOrdinal() throws Exception { + try (TsFileIOWriter writer = new TsFileIOWriter(file)) { + ChunkHeader header = + new ChunkHeader( + "s1", 0, TSDataType.INT64, CompressionType.UNCOMPRESSED, TSEncoding.PLAIN, 1, 0, 0); + Chunk chunk = new Chunk(header, ByteBuffer.allocate(0)); + assertThrows(IOException.class, () -> writer.writeChunk(chunk)); + assertThrows(IOException.class, () -> writer.writeChunk(chunk, null)); + } + } + + @Test + public void testLegacyReadPageOverloadsRejectPageAeadBeforeReading() throws Exception { + EncryptParameter encryptParameter = + TestAeadEncryptionProvider.createParameter( + new byte[16], new byte[EncryptParameter.FILE_CRYPTO_ID_LENGTH]); + try { + try (TsFileWriter writer = new TsFileWriter(file, encryptParameter)) { + writer.registerTimeseries( + new Path("d1"), + new MeasurementSchema( + "s1", TSDataType.INT64, TSEncoding.PLAIN, CompressionType.UNCOMPRESSED)); + writer.writeRecord(new TSRecord("d1", 1).addTuple(new LongDataPoint("s1", 11L))); + } + + try (TsFileSequenceReader reader = new TsFileSequenceReader(file.getPath())) { + ChunkMetadata metadata = reader.getChunkMetadataList(new Path("d1", "s1", true)).get(0); + reader.position(metadata.getOffsetOfChunkHeader()); + ChunkHeader chunkHeader = reader.readChunkHeader(reader.readMarker()); + PageHeader pageHeader = reader.readPageHeader(chunkHeader.getDataType(), false); + long pageBodyOffset = reader.position(); + + IOException noPageIndex = + assertThrows( + IOException.class, + () -> reader.readPage(pageHeader, chunkHeader.getCompressionType())); + assertTrue(noPageIndex.getMessage().contains("chunkOrdinal")); + assertEquals(pageBodyOffset, reader.position()); + + IOException noChunkOrdinal = + assertThrows( + IOException.class, + () -> reader.readPage(pageHeader, chunkHeader.getCompressionType(), 0)); + assertTrue(noChunkOrdinal.getMessage().contains("chunkOrdinal")); + assertEquals(pageBodyOffset, reader.position()); + + assertTrue( + reader + .readPage( + pageHeader, + chunkHeader.getCompressionType(), + 0, + chunkHeader.getChunkOrdinal()) + .remaining() + > 0); + } + } finally { + encryptParameter.close(); + } + } + + @Test + public void testChunkWriterReuseAllocatesDistinctOrdinals() throws Exception { + EncryptParameter encryptParameter = + TestAeadEncryptionProvider.createParameter( + new byte[16], new byte[EncryptParameter.FILE_CRYPTO_ID_LENGTH]); + try { + try (TsFileWriter writer = new TsFileWriter(file, encryptParameter)) { + writer.registerTimeseries( + new Path("d1"), new MeasurementSchema("s1", TSDataType.INT64, TSEncoding.RLE)); + writer.writeRecord(new TSRecord("d1", 1).addTuple(new LongDataPoint("s1", 11L))); + writer.flush(); + writer.writeRecord(new TSRecord("d1", 2).addTuple(new LongDataPoint("s1", 22L))); + } + + try (TsFileSequenceReader sequenceReader = new TsFileSequenceReader(file.getPath()); + TsFileReader reader = new TsFileReader(sequenceReader)) { + List chunks = + sequenceReader.getChunkMetadataList(new Path("d1", "s1", true)); + assertEquals(2, chunks.size()); + long firstOrdinal = + sequenceReader.readMemChunk(chunks.get(0)).getHeader().getChunkOrdinal(); + long secondOrdinal = + sequenceReader.readMemChunk(chunks.get(1)).getHeader().getChunkOrdinal(); + assertTrue(firstOrdinal >= 0); + assertNotEquals(firstOrdinal, secondOrdinal); + + QueryDataSet dataSet = + reader.query(QueryExpression.create(Arrays.asList(new Path("d1", "s1", true)), null)); + assertEquals(11L, dataSet.next().getFields().get(0).getLongV()); + assertEquals(22L, dataSet.next().getFields().get(0).getLongV()); + assertFalse(dataSet.hasNext()); + } + } finally { + encryptParameter.close(); + } + } + + @Test + public void testFlushAfterClearingPageWriter() throws Exception { + try (EncryptParameter parameter = + TestAeadEncryptionProvider.createParameter( + new byte[16], new byte[EncryptParameter.FILE_CRYPTO_ID_LENGTH]); + TsFileIOWriter writer = new TsFileIOWriter(file, parameter)) { + writer.startChunkGroup( + org.apache.tsfile.file.metadata.IDeviceID.Factory.DEFAULT_FACTORY.create("d1")); + ChunkWriterImpl chunkWriter = + new ChunkWriterImpl( + new MeasurementSchema("s1", TSDataType.INT64, TSEncoding.RLE), parameter); + chunkWriter.write(1L, 42L); + chunkWriter.sealCurrentPage(); + chunkWriter.clearPageWriter(); + chunkWriter.writeToFileWriter(writer); + writer.endChunkGroup(); + writer.endFile(); + } + try (TsFileReader reader = new TsFileReader(new TsFileSequenceReader(file.getPath()))) { + QueryDataSet data = + reader.query(QueryExpression.create(Arrays.asList(new Path("d1", "s1", true)), null)); + assertEquals(42L, data.next().getFields().get(0).getLongV()); + assertFalse(data.hasNext()); + } + } + + @Test + public void testAlignedFlushAfterClearingPageWriters() throws Exception { + try (EncryptParameter parameter = + TestAeadEncryptionProvider.createParameter( + new byte[16], new byte[EncryptParameter.FILE_CRYPTO_ID_LENGTH]); + TsFileIOWriter writer = new TsFileIOWriter(file, parameter)) { + writer.startChunkGroup( + org.apache.tsfile.file.metadata.IDeviceID.Factory.DEFAULT_FACTORY.create("d1")); + TimeChunkWriter time = + new TimeChunkWriter( + "", + CompressionType.UNCOMPRESSED, + TSEncoding.PLAIN, + new PlainEncoder(TSDataType.INT64, 0), + parameter); + ValueChunkWriter value = + new ValueChunkWriter( + "s1", + CompressionType.UNCOMPRESSED, + TSDataType.INT64, + TSEncoding.PLAIN, + new PlainEncoder(TSDataType.INT64, 0), + parameter); + time.write(1L); + value.write(1L, 42L, false); + time.sealCurrentPage(); + value.sealCurrentPage(); + time.clearPageWriter(); + value.clearPageWriter(); + time.writeToFileWriter(writer); + value.writeToFileWriter(writer); + writer.endChunkGroup(); + writer.endFile(); + } + try (TsFileReader reader = new TsFileReader(new TsFileSequenceReader(file.getPath()))) { + QueryDataSet data = + reader.query(QueryExpression.create(Arrays.asList(new Path("d1", "s1", true)), null)); + assertEquals(42L, data.next().getFields().get(0).getLongV()); + assertFalse(data.hasNext()); + } + } + + @Test + public void testEncryptedChunkRewriteUsesNewOrdinal() throws Exception { + EncryptParameter encryptParameter = + TestAeadEncryptionProvider.createParameter( + new byte[16], new byte[EncryptParameter.FILE_CRYPTO_ID_LENGTH]); + try { + try (TsFileWriter writer = new TsFileWriter(file, encryptParameter)) { + writer.registerTimeseries( + new Path("d1"), new MeasurementSchema("s1", TSDataType.INT64, TSEncoding.PLAIN)); + writer.writeRecord(new TSRecord("d1", 1).addTuple(new LongDataPoint("s1", 11L))); + } + + try (TsFileSequenceReader reader = new TsFileSequenceReader(file.getPath())) { + Chunk source = + reader.readMemChunk(reader.getChunkMetadataList(new Path("d1", "s1", true)).get(0)); + Chunk rewritten = source.rewrite(TSDataType.DOUBLE); + assertNotEquals( + source.getHeader().getChunkOrdinal(), rewritten.getHeader().getChunkOrdinal()); + for (IPageReader page : new ChunkReader(rewritten).loadPageReaderList()) { + BatchData data = page.getAllSatisfiedPageData(true); + IPointReader points = data.getBatchDataIterator(); + assertTrue(points.hasNextTimeValuePair()); + TimeValuePair point = points.nextTimeValuePair(); + assertEquals(1L, point.getTimestamp()); + assertEquals(11.0, point.getValue().getDouble(), 0.0); + assertFalse(points.hasNextTimeValuePair()); + } + } + } finally { + encryptParameter.close(); + } + } + + @Test + public void testSwappingEqualSizedPagesBetweenChunksFailsAuthentication() throws Exception { + EncryptParameter encryptParameter = + TestAeadEncryptionProvider.createParameter( + new byte[16], new byte[EncryptParameter.FILE_CRYPTO_ID_LENGTH]); + try { + try (TsFileWriter writer = new TsFileWriter(file, encryptParameter)) { + writer.registerTimeseries( + new Path("d1"), + new MeasurementSchema( + "s1", TSDataType.INT64, TSEncoding.PLAIN, CompressionType.UNCOMPRESSED)); + writer.writeRecord(new TSRecord("d1", 1).addTuple(new LongDataPoint("s1", 11L))); + writer.flush(); + writer.writeRecord(new TSRecord("d1", 2).addTuple(new LongDataPoint("s1", 22L))); + } + + long[] pageBodyOffsets = new long[2]; + int[] pageBodySizes = new int[2]; + try (TsFileSequenceReader reader = new TsFileSequenceReader(file.getPath())) { + List chunks = reader.getChunkMetadataList(new Path("d1", "s1", true)); + assertEquals(2, chunks.size()); + for (int i = 0; i < chunks.size(); i++) { + Chunk chunk = reader.readMemChunk(chunks.get(i)); + ByteBuffer chunkData = chunk.getData(); + PageHeader pageHeader = PageHeader.deserializeFrom(chunkData, chunk.getChunkStatistic()); + pageBodyOffsets[i] = + chunks.get(i).getOffsetOfChunkHeader() + + chunk.getHeader().getSerializedSize() + + chunkData.position(); + pageBodySizes[i] = pageHeader.getCompressedSize(); + } + } + assertEquals(pageBodySizes[0], pageBodySizes[1]); + try (RandomAccessFile data = new RandomAccessFile(file, "rw")) { + byte[] firstPage = new byte[pageBodySizes[0]]; + byte[] secondPage = new byte[pageBodySizes[1]]; + data.seek(pageBodyOffsets[0]); + data.readFully(firstPage); + data.seek(pageBodyOffsets[1]); + data.readFully(secondPage); + data.seek(pageBodyOffsets[0]); + data.write(secondPage); + data.seek(pageBodyOffsets[1]); + data.write(firstPage); + } + + try (TsFileReader reader = new TsFileReader(new TsFileSequenceReader(file.getPath()))) { + assertThrows( + Exception.class, + () -> + reader + .query(QueryExpression.create(Arrays.asList(new Path("d1", "s1", true)), null)) + .next()); + } + } finally { + encryptParameter.close(); + } + } + + @Test + public void testRejectUnsafeEncryptedChunkReuse() throws Exception { + File targetFile = new File(file.getPath() + ".target"); + byte[] sourceFileCryptoId = new byte[EncryptParameter.FILE_CRYPTO_ID_LENGTH]; + byte[] targetFileCryptoId = new byte[EncryptParameter.FILE_CRYPTO_ID_LENGTH]; + sourceFileCryptoId[0] = 1; + targetFileCryptoId[0] = 2; + EncryptParameter sourceParameter = + TestAeadEncryptionProvider.createParameter(new byte[16], sourceFileCryptoId); + EncryptParameter targetParameter = + TestAeadEncryptionProvider.createParameter(new byte[16], targetFileCryptoId); + + try (TsFileIOWriter writer = new TsFileIOWriter(targetFile, targetParameter)) { + ChunkHeader header = + new ChunkHeader( + "s1", 0, TSDataType.INT64, CompressionType.UNCOMPRESSED, TSEncoding.PLAIN, 1); + Chunk sourceChunk = new Chunk(header, ByteBuffer.allocate(0), sourceParameter); + + assertThrows(IOException.class, () -> writer.writeChunk(sourceChunk)); + Chunk sameContextChunk = new Chunk(header, ByteBuffer.allocate(0), targetParameter); + assertThrows(IOException.class, () -> writer.writeChunk(sameContextChunk)); + assertThrows(IOException.class, () -> sourceChunk.mergeChunkByAppendPage(sourceChunk)); + } finally { + sourceParameter.close(); + targetParameter.close(); + if (targetFile.exists()) { + assertTrue(targetFile.delete()); + } + } + } + + @Test + public void testFileEncryptionContextCannotChangeAfterHeader() throws Exception { + EncryptParameter first = + TestAeadEncryptionProvider.createParameter( + new byte[16], new byte[EncryptParameter.FILE_CRYPTO_ID_LENGTH]); + byte[] secondFileId = new byte[EncryptParameter.FILE_CRYPTO_ID_LENGTH]; + secondFileId[0] = 1; + EncryptParameter second = + TestAeadEncryptionProvider.createParameter(new byte[16], secondFileId); + try (TsFileIOWriter writer = new TsFileIOWriter(file, first); + EncryptParameter firstCopy = first.copy()) { + EncryptParameter owned = writer.getEncryptParameter(); + writer.setEncryptParam(owned); + assertThrows(EncryptException.class, () -> writer.setEncryptParam(second)); + assertThrows(EncryptException.class, () -> writer.setEncryptParam(firstCopy)); + assertThrows( + EncryptException.class, + () -> writer.setEncryptParam("0", "org.apache.tsfile.encrypt.UNENCRYPTED", null)); + assertTrue(writer.getEncryptParameter() == owned); + assertFalse(first.isDestroyed()); + } finally { + first.close(); + second.close(); + } + } + + @Test + public void testChunkFromDifferentFileContextIsRejectedBeforeWriting() throws Exception { + EncryptParameter target = + TestAeadEncryptionProvider.createParameter( + new byte[16], new byte[EncryptParameter.FILE_CRYPTO_ID_LENGTH]); + byte[] sourceFileId = new byte[EncryptParameter.FILE_CRYPTO_ID_LENGTH]; + sourceFileId[0] = 1; + EncryptParameter source = + TestAeadEncryptionProvider.createParameter(new byte[16], sourceFileId); + try (TsFileIOWriter writer = new TsFileIOWriter(file, target)) { + ChunkWriterImpl chunk = + new ChunkWriterImpl(new MeasurementSchema("s1", TSDataType.INT64), source); + chunk.write(1, 11L); + long position = writer.getPos(); + assertThrows(IOException.class, () -> chunk.writeToFileWriter(writer)); + assertEquals(position, writer.getPos()); + + TimeChunkWriter timeChunk = + new TimeChunkWriter( + "", + CompressionType.UNCOMPRESSED, + TSEncoding.PLAIN, + new PlainEncoder(TSDataType.INT64, 0), + source); + timeChunk.write(1); + timeChunk.sealCurrentPage(); + assertThrows(IOException.class, () -> timeChunk.writeAllPagesOfChunkToTsFile(writer)); + assertEquals(position, writer.getPos()); + + ValueChunkWriter valueChunk = + new ValueChunkWriter( + "s1", + CompressionType.UNCOMPRESSED, + TSDataType.INT64, + TSEncoding.PLAIN, + new PlainEncoder(TSDataType.INT64, 0), + source); + valueChunk.write(1, 11L, false); + valueChunk.sealCurrentPage(); + assertThrows(IOException.class, () -> valueChunk.writeAllPagesOfChunkToTsFile(writer, null)); + assertEquals(position, writer.getPos()); + + assertThrows( + IOException.class, + () -> + writer.startFlushChunk( + "s1", + CompressionType.UNCOMPRESSED, + TSDataType.INT64, + TSEncoding.PLAIN, + null, + 0, + 0, + 0, + 0, + null)); + assertEquals(position, writer.getPos()); + } finally { + source.close(); + target.close(); + } + } +} diff --git a/java/tsfile/src/test/java/org/apache/tsfile/write/TsFileIOWriterTest.java b/java/tsfile/src/test/java/org/apache/tsfile/write/TsFileIOWriterTest.java index 51a8b80a0e..2e6e06e3ef 100644 --- a/java/tsfile/src/test/java/org/apache/tsfile/write/TsFileIOWriterTest.java +++ b/java/tsfile/src/test/java/org/apache/tsfile/write/TsFileIOWriterTest.java @@ -22,6 +22,7 @@ import org.apache.tsfile.common.conf.TSFileDescriptor; import org.apache.tsfile.common.constant.TsFileConstant; import org.apache.tsfile.constant.TestConstant; +import org.apache.tsfile.encrypt.EncryptParameter; import org.apache.tsfile.enums.ColumnCategory; import org.apache.tsfile.enums.TSDataType; import org.apache.tsfile.exception.write.WriteProcessException; @@ -112,6 +113,17 @@ public void after() { } } + @Test + public void testLegacyParameterSetterPreservesSuppliedKey() throws IOException { + EncryptParameter parameter = + new EncryptParameter("org.apache.tsfile.encrypt.UNENCRYPTED", new byte[] {1, 2, 3}); + try (TsFileIOWriter writer = new TsFileIOWriter(new File(FILE_PATH))) { + writer.setEncryptParam(parameter); + Assert.assertSame(parameter, writer.getEncryptParameter()); + Assert.assertArrayEquals(new byte[] {1, 2, 3}, writer.getEncryptParameter().getKey()); + } + } + @Test public void changeTypeCompressionTest() throws IOException, WriteProcessException { TSFileConfig config = TSFileDescriptor.getInstance().getConfig(); diff --git a/java/tsfile/src/test/java/org/apache/tsfile/write/page/PageBodyEncoderTest.java b/java/tsfile/src/test/java/org/apache/tsfile/write/page/PageBodyEncoderTest.java new file mode 100644 index 0000000000..60f5e8a7f0 --- /dev/null +++ b/java/tsfile/src/test/java/org/apache/tsfile/write/page/PageBodyEncoderTest.java @@ -0,0 +1,152 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ +package org.apache.tsfile.write.page; + +import org.apache.tsfile.compress.ICompressor; +import org.apache.tsfile.compress.IUnCompressor; +import org.apache.tsfile.encrypt.EncryptParameter; +import org.apache.tsfile.encrypt.EncryptionProviderRegistry; +import org.apache.tsfile.encrypt.IDecryptor; +import org.apache.tsfile.encrypt.PageCryptoContext; +import org.apache.tsfile.encrypt.TestAeadEncryptionProvider; +import org.apache.tsfile.file.metadata.enums.CompressionType; + +import org.junit.AfterClass; +import org.junit.BeforeClass; +import org.junit.Test; + +import java.nio.ByteBuffer; + +import static org.junit.Assert.assertArrayEquals; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertSame; + +public class PageBodyEncoderTest { + + @BeforeClass + public static void registerProvider() { + EncryptionProviderRegistry.registerProvider(TestAeadEncryptionProvider.INSTANCE); + } + + @AfterClass + public static void unregisterProvider() { + EncryptionProviderRegistry.unregisterProvider(TestAeadEncryptionProvider.PROVIDER_ID); + } + + @Test + public void testUncompressedBodyWrapsSliceWithoutCopying() throws Exception { + ByteBuffer input = slicedInput(); + ByteBuffer encoded = + PageBodyEncoder.encode( + input, + input.remaining(), + ICompressor.getCompressor(CompressionType.UNCOMPRESSED), + null, + 0, + 0); + assertSame(input.array(), encoded.array()); + assertEquals(input.remaining(), encoded.remaining()); + assertEquals(input.arrayOffset() + input.position(), encoded.position()); + assertArrayEquals(remainingBytes(input), remainingBytes(encoded)); + assertEquals(2, input.position()); + } + + @Test + public void testCompressedBodyUsesOnlySliceAndValidOutputBytes() throws Exception { + for (CompressionType compression : + new CompressionType[] {CompressionType.GZIP, CompressionType.LZ4}) { + ByteBuffer input = slicedInput(); + ByteBuffer encoded = + PageBodyEncoder.encode( + input, input.remaining(), ICompressor.getCompressor(compression), null, 0, 0); + byte[] restored = new byte[input.remaining()]; + assertEquals( + restored.length, + IUnCompressor.getUnCompressor(compression) + .uncompress( + encoded.array(), + encoded.arrayOffset() + encoded.position(), + encoded.remaining(), + restored, + 0)); + assertArrayEquals(remainingBytes(input), restored); + assertEquals(2, input.position()); + } + } + + @Test + public void testEncryptedBodyRoundTripForSlicedInput() throws Exception { + encryptedBodyRoundTrip(false); + encryptedBodyRoundTrip(true); + } + + private void encryptedBodyRoundTrip(boolean padded) throws Exception { + for (CompressionType compression : + new CompressionType[] { + CompressionType.UNCOMPRESSED, CompressionType.GZIP, CompressionType.LZ4 + }) { + ByteBuffer input = slicedInput(); + try (EncryptParameter parameter = + padded + ? TestAeadEncryptionProvider.createPaddedParameter(new byte[16], new byte[16]) + : TestAeadEncryptionProvider.createParameter(new byte[16], new byte[16])) { + ByteBuffer encoded = + PageBodyEncoder.encode( + input, input.remaining(), ICompressor.getCompressor(compression), parameter, 3, 7); + PageCryptoContext context = + PageCryptoContext.forDecryption( + parameter, input.remaining(), encoded.asReadOnlyBuffer(), 3, 7); + byte[] compressed = + IDecryptor.getDecryptor(parameter) + .decryptPage( + encoded.array(), + encoded.arrayOffset() + encoded.position(), + encoded.remaining(), + context); + byte[] restored = new byte[input.remaining()]; + assertEquals( + restored.length, + IUnCompressor.getUnCompressor(compression) + .uncompress(compressed, 0, compressed.length, restored, 0)); + assertArrayEquals(remainingBytes(input), restored); + assertEquals(2, input.position()); + } + } + } + + private static ByteBuffer slicedInput() { + byte[] bytes = new byte[140]; + for (int i = 0; i < bytes.length; i++) { + bytes[i] = (byte) i; + } + ByteBuffer parent = ByteBuffer.wrap(bytes); + parent.position(3); + parent.limit(135); + ByteBuffer slice = parent.slice(); + slice.position(2); + slice.limit(130); + return slice; + } + + private static byte[] remainingBytes(ByteBuffer buffer) { + byte[] bytes = new byte[buffer.remaining()]; + buffer.duplicate().get(bytes); + return bytes; + } +} diff --git a/java/tsfile/src/test/java/org/apache/tsfile/write/writer/ForceAppendTsFileWriterTest.java b/java/tsfile/src/test/java/org/apache/tsfile/write/writer/ForceAppendTsFileWriterTest.java index 68d05cefd7..223d68a393 100644 --- a/java/tsfile/src/test/java/org/apache/tsfile/write/writer/ForceAppendTsFileWriterTest.java +++ b/java/tsfile/src/test/java/org/apache/tsfile/write/writer/ForceAppendTsFileWriterTest.java @@ -18,7 +18,11 @@ */ package org.apache.tsfile.write.writer; +import org.apache.tsfile.encrypt.EncryptParameter; +import org.apache.tsfile.encrypt.EncryptionProviderRegistry; +import org.apache.tsfile.encrypt.TestAeadEncryptionProvider; import org.apache.tsfile.enums.TSDataType; +import org.apache.tsfile.file.metadata.ChunkMetadata; import org.apache.tsfile.file.metadata.enums.TSEncoding; import org.apache.tsfile.fileSystem.FSFactoryProducer; import org.apache.tsfile.fileSystem.fsFactory.FSFactory; @@ -34,15 +38,21 @@ import org.apache.tsfile.write.record.datapoint.FloatDataPoint; import org.apache.tsfile.write.schema.MeasurementSchema; +import org.junit.AfterClass; import org.junit.Assert; +import org.junit.BeforeClass; import org.junit.Test; import java.io.File; +import java.io.IOException; import java.util.ArrayList; +import java.util.Collections; import java.util.List; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotEquals; +import static org.junit.Assert.assertThrows; import static org.junit.Assert.assertTrue; import static org.junit.Assert.fail; @@ -51,6 +61,16 @@ public class ForceAppendTsFileWriterTest { TsFileGeneratorForTest.getTestTsFilePath("root.sg1", 0, 0, 1); private static FSFactory fsFactory = FSFactoryProducer.getFSFactory(); + @BeforeClass + public static void setUpEncryptionProvider() { + EncryptionProviderRegistry.registerProvider(TestAeadEncryptionProvider.INSTANCE); + } + + @AfterClass + public static void tearDownEncryptionProvider() { + EncryptionProviderRegistry.unregisterProvider(TestAeadEncryptionProvider.PROVIDER_ID); + } + @Test public void test() throws Exception { File file = fsFactory.getFile(FILE_NAME); @@ -120,4 +140,90 @@ record = dataSet.next(); assertTrue(file.delete()); } + + @Test + public void testHeaderlessForceAppendRejectsPageAeadParameter() throws Exception { + File file = fsFactory.getFile(FILE_NAME + ".headerless"); + if (!file.getParentFile().exists()) { + assertTrue(file.getParentFile().mkdirs()); + } + try { + try (TsFileWriter writer = new TsFileWriter(file)) { + writer.registerTimeseries( + new Path("d1"), new MeasurementSchema("s1", TSDataType.FLOAT, TSEncoding.RLE)); + writer.writeRecord(new TSRecord("d1", 1).addTuple(new FloatDataPoint("s1", 5))); + } + long originalLength = file.length(); + try (EncryptParameter parameter = + TestAeadEncryptionProvider.createParameter( + new byte[16], new byte[EncryptParameter.FILE_CRYPTO_ID_LENGTH])) { + assertThrows(IOException.class, () -> new ForceAppendTsFileWriter(file, parameter)); + } + assertEquals(originalLength, file.length()); + } finally { + if (file.exists()) { + assertTrue(file.delete()); + } + } + } + + @Test + public void testEncryptedForceAppend() throws Exception { + File file = fsFactory.getFile(FILE_NAME + ".encrypted"); + if (!file.getParentFile().exists()) { + Assert.assertTrue(file.getParentFile().mkdirs()); + } + byte[] fileCryptoId = new byte[EncryptParameter.FILE_CRYPTO_ID_LENGTH]; + fileCryptoId[0] = 1; + EncryptParameter encryptParameter = + TestAeadEncryptionProvider.createParameter(new byte[16], fileCryptoId); + EncryptParameter appendParameter = + TestAeadEncryptionProvider.createParameter(new byte[16], fileCryptoId); + + try { + try (TsFileWriter writer = new TsFileWriter(file, encryptParameter)) { + writer.registerTimeseries( + new Path("d1"), new MeasurementSchema("s1", TSDataType.FLOAT, TSEncoding.RLE)); + writer.writeRecord(new TSRecord("d1", 1).addTuple(new FloatDataPoint("s1", 5))); + } + + ForceAppendTsFileWriter appendWriter = new ForceAppendTsFileWriter(file, appendParameter); + EncryptParameter ownedParameter = appendWriter.getEncryptParameter(); + Assert.assertNotSame(encryptParameter, ownedParameter); + appendWriter.doTruncate(); + try (TsFileWriter writer = new TsFileWriter(appendWriter, appendParameter)) { + writer.registerTimeseries( + new Path("d1"), new MeasurementSchema("s1", TSDataType.FLOAT, TSEncoding.RLE)); + writer.writeRecord(new TSRecord("d1", 2).addTuple(new FloatDataPoint("s1", 6))); + } + assertTrue(ownedParameter.isDestroyed()); + assertFalse(encryptParameter.isDestroyed()); + assertFalse(appendParameter.isDestroyed()); + + try (TsFileSequenceReader sequenceReader = new TsFileSequenceReader(file.getPath()); + TsFileReader reader = new TsFileReader(sequenceReader)) { + List chunks = + sequenceReader.getChunkMetadataList(new Path("d1", "s1", true)); + assertEquals(2, chunks.size()); + assertNotEquals( + sequenceReader.readMemChunk(chunks.get(0)).getHeader().getChunkOrdinal(), + sequenceReader.readMemChunk(chunks.get(1)).getHeader().getChunkOrdinal()); + QueryDataSet dataSet = + reader.query( + QueryExpression.create( + Collections.singletonList(new Path("d1", "s1", true)), null)); + RowRecord first = dataSet.next(); + assertEquals(1, first.getTimestamp()); + assertEquals(5.0f, first.getFields().get(0).getFloatV(), 0.001); + RowRecord second = dataSet.next(); + assertEquals(2, second.getTimestamp()); + assertEquals(6.0f, second.getFields().get(0).getFloatV(), 0.001); + assertFalse(dataSet.hasNext()); + } + } finally { + encryptParameter.close(); + appendParameter.close(); + file.delete(); + } + } } diff --git a/java/tsfile/src/test/java/org/apache/tsfile/write/writer/RestorableTsFileIOWriterTest.java b/java/tsfile/src/test/java/org/apache/tsfile/write/writer/RestorableTsFileIOWriterTest.java index 0e1279faeb..e1ffa067cf 100644 --- a/java/tsfile/src/test/java/org/apache/tsfile/write/writer/RestorableTsFileIOWriterTest.java +++ b/java/tsfile/src/test/java/org/apache/tsfile/write/writer/RestorableTsFileIOWriterTest.java @@ -20,6 +20,9 @@ package org.apache.tsfile.write.writer; import org.apache.tsfile.common.conf.TSFileConfig; +import org.apache.tsfile.encrypt.EncryptParameter; +import org.apache.tsfile.encrypt.EncryptionProviderRegistry; +import org.apache.tsfile.encrypt.TestAeadEncryptionProvider; import org.apache.tsfile.enums.TSDataType; import org.apache.tsfile.exception.NotCompatibleTsFileException; import org.apache.tsfile.file.MetaMarker; @@ -44,8 +47,10 @@ import org.apache.tsfile.write.schema.MeasurementSchema; import org.junit.After; +import org.junit.AfterClass; import org.junit.Assert; import org.junit.Before; +import org.junit.BeforeClass; import org.junit.Test; import java.io.File; @@ -58,6 +63,7 @@ import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertNotEquals; import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertTrue; @SuppressWarnings("squid:S4042") // Suppress use java.nio.Files#delete warning public class RestorableTsFileIOWriterTest { @@ -67,6 +73,16 @@ public class RestorableTsFileIOWriterTest { private static final FSFactory fsFactory = FSFactoryProducer.getFSFactory(); File file = fsFactory.getFile(FILE_NAME); + @BeforeClass + public static void setUpEncryptionProvider() { + EncryptionProviderRegistry.registerProvider(TestAeadEncryptionProvider.INSTANCE); + } + + @AfterClass + public static void tearDownEncryptionProvider() { + EncryptionProviderRegistry.unregisterProvider(TestAeadEncryptionProvider.PROVIDER_ID); + } + @Before public void setUp() throws IOException { if (!file.getParentFile().exists()) { @@ -105,6 +121,155 @@ public void testOnlyHeadMagic() throws Exception { rWriter.close(); } + @Test + public void testHeaderlessRecoveryKeepsOnDiskFormatWithPageAeadParameter() throws Exception { + TsFileWriter writer = new TsFileWriter(file); + writer.getIOWriter().close(); + long originalLength = file.length(); + try (EncryptParameter parameter = + TestAeadEncryptionProvider.createParameter( + new byte[16], new byte[EncryptParameter.FILE_CRYPTO_ID_LENGTH])) { + try (RestorableTsFileIOWriter recovered = + new RestorableTsFileIOWriter(file, false, parameter)) { + assertFalse(recovered.getEncryptParameter().isTdePageAead()); + assertTrue(recovered.canWrite()); + } + assertFalse(parameter.isDestroyed()); + } + assertEquals(originalLength, file.length()); + try (TsFileSequenceReader reader = new TsFileSequenceReader(file.getPath(), false)) { + assertFalse(reader.hasFileEncryptionHeader()); + } + } + + @Test + public void testPlaintextRecoveryAndAppendDoNotAdoptCallerAeadContext() throws Exception { + TsFileWriter original = new TsFileWriter(file); + original.registerTimeseries(new Path("d1"), new MeasurementSchema("s1", TSDataType.FLOAT)); + original.writeRecord(new TSRecord("d1", 1).addTuple(new FloatDataPoint("s1", 11))); + original.flush(); + original.getIOWriter().writePlanIndices(); + original.getIOWriter().close(); + try (EncryptParameter parameter = + TestAeadEncryptionProvider.createParameter( + new byte[16], new byte[EncryptParameter.FILE_CRYPTO_ID_LENGTH])) { + try (org.apache.tsfile.read.UnClosedTsFileReader reader = + new org.apache.tsfile.read.UnClosedTsFileReader(file.getPath(), parameter, null)) { + assertFalse(reader.getEncryptParam().isTdePageAead()); + } + RestorableTsFileIOWriter recovered = new RestorableTsFileIOWriter(file, parameter); + try (TsFileWriter writer = new TsFileWriter(recovered, parameter)) { + writer.writeRecord(new TSRecord("d1", 2).addTuple(new FloatDataPoint("s1", 22))); + } + assertFalse(parameter.isDestroyed()); + } + try (TsFileSequenceReader sequence = new TsFileSequenceReader(file.getPath()); + TsFileReader reader = new TsFileReader(sequence)) { + assertFalse(sequence.hasFileEncryptionHeader()); + QueryDataSet query = + reader.query( + QueryExpression.create( + java.util.Collections.singletonList(new Path("d1", "s1", true)), null)); + assertEquals(11.0f, query.next().getFields().get(0).getFloatV(), 0.0f); + assertEquals(22.0f, query.next().getFields().get(0).getFloatV(), 0.0f); + assertFalse(query.hasNext()); + } + } + + @Test + public void testEncryptedHeaderSurvivesUnclosedFile() throws Exception { + byte[] fileCryptoId = new byte[EncryptParameter.FILE_CRYPTO_ID_LENGTH]; + for (int i = 0; i < fileCryptoId.length; i++) { + fileCryptoId[i] = (byte) (i + 1); + } + EncryptParameter parameter = + TestAeadEncryptionProvider.createParameter(new byte[16], fileCryptoId); + TsFileWriter writer = new TsFileWriter(file, parameter); + long expectedDataStartOffset; + try (TsFileSequenceReader reader = new TsFileSequenceReader(file.getPath(), false)) { + expectedDataStartOffset = reader.getDataStartOffset(); + assertNotEquals( + TSFileConfig.MAGIC_STRING.getBytes().length + Byte.BYTES, expectedDataStartOffset); + assertEquals("test-key", reader.getEncryptParam().getKeyId()); + Assert.assertArrayEquals(fileCryptoId, reader.getEncryptParam().getFileCryptoId()); + } + writer.getIOWriter().close(); + writer.getIOWriter().getEncryptParameter().close(); + + RestorableTsFileIOWriter restorableWriter = new RestorableTsFileIOWriter(file, parameter); + assertEquals(expectedDataStartOffset, restorableWriter.getTruncatedSize()); + assertEquals("test-key", restorableWriter.getEncryptParameter().getKeyId()); + EncryptParameter ownedParameter = restorableWriter.getEncryptParameter(); + Assert.assertNotSame(parameter, ownedParameter); + try (TsFileWriter recoveredWriter = new TsFileWriter(restorableWriter, parameter)) { + // Closing writes a valid footer without replacing the file encryption header. + } + Assert.assertTrue(ownedParameter.isDestroyed()); + Assert.assertFalse(parameter.isDestroyed()); + parameter.close(); + } + + @Test + public void testEncryptedChunkOrdinalContinuesAfterRecovery() throws Exception { + EncryptParameter parameter = + TestAeadEncryptionProvider.createParameter( + new byte[16], new byte[EncryptParameter.FILE_CRYPTO_ID_LENGTH]); + try { + TsFileWriter writer = new TsFileWriter(file, parameter); + writer.registerTimeseries( + new Path("d1"), new MeasurementSchema("s1", TSDataType.FLOAT, TSEncoding.RLE)); + writer.writeRecord(new TSRecord("d1", 1).addTuple(new FloatDataPoint("s1", 1.0f))); + writer.flush(); + writer.getIOWriter().writePlanIndices(); + ChunkMetadata firstMetadata = + writer.getIOWriter().getChunkGroupMetadataList().get(0).getChunkMetadataList().get(0); + long firstOrdinal; + try (TsFileSequenceReader reader = new TsFileSequenceReader(file.getPath(), false)) { + firstOrdinal = reader.readMemChunk(firstMetadata).getHeader().getChunkOrdinal(); + } + writer.getIOWriter().close(); + writer.getIOWriter().getEncryptParameter().close(); + + try (EncryptParameter recoveryParameter = + TestAeadEncryptionProvider.createParameter( + new byte[16], new byte[EncryptParameter.FILE_CRYPTO_ID_LENGTH])) { + RestorableTsFileIOWriter restorableWriter = + new RestorableTsFileIOWriter(file, recoveryParameter); + try (TsFileWriter recoveredWriter = new TsFileWriter(restorableWriter, recoveryParameter)) { + recoveredWriter.writeRecord( + new TSRecord("d1", 2).addTuple(new FloatDataPoint("s1", 2.0f))); + } + } + + try (TsFileSequenceReader reader = new TsFileSequenceReader(file.getPath())) { + List chunks = reader.getChunkMetadataList(new Path("d1", "s1", true)); + assertEquals(2, chunks.size()); + assertEquals( + firstOrdinal, reader.readMemChunk(chunks.get(0)).getHeader().getChunkOrdinal()); + Assert.assertTrue( + reader.readMemChunk(chunks.get(1)).getHeader().getChunkOrdinal() > firstOrdinal); + } + } finally { + parameter.close(); + } + } + + @Test + public void testCloseDestroysOwnedTdeParameter() throws Exception { + EncryptParameter parameter = + TestAeadEncryptionProvider.createParameter( + new byte[16], new byte[EncryptParameter.FILE_CRYPTO_ID_LENGTH]); + RestorableTsFileIOWriter writer = new RestorableTsFileIOWriter(file, parameter); + EncryptParameter ownedParameter = writer.getEncryptParameter(); + + writer.close(); + + Assert.assertNotSame(parameter, ownedParameter); + Assert.assertTrue(ownedParameter.isDestroyed()); + Assert.assertFalse(parameter.isDestroyed()); + parameter.close(); + } + @Test public void testOnlyFirstMask() throws Exception { TsFileWriter writer = new TsFileWriter(file); diff --git a/java/tsfile/src/test/java/org/apache/tsfile/write/writer/TimePageWriterTest.java b/java/tsfile/src/test/java/org/apache/tsfile/write/writer/TimePageWriterTest.java index 878b9d7a76..c09ccf28fa 100644 --- a/java/tsfile/src/test/java/org/apache/tsfile/write/writer/TimePageWriterTest.java +++ b/java/tsfile/src/test/java/org/apache/tsfile/write/writer/TimePageWriterTest.java @@ -23,14 +23,24 @@ import org.apache.tsfile.encoding.decoder.PlainDecoder; import org.apache.tsfile.encoding.encoder.Encoder; import org.apache.tsfile.encoding.encoder.PlainEncoder; +import org.apache.tsfile.encrypt.EncryptParameter; +import org.apache.tsfile.encrypt.EncryptionProviderRegistry; +import org.apache.tsfile.encrypt.IDecryptor; +import org.apache.tsfile.encrypt.TestAeadEncryptionProvider; import org.apache.tsfile.enums.TSDataType; +import org.apache.tsfile.file.header.ChunkHeader; +import org.apache.tsfile.file.header.PageHeader; import org.apache.tsfile.file.metadata.enums.CompressionType; +import org.apache.tsfile.file.metadata.enums.TSEncoding; import org.apache.tsfile.file.metadata.statistics.TimeStatistics; +import org.apache.tsfile.read.reader.chunk.ChunkReader; import org.apache.tsfile.utils.PublicBAOS; import org.apache.tsfile.utils.ReadWriteForEncodingUtils; import org.apache.tsfile.utils.ReadWriteIOUtils; import org.apache.tsfile.write.page.TimePageWriter; +import org.junit.AfterClass; +import org.junit.BeforeClass; import org.junit.Test; import java.io.IOException; @@ -41,6 +51,16 @@ public class TimePageWriterTest { + @BeforeClass + public static void setUpEncryptionProvider() { + EncryptionProviderRegistry.registerProvider(TestAeadEncryptionProvider.INSTANCE); + } + + @AfterClass + public static void tearDownEncryptionProvider() { + EncryptionProviderRegistry.unregisterProvider(TestAeadEncryptionProvider.PROVIDER_ID); + } + @Test public void testWrite() { Encoder timeEncoder = new PlainEncoder(TSDataType.INT64, 0); @@ -168,4 +188,50 @@ public void testWritePageHeaderAndDataIntoBuffWithSnappy() { fail(); } } + + @Test + public void testWriteAndReadPageWithAeadEncryption() throws IOException { + Encoder timeEncoder = new PlainEncoder(TSDataType.INT64, 0); + ICompressor compressor = ICompressor.getCompressor(CompressionType.UNCOMPRESSED); + EncryptParameter encryptParameter = + TestAeadEncryptionProvider.createParameter(new byte[16], new byte[16]); + TimePageWriter pageWriter = new TimePageWriter(timeEncoder, compressor, encryptParameter); + long chunkOrdinal = encryptParameter.nextChunkOrdinal(); + pageWriter.setChunkOrdinal(chunkOrdinal); + PublicBAOS publicBAOS = new PublicBAOS(); + + try { + pageWriter.write(1L); + pageWriter.write(2L); + pageWriter.write(3L); + pageWriter.writePageHeaderAndDataIntoBuff(publicBAOS, false, 0); + + ByteBuffer pageBuffer = ByteBuffer.wrap(publicBAOS.getBuf(), 0, publicBAOS.size()); + PageHeader pageHeader = PageHeader.deserializeFrom(pageBuffer, TSDataType.VECTOR); + assertEquals(24, pageHeader.getUncompressedSize()); + assertEquals(52, pageHeader.getCompressedSize()); + + ChunkHeader chunkHeader = + new ChunkHeader( + "time", + pageHeader.getCompressedSize(), + TSDataType.INT64, + CompressionType.UNCOMPRESSED, + TSEncoding.PLAIN, + 1, + 0, + chunkOrdinal); + IDecryptor decryptor = IDecryptor.getDecryptor(encryptParameter); + ByteBuffer uncompressedPageData = + ChunkReader.deserializePageData( + pageHeader, pageBuffer, chunkHeader, decryptor, encryptParameter, 0); + + assertEquals(1L, ReadWriteIOUtils.readLong(uncompressedPageData)); + assertEquals(2L, ReadWriteIOUtils.readLong(uncompressedPageData)); + assertEquals(3L, ReadWriteIOUtils.readLong(uncompressedPageData)); + assertEquals(0, uncompressedPageData.remaining()); + } finally { + encryptParameter.close(); + } + } }