/*
 * 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.io.sstable;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.channels.FileChannel;
import java.util.Random;
import java.util.function.Consumer;

import org.junit.AfterClass;
import org.junit.BeforeClass;
import org.junit.Test;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import org.apache.cassandra.SchemaLoader;
import org.apache.cassandra.UpdateBuilder;
import org.apache.cassandra.Util;
import org.apache.cassandra.cache.ChunkCache;
import org.apache.cassandra.config.Config;
import org.apache.cassandra.config.DatabaseDescriptor;
import org.apache.cassandra.db.ColumnFamilyStore;
import org.apache.cassandra.db.DataRange;
import org.apache.cassandra.db.DecoratedKey;
import org.apache.cassandra.db.Keyspace;
import org.apache.cassandra.db.Slices;
import org.apache.cassandra.db.compaction.OperationType;
import org.apache.cassandra.db.filter.ColumnFilter;
import org.apache.cassandra.db.lifecycle.LifecycleTransaction;
import org.apache.cassandra.db.marshal.AsciiType;
import org.apache.cassandra.db.marshal.BytesType;
import org.apache.cassandra.db.rows.Row;
import org.apache.cassandra.db.rows.Unfiltered;
import org.apache.cassandra.db.rows.UnfilteredRowIterator;
import org.apache.cassandra.io.sstable.format.SSTableReader;
import org.apache.cassandra.io.sstable.format.SSTableWriter;
import org.apache.cassandra.io.util.File;
import org.apache.cassandra.io.util.FileUtils;
import org.apache.cassandra.schema.CompressionParams;
import org.apache.cassandra.schema.KeyspaceParams;
import org.apache.cassandra.schema.TableMetadata;

import static org.apache.cassandra.utils.Clock.Global.nanoTime;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertTrue;

public class SSTableCorruptionDetectionTest extends SSTableWriterTestBase
{
    private static final Logger logger = LoggerFactory.getLogger(SSTableCorruptionDetectionTest.class);

    private static final int numberOfPks = 1000;
    private static final int numberOfRuns = 100;
    private static final int valueSize = 512 * 1024;
    // Set corruption size larger or in comparable size to value size, otherwise
    // chance for corruption to land in the middle of value is quite high.
    private static final int maxCorruptionSize = 2 * 1024 * 1024;

    private static final String keyspace = "SSTableCorruptionDetectionTest";
    private static final String table = "corrupted_table";

    private static int maxValueSize;
    private static Random random;
    private static SSTableWriter writer;
    private static LifecycleTransaction txn;
    private static ColumnFamilyStore cfs;
    private static SSTableReader ssTableReader;
    private static Config.CorruptedTombstoneStrategy original;

    @BeforeClass
    public static void setUp()
    {
        // this test writes corrupted data on purpose, disable corrupted tombstone detection
        original = DatabaseDescriptor.getCorruptedTombstoneStrategy();
        DatabaseDescriptor.setCorruptedTombstoneStrategy(Config.CorruptedTombstoneStrategy.disabled);
        TableMetadata.Builder cfm =
            TableMetadata.builder(keyspace, table)
                         .addPartitionKeyColumn("pk", AsciiType.instance)
                         .addClusteringColumn("ck1", AsciiType.instance)
                         .addClusteringColumn("ck2", AsciiType.instance)
                         .addRegularColumn("reg1", BytesType.instance)
                         .addRegularColumn("reg2", BytesType.instance)
                         .compression(CompressionParams.noCompression());

        SchemaLoader.createKeyspace(keyspace, KeyspaceParams.simple(1), cfm);

        cfs = Keyspace.open(keyspace).getColumnFamilyStore(table);
        cfs.disableAutoCompaction();

        maxValueSize = DatabaseDescriptor.getMaxValueSize();
        DatabaseDescriptor.setMaxValueSize(1024 * 1024);

        long seed = nanoTime();
        logger.info("Seed {}", seed);
        random = new Random(seed);

        truncate(cfs);
        File dir = cfs.getDirectories().getDirectoryForNewSSTables();
        txn = LifecycleTransaction.offline(OperationType.WRITE);

        // Setting up/writing large values is an expensive operation, we only want to do it once per run
        writer = getWriter(cfs, dir, txn);
        for (int i = 0; i < numberOfPks; i++)
        {
            UpdateBuilder builder = UpdateBuilder.create(cfs.metadata(), String.format("pkvalue_%07d", i)).withTimestamp(1);
            byte[] reg1 = new byte[valueSize];
            random.nextBytes(reg1);
            byte[] reg2 = new byte[valueSize];
            random.nextBytes(reg2);
            builder.newRow("clustering_" + i, "clustering_" + (i + 1))
                   .add("reg1", ByteBuffer.wrap(reg1))
                   .add("reg2", ByteBuffer.wrap(reg2));
            writer.append(builder.build().unfilteredIterator());
        }
        Util.flush(cfs);

        ssTableReader = writer.finish(true);
        txn.update(ssTableReader, false);
        LifecycleTransaction.waitForDeletions();
    }

    @AfterClass
    public static void tearDown()
    {
        DatabaseDescriptor.setMaxValueSize(maxValueSize);

        txn.abort();
        writer.close();
        DatabaseDescriptor.setCorruptedTombstoneStrategy(original);
    }

    @Test
    public void testSinglePartitionIterator() throws Throwable
    {
        bruteForceCorruptionTest(ssTableReader, partitionIterator());
    }

    @Test
    public void testSSTableScanner() throws Throwable
    {
        bruteForceCorruptionTest(ssTableReader, sstableScanner());
    }

    @Test
    public void testSSTableSimpleScanner() throws Throwable
    {
        bruteForceCorruptionTest(ssTableReader, sstableSimpleScanner());
    }

    private void bruteForceCorruptionTest(SSTableReader ssTableReader, Consumer<SSTableReader> walker) throws Throwable
    {
        FileChannel fc = new File(ssTableReader.getFilename()).newReadWriteChannel();

        int corruptedCounter = 0;

        int fileLength = (int)fc.size(); // in current test, it does fit into int
        for (int i = 0; i < numberOfRuns; i++)
        {
            final int corruptionPosition = random.nextInt(fileLength - 1); //ensure at least one byte will be corrupted
            // corrupt max from position to end of file
            final int corruptionSize = Math.min(maxCorruptionSize, random.nextInt(fileLength - corruptionPosition));

            byte[] backup = corruptSstable(fc, corruptionPosition, corruptionSize);

            try
            {
                walker.accept(ssTableReader);
            }
            catch (CorruptSSTableException t)
            {
                corruptedCounter++;
            }
            finally
            {
                if (ChunkCache.instance != null)
                    ChunkCache.instance.invalidateFile(ssTableReader.getFilename());

                restore(fc, corruptionPosition, backup);
            }
        }

        assertTrue(corruptedCounter > 0);
        FileUtils.closeQuietly(fc);
    }

    private Consumer<SSTableReader> sstableScanner()
    {
        return (SSTableReader sstable) -> {
            try (var scanner = sstable.partitionIterator(ColumnFilter.NONE, DataRange.allData(sstable.getPartitioner()), SSTableReadsListener.NOOP_LISTENER))
            {
                while (scanner.hasNext())
                {
                    try (UnfilteredRowIterator rowIter = scanner.next())
                    {
                        if (rowIter.hasNext())
                        {
                            Unfiltered unfiltered = rowIter.next();
                            if (unfiltered.isRow())
                            {
                                Row row = (Row) unfiltered;
                                assertEquals(2, row.clustering().size());
                                // no-op read
                            }
                        }
                    }

                }
            }
        };
    }

    private Consumer<SSTableReader> sstableSimpleScanner()
    {
        return (SSTableReader sstable) -> {
            try (ISSTableScanner scanner = sstable.getScanner())
            {
                while (scanner.hasNext())
                {
                    try (UnfilteredRowIterator rowIter = scanner.next())
                    {
                        if (rowIter.hasNext())
                        {
                            Unfiltered unfiltered = rowIter.next();
                            if (unfiltered.isRow())
                            {
                                Row row = (Row) unfiltered;
                                assertEquals(2, row.clustering().size());
                                // no-op read
                            }
                        }
                    }

                }
            }
        };
    }

    private Consumer<SSTableReader> partitionIterator()
    {
        return (SSTableReader sstable) -> {
            for (int i = 0; i < numberOfPks; i++)
            {
                DecoratedKey dk = Util.dk(String.format("pkvalue_%07d", i));
                try (UnfilteredRowIterator rowIter = sstable.rowIterator(dk,
                                                                         Slices.ALL,
                                                                         ColumnFilter.all(cfs.metadata()),
                                                                         false,
                                                                         SSTableReadsListener.NOOP_LISTENER))
                {
                    while (rowIter.hasNext())
                    {
                        Unfiltered unfiltered = rowIter.next();
                        if (unfiltered.isRow())
                        {
                            Row row = (Row) unfiltered;
                            assertEquals(2, row.clustering().size());
                            // no-op read
                        }
                    }
                }
            }
        };
    }

    private byte[] corruptSstable(FileChannel fc, int position, int corruptionSize) throws IOException
    {
        byte[] backup = new byte[corruptionSize];
        fc.position(position);
        fc.read(ByteBuffer.wrap(backup));

        fc.position(position);
        byte[] corruption = new byte[corruptionSize];
        random.nextBytes(corruption);
        fc.write(ByteBuffer.wrap(corruption));

        return backup;
    }

    private void restore(FileChannel fc, int position, byte[] backup) throws IOException
    {
        fc.position(position);
        fc.write(ByteBuffer.wrap(backup));
    }
}
