content: /*
   * 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.bookkeeper.proto.checksum;
  
  import io.netty.buffer.ByteBuf;
  import io.netty.buffer.ByteBufAllocator;
  import io.netty.buffer.ByteBufUtil;
  import io.netty.buffer.PooledByteBufAllocator;
  import io.netty.buffer.Unpooled;
  import io.netty.util.ReferenceCounted;
  import io.netty.util.concurrent.FastThreadLocal;
  import java.security.GeneralSecurityException;
  import java.security.NoSuchAlgorithmException;
  import org.apache.bookkeeper.client.BKException.BKDigestMatchException;
  import org.apache.bookkeeper.client.LedgerHandle;
  import org.apache.bookkeeper.proto.BookieProtoEncoding;
  import org.apache.bookkeeper.proto.BookieProtocol;
  import org.apache.bookkeeper.proto.DataFormats.LedgerMetadataFormat.DigestType;
  import org.apache.bookkeeper.util.ByteBufList;
  import org.apache.bookkeeper.util.ByteBufVisitor;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  /**
   * This class takes an entry, attaches a digest to it and packages it with relevant
   * data so that it can be shipped to the bookie. On the return side, it also
   * gets a packet, checks that the digest matches, and extracts the original entry
   * for the packet. Currently 3 types of digests are supported: MAC (based on SHA-1) and CRC32 and CRC32C.
   */
  
  public abstract class DigestManager {
      private static final Logger logger = LoggerFactory.getLogger(DigestManager.class);
  
      public static final int METADATA_LENGTH = 32;
      public static final int LAC_METADATA_LENGTH = 16;
  
      final long ledgerId;
      final boolean useV2Protocol;
      private final ByteBufAllocator allocator;
      private final DigestUpdaterByteBufVisitorCallback byteBufVisitorCallback;
  
      abstract int getMacCodeLength();
  
      abstract int internalUpdate(int digest, ByteBuf buffer, int offset, int len);
  
      abstract int internalUpdate(int digest, byte[] buffer, int offset, int len);
  
      final int update(int digest, ByteBuf buffer, int offset, int len) {
          if (buffer.hasMemoryAddress() && acceptsMemoryAddressBuffer()) {
              return internalUpdate(digest, buffer, offset, len);
          } else if (buffer.hasArray()) {
              return internalUpdate(digest, buffer.array(), buffer.arrayOffset() + offset, len);
          } else {
              UpdateContext updateContext = new UpdateContext(digest);
              ByteBufVisitor.visitBuffers(buffer, offset, len, byteBufVisitorCallback, updateContext);
              return updateContext.digest;
          }
      }
  
      abstract void populateValueAndReset(int digest, ByteBuf buffer);
  
      abstract boolean isInt32Digest();
  
      final int macCodeLength;
  
      public DigestManager(long ledgerId, boolean useV2Protocol, ByteBufAllocator allocator) {
          this.ledgerId = ledgerId;
          this.useV2Protocol = useV2Protocol;
          this.macCodeLength = getMacCodeLength();
          this.allocator = allocator;
          this.byteBufVisitorCallback = new DigestUpdaterByteBufVisitorCallback();
      }
  
      public static DigestManager instantiate(long ledgerId, byte[] passwd, DigestType digestType,
              ByteBufAllocator allocator, boolean useV2Protocol) throws GeneralSecurityException {
          switch(digestType) {
          case HMAC:
              return new MacDigestManager(ledgerId, passwd, useV2Protocol, allocator);
          case CRC32:
              return new CRC32DigestManager(ledgerId, useV2Protocol, allocator);
          case CRC32C:
              return new CRC32CDigestManager(ledgerId, useV2Protocol, allocator);
          case DUMMY:
              return new DummyDigestManager(ledgerId, useV2Protocol, allocator);
          default:
              throw new GeneralSecurityException("Unknown checksum type: " + digestType);
          }
      }
  
      public static byte[] generateMasterKey(byte[] password) throws NoSuchAlgorithmException {
          return password.length > 0 ? MacDigestManager.genDigest("ledger", password) : MacDigestManager.EMPTY_LEDGER_KEY;
      }
  
      /**
       * Computes the digest for an entry and put bytes together for sending.
       *
       * @param entryId
       * @param lastAddConfirmed
       * @param length
       * @param data
       * @return
       */
      public ReferenceCounted computeDigestAndPackageForSending(long entryId, long lastAddConfirmed, long length,
                                                                ByteBuf data, byte[] masterKey, int flags) {
          if (this.useV2Protocol) {
              return computeDigestAndPackageForSendingV2(entryId, lastAddConfirmed, length, data, masterKey, flags);
          } else {
              return computeDigestAndPackageForSendingV3(entryId, lastAddConfirmed, length, data);
          }
      }
  
      private ReferenceCounted computeDigestAndPackageForSendingV2(long entryId, long lastAddConfirmed, long length,
                                                                   ByteBuf data, byte[] masterKey, int flags) {
          boolean isSmallEntry = data.readableBytes() < BookieProtoEncoding.SMALL_ENTRY_SIZE_THRESHOLD;
  
          int headersSize = 4 // Request header
                          + BookieProtocol.MASTER_KEY_LENGTH // for the master key
                          + METADATA_LENGTH  //
                          + macCodeLength;
          int payloadSize = data.readableBytes();
          int bufferSize = 4 + headersSize + (isSmallEntry ? payloadSize : 0);
  
          ByteBuf buf = allocator.buffer(bufferSize, bufferSize);
          buf.writeInt(headersSize + payloadSize);
          buf.writeInt(
                  BookieProtocol.PacketHeader.toInt(
                          BookieProtocol.CURRENT_PROTOCOL_VERSION, BookieProtocol.ADDENTRY, (short) flags));
          buf.writeBytes(masterKey, 0, BookieProtocol.MASTER_KEY_LENGTH);
  
          // The checksum is computed on the next part of the buffer only
          buf.readerIndex(buf.writerIndex());
          buf.writeLong(ledgerId);
          buf.writeLong(entryId);
          buf.writeLong(lastAddConfirmed);
          buf.writeLong(length);
  
          // Compute checksum over the headers
          int digest = update(0, buf, buf.readerIndex(), buf.readableBytes());
          digest = update(digest, data, data.readerIndex(), data.readableBytes());
  
          populateValueAndReset(digest, buf);
  
          // Reset the reader index to the beginning
          buf.readerIndex(0);
  
          if (isSmallEntry) {
              buf.writeBytes(data, data.readerIndex(), data.readableBytes());
              data.release();
              return buf;
          } else {
              return ByteBufList.get(buf, data);
          }
      }
  
      private ByteBufList computeDigestAndPackageForSendingV3(long entryId, long lastAddConfirmed, long length,
                                                              ByteBuf data) {
          ByteBuf headersBuffer = Unpooled.buffer(METADATA_LENGTH + macCodeLength);
          headersBuffer.writeLong(ledgerId);
          headersBuffer.writeLong(entryId);
          headersBuffer.writeLong(lastAddConfirmed);
          headersBuffer.writeLong(length);
  
          int digest = update(0, headersBuffer, 0, METADATA_LENGTH);
          digest = update(digest, data, data.readerIndex(), data.readableBytes());
          populateValueAndReset(digest, headersBuffer);
          return ByteBufList.get(headersBuffer, data);
      }
  
      /**
       * Computes the digest for writeLac for sending.
       *
       * @param lac
       * @return
       */
  
      public ByteBufList computeDigestAndPackageForSendingLac(long lac) {
          ByteBuf headersBuffer;
          if (this.useV2Protocol) {
              headersBuffer = allocator.buffer(LAC_METADATA_LENGTH + macCodeLength);
          } else {
              headersBuffer = Unpooled.buffer(LAC_METADATA_LENGTH + macCodeLength);
          }
          headersBuffer.writeLong(ledgerId);
          headersBuffer.writeLong(lac);
  
          int digest = update(0, headersBuffer, 0, LAC_METADATA_LENGTH);
          populateValueAndReset(digest, headersBuffer);
  
          return ByteBufList.get(headersBuffer);
      }
  
      private void verifyDigest(ByteBuf dataReceived) throws BKDigestMatchException {
          verifyDigest(LedgerHandle.INVALID_ENTRY_ID, dataReceived, true);
      }
  
      private void verifyDigest(long entryId, ByteBuf dataReceived) throws BKDigestMatchException {
          verifyDigest(entryId, dataReceived, false);
      }
  
      private static final FastThreadLocal<ByteBuf> DIGEST_BUFFER = new FastThreadLocal<ByteBuf>() {
          @Override
          protected ByteBuf initialValue() throws Exception {
              return PooledByteBufAllocator.DEFAULT.directBuffer(1024);
          }
  
          @Override
          protected void onRemoval(ByteBuf value) throws Exception {
              value.release();
          }
      };
  
      private void verifyDigest(long entryId, ByteBuf dataReceived, boolean skipEntryIdCheck)
              throws BKDigestMatchException {
  
          if ((METADATA_LENGTH + macCodeLength) > dataReceived.readableBytes()) {
              logger.error("Data received is smaller than the minimum for this digest type. "
                      + " Either the packet it corrupt, or the wrong digest is configured. "
                      + " Digest type: {}, Packet Length: {}",
                      this.getClass().getName(), dataReceived.readableBytes());
              throw new BKDigestMatchException();
          }
          int digest = update(0, dataReceived, 0, METADATA_LENGTH);
  
          int offset = METADATA_LENGTH + macCodeLength;
          digest = update(digest, dataReceived, offset, dataReceived.readableBytes() - offset);
  
          if (isInt32Digest()) {
              int receivedDigest = dataReceived.getInt(METADATA_LENGTH);
              if (receivedDigest != digest) {
                  logger.error("Digest mismatch for ledger-id: " + ledgerId + ", entry-id: " + entryId);
                  throw new BKDigestMatchException();
              }
          } else {
              ByteBuf digestBuf = DIGEST_BUFFER.get();
              digestBuf.clear();
              populateValueAndReset(digest, digestBuf);
  
              if (!ByteBufUtil.equals(digestBuf, 0, dataReceived, METADATA_LENGTH, macCodeLength)) {
                  logger.error("Mac mismatch for ledger-id: " + ledgerId + ", entry-id: " + entryId);
                  throw new BKDigestMatchException();
              }
          }
  
          long actualLedgerId = dataReceived.readLong();
          long actualEntryId = dataReceived.readLong();
  
          if (actualLedgerId != ledgerId) {
              logger.error("Ledger-id mismatch in authenticated message, expected: " + ledgerId + " , actual: "
                           + actualLedgerId);
              throw new BKDigestMatchException();
          }
  
          if (!skipEntryIdCheck && actualEntryId != entryId) {
              logger.error("Entry-id mismatch in authenticated message, expected: " + entryId + " , actual: "
                           + actualEntryId);
              throw new BKDigestMatchException();
          }
  
      }
  
      public long verifyDigestAndReturnLac(ByteBuf dataReceived) throws BKDigestMatchException{
          if ((LAC_METADATA_LENGTH + macCodeLength) > dataReceived.readableBytes()) {
              logger.error("Data received is smaller than the minimum for this digest type."
                      + " Either the packet it corrupt, or the wrong digest is configured. "
                      + " Digest type: {}, Packet Length: {}",
                      this.getClass().getName(), dataReceived.readableBytes());
              throw new BKDigestMatchException();
          }
  
          int digest = update(0, dataReceived, 0, LAC_METADATA_LENGTH);
  
          if (isInt32Digest()) {
              int receivedDigest = dataReceived.getInt(LAC_METADATA_LENGTH);
              if (receivedDigest != digest) {
                  logger.error("Digest mismatch for ledger-id LAC: " + ledgerId);
                  throw new BKDigestMatchException();
              }
          } else {
              ByteBuf digestBuf = DIGEST_BUFFER.get();
              digestBuf.clear();
              populateValueAndReset(digest, digestBuf);
  
              if (!ByteBufUtil.equals(digestBuf, 0, dataReceived, LAC_METADATA_LENGTH, macCodeLength)) {
                  logger.error("Mac mismatch for ledger-id LAC: " + ledgerId);
                  throw new BKDigestMatchException();
              }
          }
  
          long actualLedgerId = dataReceived.readLong();
          long lac = dataReceived.readLong();
          if (actualLedgerId != ledgerId) {
              logger.error("Ledger-id mismatch in authenticated message, expected: " + ledgerId + " , actual: "
                           + actualLedgerId);
              throw new BKDigestMatchException();
          }
          return lac;
      }
  
      /**
       * Verify that the digest matches and returns the data in the entry.
       *
       * @param entryId
       * @param dataReceived
       * @return
       * @throws BKDigestMatchException
       */
      public ByteBuf verifyDigestAndReturnData(long entryId, ByteBuf dataReceived)
              throws BKDigestMatchException {
          verifyDigest(entryId, dataReceived);
          dataReceived.readerIndex(METADATA_LENGTH + macCodeLength);
          return dataReceived;
      }
  
      /**
       * A representation of RecoveryData.
       */
      public static final class RecoveryData {
          final long lastAddConfirmed;
          final long length;
  
          public RecoveryData(long lastAddConfirmed, long length) {
              this.lastAddConfirmed = lastAddConfirmed;
              this.length = length;
          }
  
          public long getLastAddConfirmed() {
              return lastAddConfirmed;
          }
  
          public long getLength() {
              return length;
          }
  
      }
  
      public RecoveryData verifyDigestAndReturnLastConfirmed(ByteBuf dataReceived) throws BKDigestMatchException {
          verifyDigest(dataReceived);
          dataReceived.readerIndex(8);
  
          dataReceived.readLong(); // skip unused entryId
          long lastAddConfirmed = dataReceived.readLong();
          long length = dataReceived.readLong();
          return new RecoveryData(lastAddConfirmed, length);
      }
  
      private static class UpdateContext {
          int digest;
  
          UpdateContext(int digest) {
              this.digest = digest;
          }
      }
  
      private class DigestUpdaterByteBufVisitorCallback implements ByteBufVisitor.ByteBufVisitorCallback<UpdateContext> {
  
          @Override
          public void visitBuffer(UpdateContext context, ByteBuf visitBuffer, int visitIndex, int visitLength) {
              // recursively visit the sub buffer and update the digest
              context.digest = internalUpdate(context.digest, visitBuffer, visitIndex, visitLength);
          }
  
          @Override
          public void visitArray(UpdateContext context, byte[] visitArray, int visitIndex, int visitLength) {
              // update the digest with the array
              context.digest = internalUpdate(context.digest, visitArray, visitIndex, visitLength);
          }
  
          @Override
          public boolean acceptsMemoryAddress(UpdateContext context) {
              return DigestManager.this.acceptsMemoryAddressBuffer();
          }
      }
  
      abstract boolean acceptsMemoryAddressBuffer();
  }
  