code: /*
   * Copyright (c) 2021 the Eclipse Milo Authors
   *
   * This program and the accompanying materials are made
   * available under the terms of the Eclipse Public License 2.0
   * which is available at https://www.eclipse.org/legal/epl-2.0/
   *
   * SPDX-License-Identifier: EPL-2.0
   */
  
  package org.eclipse.milo.opcua.stack.core.channel;
  
  import java.nio.Buffer;
  import java.nio.ByteBuffer;
  import java.security.GeneralSecurityException;
  import java.security.InvalidKeyException;
  import java.security.MessageDigest;
  import java.security.NoSuchAlgorithmException;
  import java.security.Signature;
  import java.security.SignatureException;
  import java.util.List;
  import javax.crypto.Cipher;
  import javax.crypto.spec.IvParameterSpec;
  import javax.crypto.spec.SecretKeySpec;
  
  import io.netty.buffer.ByteBuf;
  import io.netty.buffer.CompositeByteBuf;
  import io.netty.util.ReferenceCountUtil;
  import org.eclipse.milo.opcua.stack.core.StatusCodes;
  import org.eclipse.milo.opcua.stack.core.UaException;
  import org.eclipse.milo.opcua.stack.core.channel.headers.AsymmetricSecurityHeader;
  import org.eclipse.milo.opcua.stack.core.channel.headers.SecureMessageHeader;
  import org.eclipse.milo.opcua.stack.core.channel.headers.SequenceHeader;
  import org.eclipse.milo.opcua.stack.core.channel.headers.SymmetricSecurityHeader;
  import org.eclipse.milo.opcua.stack.core.channel.messages.ErrorMessage;
  import org.eclipse.milo.opcua.stack.core.security.SecurityAlgorithm;
  import org.eclipse.milo.opcua.stack.core.util.BufferUtil;
  import org.eclipse.milo.opcua.stack.core.util.SignatureUtil;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  public final class ChunkDecoder {
  
      private final AsymmetricDecoder asymmetricDecoder = new AsymmetricDecoder();
      private final SymmetricDecoder symmetricDecoder = new SymmetricDecoder();
  
      private volatile long lastSequenceNumber = -1L;
  
      private final ChannelParameters parameters;
      private final EncodingLimits encodingLimits;
  
      public ChunkDecoder(ChannelParameters parameters, EncodingLimits encodingLimits) {
          this.parameters = parameters;
          this.encodingLimits = encodingLimits;
      }
  
      public DecodedMessage decodeAsymmetric(
          SecureChannel channel,
          List<ByteBuf> chunkBuffers
      ) throws MessageAbortException, MessageDecodeException {
  
          return decode(asymmetricDecoder, channel, chunkBuffers);
      }
  
      public DecodedMessage decodeSymmetric(
          SecureChannel channel,
          List<ByteBuf> chunkBuffers
      ) throws MessageAbortException, MessageDecodeException {
  
          try {
              validateSymmetricSecurityHeaders(channel, chunkBuffers);
          } catch (UaException e) {
              chunkBuffers.forEach(ReferenceCountUtil::safeRelease);
              throw new MessageDecodeException(e);
          }
  
          return decode(symmetricDecoder, channel, chunkBuffers);
      }
  
      private static DecodedMessage decode(
          AbstractDecoder decoder,
          SecureChannel channel,
          List<ByteBuf> chunkBuffers
      ) throws MessageAbortException, MessageDecodeException {
  
          CompositeByteBuf composite = BufferUtil.compositeBuffer();
  
          try {
              return decoder.decode(channel, composite, chunkBuffers);
          } catch (MessageAbortException e) {
              ReferenceCountUtil.safeRelease(composite);
              chunkBuffers.forEach(ReferenceCountUtil::safeRelease);
              throw e;
          } catch (UaException e) {
              ReferenceCountUtil.safeRelease(composite);
              chunkBuffers.forEach(ReferenceCountUtil::safeRelease);
              throw new MessageDecodeException(e);
          }
      }
  
      private static void validateSymmetricSecurityHeaders(
          SecureChannel secureChannel,
          List<ByteBuf> chunkBuffers
      ) throws UaException {
  
          ChannelSecurity channelSecurity = secureChannel.getChannelSecurity();
          long currentTokenId = channelSecurity.getCurrentToken().getTokenId().longValue();
          long previousTokenId = channelSecurity.getPreviousToken()
              .map(t -> t.getTokenId().longValue())
              .orElse(-1L);
  
          for (ByteBuf chunkBuffer : chunkBuffers) {
              // tokenId starts after messageType + chunkType + messageSize + secureChannelId
              long tokenId = chunkBuffer.getUnsignedIntLE(3 + 1 + 4 + 4);
  
              if (tokenId != currentTokenId && tokenId != previousTokenId) {
                  String message = String.format(
                      "received unknown secure channel token: " +
                          "tokenId=%s currentTokenId=%s previousTokenId=%s",
                      tokenId, currentTokenId, previousTokenId
                  );
  
                  throw new UaException(StatusCodes.Bad_SecureChannelTokenUnknown, message);
              }
          }
      }
  
      /**
       * A full decoded message, assembled from one or more successfully decoded chunks.
       */
      public static class DecodedMessage {
  
          private final ByteBuf message;
          private final long requestId;
  
          private DecodedMessage(ByteBuf message, long requestId) {
              this.message = message;
              this.requestId = requestId;
          }
  
          public ByteBuf getMessage() {
              return message;
          }
  
          public long getRequestId() {
              return requestId;
          }
  
      }
  
      private abstract class AbstractDecoder {
  
          protected final Logger logger = LoggerFactory.getLogger(getClass());
  
          DecodedMessage decode(
              SecureChannel channel,
              CompositeByteBuf composite,
              List<ByteBuf> chunkBuffers
          ) throws MessageAbortException, UaException {
  
              int signatureSize = getSignatureSize(channel);
              int cipherTextBlockSize = getCipherTextBlockSize(channel);
  
              boolean encrypted = isEncryptionEnabled(channel);
              boolean signed = isSigningEnabled(channel);
  
              long requestId = -1L;
  
              for (ByteBuf chunkBuffer : chunkBuffers) {
                  final char chunkType = (char) chunkBuffer.getByte(3);
  
                  chunkBuffer.skipBytes(SecureMessageHeader.SECURE_MESSAGE_HEADER_SIZE);
  
                  readSecurityHeader(channel, chunkBuffer);
  
                  if (encrypted) {
                      decryptChunk(channel, chunkBuffer);
                  }
  
                  int encryptedStart = chunkBuffer.readerIndex();
                  chunkBuffer.readerIndex(0);
  
                  if (signed) {
                      verifyChunk(channel, chunkBuffer);
                  }
  
                  final int paddingOverhead = encrypted ? (cipherTextBlockSize > 256 ? 2 : 1) : 0;
                  final int paddingSize = encrypted ? getPaddingSize(cipherTextBlockSize, signatureSize, chunkBuffer) : 0;
                  final int bodyEnd = chunkBuffer.readableBytes() - signatureSize - paddingOverhead - paddingSize;
  
                  chunkBuffer.readerIndex(encryptedStart);
  
                  SequenceHeader sequenceHeader = SequenceHeader.decode(chunkBuffer);
                  long sequenceNumber = sequenceHeader.getSequenceNumber();
                  requestId = sequenceHeader.getRequestId();
  
                  if (lastSequenceNumber != -1) {
                      if (lastSequenceNumber + 1 != sequenceNumber) {
                          String message = String.format(
                              "expected sequence number %s but received %s",
                              lastSequenceNumber + 1, sequenceNumber);
  
                          throw new UaException(StatusCodes.Bad_SecurityChecksFailed, message);
                      }
                  }
  
                  lastSequenceNumber = sequenceNumber;
  
                  ByteBuf bodyBuffer = chunkBuffer.readSlice(bodyEnd - chunkBuffer.readerIndex());
  
                  if (encrypted) {
                      int expectedPaddingSize = chunkBuffer.readableBytes() - signatureSize - paddingOverhead;
                      if (paddingSize != expectedPaddingSize) {
                          throw new UaException(StatusCodes.Bad_SecurityChecksFailed, "bad padding size");
                      }
                      byte expectedPaddingByte = (byte) (paddingSize & 0xFF);
                      for (int i = chunkBuffer.readerIndex(); i < chunkBuffer.readerIndex() + paddingSize + 1; i++) {
                          if (chunkBuffer.getByte(i) != expectedPaddingByte) {
                              throw new UaException(StatusCodes.Bad_SecurityChecksFailed, "bad padding sequence");
                          }
                      }
                  }
  
                  if (chunkType == 'A') {
                      ErrorMessage errorMessage = ErrorMessage.decode(bodyBuffer);
  
                      throw new MessageAbortException(errorMessage.getReason(), requestId, errorMessage.getError());
                  }
  
                  composite.addComponent(bodyBuffer);
                  composite.writerIndex(composite.writerIndex() + bodyBuffer.readableBytes());
              }
  
              if (parameters.getLocalMaxMessageSize() > 0 &&
                  composite.readableBytes() > parameters.getLocalMaxMessageSize()) {
  
                  String errorMessage = String.format(
                      "message size exceeds configured limit: %s > %s",
                      composite.readableBytes(), parameters.getLocalMaxMessageSize());
  
                  throw new UaException(StatusCodes.Bad_TcpMessageTooLarge, errorMessage);
              }
  
              return new DecodedMessage(composite, requestId);
          }
  
          private void decryptChunk(SecureChannel channel, ByteBuf chunkBuffer) throws UaException {
              int cipherTextBlockSize = getCipherTextBlockSize(channel);
              int blockCount = chunkBuffer.readableBytes() / cipherTextBlockSize;
  
              int plainTextBufferSize = cipherTextBlockSize * blockCount;
  
              ByteBuf plainTextBuffer = BufferUtil.pooledBuffer(plainTextBufferSize);
  
              ByteBuffer plainTextNioBuffer = plainTextBuffer
                  .writerIndex(plainTextBufferSize)
                  .nioBuffer();
  
              ByteBuffer chunkNioBuffer = chunkBuffer.nioBuffer();
  
              try {
                  Cipher cipher = getCipher(channel);
  
                  assert (chunkBuffer.readableBytes() % cipherTextBlockSize == 0);
  
                  if (isAsymmetric()) {
                      for (int blockNumber = 0; blockNumber < blockCount; blockNumber++) {
                          ((Buffer) chunkNioBuffer).limit(chunkNioBuffer.position() + cipherTextBlockSize);
  
                          cipher.doFinal(chunkNioBuffer, plainTextNioBuffer);
                      }
                  } else {
                      cipher.doFinal(chunkNioBuffer, plainTextNioBuffer);
                  }
  
                  /* Write plainTextBuffer back into the chunk buffer we decrypted from. */
                  ((Buffer) plainTextNioBuffer).flip(); // limit = pos, pos = 0
  
                  chunkBuffer.writerIndex(chunkBuffer.readerIndex());
                  chunkBuffer.writeBytes(plainTextNioBuffer);
              } catch (GeneralSecurityException e) {
                  throw new UaException(StatusCodes.Bad_SecurityChecksFailed, e);
              } finally {
                  plainTextBuffer.release();
              }
          }
  
          private int getPaddingSize(int cipherTextBlockSize, int signatureSize, ByteBuf buffer) {
              int lastPaddingByteOffset = buffer.readableBytes() - signatureSize - 1;
  
              return cipherTextBlockSize <= 256 ?
                  buffer.getUnsignedByte(lastPaddingByteOffset) :
                  buffer.getUnsignedShortLE(lastPaddingByteOffset - 1);
          }
  
          protected abstract void readSecurityHeader(SecureChannel channel, ByteBuf chunkBuffer) throws UaException;
  
          protected abstract Cipher getCipher(SecureChannel channel) throws UaException;
  
          protected abstract int getCipherTextBlockSize(SecureChannel channel);
  
          protected abstract int getSignatureSize(SecureChannel channel);
  
          protected abstract void verifyChunk(SecureChannel channel, ByteBuf chunkBuffer) throws UaException;
  
          protected abstract boolean isAsymmetric();
  
          protected abstract boolean isEncryptionEnabled(SecureChannel channel);
  
          protected abstract boolean isSigningEnabled(SecureChannel channel);
  
      }
  
      private final class AsymmetricDecoder extends AbstractDecoder {
  
          @Override
          public void readSecurityHeader(SecureChannel channel, ByteBuf chunkBuffer) {
              AsymmetricSecurityHeader.decode(chunkBuffer, encodingLimits);
          }
  
          @Override
          public Cipher getCipher(SecureChannel channel) throws UaException {
              try {
                  String transformation = channel.getSecurityPolicy()
                      .getAsymmetricEncryptionAlgorithm().getTransformation();
                  Cipher cipher = Cipher.getInstance(transformation);
                  cipher.init(Cipher.DECRYPT_MODE, channel.getKeyPair().getPrivate());
                  return cipher;
              } catch (GeneralSecurityException e) {
                  throw new UaException(StatusCodes.Bad_InternalError, e);
              }
          }
  
          @Override
          public int getCipherTextBlockSize(SecureChannel channel) {
              return channel.getLocalAsymmetricCipherTextBlockSize();
          }
  
          @Override
          public int getSignatureSize(SecureChannel channel) {
              return channel.getRemoteAsymmetricSignatureSize();
          }
  
          @Override
          public void verifyChunk(SecureChannel channel, ByteBuf chunkBuffer) throws UaException {
              String transformation = channel.getSecurityPolicy().getAsymmetricSignatureAlgorithm().getTransformation();
              int signatureSize = channel.getRemoteAsymmetricSignatureSize();
  
              ByteBuffer chunkNioBuffer = chunkBuffer.nioBuffer(0, chunkBuffer.writerIndex());
              ((Buffer) chunkNioBuffer).position(0);
              ((Buffer) chunkNioBuffer).limit(chunkBuffer.writerIndex() - signatureSize);
  
              try {
                  Signature signature = Signature.getInstance(transformation);
  
                  signature.initVerify(channel.getRemoteCertificate().getPublicKey());
                  signature.update(chunkNioBuffer);
  
                  byte[] signatureBytes = new byte[signatureSize];
                  ((Buffer) chunkNioBuffer).limit(chunkNioBuffer.position() + signatureSize);
                  chunkNioBuffer.get(signatureBytes);
  
                  if (!signature.verify(signatureBytes)) {
                      throw new UaException(StatusCodes.Bad_SecurityChecksFailed, "could not verify signature");
                  }
              } catch (NoSuchAlgorithmException e) {
                  throw new UaException(StatusCodes.Bad_InternalError, e);
              } catch (SignatureException e) {
                  throw new UaException(StatusCodes.Bad_ApplicationSignatureInvalid, e);
              } catch (InvalidKeyException e) {
                  throw new UaException(StatusCodes.Bad_CertificateInvalid, e);
              }
          }
  
          @Override
          protected boolean isAsymmetric() {
              return true;
          }
  
          @Override
          public boolean isEncryptionEnabled(SecureChannel channel) {
              return channel.isAsymmetricEncryptionEnabled();
          }
  
          @Override
          public boolean isSigningEnabled(SecureChannel channel) {
              return channel.isAsymmetricEncryptionEnabled();
          }
  
      }
  
      private final class SymmetricDecoder extends AbstractDecoder {
  
          private volatile ChannelSecurity.SecurityKeys securityKeys;
          private volatile Cipher cipher = null;
          private volatile long cipherId = -1;
  
          @Override
          public void readSecurityHeader(SecureChannel channel, ByteBuf chunkBuffer) throws UaException {
              long receivedTokenId = SymmetricSecurityHeader.decode(chunkBuffer).getTokenId();
  
              ChannelSecurity channelSecurity = channel.getChannelSecurity();
  
              if (channelSecurity == null) {
                  if (receivedTokenId != 0L) {
                      throw new UaException(StatusCodes.Bad_SecureChannelTokenUnknown,
                          "unknown secure channel token: " + receivedTokenId);
                  }
              } else {
                  long currentTokenId = channelSecurity.getCurrentToken().getTokenId().longValue();
  
                  if (receivedTokenId == currentTokenId) {
                      securityKeys = channelSecurity.getCurrentKeys();
                  } else {
                      long previousTokenId = channelSecurity.getPreviousToken()
                          .map(t -> t.getTokenId().longValue())
                          .orElse(-1L);
  
                      logger.debug("Attempting to use SecurityKeys from previousTokenId={}", previousTokenId);
  
                      if (receivedTokenId != previousTokenId) {
                          logger.warn(
                              "receivedTokenId={} did not match previousTokenId={}",
                              receivedTokenId, previousTokenId);
  
                          throw new UaException(StatusCodes.Bad_SecureChannelTokenUnknown,
                              "unknown secure channel token: " + receivedTokenId);
                      }
  
                      if (channel.isSymmetricEncryptionEnabled() && channelSecurity.getPreviousKeys().isPresent()) {
                          securityKeys = channelSecurity.getPreviousKeys().get();
                      }
                  }
  
                  if (cipherId != receivedTokenId && channel.isSymmetricEncryptionEnabled()) {
                      cipher = initCipher(channel);
                      cipherId = receivedTokenId;
                  }
              }
          }
  
          @Override
          public Cipher getCipher(SecureChannel channel) {
              assert cipher != null;
              return cipher;
          }
  
          @Override
          public int getCipherTextBlockSize(SecureChannel channel) {
              return channel.getSymmetricBlockSize();
          }
  
          @Override
          public int getSignatureSize(SecureChannel channel) {
              return channel.getSymmetricSignatureSize();
          }
  
          @Override
          public void verifyChunk(SecureChannel channel, ByteBuf chunkBuffer) throws UaException {
              SecurityAlgorithm securityAlgorithm = channel.getSecurityPolicy().getSymmetricSignatureAlgorithm();
              byte[] secretKey = channel.getDecryptionKeys(securityKeys).getSignatureKey();
              int signatureSize = channel.getSymmetricSignatureSize();
  
              ByteBuffer chunkNioBuffer = chunkBuffer.nioBuffer(0, chunkBuffer.writerIndex());
              ((Buffer) chunkNioBuffer).position(0);
              ((Buffer) chunkNioBuffer).limit(chunkBuffer.writerIndex() - signatureSize);
  
              byte[] signature = SignatureUtil.hmac(securityAlgorithm, secretKey, chunkNioBuffer);
  
              byte[] signatureBytes = new byte[signatureSize];
              ((Buffer) chunkNioBuffer).limit(chunkNioBuffer.position() + signatureSize);
              chunkNioBuffer.get(signatureBytes);
  
              if (!MessageDigest.isEqual(signature, signatureBytes)) {
                  throw new UaException(StatusCodes.Bad_SecurityChecksFailed, "could not verify signature");
              }
          }
  
          @Override
          protected boolean isAsymmetric() {
              return false;
          }
  
          @Override
          public boolean isEncryptionEnabled(SecureChannel channel) {
              return channel.isSymmetricEncryptionEnabled();
          }
  
          @Override
          public boolean isSigningEnabled(SecureChannel channel) {
              return channel.isSymmetricSigningEnabled();
          }
  
          private Cipher initCipher(SecureChannel channel) throws UaException {
              try {
                  String transformation = channel.getSecurityPolicy()
                      .getSymmetricEncryptionAlgorithm().getTransformation();
  
                  ChannelSecurity.SecretKeys decryptionKeys = channel.getDecryptionKeys(securityKeys);
  
                  SecretKeySpec keySpec = new SecretKeySpec(decryptionKeys.getEncryptionKey(), "AES");
                  IvParameterSpec ivSpec = new IvParameterSpec(decryptionKeys.getInitializationVector());
  
                  Cipher cipher = Cipher.getInstance(transformation);
                  cipher.init(Cipher.DECRYPT_MODE, keySpec, ivSpec);
  
                  return cipher;
              } catch (GeneralSecurityException e) {
                  throw new UaException(StatusCodes.Bad_InternalError, e);
              }
          }
  
      }
  
  }
  