Java Code Examples for io.netty.buffer.ByteBuf#nioBuffer()

private void run(ChannelHandlerContext ctx, ByteBuf buf) throws Exception {
    int readerIndex = buf.readerIndex() - 4; 
    buf.readerIndex(readerIndex);
    int size = buf.getIntLE(readerIndex) & 0xffffff;
    ByteBuffer niobuf = buf.nioBuffer(readerIndex, size + 4);
    niobuf.order(ByteOrder.LITTLE_ENDIAN);
    buf.readerIndex(readerIndex + size + 4);
    int result = this.mysession.run(niobuf);
    if (result == 1) {
        return;
    }
    else if (result == -1) {
        switchToSSL();
        return;
    }
    else if (result == -2) {
        ctx.close();
        return;
    }
    else {
        throw new IllegalArgumentException();
    }
}
 

@Override
public ByteBuf decode(ByteBuf encoded, int uncompressedLength) throws IOException {
    ByteBuf uncompressed = PulsarByteBufAllocator.DEFAULT.directBuffer(uncompressedLength, uncompressedLength);

    if (encoded.hasMemoryAddress()) {
        Zstd.decompressUnsafe(uncompressed.memoryAddress(), uncompressedLength,
                encoded.memoryAddress() + encoded.readerIndex(),
                encoded.readableBytes());
    } else {
        ByteBuffer uncompressedNio = uncompressed.nioBuffer(0, uncompressedLength);
        ByteBuffer encodedNio = encoded.nioBuffer(encoded.readerIndex(), encoded.readableBytes());

        Zstd.decompress(uncompressedNio, encodedNio);
    }

    uncompressed.writerIndex(uncompressedLength);
    return uncompressed;
}
 

/**
 * Transfer {@code dataSize} bytes of data from the given crc stream to a newly create {@link ByteBuffer}. The method
 * would also update the crc value in the crc stream.
 * @param crcStream The crc stream.
 * @param dataSize The number of bytes to transfer.
 * @return the newly created {@link ByteBuffer} which contains the transferred data.
 * @throws IOException Any I/O error.
 */
public static ByteBuffer readByteBufferFromCrcInputStream(CrcInputStream crcStream, int dataSize) throws IOException {
  ByteBuffer output;
  InputStream inputStream = crcStream.getUnderlyingInputStream();
  if (inputStream instanceof NettyByteBufDataInputStream) {
    // getBuffer() doesn't increase the reference count on this ByteBuf.
    ByteBuf nettyByteBuf = ((NettyByteBufDataInputStream) inputStream).getBuffer();
    // construct a java.nio.ByteBuffer to create a ByteBufferInputStream
    int startIndex = nettyByteBuf.readerIndex();
    output = nettyByteBuf.nioBuffer(startIndex, dataSize);
    crcStream.updateCrc(output.duplicate());
    nettyByteBuf.readerIndex(startIndex + dataSize);
  } else {
    output = getByteBufferFromInputStream(crcStream, dataSize);
  }
  return output;
}
 

@SuppressWarnings("BetaApi") // verify is stable in Guava
@Override
public void decrypt(ByteBuf out, ByteBuf ciphertextAndTag) throws GeneralSecurityException {
  int bytesRead = ciphertextAndTag.readableBytes();
  checkArgument(bytesRead == out.writableBytes());

  checkArgument(out.nioBufferCount() == 1);
  ByteBuffer outBuffer = out.internalNioBuffer(out.writerIndex(), out.writableBytes());

  checkArgument(ciphertextAndTag.nioBufferCount() == 1);
  ByteBuffer ciphertextAndTagBuffer =
      ciphertextAndTag.nioBuffer(ciphertextAndTag.readerIndex(), bytesRead);

  byte[] counter = incrementInCounter();
  int outPosition = outBuffer.position();
  aeadCrypter.decrypt(outBuffer, ciphertextAndTagBuffer, counter);
  int bytesWritten = outBuffer.position() - outPosition;
  out.writerIndex(out.writerIndex() + bytesWritten);
  ciphertextAndTag.readerIndex(out.readerIndex() + bytesRead);
  verify(out.writableBytes() == TAG_LENGTH);
}
 

@Override
// length parameter is ignored here because of the way the fastDecompressor works.
public ByteBuf decompress(ByteBuf compressed, int decompressedSize) {
    checkNotNull(compressed);
    checkArgument(compressed.readableBytes() >= 0);
    checkArgument(decompressedSize >= 0);

    ByteBuf uncompressed = PooledByteBufAllocator.DEFAULT.buffer(decompressedSize, decompressedSize);
    ByteBuffer uncompressedNio = uncompressed.nioBuffer(0, decompressedSize);
    ByteBuffer compressedNio = compressed.nioBuffer(compressed.readerIndex(), compressed.readableBytes());

    decompressor.decompress(
            compressedNio, compressedNio.position(),
            uncompressedNio, uncompressedNio.position(), uncompressedNio.remaining());
    uncompressed.writerIndex(decompressedSize);
    return uncompressed;
}
 

public ByteBuffer getBuffer(int len)
{
    ByteBuf byteBuf = memoryAllocator.directBuffer(len);
    ByteBuffer  byteBuffer = byteBuf.nioBuffer(0,len);
    freeMaps.put(PlatformDependent.directBufferAddress(byteBuffer),byteBuf);
    return byteBuffer;
}
 

@Override
public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise promise) {
    DatagramPacket datagramPacket = (DatagramPacket) msg;
    ByteBuf data = datagramPacket.content();
    ByteBuffer byteBuffer = data.nioBuffer(0,data.readableBytes());
    crc32.reset();
    crc32.update(byteBuffer);
    long checksum = crc32.getValue();
    ByteBuf headByteBuf = ctx.alloc().ioBuffer(4);
    headByteBuf.writeIntLE((int) checksum);
    ByteBuf newByteBuf = Unpooled.wrappedBuffer(headByteBuf,data);
    datagramPacket = datagramPacket.replace(newByteBuf);
    ctx.write(datagramPacket, promise);
}
 

public DataTransfer(RemotingHeader requestHeader, ByteBuf payload) {
    this.payload = payload;
    this.header = HeaderSerializer.serialize(requestHeader, payload.readableBytes(), 0);

    this.buffers = new ByteBuffer[2];
    this.buffers[0] = header.nioBuffer();
    this.buffers[1] = payload.nioBuffer();

    this.count = header.readableBytes() + payload.readableBytes();
}
 

public boolean appendIndex(Map<String, PullLogIndexEntry> indexMap) {
    positionOfIndex = tablet.getWrotePosition();

    for (Map.Entry<String, PullLogIndexEntry> entry : indexMap.entrySet()) {
        final byte[] consumerBytes = entry.getKey().getBytes(StandardCharsets.UTF_8);
        int size = Integer.BYTES + Short.BYTES + consumerBytes.length + Long.BYTES + Long.BYTES + Integer.BYTES + Integer.BYTES;
        PullLogIndexEntry indexEntry = entry.getValue();

        ByteBuf buffer = ByteBufAllocator.DEFAULT.ioBuffer(size);
        try {
            buffer.writeInt(MagicCode.PULL_LOG_MAGIC_V1);
            buffer.writeShort((short) consumerBytes.length);
            buffer.writeBytes(consumerBytes);
            buffer.writeLong(indexEntry.startOfPullLogSequence);
            buffer.writeLong(indexEntry.baseOfMessageSequence);
            buffer.writeInt(indexEntry.position);
            buffer.writeInt(indexEntry.num);
            ByteBuffer nioBuffer = buffer.nioBuffer();
            Checksums.update(crc, nioBuffer, nioBuffer.limit());
            boolean result = tablet.appendData(nioBuffer);
            if (!result) return false;
        } finally {
            ReferenceCountUtil.safeRelease(buffer);
        }

        ConcurrentSkipListMap<PullLogSequence, SegmentLocation> index = sortedPullLogTable.index;
        index.put(new PullLogSequence(entry.getKey(), indexEntry.startOfPullLogSequence), new SegmentLocation(indexEntry.baseOfMessageSequence, indexEntry.position, indexEntry.num, tablet));
        tablet.retain();
    }
    return true;
}
 

@Override
public Collection<Message<P>> decode(ByteBuf inputBuffer) {
	ByteBuffer nioBuffer = inputBuffer.nioBuffer();
	int start = nioBuffer.position();
	List<Message<P>> messages = decodeInternal(nioBuffer);
	inputBuffer.skipBytes(nioBuffer.position() - start);
	return messages;
}
 

@Override
public void setContent(ByteBuf buffer) throws IOException {
    if (buffer == null) {
        throw new NullPointerException("buffer");
    }
    try {
        size = buffer.readableBytes();
        if (definedSize > 0 && definedSize < size) {
            throw new IOException("Out of size: " + size + " > " + definedSize);
        }
        if (file == null) {
            file = tempFile();
        }
        if (buffer.readableBytes() == 0) {
            // empty file
            file.createNewFile();
            return;
        }
        FileOutputStream outputStream = new FileOutputStream(file);
        try {
            FileChannel localfileChannel = outputStream.getChannel();
            ByteBuffer byteBuffer = buffer.nioBuffer();
            int written = 0;
            while (written < size) {
                written += localfileChannel.write(byteBuffer);
            }
            buffer.readerIndex(buffer.readerIndex() + written);
            localfileChannel.force(false);
        } finally {
            outputStream.close();
        }
        completed = true;
    } finally {
        // Release the buffer as it was retained before and we not need a reference to it at all
        // See https://github.com/netty/netty/issues/1516
        buffer.release();
    }
}
 

public static char[] readUtf8(ByteBuf byteBuf, int length) {
  CharsetDecoder charsetDecoder = CharsetUtil.UTF_8.newDecoder();
  int en = (int) (length * (double) charsetDecoder.maxCharsPerByte());
  char[] ca = new char[en];

  CharBuffer charBuffer = CharBuffer.wrap(ca);
  ByteBuffer byteBuffer =
      byteBuf.nioBufferCount() == 1
          ? byteBuf.internalNioBuffer(byteBuf.readerIndex(), length)
          : byteBuf.nioBuffer(byteBuf.readerIndex(), length);
  byteBuffer.mark();
  try {
    CoderResult cr = charsetDecoder.decode(byteBuffer, charBuffer, true);
    if (!cr.isUnderflow()) cr.throwException();
    cr = charsetDecoder.flush(charBuffer);
    if (!cr.isUnderflow()) cr.throwException();

    byteBuffer.reset();
    byteBuf.skipBytes(length);

    return safeTrim(charBuffer.array(), charBuffer.position());
  } catch (CharacterCodingException x) {
    // Substitution is always enabled,
    // so this shouldn't happen
    throw new IllegalStateException("unable to decode char array from the given buffer", x);
  }
}
 

@Override
protected void channelRead0(ChannelHandlerContext ctx, WebSocketFrame frame) {
  if (!frame.isFinalFragment()) {
    LOG.warning("Frame is not final. Chaos may ensue");
  }

  Message message = null;

  if (frame instanceof TextWebSocketFrame) {
    message = new TextMessage(((TextWebSocketFrame) frame).text());
  } else if (frame instanceof BinaryWebSocketFrame) {
    ByteBuf buf = frame.content();
    if (buf.nioBufferCount() != -1) {
      message = new BinaryMessage(buf.nioBuffer());
    } else if (buf.hasArray()) {
      message = new BinaryMessage(ByteBuffer.wrap(buf.array()));
    } else {
      throw new IllegalStateException("Unable to handle bytebuf: " + buf);
    }
  } else if (frame instanceof CloseWebSocketFrame) {
    CloseWebSocketFrame closeFrame = (CloseWebSocketFrame) frame;
    message = new CloseMessage(closeFrame.statusCode(), closeFrame.reasonText());
  }

  if (message != null) {
    ctx.fireChannelRead(message);
  } else {
    ctx.write(frame);
  }
}
 

@Test
public void testResponseToByteBuf() throws Exception {
    int correlationId = 7777;
    String clientId = "KopClientId";

    ApiVersionsRequest apiVersionsRequest = new ApiVersionsRequest.Builder().build();
    RequestHeader requestHeader = new RequestHeader(
        ApiKeys.API_VERSIONS,
        ApiKeys.API_VERSIONS.latestVersion(),
        clientId,
        correlationId);

    KafkaHeaderAndRequest kopRequest = new KafkaHeaderAndRequest(
        requestHeader,
        apiVersionsRequest,
        Unpooled.buffer(20),
        null);

    ApiVersionsResponse apiVersionsResponse = ApiVersionsResponse.defaultApiVersionsResponse();
    KafkaHeaderAndResponse kopResponse = KafkaHeaderAndResponse.responseForRequest(
        kopRequest, apiVersionsResponse);

    // 1. serialize response into ByteBuf
    ByteBuf serializedResponse = handler.responseToByteBuf(kopResponse.getResponse(), kopRequest);

    // 2. verify responseToByteBuf works well.
    ByteBuffer byteBuffer = serializedResponse.nioBuffer();
    ResponseHeader responseHeader = ResponseHeader.parse(byteBuffer);
    assertEquals(responseHeader.correlationId(), correlationId);

    ApiVersionsResponse parsedResponse = ApiVersionsResponse.parse(
        byteBuffer, kopResponse.getApiVersion());

    assertEquals(parsedResponse.apiVersions().size(), apiVersionsResponse.apiVersions().size());
}
 

private void decryptChunk(Delegate delegate, SecureChannel channel, ByteBuf chunkBuffer) throws UaException {
    int cipherTextBlockSize = delegate.getCipherTextBlockSize(channel);
    int blockCount = chunkBuffer.readableBytes() / cipherTextBlockSize;

    int plainTextBufferSize = cipherTextBlockSize * blockCount;

    ByteBuf plainTextBuffer = BufferUtil.buffer(plainTextBufferSize);

    ByteBuffer plainTextNioBuffer = plainTextBuffer
            .writerIndex(plainTextBufferSize)
            .nioBuffer();

    ByteBuffer chunkNioBuffer = chunkBuffer.nioBuffer();

    try {
        Cipher cipher = delegate.getCipher(channel);

        assert (chunkBuffer.readableBytes() % cipherTextBlockSize == 0);

        if (delegate instanceof AsymmetricDelegate) {
            for (int blockNumber = 0; blockNumber < blockCount; blockNumber++) {
                chunkNioBuffer.limit(chunkNioBuffer.position() + cipherTextBlockSize);

                cipher.doFinal(chunkNioBuffer, plainTextNioBuffer);
            }
        } else {
            cipher.doFinal(chunkNioBuffer, plainTextNioBuffer);
        }
    } catch (GeneralSecurityException e) {
        throw new UaException(StatusCodes.Bad_SecurityChecksFailed, e);
    }

    /* Write plainTextBuffer back into the chunk buffer we decrypted from. */
    plainTextNioBuffer.flip(); // limit = pos, pos = 0

    chunkBuffer.writerIndex(chunkBuffer.readerIndex());
    chunkBuffer.writeBytes(plainTextNioBuffer);

    plainTextBuffer.release();
}
 

/**
 * Helper function to encrypt ByteBuf with the given key. When useFixedIv is true, don't use a random iv byte
 * array, using a all zero byte array instead. Only set it to be true in test.
 * @param toEncrypt {@link ByteBuf} that needs to be encrypted
 * @param key the secret key (of type T) to use to encrypt
 * @param iv If null, will create a random byte array serve as iv bytes.
 * @return the {@link ByteBuf} containing the encrypted content. Ensure the result has all
 * the information like the IV along with the encrypted content, in order to decrypt the content with a given key
 * @throws {@link GeneralSecurityException} on any exception with encryption
 */
public ByteBuf encrypt(ByteBuf toEncrypt, SecretKeySpec key, byte[] iv) throws GeneralSecurityException {
  try {
    Cipher encrypter = Cipher.getInstance(GCM_CRYPTO_INSTANCE, "BC");
    if (iv == null) {
      iv = new byte[ivValSize];
      random.nextBytes(iv);
    }
    encrypter.init(Cipher.ENCRYPT_MODE, key, new IvParameterSpec(iv));
    int outputSize = encrypter.getOutputSize(toEncrypt.readableBytes());

    // stick with heap memory for now so to compare with the java.nio.ByteBuffer.
    ByteBuf encryptedContent =
        ByteBufAllocator.DEFAULT.heapBuffer(IVRecord_Format_V1.getIVRecordSize(iv) + outputSize);
    IVRecord_Format_V1.serializeIVRecord(encryptedContent, iv);

    boolean toRelease = false;
    if (toEncrypt.nioBufferCount() != 1) {
      toRelease = true;
      ByteBuf temp = ByteBufAllocator.DEFAULT.heapBuffer(toEncrypt.readableBytes());
      temp.writeBytes(toEncrypt);
      toEncrypt = temp;
    }
    try {
      ByteBuffer toEncryptBuffer = toEncrypt.nioBuffer();
      ByteBuffer encryptedContentBuffer = encryptedContent.nioBuffer(encryptedContent.writerIndex(),
          encryptedContent.capacity() - encryptedContent.writerIndex());
      int n = encrypter.doFinal(toEncryptBuffer, encryptedContentBuffer);
      encryptedContent.writerIndex(encryptedContent.writerIndex() + n);
      return encryptedContent;
    } finally {
      if (toRelease) {
        toEncrypt.release();
      } else {
        toEncrypt.skipBytes(toEncrypt.readableBytes());
      }
    }
  } catch (Exception e) {
    throw new GeneralSecurityException("Exception thrown while encrypting data", e);
  }
}
 

@Override
public synchronized ByteBuf encrypt(Set<String> encKeys, CryptoKeyReader keyReader,
                                    Supplier<Builder> messageMetadataBuilderSupplier, ByteBuf payload) throws PulsarClientException {

    Builder msgMetadata = messageMetadataBuilderSupplier.get();

    if (encKeys.isEmpty()) {
        return payload;
    }

    // Update message metadata with encrypted data key
    for (String keyName : encKeys) {
        if (encryptedDataKeyMap.get(keyName) == null) {
            // Attempt to load the key. This will allow us to load keys as soon as
            // a new key is added to producer config
            addPublicKeyCipher(keyName, keyReader);
        }
        EncryptionKeyInfo keyInfo = encryptedDataKeyMap.get(keyName);
        if (keyInfo != null) {
            if (keyInfo.getMetadata() != null && !keyInfo.getMetadata().isEmpty()) {
                List<KeyValue> kvList = new ArrayList<KeyValue>();
                keyInfo.getMetadata().forEach((key, value) -> {
                    kvList.add(KeyValue.newBuilder().setKey(key).setValue(value).build());
                });
                msgMetadata.addEncryptionKeys(EncryptionKeys.newBuilder().setKey(keyName)
                        .setValue(ByteString.copyFrom(keyInfo.getKey())).addAllMetadata(kvList).build());
            } else {
                msgMetadata.addEncryptionKeys(EncryptionKeys.newBuilder().setKey(keyName)
                        .setValue(ByteString.copyFrom(keyInfo.getKey())).build());
            }
        } else {
            // We should never reach here.
            log.error("{} Failed to find encrypted Data key for key {}.", logCtx, keyName);
        }

    }

    // Create gcm param
    // TODO: Replace random with counter and periodic refreshing based on timer/counter value
    secureRandom.nextBytes(iv);
    GCMParameterSpec gcmParam = new GCMParameterSpec(tagLen, iv);

    // Update message metadata with encryption param
    msgMetadata.setEncryptionParam(ByteString.copyFrom(iv));

    ByteBuf targetBuf = null;
    try {
        // Encrypt the data
        cipher.init(Cipher.ENCRYPT_MODE, dataKey, gcmParam);

        ByteBuffer sourceNioBuf = payload.nioBuffer(payload.readerIndex(), payload.readableBytes());

        int maxLength = cipher.getOutputSize(payload.readableBytes());
        targetBuf = PulsarByteBufAllocator.DEFAULT.buffer(maxLength, maxLength);
        ByteBuffer targetNioBuf = targetBuf.nioBuffer(0, maxLength);

        int bytesStored = cipher.doFinal(sourceNioBuf, targetNioBuf);
        targetBuf.writerIndex(bytesStored);

    } catch (IllegalBlockSizeException | BadPaddingException | InvalidKeyException
            | InvalidAlgorithmParameterException | ShortBufferException e) {

        targetBuf.release();
        log.error("{} Failed to encrypt message. {}", logCtx, e);
        throw new PulsarClientException.CryptoException(e.getMessage());

    }

    payload.release();
    return targetBuf;
}
 

private SSLEngineResult wrap(ByteBufAllocator alloc, SSLEngine engine, ByteBuf in, ByteBuf out)
        throws SSLException {
    ByteBuf newDirectIn = null;
    try {
        int readerIndex = in.readerIndex();
        int readableBytes = in.readableBytes();

        // We will call SslEngine.wrap(ByteBuffer[], ByteBuffer) to allow efficient handling of
        // CompositeByteBuf without force an extra memory copy when CompositeByteBuffer.nioBuffer() is called.
        final ByteBuffer[] in0;
        if (in.isDirect() || !wantsDirectBuffer) {
            // As CompositeByteBuf.nioBufferCount() can be expensive (as it needs to check all composed ByteBuf
            // to calculate the count) we will just assume a CompositeByteBuf contains more then 1 ByteBuf.
            // The worst that can happen is that we allocate an extra ByteBuffer[] in CompositeByteBuf.nioBuffers()
            // which is better then walking the composed ByteBuf in most cases.
            if (!(in instanceof CompositeByteBuf) && in.nioBufferCount() == 1) {
                in0 = singleBuffer;
                // We know its only backed by 1 ByteBuffer so use internalNioBuffer to keep object allocation
                // to a minimum.
                in0[0] = in.internalNioBuffer(readerIndex, readableBytes);
            } else {
                in0 = in.nioBuffers();
            }
        } else {
            // We could even go further here and check if its a CompositeByteBuf and if so try to decompose it and
            // only replace the ByteBuffer that are not direct. At the moment we just will replace the whole
            // CompositeByteBuf to keep the complexity to a minimum
            newDirectIn = alloc.directBuffer(readableBytes);
            newDirectIn.writeBytes(in, readerIndex, readableBytes);
            in0 = singleBuffer;
            in0[0] = newDirectIn.internalNioBuffer(0, readableBytes);
        }

        for (;;) {
            ByteBuffer out0 = out.nioBuffer(out.writerIndex(), out.writableBytes());
            SSLEngineResult result = engine.wrap(in0, out0);
            in.skipBytes(result.bytesConsumed());
            out.writerIndex(out.writerIndex() + result.bytesProduced());

            switch (result.getStatus()) {
            case BUFFER_OVERFLOW:
                out.ensureWritable(maxPacketBufferSize);
                break;
            default:
                return result;
            }
        }
    } finally {
        // Null out to allow GC of ByteBuffer
        singleBuffer[0] = null;

        if (newDirectIn != null) {
            newDirectIn.release();
        }
    }
}
 

@Override
public ByteBuffer encodeValue(V value) {
  ByteBuf buf = Unpooled.buffer(estimateSize(value));
  encodeValue(value, buf);
  return buf.nioBuffer();
}
 

ByteBufSegmentBuffer(long startOffset, LogSegment segment, ByteBuf buffer, int size) {
    super(startOffset, buffer.nioBuffer(), size, segment);
    this.byteBuf = buffer;
}