package org.torproject.jtor.circuits.impl;

import java.io.InputStream;
import java.io.OutputStream;
import java.util.Date;

import org.torproject.jtor.TorException;
import org.torproject.jtor.TorTimeoutException;
import org.torproject.jtor.circuits.Circuit;
import org.torproject.jtor.circuits.CircuitNode;
import org.torproject.jtor.circuits.Stream;
import org.torproject.jtor.circuits.cells.RelayCell;

public class StreamImpl implements Stream {
	private final static int STREAM_CONNECT_TIMEOUT = 20 * 1000;
	private final CircuitImpl circuit;
	private final int streamId;
	private final CircuitNode targetNode;
	private final TorInputStream inputStream;
	private final TorOutputStream outputStream;
	private boolean isClosed;
	private boolean relayEndReceived;
	private boolean relayConnectedReceived;
	private final Object waitConnectLock = new Object();

	StreamImpl(CircuitImpl circuit, CircuitNode targetNode, int streamId) {
		this.circuit = circuit;
		this.targetNode = targetNode;
		this.streamId = streamId;
		this.inputStream = new TorInputStream(this);
		this.outputStream = new TorOutputStream(this);
	}

	void addInputCell(RelayCell cell) {
		if(isClosed)
			return;
		if(cell.getRelayCommand() == RelayCell.RELAY_END) {
			synchronized(waitConnectLock) {
				relayEndReceived = true;
				inputStream.addEndCell(cell);
				waitConnectLock.notifyAll();
			}
		} else if(cell.getRelayCommand() == RelayCell.RELAY_CONNECTED) {
			synchronized(waitConnectLock) {
				relayConnectedReceived = true;
				waitConnectLock.notifyAll();
			}
		}
		else
			inputStream.addInputCell(cell);
	}

	public int getStreamId() {
		return streamId;
	}

	public Circuit getCircuit() {
		return circuit;
	}

	CircuitNode getTargetNode() {
		return targetNode;
	}

	public void close() {
		if(isClosed)
			return;
		isClosed = true;
		inputStream.close();
		outputStream.close();
		if(!relayEndReceived) {
			final RelayCell cell = new RelayCellImpl(circuit.getFinalCircuitNode(), circuit.getCircuitId(), streamId, RelayCell.RELAY_END);
			cell.putByte(RelayCell.REASON_DONE);
			circuit.sendRelayCellToFinalNode(cell);
		}
		circuit.removeStream(this);
	}

	void openDirectory() {
		final RelayCell cell = new RelayCellImpl(circuit.getFinalCircuitNode(), circuit.getCircuitId(), streamId, RelayCell.RELAY_BEGIN_DIR);
		circuit.sendRelayCellToFinalNode(cell);
		waitForRelayConnected();
	}

	void openExit(String target, int port) {
		final RelayCell cell = new RelayCellImpl(circuit.getFinalCircuitNode(), circuit.getCircuitId(), streamId, RelayCell.RELAY_BEGIN);
		cell.putString(target + ":"+ port);
		circuit.sendRelayCellToFinalNode(cell);
		waitForRelayConnected();
	}

	private void waitForRelayConnected() {
		final Date startWait = new Date();
		synchronized(waitConnectLock) {
			while(!relayConnectedReceived && !relayEndReceived)
				try {
					waitConnectLock.wait(1000);
					final Date now = new Date();
					if(now.getTime() - startWait.getTime() >= STREAM_CONNECT_TIMEOUT)
						throw new TorTimeoutException("Timeout waiting for RELAY_CONNECTED cell on stream id="+ streamId);
				} catch (InterruptedException e) {
					Thread.currentThread().interrupt();
					throw new TorTimeoutException("Interrupted waiting for RELAY_CONNECTED cell on stream id="+ streamId);
				}
				if(relayEndReceived)
					throw new TorException("RELAY_END cell received while waiting for RELAY_CONNECTED cell on stream id="+ streamId);
				if(relayConnectedReceived)
					return;
		}
	}

	public InputStream getInputStream() {
		return inputStream;
	}

	public OutputStream getOutputStream() {
		return outputStream;
	}

	public String toString() {
		return "[Stream stream_id="+ streamId + " circuit="+ circuit +" ]";
	}
}

@@ -1867,16 +1867,47 @@
 lose();%0A
+%09%09circuit.removeStream(this);%0A%0A
 %09%09if(!re
@@ -2146,38 +2146,8 @@
 %09%09%7D%0A
-%09%09circuit.removeStream(this);%0A
 %09%7D%0A%0A
@@ -2811,13 +2811,20 @@
 ived
- && !
+) %7B%0A%0A%09%09%09%09if(
 rela
@@ -2845,153 +2845,168 @@
 %09%09%09%09
-try %7B%0A%09%09%09%09%09waitConnectLock.wait(1000);%0A%09%09%09%09%09final Date now = new Date();%0A%09
+%09throw new TorException(%22RELAY_END cell received while waiting for RELAY_CONNECTED cell on stream id=%22+ streamId);%0A%0A
 %09%09%09%09if(
-now.getTime() - startWait.getTime() %3E= STREAM_CONNECT_TIMEOUT
+hasStreamConnectTimedOut(startWait)
 )%0A
-%09
 %09%09%09%09
@@ -3100,24 +3100,68 @@
 + streamId);
+%0A%0A%09%09%09%09try %7B%0A%09%09%09%09%09waitConnectLock.wait(1000);
 %0A%09%09%09%09%7D catch
@@ -3350,196 +3350,213 @@
 %0A%09%09%09
-%09if(relayEndReceived)%0A%09%09%09%09%09throw new TorException(%22RELAY_END cell received while waiting for RELAY_CONNECTED cell on stream id=%22+ streamId);%0A%09%09%09%09if(relayConnectedReceived)%0A%09%09%09%09%09return;%0A%09%09%7D
+%7D%0A%09%09%7D%0A%09%7D%0A%0A%09private static boolean hasStreamConnectTimedOut(Date startTime) %7B%0A%09%09final Date now = new Date();%0A%09%09final long diff = now.getTime() - startTime.getTime();%0A%09%09return diff %3E= STREAM_CONNECT_TIMEOUT;
 %0A%09%7D%0A

Refactor the loop that waits for RELAY_CONNECTED cells a bit.