| /* |
| * 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.cassandra.db.commitlog; |
| |
| import org.apache.cassandra.io.util.*; |
| import java.io.IOException; |
| import java.nio.ByteBuffer; |
| import java.nio.channels.FileChannel; |
| import java.util.Collections; |
| import java.util.Random; |
| import java.util.function.BiFunction; |
| |
| import javax.crypto.Cipher; |
| |
| import org.junit.Assert; |
| import org.junit.BeforeClass; |
| import org.junit.Test; |
| |
| import org.apache.cassandra.config.DatabaseDescriptor; |
| import org.apache.cassandra.db.commitlog.CommitLogSegmentReader.CompressedSegmenter; |
| import org.apache.cassandra.db.commitlog.CommitLogSegmentReader.EncryptedSegmenter; |
| import org.apache.cassandra.db.commitlog.CommitLogSegmentReader.SyncSegment; |
| import org.apache.cassandra.io.compress.DeflateCompressor; |
| import org.apache.cassandra.io.compress.ICompressor; |
| import org.apache.cassandra.io.compress.LZ4Compressor; |
| import org.apache.cassandra.io.compress.SnappyCompressor; |
| import org.apache.cassandra.io.compress.ZstdCompressor; |
| import org.apache.cassandra.security.CipherFactory; |
| import org.apache.cassandra.security.EncryptionUtils; |
| import org.apache.cassandra.security.EncryptionContext; |
| import org.apache.cassandra.security.EncryptionContextGenerator; |
| import org.apache.cassandra.utils.ByteBufferUtil; |
| |
| public class SegmentReaderTest |
| { |
| static final Random random = new Random(); |
| |
| @BeforeClass |
| public static void setupDD() |
| { |
| DatabaseDescriptor.daemonInitialization(); |
| } |
| |
| @Test |
| public void compressedSegmenter_LZ4() throws IOException |
| { |
| compressedSegmenter(LZ4Compressor.create(Collections.emptyMap())); |
| } |
| |
| @Test |
| public void compressedSegmenter_Snappy() throws IOException |
| { |
| compressedSegmenter(SnappyCompressor.create(null)); |
| } |
| |
| @Test |
| public void compressedSegmenter_Deflate() throws IOException |
| { |
| compressedSegmenter(DeflateCompressor.create(null)); |
| } |
| |
| @Test |
| public void compressedSegmenter_Zstd() throws IOException |
| { |
| compressedSegmenter(ZstdCompressor.create(Collections.emptyMap())); |
| } |
| |
| private void compressedSegmenter(ICompressor compressor) throws IOException |
| { |
| int rawSize = (1 << 15) - 137; |
| ByteBuffer plainTextBuffer = compressor.preferredBufferType().allocate(rawSize); |
| byte[] b = new byte[rawSize]; |
| random.nextBytes(b); |
| plainTextBuffer.put(b); |
| plainTextBuffer.flip(); |
| |
| int uncompressedHeaderSize = 4; // need to add in the plain text size to the block we write out |
| int length = compressor.initialCompressedBufferLength(rawSize); |
| ByteBuffer compBuffer = ByteBufferUtil.ensureCapacity(null, length + uncompressedHeaderSize, true, compressor.preferredBufferType()); |
| compBuffer.putInt(rawSize); |
| compressor.compress(plainTextBuffer, compBuffer); |
| compBuffer.flip(); |
| |
| File compressedFile = FileUtils.createTempFile("compressed-segment-", ".log"); |
| compressedFile.deleteOnExit(); |
| FileOutputStreamPlus fos = new FileOutputStreamPlus(compressedFile); |
| fos.getChannel().write(compBuffer); |
| fos.close(); |
| |
| try (RandomAccessReader reader = RandomAccessReader.open(compressedFile)) |
| { |
| CompressedSegmenter segmenter = new CompressedSegmenter(compressor, reader); |
| int fileLength = (int) compressedFile.length(); |
| SyncSegment syncSegment = segmenter.nextSegment(0, fileLength); |
| FileDataInput fileDataInput = syncSegment.input; |
| ByteBuffer fileBuffer = readBytes(fileDataInput, rawSize); |
| |
| plainTextBuffer.flip(); |
| Assert.assertEquals(plainTextBuffer, fileBuffer); |
| |
| // CompressedSegmenter includes the Sync header length in the syncSegment.endPosition (value) |
| Assert.assertEquals(rawSize, syncSegment.endPosition - CommitLogSegment.SYNC_MARKER_SIZE); |
| } |
| } |
| |
| private ByteBuffer readBytes(FileDataInput input, int len) |
| { |
| byte[] buf = new byte[len]; |
| try |
| { |
| input.readFully(buf); |
| } |
| catch (IOException e) |
| { |
| throw new RuntimeException(e); |
| } |
| return ByteBuffer.wrap(buf); |
| } |
| |
| private ByteBuffer readBytesSeek(FileDataInput input, int len) |
| { |
| byte[] buf = new byte[len]; |
| |
| /// divide output buffer into 5 |
| int[] offsets = new int[] { 0, len / 5, 2 * len / 5, 3 * len / 5, 4 * len / 5, len }; |
| |
| //seek offset |
| long inputStart = input.getFilePointer(); |
| |
| for (int i = 0; i < offsets.length - 1; i++) |
| { |
| try |
| { |
| // seek to beginning of offet |
| input.seek(inputStart + offsets[i]); |
| //read this segment |
| input.readFully(buf, offsets[i], offsets[i + 1] - offsets[i]); |
| } |
| catch (IOException e) |
| { |
| throw new RuntimeException(e); |
| } |
| } |
| return ByteBuffer.wrap(buf); |
| } |
| |
| @Test |
| public void encryptedSegmenterRead() throws IOException |
| { |
| underlyingEncryptedSegmenterTest((s, t) -> readBytes(s, t)); |
| } |
| |
| @Test |
| public void encryptedSegmenterSeek() throws IOException |
| { |
| underlyingEncryptedSegmenterTest((s, t) -> readBytesSeek(s, t)); |
| } |
| |
| public void underlyingEncryptedSegmenterTest(BiFunction<FileDataInput, Integer, ByteBuffer> readFun) |
| throws IOException |
| { |
| EncryptionContext context = EncryptionContextGenerator.createContext(true); |
| CipherFactory cipherFactory = new CipherFactory(context.getTransparentDataEncryptionOptions()); |
| |
| int plainTextLength = (1 << 13) - 137; |
| ByteBuffer plainTextBuffer = ByteBuffer.allocate(plainTextLength); |
| random.nextBytes(plainTextBuffer.array()); |
| |
| ByteBuffer compressedBuffer = EncryptionUtils.compress(plainTextBuffer, null, true, context.getCompressor()); |
| Cipher cipher = cipherFactory.getEncryptor(context.getTransparentDataEncryptionOptions().cipher, context.getTransparentDataEncryptionOptions().key_alias); |
| File encryptedFile = FileUtils.createTempFile("encrypted-segment-", ".log"); |
| encryptedFile.deleteOnExit(); |
| FileChannel channel = encryptedFile.newReadWriteChannel(); |
| channel.write(ByteBufferUtil.bytes(plainTextLength)); |
| EncryptionUtils.encryptAndWrite(compressedBuffer, channel, true, cipher); |
| channel.close(); |
| |
| try (RandomAccessReader reader = RandomAccessReader.open(encryptedFile)) |
| { |
| context = EncryptionContextGenerator.createContext(cipher.getIV(), true); |
| EncryptedSegmenter segmenter = new EncryptedSegmenter(reader, context); |
| SyncSegment syncSegment = segmenter.nextSegment(0, (int) reader.length()); |
| |
| // EncryptedSegmenter includes the Sync header length in the syncSegment.endPosition (value) |
| Assert.assertEquals(plainTextLength, syncSegment.endPosition - CommitLogSegment.SYNC_MARKER_SIZE); |
| ByteBuffer fileBuffer = readFun.apply(syncSegment.input, plainTextLength); |
| plainTextBuffer.position(0); |
| Assert.assertEquals(plainTextBuffer, fileBuffer); |
| } |
| } |
| } |