/*
 * 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.hadoop.hbase.io.hfile;

import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.Matchers.allOf;
import static org.hamcrest.Matchers.hasProperty;
import static org.hamcrest.Matchers.instanceOf;
import static org.hamcrest.Matchers.startsWith;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertNotEquals;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.junit.jupiter.api.Assertions.fail;

import java.io.ByteArrayOutputStream;
import java.io.Closeable;
import java.io.IOException;
import java.io.PrintStream;
import java.nio.ByteBuffer;
import java.nio.channels.SeekableByteChannel;
import java.nio.charset.StandardCharsets;
import java.nio.file.FileSystems;
import java.nio.file.Files;
import java.nio.file.StandardOpenOption;
import java.time.Instant;
import java.util.LinkedList;
import java.util.List;
import java.util.NoSuchElementException;
import java.util.Random;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.hbase.Cell;
import org.apache.hadoop.hbase.CellBuilderType;
import org.apache.hadoop.hbase.ExtendedCell;
import org.apache.hadoop.hbase.ExtendedCellBuilder;
import org.apache.hadoop.hbase.ExtendedCellBuilderFactory;
import org.apache.hadoop.hbase.HBaseTestingUtil;
import org.apache.hadoop.hbase.fs.HFileSystem;
import org.apache.hadoop.hbase.nio.ByteBuff;
import org.apache.hadoop.hbase.testclassification.IOTests;
import org.apache.hadoop.hbase.testclassification.SmallTests;
import org.apache.hadoop.hbase.util.Bytes;
import org.apache.hadoop.hbase.util.ChecksumType;
import org.hamcrest.Description;
import org.hamcrest.Matcher;
import org.hamcrest.TypeSafeMatcher;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Tag;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.TestInfo;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * This test provides coverage for HFileHeader block fields that are read and interpreted before
 * HBase checksum validation can be applied. As of now, this is just
 * {@code onDiskSizeWithoutHeader}.
 */
@Tag(IOTests.TAG)
@Tag(SmallTests.TAG)
public class TestHFileBlockHeaderCorruption {

  private static final Logger LOG = LoggerFactory.getLogger(TestHFileBlockHeaderCorruption.class);

  private static final HBaseTestingUtil UTIL = new HBaseTestingUtil();
  private static HFileSystem HFS;
  private HFileContext hfileCtx;
  private Path path;

  @BeforeAll
  public static void setUpBeforeAll() throws IOException {
    HFS = (HFileSystem) HFileSystem.get(UTIL.getConfiguration());
  }

  @BeforeEach
  public void setUp(TestInfo testInfo) throws IOException {
    path = new Path(UTIL.getDataTestDirOnTestFS(), testInfo.getTestMethod().get().getName());
    hfileCtx = new HFileContextBuilder().withBlockSize(4 * 1024).withHBaseCheckSum(true).build();
    HFile.WriterFactory factory =
      HFile.getWriterFactory(UTIL.getConfiguration(), CacheConfig.DISABLED).withPath(HFS, path)
        .withFileContext(hfileCtx);

    ExtendedCellBuilder cellBuilder = ExtendedCellBuilderFactory.create(CellBuilderType.DEEP_COPY);
    Random rand = new Random(Instant.now().toEpochMilli());
    byte[] family = Bytes.toBytes("f");
    try (HFile.Writer writer = factory.create()) {
      for (int i = 0; i < 40; i++) {
        byte[] row = RandomKeyValueUtil.randomOrderedFixedLengthKey(rand, i, 100);
        byte[] qualifier = RandomKeyValueUtil.randomRowOrQualifier(rand);
        byte[] value = RandomKeyValueUtil.randomValue(rand);
        ExtendedCell cell = cellBuilder.setType(Cell.Type.Put).setRow(row).setFamily(family)
          .setQualifier(qualifier).setValue(value).build();
        writer.append(cell);
        cellBuilder.clear();
      }
    }
  }

  @Test
  public void testChecksumTypeCorruptionFirstBlock() throws Exception {
    HFileBlockChannelPosition firstBlock = null;
    try {
      try (
        HFileBlockChannelPositionIterator it = new HFileBlockChannelPositionIterator(HFS, path)) {
        assertTrue(it.hasNext());
        firstBlock = it.next();
      }

      Corrupter c = new Corrupter(firstBlock);

      logHeader(firstBlock);

      // test corrupted HFileBlock with unknown checksumType code -1
      c.write(HFileBlock.Header.CHECKSUM_TYPE_INDEX, ByteBuffer.wrap(new byte[] { -1 }));
      logHeader(firstBlock);
      try (
        HFileBlockChannelPositionIterator it = new HFileBlockChannelPositionIterator(HFS, path)) {
        CountingConsumer consumer = new CountingConsumer(it);
        try {
          consumer.readFully();
          fail();
        } catch (Exception e) {
          assertThat(e, new IsThrowableMatching().withInstanceOf(IOException.class)
            .withMessage(startsWith("Unknown checksum type code")));
        }
        assertEquals(0, consumer.getItemsRead());
      }

      // valid checksumType code test
      for (ChecksumType t : ChecksumType.values()) {
        testValidChecksumTypeReadBlock(t.getCode(), c, firstBlock);
      }

      c.restore();
      // test corrupted HFileBlock with unknown checksumType code 3
      c.write(HFileBlock.Header.CHECKSUM_TYPE_INDEX, ByteBuffer.wrap(new byte[] { 3 }));
      logHeader(firstBlock);
      try (
        HFileBlockChannelPositionIterator it = new HFileBlockChannelPositionIterator(HFS, path)) {
        CountingConsumer consumer = new CountingConsumer(it);
        try {
          consumer.readFully();
          fail();
        } catch (Exception e) {
          assertThat(e, new IsThrowableMatching().withInstanceOf(IOException.class)
            .withMessage(startsWith("Unknown checksum type code")));
        }
        assertEquals(0, consumer.getItemsRead());
      }
    } finally {
      if (firstBlock != null) {
        firstBlock.close();
      }
    }
  }

  @Test
  public void testChecksumTypeCorruptionSecondBlock() throws Exception {
    HFileBlockChannelPosition secondBlock = null;
    try {
      try (
        HFileBlockChannelPositionIterator it = new HFileBlockChannelPositionIterator(HFS, path)) {
        assertTrue(it.hasNext());
        it.next();
        assertTrue(it.hasNext());
        secondBlock = it.next();
      }

      Corrupter c = new Corrupter(secondBlock);

      logHeader(secondBlock);
      // test corrupted HFileBlock with unknown checksumType code -1
      c.write(HFileBlock.Header.CHECKSUM_TYPE_INDEX, ByteBuffer.wrap(new byte[] { -1 }));
      logHeader(secondBlock);
      try (
        HFileBlockChannelPositionIterator it = new HFileBlockChannelPositionIterator(HFS, path)) {
        CountingConsumer consumer = new CountingConsumer(it);
        try {
          consumer.readFully();
          fail();
        } catch (Exception e) {
          assertThat(e, new IsThrowableMatching().withInstanceOf(RuntimeException.class)
            .withMessage(startsWith("Unknown checksum type code")));
        }
        assertEquals(1, consumer.getItemsRead());
      }

      // valid checksumType code test
      for (ChecksumType t : ChecksumType.values()) {
        testValidChecksumTypeReadBlock(t.getCode(), c, secondBlock);
      }

      c.restore();
      // test corrupted HFileBlock with unknown checksumType code 3
      c.write(HFileBlock.Header.CHECKSUM_TYPE_INDEX, ByteBuffer.wrap(new byte[] { 3 }));
      logHeader(secondBlock);
      try (
        HFileBlockChannelPositionIterator it = new HFileBlockChannelPositionIterator(HFS, path)) {
        CountingConsumer consumer = new CountingConsumer(it);
        try {
          consumer.readFully();
          fail();
        } catch (Exception e) {
          assertThat(e, new IsThrowableMatching().withInstanceOf(RuntimeException.class)
            .withMessage(startsWith("Unknown checksum type code")));
        }
        assertEquals(1, consumer.getItemsRead());
      }
    } finally {
      if (secondBlock != null) {
        secondBlock.close();
      }
    }
  }

  public void testValidChecksumTypeReadBlock(byte checksumTypeCode, Corrupter c,
    HFileBlockChannelPosition testBlock) throws IOException {
    c.restore();
    c.write(HFileBlock.Header.CHECKSUM_TYPE_INDEX,
      ByteBuffer.wrap(new byte[] { checksumTypeCode }));
    logHeader(testBlock);
    try (HFileBlockChannelPositionIterator it = new HFileBlockChannelPositionIterator(HFS, path)) {
      CountingConsumer consumer = new CountingConsumer(it);
      try {
        consumer.readFully();
      } catch (Exception e) {
        fail("test fail: valid checksumType are not executing properly");
      }
      assertNotEquals(0, consumer.getItemsRead());
    }
  }

  @Test
  public void testOnDiskSizeWithoutHeaderCorruptionFirstBlock() throws Exception {
    HFileBlockChannelPosition firstBlock = null;
    try {
      try (
        HFileBlockChannelPositionIterator it = new HFileBlockChannelPositionIterator(HFS, path)) {
        assertTrue(it.hasNext());
        firstBlock = it.next();
      }

      Corrupter c = new Corrupter(firstBlock);

      logHeader(firstBlock);
      c.write(HFileBlock.Header.ON_DISK_SIZE_WITHOUT_HEADER_INDEX,
        ByteBuffer.wrap(Bytes.toBytes(Integer.MIN_VALUE)));
      logHeader(firstBlock);
      try (
        HFileBlockChannelPositionIterator it = new HFileBlockChannelPositionIterator(HFS, path)) {
        CountingConsumer consumer = new CountingConsumer(it);
        try {
          consumer.readFully();
          fail();
        } catch (Exception e) {
          assertThat(e, new IsThrowableMatching().withInstanceOf(IOException.class)
            .withMessage(startsWith("Invalid onDiskSizeWithHeader=")));
        }
        assertEquals(0, consumer.getItemsRead());
      }

      c.restore();
      c.write(HFileBlock.Header.ON_DISK_SIZE_WITHOUT_HEADER_INDEX,
        ByteBuffer.wrap(Bytes.toBytes(0)));
      logHeader(firstBlock);
      try (
        HFileBlockChannelPositionIterator it = new HFileBlockChannelPositionIterator(HFS, path)) {
        CountingConsumer consumer = new CountingConsumer(it);
        try {
          consumer.readFully();
          fail();
        } catch (Exception e) {
          assertThat(e, new IsThrowableMatching().withInstanceOf(IllegalArgumentException.class));
        }
        assertEquals(0, consumer.getItemsRead());
      }

      c.restore();
      c.write(HFileBlock.Header.ON_DISK_SIZE_WITHOUT_HEADER_INDEX,
        ByteBuffer.wrap(Bytes.toBytes(Integer.MAX_VALUE)));
      logHeader(firstBlock);
      try (
        HFileBlockChannelPositionIterator it = new HFileBlockChannelPositionIterator(HFS, path)) {
        CountingConsumer consumer = new CountingConsumer(it);
        try {
          consumer.readFully();
          fail();
        } catch (Exception e) {
          assertThat(e, new IsThrowableMatching().withInstanceOf(IOException.class)
            .withMessage(startsWith("Invalid onDiskSizeWithHeader=")));
        }
        assertEquals(0, consumer.getItemsRead());
      }
    } finally {
      if (firstBlock != null) {
        firstBlock.close();
      }
    }
  }

  @Test
  public void testOnDiskSizeWithoutHeaderCorruptionSecondBlock() throws Exception {
    HFileBlockChannelPosition secondBlock = null;
    try {
      try (
        HFileBlockChannelPositionIterator it = new HFileBlockChannelPositionIterator(HFS, path)) {
        assertTrue(it.hasNext());
        it.next();
        assertTrue(it.hasNext());
        secondBlock = it.next();
      }

      Corrupter c = new Corrupter(secondBlock);

      logHeader(secondBlock);
      c.write(HFileBlock.Header.ON_DISK_SIZE_WITHOUT_HEADER_INDEX,
        ByteBuffer.wrap(Bytes.toBytes(Integer.MIN_VALUE)));
      logHeader(secondBlock);
      try (
        HFileBlockChannelPositionIterator it = new HFileBlockChannelPositionIterator(HFS, path)) {
        CountingConsumer consumer = new CountingConsumer(it);
        try {
          consumer.readFully();
          fail();
        } catch (Exception e) {
          assertThat(e, new IsThrowableMatching().withInstanceOf(IOException.class)
            .withMessage(startsWith("Invalid onDiskSizeWithHeader=")));
        }
        assertEquals(1, consumer.getItemsRead());
      }

      c.restore();
      c.write(HFileBlock.Header.ON_DISK_SIZE_WITHOUT_HEADER_INDEX,
        ByteBuffer.wrap(Bytes.toBytes(0)));
      logHeader(secondBlock);
      try (
        HFileBlockChannelPositionIterator it = new HFileBlockChannelPositionIterator(HFS, path)) {
        CountingConsumer consumer = new CountingConsumer(it);
        try {
          consumer.readFully();
          fail();
        } catch (Exception e) {
          assertThat(e, new IsThrowableMatching().withInstanceOf(IllegalArgumentException.class));
        }
        assertEquals(1, consumer.getItemsRead());
      }

      c.restore();
      c.write(HFileBlock.Header.ON_DISK_SIZE_WITHOUT_HEADER_INDEX,
        ByteBuffer.wrap(Bytes.toBytes(Integer.MAX_VALUE)));
      logHeader(secondBlock);
      try (
        HFileBlockChannelPositionIterator it = new HFileBlockChannelPositionIterator(HFS, path)) {
        CountingConsumer consumer = new CountingConsumer(it);
        try {
          consumer.readFully();
          fail();
        } catch (Exception e) {
          assertThat(e, new IsThrowableMatching().withInstanceOf(IOException.class)
            .withMessage(startsWith("Invalid onDiskSizeWithHeader=")));
        }
        assertEquals(1, consumer.getItemsRead());
      }
    } finally {
      if (secondBlock != null) {
        secondBlock.close();
      }
    }
  }

  private static void logHeader(HFileBlockChannelPosition hbcp) throws IOException {
    ByteBuff buf = ByteBuff.wrap(ByteBuffer.allocate(HFileBlock.headerSize(true)));
    hbcp.rewind();
    assertEquals(buf.capacity(), buf.read(hbcp.getChannel()));
    buf.rewind();
    hbcp.rewind();
    logHeader(buf);
  }

  private static void logHeader(ByteBuff buf) {
    byte[] blockMagic = new byte[8];
    buf.get(blockMagic);
    int onDiskSizeWithoutHeader = buf.getInt();
    int uncompressedSizeWithoutHeader = buf.getInt();
    long prevBlockOffset = buf.getLong();
    byte checksumType = buf.get();
    int bytesPerChecksum = buf.getInt();
    int onDiskDataSizeWithHeader = buf.getInt();
    LOG.debug(
      "blockMagic={}, onDiskSizeWithoutHeader={}, uncompressedSizeWithoutHeader={}, "
        + "prevBlockOffset={}, checksumType={}, bytesPerChecksum={}, onDiskDataSizeWithHeader={}",
      Bytes.toStringBinary(blockMagic), onDiskSizeWithoutHeader, uncompressedSizeWithoutHeader,
      prevBlockOffset, checksumType, bytesPerChecksum, onDiskDataSizeWithHeader);
  }

  /**
   * Data class to enabled messing with the bytes behind an {@link HFileBlock}.
   */
  public static class HFileBlockChannelPosition implements Closeable {
    private final SeekableByteChannel channel;
    private final long position;

    public HFileBlockChannelPosition(SeekableByteChannel channel, long position) {
      this.channel = channel;
      this.position = position;
    }

    public SeekableByteChannel getChannel() {
      return channel;
    }

    public long getPosition() {
      return position;
    }

    public void rewind() throws IOException {
      channel.position(position);
    }

    @Override
    public void close() throws IOException {
      channel.close();
    }
  }

  /**
   * Reads blocks off of an {@link HFileBlockChannelPositionIterator}, counting them as it does.
   */
  public static class CountingConsumer {
    private final HFileBlockChannelPositionIterator iterator;
    private int itemsRead = 0;

    public CountingConsumer(HFileBlockChannelPositionIterator iterator) {
      this.iterator = iterator;
    }

    public int getItemsRead() {
      return itemsRead;
    }

    public Object readFully() throws IOException {
      Object val = null;
      for (itemsRead = 0; iterator.hasNext(); itemsRead++) {
        val = iterator.next();
      }
      return val;
    }
  }

  /**
   * A simplified wrapper over an {@link HFileBlock.BlockIterator} that looks a lot like an
   * {@link java.util.Iterator}.
   */
  public static class HFileBlockChannelPositionIterator implements Closeable {

    private final Path hfsPath;
    private final HFile.Reader reader;
    private final HFileBlock.BlockIterator iter;
    private HFileBlockChannelPosition current = null;

    public HFileBlockChannelPositionIterator(HFileSystem hfs, Path hfsPath) throws IOException {
      this.hfsPath = hfsPath;
      this.reader = HFile.createReader(hfs, hfsPath, CacheConfig.DISABLED, true, hfs.getConf());
      HFileBlock.FSReader fsreader = reader.getUncachedBlockReader();
      // The read block offset cannot out of the range:0,loadOnOpenDataOffset
      this.iter = fsreader.blockRange(0, reader.getTrailer().getLoadOnOpenDataOffset());
    }

    public boolean hasNext() throws IOException {
      HFileBlock next = iter.nextBlock();
      if (next != null) {
        java.nio.file.Path p = FileSystems.getDefault().getPath(hfsPath.toString());
        SeekableByteChannel channel = Files.newByteChannel(p, StandardOpenOption.READ,
          StandardOpenOption.WRITE, StandardOpenOption.DSYNC);
        current = new HFileBlockChannelPosition(channel, next.getOffset());
      }
      return next != null;
    }

    public HFileBlockChannelPosition next() {
      if (current == null) {
        throw new NoSuchElementException();
      }
      HFileBlockChannelPosition ret = current;
      current = null;
      return ret;
    }

    @Override
    public void close() throws IOException {
      if (current != null) {
        closeQuietly(current::close);
      }
      closeQuietly(reader::close);
    }

    @FunctionalInterface
    private interface CloseMethod {
      void run() throws IOException;
    }

    private static void closeQuietly(CloseMethod closeMethod) {
      try {
        closeMethod.run();
      } catch (Throwable e) {
        LOG.debug("Ignoring thrown exception.", e);
      }
    }
  }

  /**
   * Enables writing and rewriting portions of the file backing an {@link HFileBlock}.
   */
  public static class Corrupter {

    private final HFileBlockChannelPosition channelAndPosition;
    private final ByteBuffer originalHeader;

    public Corrupter(HFileBlockChannelPosition channelAndPosition) throws IOException {
      this.channelAndPosition = channelAndPosition;
      this.originalHeader = readHeaderData(channelAndPosition);
    }

    private static ByteBuffer readHeaderData(HFileBlockChannelPosition channelAndPosition)
      throws IOException {
      SeekableByteChannel channel = channelAndPosition.getChannel();
      ByteBuffer originalHeader = ByteBuffer.allocate(HFileBlock.headerSize(true));
      channelAndPosition.rewind();
      channel.read(originalHeader);
      return originalHeader;
    }

    public void write(int offset, ByteBuffer src) throws IOException {
      SeekableByteChannel channel = channelAndPosition.getChannel();
      long position = channelAndPosition.getPosition();
      channel.position(position + offset);
      channel.write(src);
    }

    public void restore() throws IOException {
      SeekableByteChannel channel = channelAndPosition.getChannel();
      originalHeader.rewind();
      channelAndPosition.rewind();
      assertEquals(originalHeader.capacity(), channel.write(originalHeader));
    }
  }

  /**
   * A Matcher implementation that can make basic assertions over a provided {@link Throwable}.
   * Assertion failures include the full stacktrace in their description.
   */
  private static final class IsThrowableMatching extends TypeSafeMatcher<Throwable> {

    private final List<Matcher<? super Throwable>> requirements = new LinkedList<>();

    public IsThrowableMatching withInstanceOf(Class<?> type) {
      requirements.add(instanceOf(type));
      return this;
    }

    public IsThrowableMatching withMessage(Matcher<String> matcher) {
      requirements.add(hasProperty("message", matcher));
      return this;
    }

    @Override
    protected boolean matchesSafely(Throwable throwable) {
      return allOf(requirements).matches(throwable);
    }

    @Override
    protected void describeMismatchSafely(Throwable item, Description mismatchDescription) {
      allOf(requirements).describeMismatch(item, mismatchDescription);
      // would be nice if `item` could be provided as the cause of the AssertionError instead.
      mismatchDescription.appendText(String.format("%nProvided: "));
      try (ByteArrayOutputStream baos = new ByteArrayOutputStream()) {
        try (PrintStream ps = new PrintStream(baos, false, StandardCharsets.UTF_8.name())) {
          item.printStackTrace(ps);
          ps.flush();
        }
        mismatchDescription.appendText(baos.toString(StandardCharsets.UTF_8.name()));
      } catch (Exception e) {
        throw new RuntimeException(e);
      }
    }

    @Override
    public void describeTo(Description description) {
      description.appendDescriptionOf(allOf(requirements));
    }
  }
}