package org.infinispan.client.hotrod.impl.transport.tcp;
import java.io.BufferedOutputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import javax.security.sasl.Sasl;
import javax.security.sasl.SaslClient;
import javax.security.sasl.SaslException;
/**
* SaslOutputStream.
*
* @author Tristan Tarrant
* @since 7.0
*/
public class SaslOutputStream extends OutputStream {
private static final int BUFFER_SIZE = 64 * 1024;
private final OutputStream outStream;
private final SaslClient saslClient;
private final int bufferSize;
private final ByteArrayOutputStream buffer;
public SaslOutputStream(OutputStream outStream, SaslClient saslClient) {
this.saslClient = saslClient;
this.outStream = new BufferedOutputStream(outStream, BUFFER_SIZE);
String maxSendBuf = (String) saslClient.getNegotiatedProperty(Sasl.RAW_SEND_SIZE);
if (maxSendBuf != null) {
bufferSize = Integer.parseInt(maxSendBuf);
} else {
bufferSize = BUFFER_SIZE;
}
buffer = new ByteArrayOutputStream(bufferSize);
}
@Override
public void write(int b) throws IOException {
checkCapacity(1);
buffer.write(b);
}
@Override
public void write(byte[] b) throws IOException {
write(b, 0, b.length);
}
@Override
public void write(byte[] b, int off, int len) throws IOException {
checkCapacity(len);
buffer.write(b, off, len);
}
private void writeInt(int i) throws IOException {
outStream.write((i >>> 24) & 0xFF);
outStream.write((i >>> 16) & 0xFF);
outStream.write((i >>> 8) & 0xFF);
outStream.write((i >>> 0) & 0xFF);
}
private void checkCapacity(int capacity) throws IOException {
if (buffer.size() + capacity >= bufferSize) {
flush();
}
}
private void wrapAndWrite() throws IOException {
try {
byte[] saslToken = saslClient.wrap(buffer.toByteArray(), 0, buffer.size());
writeInt(saslToken.length);
outStream.write(saslToken);
buffer.reset();
} catch (SaslException se) {
try {
saslClient.dispose();
} catch (SaslException ignored) {
}
throw se;
}
}
/**
* Flushes this output stream
*
* @exception IOException
* if an I/O error occurs.
*/
@Override
public void flush() throws IOException {
wrapAndWrite();
outStream.flush();
}
/**
* Closes this output stream and releases any system resources associated with this stream.
*
* @exception IOException
* if an I/O error occurs.
*/
@Override
public void close() throws IOException {
saslClient.dispose();
outStream.close();
}
}