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

@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
public void flushBuffer(ByteBuf byteBuf, boolean requestedSync, List<IOCallback> callbacks) {
   //maybe no need to perform any copy
   final int bytes = byteBuf.readableBytes();
   if (bytes == 0) {
      IOCallback.done(callbacks);
   } else {
      //enable zero copy case
      if (byteBuf.nioBufferCount() == 1 && byteBuf.isDirect()) {
         final ByteBuffer buffer = byteBuf.internalNioBuffer(byteBuf.readerIndex(), bytes);
         final IOCallback callback = new DelegateCallback(callbacks);
         try {
            //no need to pool the buffer and don't care if the NIO buffer got modified
            internalWrite(buffer, requestedSync, callback, false);
         } catch (Exception e) {
            callback.onError(ActiveMQExceptionType.GENERIC_EXCEPTION.getCode(), e.getMessage());
         }
      } else {
         super.flushBuffer(byteBuf, requestedSync, callbacks);
      }
   }
}
 

/**
 * Optimized accessor for obtaining the underlying NIO buffers for a Netty {@link ByteBuf}. Based
 * on code from Netty's {@code SslHandler}. This method returns NIO buffers that span the readable
 * region of the {@link ByteBuf}.
 */
private static ByteBuffer[] nioBuffers(ByteBuf buf, ByteBuffer[] singleBuffer) {
  // 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 than 1
  // ByteBuf. The worst that can happen is that we allocate an extra ByteBuffer[] in
  // CompositeByteBuf.nioBuffers() which is better than walking the composed ByteBuf in most
  // cases.
  if (!(buf instanceof CompositeByteBuf) && buf.nioBufferCount() == 1) {
    // We know its only backed by 1 ByteBuffer so use internalNioBuffer to keep object
    // allocation to a minimum.
    singleBuffer[0] = buf.internalNioBuffer(buf.readerIndex(), buf.readableBytes());
    return singleBuffer;
  }

  return buf.nioBuffers();
}
 

/**
 * Read bytes into the given {@link ByteBuf} and return the amount.将字节读入给定的ByteBuf并返回数量。
 */
protected final int doReadBytes(ByteBuf byteBuf) throws Exception {
    int writerIndex = byteBuf.writerIndex();
    int localReadAmount;
    unsafe().recvBufAllocHandle().attemptedBytesRead(byteBuf.writableBytes());
    if (byteBuf.hasMemoryAddress()) {
        localReadAmount = socket.readAddress(byteBuf.memoryAddress(), writerIndex, byteBuf.capacity());
    } else {
        ByteBuffer buf = byteBuf.internalNioBuffer(writerIndex, byteBuf.writableBytes());
        localReadAmount = socket.read(buf, buf.position(), buf.limit());
    }
    if (localReadAmount > 0) {
        byteBuf.writerIndex(writerIndex + localReadAmount);
    }
    return localReadAmount;
}
 

/**
 * Read bytes into the given {@link ByteBuf} and return the amount.
 */
protected final int doReadBytes(ByteBuf byteBuf) throws Exception {
    int writerIndex = byteBuf.writerIndex();
    int localReadAmount;
    unsafe().recvBufAllocHandle().attemptedBytesRead(byteBuf.writableBytes());
    if (byteBuf.hasMemoryAddress()) {
        localReadAmount = socket.readAddress(byteBuf.memoryAddress(), writerIndex, byteBuf.capacity());
    } else {
        ByteBuffer buf = byteBuf.internalNioBuffer(writerIndex, byteBuf.writableBytes());
        localReadAmount = socket.read(buf, buf.position(), buf.limit());
    }
    if (localReadAmount > 0) {
        byteBuf.writerIndex(writerIndex + localReadAmount);
    }
    return localReadAmount;
}
 

@Test
public void testNioBuffersSingleBacked() {
    TestChannel channel = new TestChannel();

    ChannelOutboundBuffer buffer = new ChannelOutboundBuffer(channel);
    assertEquals(0, buffer.nioBufferCount());

    ByteBuf buf = copiedBuffer("buf1", CharsetUtil.US_ASCII);
    ByteBuffer nioBuf = buf.internalNioBuffer(buf.readerIndex(), buf.readableBytes());
    buffer.addMessage(buf, buf.readableBytes(), channel.voidPromise());
    assertEquals("Should still be 0 as not flushed yet", 0, buffer.nioBufferCount());
    buffer.addFlush();
    ByteBuffer[] buffers = buffer.nioBuffers();
    assertNotNull(buffers);
    assertEquals("Should still be 0 as not flushed yet", 1, buffer.nioBufferCount());
    for (int i = 0;  i < buffer.nioBufferCount(); i++) {
        if (i == 0) {
            assertEquals(buffers[i], nioBuf);
        } else {
            assertNull(buffers[i]);
        }
    }
    release(buffer);
}
 

@Test
public void testNioBuffersSingleBacked() {
    TestChannel channel = new TestChannel();

    ChannelOutboundBuffer buffer = new ChannelOutboundBuffer(channel);
    assertEquals(0, buffer.nioBufferCount());

    ByteBuf buf = copiedBuffer("buf1", CharsetUtil.US_ASCII);
    ByteBuffer nioBuf = buf.internalNioBuffer(0, buf.readableBytes());
    buffer.addMessage(buf, buf.readableBytes(), channel.voidPromise());
    assertEquals("Should still be 0 as not flushed yet", 0, buffer.nioBufferCount());
    buffer.addFlush();
    ByteBuffer[] buffers = buffer.nioBuffers();
    assertNotNull(buffers);
    assertEquals("Should still be 0 as not flushed yet", 1, buffer.nioBufferCount());
    for (int i = 0;  i < buffer.nioBufferCount(); i++) {
        if (i == 0) {
            assertEquals(buffers[i], nioBuf);
        } else {
            assertNull(buffers[i]);
        }
    }
    release(buffer);
}
 

/**
 * Optimized accessor for obtaining the underlying NIO buffers for a Netty {@link ByteBuf}. Based
 * on code from Netty's {@code SslHandler}. This method returns NIO buffers that span the readable
 * region of the {@link ByteBuf}.
 */
private static ByteBuffer[] nioBuffers(ByteBuf buf, ByteBuffer[] singleBuffer) {
  // 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 than 1
  // ByteBuf. The worst that can happen is that we allocate an extra ByteBuffer[] in
  // CompositeByteBuf.nioBuffers() which is better than walking the composed ByteBuf in most
  // cases.
  if (!(buf instanceof CompositeByteBuf) && buf.nioBufferCount() == 1) {
    // We know its only backed by 1 ByteBuffer so use internalNioBuffer to keep object
    // allocation to a minimum.
    singleBuffer[0] = buf.internalNioBuffer(buf.readerIndex(), buf.readableBytes());
    return singleBuffer;
  }

  return buf.nioBuffers();
}
 

@Override
protected int doReadMessages(List<Object> buf) throws Exception {
    SctpChannel ch = javaChannel();

    RecvByteBufAllocator.Handle allocHandle = unsafe().recvBufAllocHandle();
    ByteBuf buffer = allocHandle.allocate(config().getAllocator());
    boolean free = true;
    try {
        ByteBuffer data = buffer.internalNioBuffer(buffer.writerIndex(), buffer.writableBytes());
        int pos = data.position();

        MessageInfo messageInfo = ch.receive(data, null, notificationHandler);
        if (messageInfo == null) {
            return 0;
        }

        allocHandle.lastBytesRead(data.position() - pos);
        buf.add(new SctpMessage(messageInfo,
                buffer.writerIndex(buffer.writerIndex() + allocHandle.lastBytesRead())));
        free = false;
        return 1;
    } catch (Throwable cause) {
        PlatformDependent.throwException(cause);
        return -1;
    }  finally {
        if (free) {
            buffer.release();
        }
    }
}
 

private void flushBufferedData(ByteBuf out) {
    int flushableBytes = buffer.readableBytes();
    if (flushableBytes == 0) {
        return;
    }
    checksum.reset();
    checksum.update(buffer, buffer.readerIndex(), flushableBytes);
    final int check = (int) checksum.getValue();

    final int bufSize = compressor.maxCompressedLength(flushableBytes) + HEADER_LENGTH;
    out.ensureWritable(bufSize);
    final int idx = out.writerIndex();
    int compressedLength;
    try {
        ByteBuffer outNioBuffer = out.internalNioBuffer(idx + HEADER_LENGTH, out.writableBytes() - HEADER_LENGTH);
        int pos = outNioBuffer.position();
        // We always want to start at position 0 as we take care of reusing the buffer in the encode(...) loop.
        compressor.compress(buffer.internalNioBuffer(buffer.readerIndex(), flushableBytes), outNioBuffer);
        compressedLength = outNioBuffer.position() - pos;
    } catch (LZ4Exception e) {
        throw new CompressionException(e);
    }
    final int blockType;
    if (compressedLength >= flushableBytes) {
        blockType = BLOCK_TYPE_NON_COMPRESSED;
        compressedLength = flushableBytes;
        out.setBytes(idx + HEADER_LENGTH, buffer, 0, flushableBytes);
    } else {
        blockType = BLOCK_TYPE_COMPRESSED;
    }

    out.setLong(idx, MAGIC_NUMBER);
    out.setByte(idx + TOKEN_OFFSET, (byte) (blockType | compressionLevel));
    out.setIntLE(idx + COMPRESSED_LENGTH_OFFSET, compressedLength);
    out.setIntLE(idx + DECOMPRESSED_LENGTH_OFFSET, flushableBytes);
    out.setIntLE(idx + CHECKSUM_OFFSET, check);
    out.writerIndex(idx + HEADER_LENGTH + compressedLength);
    buffer.clear();
}
 

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 int doReadMessages(List<Object> buf) throws Exception {
    DatagramChannel ch = javaChannel();
    ChannelConfig config = config();
    RecvByteBufAllocator.Handle allocHandle = unsafe().recvBufAllocHandle();

    ByteBuf data = allocHandle.allocate(config.getAllocator());
    allocHandle.attemptedBytesRead(data.writableBytes());
    boolean free = true;
    try {
        ByteBuffer nioData = data.internalNioBuffer(data.writerIndex(), data.writableBytes());
        int pos = nioData.position();
        int read = ch.read(nioData);
        if (read <= 0) {
            return read;
        }

        allocHandle.lastBytesRead(nioData.position() - pos);
        buf.add(data.writerIndex(data.writerIndex() + allocHandle.lastBytesRead()));
        free = false;
        return 1;
    } catch (Throwable cause) {
        PlatformDependent.throwException(cause);
        return -1;
    } finally {
        if (free) {
            data.release();
        }
    }
}
 

@Override
protected boolean doWriteMessage(Object msg, ChannelOutboundBuffer in) throws Exception {
    ByteBuf data = (ByteBuf) msg;

    final int dataLen = data.readableBytes();
    if (dataLen == 0) {
        return true;
    }

    final ByteBuffer nioData = data.internalNioBuffer(data.readerIndex(), dataLen);
    final int writtenBytes;
    writtenBytes = javaChannel().write(nioData);
    return writtenBytes > 0;
}
 

@Override
protected int doReadMessages(List<Object> buf) throws Exception {
    DatagramChannel ch = javaChannel();
    UkcpServerChannelConfig config = config();
    RecvByteBufAllocator.Handle allocHandle = unsafe().recvBufAllocHandle();

    ByteBuf data = allocHandle.allocate(config.getAllocator());
    allocHandle.attemptedBytesRead(data.writableBytes());
    boolean free = true;
    try {
        ByteBuffer nioData = data.internalNioBuffer(data.writerIndex(), data.writableBytes());
        int pos = nioData.position();
        InetSocketAddress remoteAddress = (InetSocketAddress) ch.receive(nioData);
        if (remoteAddress == null) {
            return 0;
        }

        allocHandle.lastBytesRead(nioData.position() - pos);
        buf.add(UkcpPacket.newInstance(data.writerIndex(data.writerIndex() + allocHandle.lastBytesRead()),
                remoteAddress));
        free = false;
        return 1;
    } catch (Throwable cause) {
        PlatformDependent.throwException(cause);
        return -1;
    } finally {
        if (free) {
            data.release();
        }
    }
}
 

@Override
protected int doReadMessages(List<Object> buf) throws Exception {
    SctpChannel ch = javaChannel();

    RecvByteBufAllocator.Handle allocHandle = this.allocHandle;
    if (allocHandle == null) {
        this.allocHandle = allocHandle = config().getRecvByteBufAllocator().newHandle();
    }
    ByteBuf buffer = allocHandle.allocate(config().getAllocator());
    boolean free = true;
    try {
        ByteBuffer data = buffer.internalNioBuffer(buffer.writerIndex(), buffer.writableBytes());
        int pos = data.position();

        MessageInfo messageInfo = ch.receive(data, null, notificationHandler);
        if (messageInfo == null) {
            return 0;
        }
        buf.add(new SctpMessage(messageInfo, buffer.writerIndex(buffer.writerIndex() + data.position() - pos)));
        free = false;
        return 1;
    } catch (Throwable cause) {
        PlatformDependent.throwException(cause);
        return -1;
    }  finally {
        int bytesRead = buffer.readableBytes();
        allocHandle.record(bytesRead);
        if (free) {
            buffer.release();
        }
    }
}
 

private void flushBufferedData(ByteBuf out) {
    final int flushableBytes = buffer.readableBytes();
    if (flushableBytes == 0) {
        return;
    }
    checksum.reset();
    checksum.update(buffer.internalNioBuffer(buffer.readerIndex(), flushableBytes));
    final int check = (int) checksum.getValue();

    final int bufSize = (int) Zstd.compressBound(flushableBytes) + HEADER_LENGTH;
    out.ensureWritable(bufSize);
    final int idx = out.writerIndex();
    int compressedLength;
    try {
        ByteBuffer outNioBuffer = out.internalNioBuffer(idx + HEADER_LENGTH, out.writableBytes() - HEADER_LENGTH);
        compressedLength = Zstd.compress(
                outNioBuffer,
                buffer.internalNioBuffer(buffer.readerIndex(), flushableBytes),
                DEFAULT_COMPRESS_LEVEL);
    } catch (Exception e) {
        throw new CompressionException(e);
    }
    final int blockType;
    if (compressedLength >= flushableBytes) {
        blockType = BLOCK_TYPE_NON_COMPRESSED;
        compressedLength = flushableBytes;
        out.setBytes(idx + HEADER_LENGTH, buffer, 0, flushableBytes);
    } else {
        blockType = BLOCK_TYPE_COMPRESSED;
    }

    out.setInt(idx, MAGIC_NUMBER);
    out.setByte(idx + TOKEN_OFFSET, (byte) (blockType | compressionLevel));
    out.setIntLE(idx + COMPRESSED_LENGTH_OFFSET, compressedLength);
    out.setIntLE(idx + DECOMPRESSED_LENGTH_OFFSET, flushableBytes);
    out.setIntLE(idx + CHECKSUM_OFFSET, check);
    out.writerIndex(idx + HEADER_LENGTH + compressedLength);
    buffer.clear();
}
 

@Test
public void testZstdCompressingDecompressingStream() throws IOException {

    String sample = randomString(DEFAULT_BLOCK_SIZE);
    ByteBuf source = Unpooled.directBuffer(DEFAULT_BLOCK_SIZE);
    source.writeBytes(sample.getBytes());

    ByteBuf target = Unpooled.directBuffer(DEFAULT_BLOCK_SIZE);
    int idx = target.writerIndex();
    ByteBuffer buffer = target.internalNioBuffer(idx + HEADER_LENGTH, target.writableBytes() - HEADER_LENGTH);

    int compressLength = Zstd.compress(
            buffer,
            source.internalNioBuffer(source.readerIndex(), source.readableBytes()),
            0);
    target.writerIndex(idx + HEADER_LENGTH + compressLength);
    logger.info("{}", target.readerIndex());
    logger.info("{}", target.readableBytes());
    logger.info("{}", target.writerIndex());

    ByteBuf output = Unpooled.directBuffer(DEFAULT_BLOCK_SIZE);
    ByteBuffer outputBuffer = output.internalNioBuffer(output.writerIndex(), output.writableBytes());

    ByteBuffer inputBuffer = target.internalNioBuffer(target.readerIndex() + HEADER_LENGTH, target.readableBytes() - HEADER_LENGTH);

    logger.info("{}", inputBuffer.position());
    logger.info("{}", inputBuffer.limit());
    Zstd.decompress(outputBuffer, inputBuffer);

    outputBuffer.flip();
    output = Unpooled.wrappedBuffer(outputBuffer);

    String val = output.toString(Charset.defaultCharset());

    Assert.assertEquals(sample, val);
}
 

static ByteBuffer safeNioBuffer(ByteBuf buffer) {
    return buffer.nioBufferCount() == 1 ? buffer.internalNioBuffer(buffer.readerIndex(), buffer.readableBytes())
            : buffer.nioBuffer();
}
 

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() || !engineType.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(newDirectIn.readerIndex(), 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(engine.getSession().getPacketBufferSize());
                break;
            default:
                return result;
            }
        }
    } finally {
        // Null out to allow GC of ByteBuffer
        singleBuffer[0] = null;

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

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();
        }
    }
}