protected BufferChain wrapWithSasl(BufferChain bc) throws IOException { if (!this.connection.useSasl) return bc; // Looks like no way around this; saslserver wants a byte array. I have to make it one. // THIS IS A BIG UGLY COPY. byte [] responseBytes = bc.getBytes(); byte [] token; // synchronization may be needed since there can be multiple Handler // threads using saslServer or Crypto AES to wrap responses. if (connection.useCryptoAesWrap) { // wrap with Crypto AES synchronized (connection.cryptoAES) { token = connection.cryptoAES.wrap(responseBytes, 0, responseBytes.length); } } else { synchronized (connection.saslServer) { token = connection.saslServer.wrap(responseBytes, 0, responseBytes.length); } } if (RpcServer.LOG.isTraceEnabled()) { RpcServer.LOG.trace("Adding saslServer wrapped token of size " + token.length + " as call response."); } ByteBuffer[] responseBufs = new ByteBuffer[2]; responseBufs[0] = ByteBuffer.wrap(Bytes.toBytes(token.length)); responseBufs[1] = ByteBuffer.wrap(token); return new BufferChain(responseBufs); }
private void writeAndVerify(BufferChain chain, String string, int chunkSize) throws IOException { FileOutputStream fos = new FileOutputStream(tmpFile); FileChannel ch = fos.getChannel(); try { long remaining = string.length(); while (chain.hasRemaining()) { long n = chain.write(ch, chunkSize); assertTrue(n == chunkSize || n == remaining); remaining -= n; } assertEquals(0, remaining); } finally { fos.close(); } assertFalse(chain.hasRemaining()); assertEquals(string, Files.toString(tmpFile, Charsets.UTF_8)); } }
@Override public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise promise) throws Exception { if (msg instanceof RpcResponse) { RpcResponse resp = (RpcResponse) msg; BufferChain buf = resp.getResponse(); ctx.write(Unpooled.wrappedBuffer(buf.getBuffers()), promise).addListener(f -> { resp.done(); if (f.isSuccess()) { metrics.sentBytes(buf.size()); } }); } else { ctx.write(msg, promise); } } }
/** * No protobuf encoding of raw sasl messages */ protected final void doRawSaslReply(SaslStatus status, Writable rv, String errorClass, String error) throws IOException { BufferChain bc; // In my testing, have noticed that sasl messages are usually // in the ballpark of 100-200. That's why the initial capacity is 256. try (ByteBufferOutputStream saslResponse = new ByteBufferOutputStream(256); DataOutputStream out = new DataOutputStream(saslResponse)) { out.writeInt(status.state); // write status if (status == SaslStatus.SUCCESS) { rv.write(out); } else { WritableUtils.writeString(out, errorClass); WritableUtils.writeString(out, error); } bc = new BufferChain(saslResponse.getByteBuffer()); } doRespond(() -> bc); }
@Test public void testWithSpy() throws IOException { ByteBuffer[] bufs = new ByteBuffer[] { stringBuf("XXXhelloYYY", 3, 5), stringBuf(" ", 0, 1), stringBuf("XXXXworldY", 4, 5) }; BufferChain chain = new BufferChain(bufs); FileOutputStream fos = new FileOutputStream(tmpFile); FileChannel ch = Mockito.spy(fos.getChannel()); try { chain.write(ch, 2); assertEquals("he", Files.toString(tmpFile, Charsets.UTF_8)); chain.write(ch, 2); assertEquals("hell", Files.toString(tmpFile, Charsets.UTF_8)); chain.write(ch, 3); assertEquals("hello w", Files.toString(tmpFile, Charsets.UTF_8)); chain.write(ch, 8); assertEquals("hello world", Files.toString(tmpFile, Charsets.UTF_8)); } finally { ch.close(); fos.close(); } }
/** * Expensive. Makes a new buffer to hold a copy of what is in contained ByteBuffers. This * call drains this instance; it cannot be used subsequent to the call. * @return A new byte buffer with the content of all contained ByteBuffers. */ byte [] getBytes() { if (!hasRemaining()) throw new IllegalAccessError(); byte [] bytes = new byte [this.remaining]; int offset = 0; for (ByteBuffer bb: this.buffers) { int length = bb.remaining(); bb.get(bytes, offset, length); offset += length; } return bytes; }
/** * This is a wrapper around {@link java.nio.channels.WritableByteChannel#write(java.nio.ByteBuffer)}. * If the amount of data is large, it writes to channel in smaller chunks. * This is to avoid jdk from creating many direct buffers as the size of * buffer increases. This also minimizes extra copies in NIO layer * as a result of multiple write operations required to write a large * buffer. * * @param channel writable byte channel to write to * @param bufferChain Chain of buffers to write * @return number of bytes written * @throws java.io.IOException e * @see java.nio.channels.WritableByteChannel#write(java.nio.ByteBuffer) */ protected long channelWrite(GatheringByteChannel channel, BufferChain bufferChain) throws IOException { long count = bufferChain.write(channel, NIO_BUFFER_LIMIT); if (count > 0) this.metrics.sentBytes(count); return count; }
protected final RpcResponse getErrorResponse(String msg, Exception e) throws IOException { ResponseHeader.Builder headerBuilder = ResponseHeader.newBuilder().setCallId(-1); ServerCall.setExceptionResponse(e, msg, headerBuilder); ByteBuffer headerBuf = ServerCall.createHeaderAndMessageBytes(null, headerBuilder.build(), 0, null); BufferChain buf = new BufferChain(headerBuf); return () -> buf; }
@Test public void testWithSpy() throws IOException { ByteBuffer[] bufs = new ByteBuffer[] { stringBuf("XXXhelloYYY", 3, 5), stringBuf(" ", 0, 1), stringBuf("XXXXworldY", 4, 5) }; BufferChain chain = new BufferChain(bufs); FileOutputStream fos = new FileOutputStream(tmpFile); FileChannel ch = Mockito.spy(fos.getChannel()); try { chain.write(ch, 2); assertEquals("he", Files.toString(tmpFile, Charsets.UTF_8)); chain.write(ch, 2); assertEquals("hell", Files.toString(tmpFile, Charsets.UTF_8)); chain.write(ch, 3); assertEquals("hello w", Files.toString(tmpFile, Charsets.UTF_8)); chain.write(ch, 8); assertEquals("hello world", Files.toString(tmpFile, Charsets.UTF_8)); } finally { ch.close(); fos.close(); } }
assert !hasRemaining();
/** * This is a wrapper around {@link java.nio.channels.WritableByteChannel#write(java.nio.ByteBuffer)}. * If the amount of data is large, it writes to channel in smaller chunks. * This is to avoid jdk from creating many direct buffers as the size of * buffer increases. This also minimizes extra copies in NIO layer * as a result of multiple write operations required to write a large * buffer. * * @param channel writable byte channel to write to * @param bufferChain Chain of buffers to write * @return number of bytes written * @throws java.io.IOException e * @see java.nio.channels.WritableByteChannel#write(java.nio.ByteBuffer) */ protected long channelWrite(GatheringByteChannel channel, BufferChain bufferChain) throws IOException { long count = bufferChain.write(channel, NIO_BUFFER_LIMIT); if (count > 0) this.metrics.sentBytes(count); return count; }
@Test public void testGetBackBytesWePutIn() { ByteBuffer[] bufs = wrapArrays(HELLO_WORLD_CHUNKS); BufferChain chain = new BufferChain(bufs); assertTrue(Bytes.equals(Bytes.toBytes("hello world"), chain.getBytes())); }
/** * Send the response for connection header */ private void responseConnectionHeader(RPCProtos.ConnectionHeaderResponse.Builder chrBuilder) throws FatalConnectionException { // Response the connection header if Crypto AES is enabled if (!chrBuilder.hasCryptoCipherMeta()) return; try { byte[] connectionHeaderResBytes = chrBuilder.build().toByteArray(); // encrypt the Crypto AES cipher meta data with sasl server, and send to client byte[] unwrapped = new byte[connectionHeaderResBytes.length + 4]; Bytes.putBytes(unwrapped, 0, Bytes.toBytes(connectionHeaderResBytes.length), 0, 4); Bytes.putBytes(unwrapped, 4, connectionHeaderResBytes, 0, connectionHeaderResBytes.length); byte[] wrapped = saslServer.wrap(unwrapped, 0, unwrapped.length); BufferChain bc; try (ByteBufferOutputStream response = new ByteBufferOutputStream(wrapped.length + 4); DataOutputStream out = new DataOutputStream(response)) { out.writeInt(wrapped.length); out.write(wrapped); bc = new BufferChain(response.getByteBuffer()); } doRespond(() -> bc); } catch (IOException ex) { throw new UnsupportedCryptoException(ex.getMessage(), ex); } }
private void writeAndVerify(BufferChain chain, String string, int chunkSize) throws IOException { FileOutputStream fos = new FileOutputStream(tmpFile); FileChannel ch = fos.getChannel(); try { long remaining = string.length(); while (chain.hasRemaining()) { long n = chain.write(ch, chunkSize); assertTrue(n == chunkSize || n == remaining); remaining -= n; } assertEquals(0, remaining); } finally { fos.close(); } assertFalse(chain.hasRemaining()); assertEquals(string, Files.toString(tmpFile, Charsets.UTF_8)); } }
if (!buf.hasRemaining()) { resp.done(); return true;
private BufferChain wrapWithSasl(BufferChain bc) throws IOException { if (!this.connection.useSasl) return bc; // Looks like no way around this; saslserver wants a byte array. I have to make it one. // THIS IS A BIG UGLY COPY. byte [] responseBytes = bc.getBytes(); byte [] token; // synchronization may be needed since there can be multiple Handler // threads using saslServer to wrap responses. synchronized (connection.saslServer) { token = connection.saslServer.wrap(responseBytes, 0, responseBytes.length); } if (LOG.isTraceEnabled()) { LOG.trace("Adding saslServer wrapped token of size " + token.length + " as call response."); } ByteBuffer bbTokenLength = ByteBuffer.wrap(Bytes.toBytes(token.length)); ByteBuffer bbTokenBytes = ByteBuffer.wrap(token); return new BufferChain(bbTokenLength, bbTokenBytes); }
bc = new BufferChain(responseBufs); if (connection.useWrap) { bc = wrapWithSasl(bc);
/** * Expensive. Makes a new buffer to hold a copy of what is in contained ByteBuffers. This * call drains this instance; it cannot be used subsequent to the call. * @return A new byte buffer with the content of all contained ByteBuffers. */ byte [] getBytes() { if (!hasRemaining()) throw new IllegalAccessError(); byte [] bytes = new byte [this.remaining]; int offset = 0; for (ByteBuffer bb: this.buffers) { System.arraycopy(bb.array(), bb.arrayOffset(), bytes, offset, bb.limit()); offset += bb.capacity(); } return bytes; }
@Test public void testGetBackBytesWePutIn() { ByteBuffer[] bufs = wrapArrays(HELLO_WORLD_CHUNKS); BufferChain chain = new BufferChain(bufs); assertTrue(Bytes.equals(Bytes.toBytes("hello world"), chain.getBytes())); }
@Test public void testChainChunkBiggerThanWholeArray() throws IOException { ByteBuffer[] bufs = wrapArrays(HELLO_WORLD_CHUNKS); BufferChain chain = new BufferChain(bufs); writeAndVerify(chain, "hello world", 8192); assertNoRemaining(bufs); }