Examples with CompositeByteBuf io.netty.buffer.CompositeByteBuf used on opensource projects

Example 76 with CompositeByteBuf

use of io.netty.buffer.CompositeByteBuf in project drill by apache.

the class SaslEncryptionHandler method encode.

public void encode(ChannelHandlerContext ctx, ByteBuf msg, List<Object> out) throws IOException {
    if (!ctx.channel().isOpen()) {
        logger.debug("In " + RpcConstants.SASL_ENCRYPTION_HANDLER + " and channel is not open. " + "So releasing msg memory before encryption.");
        msg.release();
        return;
    }
    try {
        // If encryption is enabled then this handler will always get ByteBuf of type Composite ByteBuf
        assert (msg instanceof CompositeByteBuf);
        final CompositeByteBuf cbb = (CompositeByteBuf) msg;
        final int numComponents = cbb.numComponents();
        // Get all the components inside the Composite ByteBuf for encryption
        for (int currentIndex = 0; currentIndex < numComponents; ++currentIndex) {
            final ByteBuf component = cbb.component(currentIndex);
            // will break the RPC message into chunks of wrapSizeLimit.
            if (component.readableBytes() > wrapSizeLimit) {
                throw new RpcException(String.format("Component Chunk size: %d is greater than the wrapSizeLimit: %d", component.readableBytes(), wrapSizeLimit));
            }
            // Uncomment the below code if msg can contain both of Direct and Heap ByteBuf. Currently Drill only supports
            // DirectByteBuf so the below condition will always be false. If the msg are always HeapByteBuf then in
            // addition also remove the allocation of origMsgBuffer from constructor.
            /*if (component.hasArray()) {
          origMsg = component.array();
        } else {

        if (RpcConstants.EXTRA_DEBUGGING) {
          logger.trace("The input bytebuf is not backed by a byte array so allocating a new one");
        }*/
            final byte[] origMsg = origMsgBuffer;
            component.getBytes(component.readerIndex(), origMsg, 0, component.readableBytes());
            if (logger.isTraceEnabled()) {
                logger.trace("Trying to encrypt chunk of size:{} with wrapSizeLimit:{}", component.readableBytes(), wrapSizeLimit);
            }
            // Length to encrypt will be component length not origMsg length since that can be greater.
            final byte[] wrappedMsg = saslCodec.wrap(origMsg, 0, component.readableBytes());
            if (logger.isTraceEnabled()) {
                logger.trace("Successfully encrypted message, original size: {} Final Size: {}", component.readableBytes(), wrappedMsg.length);
            }
            // Allocate the buffer (directByteBuff) for copying the encrypted byte array and 4 octets for length of the
            // encrypted message. This is preferred since later on if the passed buffer is not in direct memory then it
            // will be copied by the channel into a temporary direct memory which will be cached to the thread. The size
            // of that temporary direct memory will be size of largest message send.
            final ByteBuf encryptedBuf = ctx.alloc().buffer(wrappedMsg.length + RpcConstants.LENGTH_FIELD_LENGTH);
            // Based on SASL RFC 2222/4422 we should have starting 4 octet as the length of the encrypted buffer in network
            // byte order. SASL framework provided by JDK doesn't do that by default and leaves it upto application. Whereas
            // Cyrus SASL implementation of sasl_encode does take care of this.
            lengthOctets.putInt(wrappedMsg.length);
            encryptedBuf.writeBytes(lengthOctets.array());
            // reset the position for re-use in next round
            lengthOctets.rewind();
            // Write the encrypted bytes inside the buffer
            encryptedBuf.writeBytes(wrappedMsg);
            // Update the msg and component reader index
            msg.skipBytes(component.readableBytes());
            component.skipBytes(component.readableBytes());
            // Add the encrypted buffer into the output to send it on wire.
            out.add(encryptedBuf);
        }
    } catch (OutOfMemoryException e) {
        logger.warn("Failure allocating buffer on incoming stream due to memory limits.");
        msg.resetReaderIndex();
        outOfMemoryHandler.handle();
    } catch (IOException e) {
        logger.error("Something went wrong while wrapping the message: {} with MaxRawWrapSize: {}, " + "and error: {}", msg, wrapSizeLimit, e.getMessage());
        throw e;
    }
}

Example 77 with CompositeByteBuf

use of io.netty.buffer.CompositeByteBuf in project mongo-java-driver by mongodb.

the class NettyStream method readAsync.

/**
 * @param numBytes Must be equal to {@link #pendingReader}{@code .numBytes} when called by a Netty channel handler.
 * @param handler Must be equal to {@link #pendingReader}{@code .handler} when called by a Netty channel handler.
 * @param readTimeoutMillis Must be equal to {@link #NO_SCHEDULE_TIME} when called by a Netty channel handler.
 *                          Timeouts may be scheduled only by the public read methods. Taking into account that concurrent pending
 *                          readers are not allowed, there must not be a situation when threads attempt to schedule a timeout
 *                          before the previous one is either cancelled or completed.
 */
private void readAsync(final int numBytes, final AsyncCompletionHandler<ByteBuf> handler, final long readTimeoutMillis) {
    ByteBuf buffer = null;
    Throwable exceptionResult = null;
    synchronized (this) {
        exceptionResult = pendingException;
        if (exceptionResult == null) {
            if (!hasBytesAvailable(numBytes)) {
                if (pendingReader == null) {
                    // called by a public read method
                    pendingReader = new PendingReader(numBytes, handler, scheduleReadTimeout(readTimeoutTask, readTimeoutMillis));
                }
            } else {
                CompositeByteBuf composite = allocator.compositeBuffer(pendingInboundBuffers.size());
                int bytesNeeded = numBytes;
                for (Iterator<io.netty.buffer.ByteBuf> iter = pendingInboundBuffers.iterator(); iter.hasNext(); ) {
                    io.netty.buffer.ByteBuf next = iter.next();
                    int bytesNeededFromCurrentBuffer = Math.min(next.readableBytes(), bytesNeeded);
                    if (bytesNeededFromCurrentBuffer == next.readableBytes()) {
                        composite.addComponent(next);
                        iter.remove();
                    } else {
                        next.retain();
                        composite.addComponent(next.readSlice(bytesNeededFromCurrentBuffer));
                    }
                    composite.writerIndex(composite.writerIndex() + bytesNeededFromCurrentBuffer);
                    bytesNeeded -= bytesNeededFromCurrentBuffer;
                    if (bytesNeeded == 0) {
                        break;
                    }
                }
                buffer = new NettyByteBuf(composite).flip();
            }
        }
        if (// the read operation has completed
        !(exceptionResult == null && buffer == null) && pendingReader != null) {
            // we need to clear the pending reader
            cancel(pendingReader.timeout);
            this.pendingReader = null;
        }
    }
    if (exceptionResult != null) {
        handler.failed(exceptionResult);
    }
    if (buffer != null) {
        handler.completed(buffer);
    }
}

Example 78 with CompositeByteBuf

use of io.netty.buffer.CompositeByteBuf in project grpc-java by grpc.

the class NettyHandlerTestBase method captureWrite.

protected final ByteBuf captureWrite(ChannelHandlerContext ctx) {
    ArgumentCaptor<ByteBuf> captor = ArgumentCaptor.forClass(ByteBuf.class);
    verify(ctx, atLeastOnce()).write(captor.capture(), any(ChannelPromise.class));
    CompositeByteBuf composite = Unpooled.compositeBuffer();
    for (ByteBuf buf : captor.getAllValues()) {
        composite.addComponent(buf);
        composite.writerIndex(composite.writerIndex() + buf.readableBytes());
    }
    return composite;
}

Example 79 with CompositeByteBuf

use of io.netty.buffer.CompositeByteBuf in project grpc-java by grpc.

the class AltsProtocolNegotiatorTest method protectShouldRoundtrip.

@Test
// List cast
@SuppressWarnings("unchecked")
public void protectShouldRoundtrip() throws Exception {
    doHandshake();
    // Write the message 1 character at a time. The message should be buffered
    // and not interfere with the handshake.
    final AtomicInteger writeCount = new AtomicInteger();
    String message = "hello";
    for (int ix = 0; ix < message.length(); ++ix) {
        ByteBuf in = Unpooled.copiedBuffer(message, ix, 1, UTF_8);
        // go/futurereturn-lsc
        @SuppressWarnings("unused") Future<?> possiblyIgnoredError = channel.write(in).addListener(new ChannelFutureListener() {

            @Override
            public void operationComplete(ChannelFuture future) throws Exception {
                if (future.isSuccess()) {
                    writeCount.incrementAndGet();
                }
            }
        });
    }
    channel.flush();
    // Capture the protected data written to the wire.
    assertEquals(1, channel.outboundMessages().size());
    ByteBuf protectedData = channel.readOutbound();
    assertEquals(message.length(), writeCount.get());
    // Read the protected message at the server and verify it matches the original message.
    TsiFrameProtector serverProtector = serverHandshaker.createFrameProtector(channel.alloc());
    List<ByteBuf> unprotected = new ArrayList<>();
    serverProtector.unprotect(protectedData, (List<Object>) (List<?>) unprotected, channel.alloc());
    // We try our best to remove the HTTP2 handler as soon as possible, but just by constructing it
    // a settings frame is written (and an HTTP2 preface).  This is hard coded into Netty, so we
    // have to remove it here.  See {@code Http2ConnectionHandler.PrefaceDecode.sendPreface}.
    int settingsFrameLength = 9;
    CompositeByteBuf unprotectedAll = new CompositeByteBuf(channel.alloc(), false, unprotected.size() + 1, unprotected);
    ByteBuf unprotectedData = unprotectedAll.slice(settingsFrameLength, message.length());
    assertEquals(message, unprotectedData.toString(UTF_8));
    // Protect the same message at the server.
    final AtomicReference<ByteBuf> newlyProtectedData = new AtomicReference<>();
    serverProtector.protectFlush(Collections.singletonList(unprotectedData), new Consumer<ByteBuf>() {

        @Override
        public void accept(ByteBuf buf) {
            newlyProtectedData.set(buf);
        }
    }, channel.alloc());
    // Read the protected message at the client and verify that it matches the original message.
    channel.writeInbound(newlyProtectedData.get());
    assertEquals(1, channel.inboundMessages().size());
    assertEquals(message, channel.<ByteBuf>readInbound().toString(UTF_8));
}

Example 80 with CompositeByteBuf

use of io.netty.buffer.CompositeByteBuf in project ambry by linkedin.

the class WriteCallback method writeToChannel.

/**
 * Writes {@code writeCount} number of random chunks to the given {@link ByteBufferAsyncWritableChannel}.
 * @param writeCount the number of chunks to write.
 */
public void writeToChannel(int writeCount) {
    for (int i = 0; i < writeCount; i++) {
        WriteCallback writeCallback = new WriteCallback(i);
        byte[] data = new byte[100];
        random.nextBytes(data);
        Future<Long> future = null;
        if (useNettyByteBuf) {
            ByteBuf chunk = null;
            if (!useCompositeByteBuf) {
                chunk = ByteBufAllocator.DEFAULT.heapBuffer(data.length);
                chunk.writeBytes(data);
            } else {
                CompositeByteBuf composite = ByteBufAllocator.DEFAULT.compositeHeapBuffer(100);
                ByteBuf c = ByteBufAllocator.DEFAULT.heapBuffer(50);
                c.writeBytes(data, 0, 50);
                composite.addComponent(true, c);
                c = ByteBufAllocator.DEFAULT.heapBuffer(50);
                c.writeBytes(data, 50, 50);
                composite.addComponent(true, c);
                chunk = composite;
            }
            final ByteBuf finalByteBuf = chunk;
            future = channel.write(finalByteBuf, (result, exception) -> {
                finalByteBuf.release();
                writeCallback.onCompletion(result, exception);
            });
        } else {
            ByteBuffer chunk = ByteBuffer.wrap(data);
            future = channel.write(chunk, writeCallback);
        }
        writes.add(new WriteData(data, future, writeCallback));
    }
}