1 /*
2  * Copyright 2018 The gRPC Authors
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *     http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 package io.grpc.alts.internal;
18 
19 import io.netty.buffer.ByteBuf;
20 import io.netty.buffer.CompositeByteBuf;
21 import java.nio.ByteBuffer;
22 
23 /** Unwraps {@link ByteBuf}s into {@link ByteBuffer}s. */
24 final class BufUnwrapper implements AutoCloseable {
25 
26   private final ByteBuffer[] singleReadBuffer = new ByteBuffer[1];
27   private final ByteBuffer[] singleWriteBuffer = new ByteBuffer[1];
28 
29   /**
30    * Called to get access to the underlying NIO buffers for a {@link ByteBuf} that will be used for
31    * writing.
32    */
writableNioBuffers(ByteBuf buf)33   ByteBuffer[] writableNioBuffers(ByteBuf buf) {
34     // Set the writer index to the capacity to guarantee that the returned NIO buffers will have
35     // the capacity available.
36     int readerIndex = buf.readerIndex();
37     int writerIndex = buf.writerIndex();
38     buf.readerIndex(writerIndex);
39     buf.writerIndex(buf.capacity());
40 
41     try {
42       return nioBuffers(buf, singleWriteBuffer);
43     } finally {
44       // Restore the writer index before returning.
45       buf.readerIndex(readerIndex);
46       buf.writerIndex(writerIndex);
47     }
48   }
49 
50   /**
51    * Called to get access to the underlying NIO buffers for a {@link ByteBuf} that will be used for
52    * reading.
53    */
readableNioBuffers(ByteBuf buf)54   ByteBuffer[] readableNioBuffers(ByteBuf buf) {
55     return nioBuffers(buf, singleReadBuffer);
56   }
57 
58   @Override
close()59   public void close() {
60     singleReadBuffer[0] = null;
61     singleWriteBuffer[0] = null;
62   }
63 
64   /**
65    * Optimized accessor for obtaining the underlying NIO buffers for a Netty {@link ByteBuf}. Based
66    * on code from Netty's {@code SslHandler}. This method returns NIO buffers that span the readable
67    * region of the {@link ByteBuf}.
68    */
nioBuffers(ByteBuf buf, ByteBuffer[] singleBuffer)69   private static ByteBuffer[] nioBuffers(ByteBuf buf, ByteBuffer[] singleBuffer) {
70     // As CompositeByteBuf.nioBufferCount() can be expensive (as it needs to check all composed
71     // ByteBuf to calculate the count) we will just assume a CompositeByteBuf contains more than 1
72     // ByteBuf. The worst that can happen is that we allocate an extra ByteBuffer[] in
73     // CompositeByteBuf.nioBuffers() which is better than walking the composed ByteBuf in most
74     // cases.
75     if (!(buf instanceof CompositeByteBuf) && buf.nioBufferCount() == 1) {
76       // We know its only backed by 1 ByteBuffer so use internalNioBuffer to keep object
77       // allocation to a minimum.
78       singleBuffer[0] = buf.internalNioBuffer(buf.readerIndex(), buf.readableBytes());
79       return singleBuffer;
80     }
81 
82     return buf.nioBuffers();
83   }
84 }
85