Commit 9300692a authored by Alex Wenckus's avatar Alex Wenckus Committed by alex

Implemented blocking NIO for file transfer proxy. JM-782

git-svn-id: http://svn.igniterealtime.org/svn/repos/wildfire/trunk@4535 b35dd754-fafc-0310-a699-88a17e54d16e
parent 25653287
......@@ -11,8 +11,8 @@
package org.jivesoftware.wildfire.filetransfer;
import java.util.concurrent.Future;
import java.io.InputStream;
import java.io.OutputStream;
import java.nio.channels.ReadableByteChannel;
import java.nio.channels.WritableByteChannel;
/**
* An interface to track the progress of a file transfer through the server. This interface is used
......@@ -58,11 +58,11 @@ public interface FileTransferProgress {
*/
public void setTransferFuture(Future<?> future);
public void setInputStream(InputStream initiatorInputStream);
public void setInputChannel(ReadableByteChannel inputChannel);
public InputStream getInputStream();
public ReadableByteChannel getInputChannel();
public void setOutputStream(OutputStream targetOutputStream);
public void setOutputChannel(WritableByteChannel outputChannel);
public OutputStream getOutputStream();
public WritableByteChannel getOutputChannel();
}
......@@ -14,8 +14,9 @@ import org.jivesoftware.util.CacheSizes;
import java.util.concurrent.Future;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.nio.channels.ReadableByteChannel;
import java.nio.channels.WritableByteChannel;
import java.nio.ByteBuffer;
/**
* Tracks the different connections related to a file transfer. There are two connections, the
......@@ -25,9 +26,9 @@ public class DefaultProxyTransfer implements ProxyTransfer {
private String initiator;
private InputStream inputStream;
private ReadableByteChannel inputStream;
private OutputStream outputStream;
private WritableByteChannel outputStream;
private String target;
......@@ -52,20 +53,20 @@ public class DefaultProxyTransfer implements ProxyTransfer {
this.initiator = initiator;
}
public InputStream getInputStream() {
public ReadableByteChannel getInputChannel() {
return inputStream;
}
public void setInputStream(InputStream initiatorInputStream) {
this.inputStream = initiatorInputStream;
public void setInputChannel(ReadableByteChannel inputChannel) {
this.inputStream = inputChannel;
}
public OutputStream getOutputStream() {
public WritableByteChannel getOutputChannel() {
return outputStream;
}
public void setOutputStream(OutputStream outputStream) {
this.outputStream = outputStream;
public void setOutputChannel(WritableByteChannel outputChannel) {
this.outputStream = outputChannel;
}
public String getTarget() {
......@@ -112,25 +113,27 @@ public class DefaultProxyTransfer implements ProxyTransfer {
if(!isActivatable()) {
throw new IOException("Transfer missing party");
}
InputStream in = getInputStream();
OutputStream out = new ProxyOutputStream(getOutputStream());
ReadableByteChannel in = getInputChannel();
WritableByteChannel out = new ProxyOutputChannel(getOutputChannel());
final byte[] b = new byte[BUFFER_SIZE];
ByteBuffer buffer = ByteBuffer.allocate(BUFFER_SIZE);
int count = 0;
amountWritten = 0;
do {
// write to the output stream
out.write(b, 0, count);
// write to the output channel
out.write(buffer);
amountWritten += count;
// read more bytes from the input stream
count = in.read(b);
// read more bytes from the input channel
buffer.clear();
count = in.read(buffer);
buffer.flip();
} while (count >= 0);
getInputStream().close();
getOutputStream().close();
in.close();
out.close();
}
public int getCachedSize() {
......
......@@ -20,13 +20,15 @@ import org.jivesoftware.wildfire.stats.i18nStatistic;
import org.xmpp.packet.JID;
import java.io.*;
import java.net.ServerSocket;
import java.net.Socket;
import java.net.InetAddress;
import java.net.*;
import java.util.Map;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.nio.channels.ServerSocketChannel;
import java.nio.channels.SocketChannel;
import java.nio.channels.WritableByteChannel;
import java.nio.ByteBuffer;
/**
* Manages the connections to the proxy server. The connections go through two stages before
......@@ -48,14 +50,14 @@ public class ProxyConnectionManager {
private Future<?> socketProcess;
private ServerSocket serverSocket;
private int proxyPort;
private FileTransferManager transferManager;
private String className;
private ServerSocketChannel serverChannel;
public ProxyConnectionManager(FileTransferManager manager) {
String cacheName = "File Transfer";
connectionMap = new Cache<String, ProxyTransfer>(cacheName, -1, 1000 * 60 * 10);
......@@ -80,20 +82,23 @@ public class ProxyConnectionManager {
reset();
socketProcess = executor.submit(new Runnable() {
public void run() {
ServerSocket serverSocket;
try {
serverSocket = new ServerSocket(port, -1, bindInterface);
serverChannel = ServerSocketChannel.open();
serverSocket = serverChannel.socket();
serverSocket.bind(new InetSocketAddress(bindInterface, port));
}
catch (IOException e) {
Log.error("Error creating server socket", e);
Log.error("Error binding server socket", e);
return;
}
while (serverSocket.isBound()) {
final Socket socket;
final SocketChannel channel;
try {
socket = serverSocket.accept();
channel = serverChannel.accept();
}
catch (IOException e) {
if (!serverSocket.isClosed()) {
if (serverChannel.isOpen()) {
Log.error("Error accepting proxy connection", e);
continue;
}
......@@ -104,13 +109,13 @@ public class ProxyConnectionManager {
executor.submit(new Runnable() {
public void run() {
try {
processConnection(socket);
processConnection(channel);
}
catch (IOException ie) {
Log.error("Error processing file transfer proxy connection",
ie);
try {
socket.close();
channel.close();
}
catch (IOException e) {
/* Do Nothing */
......@@ -128,21 +133,21 @@ public class ProxyConnectionManager {
return proxyPort;
}
private void processConnection(Socket connection) throws IOException {
OutputStream out = new DataOutputStream(connection.getOutputStream());
InputStream in = new DataInputStream(connection.getInputStream());
private void processConnection(SocketChannel connection) throws IOException {
ByteBuffer buffer = ByteBuffer.allocate(512);
connection.read(buffer);
buffer.flip();
// first byte is version should be 5
int b = in.read();
int b = buffer.get();
if (b != 5) {
throw new IOException("Only SOCKS5 supported");
}
// second byte number of authentication methods supported
b = in.read();
b = buffer.get();
int[] auth = new int[b];
for (int i = 0; i < b; i++) {
auth[i] = in.read();
auth[i] = buffer.get();
}
int authMethod = -1;
......@@ -163,9 +168,11 @@ public class ProxyConnectionManager {
byte[] cmd = new byte[2];
cmd[0] = (byte) 0x05;
cmd[1] = (byte) 0x00;
out.write(cmd);
buffer.clear();
buffer.put(cmd).flip();
connection.write(buffer);
String responseDigest = processIncomingSocks5Message(in);
String responseDigest = processIncomingSocks5Message(connection, buffer);
try {
synchronized (connectionLock) {
ProxyTransfer transfer = connectionMap.get(responseDigest);
......@@ -175,21 +182,26 @@ public class ProxyConnectionManager {
connectionMap.put(responseDigest, transfer);
}
else {
transfer.setInputStream(connection.getInputStream());
transfer.setInputChannel(connection);
}
}
cmd = createOutgoingSocks5Message(0, responseDigest);
out.write(cmd);
buffer.clear();
buffer.put(cmd).flip();
connection.write(buffer);
}
catch (UnauthorizedException eu) {
cmd = createOutgoingSocks5Message(2, responseDigest);
out.write(cmd);
buffer.clear();
buffer.put(cmd).flip();
connection.write(buffer);
throw new IOException("Illegal proxy transfer");
}
}
private ProxyTransfer createProxyTransfer(String transferDigest, Socket targetSocket)
throws IOException {
private ProxyTransfer createProxyTransfer(String transferDigest,
WritableByteChannel targetSocket)
{
ProxyTransfer provider;
try {
Class c = ClassUtils.forName(className);
......@@ -201,32 +213,26 @@ public class ProxyConnectionManager {
}
provider.setTransferDigest(transferDigest);
provider.setOutputStream(targetSocket.getOutputStream());
provider.setOutputChannel(targetSocket);
return provider;
}
@SuppressWarnings({"ResultOfMethodCallIgnored"})
private static String processIncomingSocks5Message(InputStream in)
private static String processIncomingSocks5Message(SocketChannel in, ByteBuffer buffer)
throws IOException {
buffer.clear();
// read the version and command
byte[] cmd = new byte[5];
int read = in.read(cmd, 0, 5);
if (read != 5) {
int read = in.read(buffer);
if (read < 5) {
throw new IOException("Error reading Socks5 version and command");
}
buffer.position(5);
// read the digest
byte[] addr = new byte[cmd[4]];
read = in.read(addr, 0, addr.length);
if (read != addr.length) {
throw new IOException("Error reading provided address");
}
String digest = new String(addr);
in.read();
in.read();
byte[] addr = new byte[buffer.get(4)];
buffer.get(addr, 0, addr.length);
return digest;
return new String(addr);
}
private static byte[] createOutgoingSocks5Message(int cmd, String digest) {
......@@ -338,9 +344,9 @@ public class ProxyConnectionManager {
socketProcess.cancel(true);
socketProcess = null;
}
if (serverSocket != null) {
if (serverChannel != null) {
try {
serverSocket.close();
serverChannel.close();
}
catch (IOException e) {
Log.warn("Error closing proxy listening socket", e);
......@@ -354,7 +360,7 @@ public class ProxyConnectionManager {
}
public double sample() {
return (ProxyOutputStream.amountTransfered.getAndSet(0) / 1000);
return (ProxyOutputChannel.amountTransfered.getAndSet(0) / 1000);
}
}
}
\ No newline at end of file
......@@ -8,23 +8,33 @@
*/
package org.jivesoftware.wildfire.filetransfer.proxy;
import java.io.OutputStream;
import java.io.IOException;
import java.io.DataOutputStream;
import java.util.concurrent.atomic.AtomicLong;
import java.nio.channels.WritableByteChannel;
import java.nio.ByteBuffer;
/**
* An output stream which tracks the amount of bytes transfered by proxy sockets.
*/
public class ProxyOutputStream extends DataOutputStream {
public class ProxyOutputChannel implements WritableByteChannel {
static AtomicLong amountTransfered = new AtomicLong(0);
private WritableByteChannel channel;
public ProxyOutputStream(OutputStream out) {
super(out);
public ProxyOutputChannel(WritableByteChannel channel) {
this.channel = channel;
}
public synchronized void write(byte b[], int off, int len) throws IOException {
super.write(b, off, len);
amountTransfered.addAndGet(len);
public int write(ByteBuffer src) throws IOException {
int bytesWritten = channel.write(src);
amountTransfered.addAndGet(bytesWritten);
return bytesWritten;
}
public boolean isOpen() {
return channel.isOpen();
}
public void close() throws IOException {
channel.close();
}
}
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment