Get at least SOCKS4A working

1 files changed, 26 insertions, 20 deletions

diff --git a/common.c b/common.c
index cae9248..7e0b0f1 100644
--- a/common.c
+++ b/common.c
@@ -390,6 +390,7 @@ struct socks_connector
 
 	struct poller_timer timeout;        ///< Timeout timer
 
+	bool done;                          ///< We're connected
 	uint8_t bound_address_len;          ///< Length of domain name
 	struct socks_addr bound_address;    ///< Bound address at the server
 	uint16_t bound_port;                ///< Bound port at the server
@@ -435,16 +436,6 @@ struct socks_connector
 		return false;                                                          \
 	BLOCK_END
 
-// FIXME: we need to cancel all events
-#define SOCKS_DONE()                                                           \
-	BLOCK_START                                                                \
-		int fd = self->socket_fd;                                              \
-		set_blocking (fd, true);                                               \
-		self->socket_fd = -1;                                                  \
-		self->on_connected (self->user_data, fd);                              \
-		return true;                                                           \
-	BLOCK_END
-
 #define SOCKS_NEED_DATA(n)                                                     \
 	if (!socks_try_fill_read_buffer (self, (n)))                               \
 		return false;                                                          \
@@ -521,7 +512,8 @@ socks_4a_finish (struct socks_connector *self)
 	switch (status)
 	{
 	case 90:
-		SOCKS_DONE ();
+		self->done = true;
+		return false;
 	case 91:
 		SOCKS_FAIL ("request rejected or failed");
 	case 92:
@@ -593,7 +585,8 @@ socks_5_request_port (struct socks_connector *self)
 	hard_assert (msg_unpacker_u16 (&unpacker, &self->bound_port));
 	str_remove_slice (&self->read_buffer, 0, unpacker.offset);
 
-	SOCKS_DONE ();
+	self->done = true;
+	return false;
 }
 
 static bool
@@ -890,11 +883,20 @@ socks_connector_start (struct socks_connector *self)
 
 	connector_step (connector);
 	poller_timer_set (&self->timeout, 60 * 1000);
+	self->done = false;
 }
 
 static void
 socks_connector_step (struct socks_connector *self)
 {
+	// Close the socket if needed
+	if (self->socket_fd != -1)
+	{
+		poller_fd_reset (&self->socket_event);
+		xclose (self->socket_fd);
+		self->socket_fd = -1;
+	}
+
 	// Destroy current connector if needed
 	if (self->connector)
 	{
@@ -931,18 +933,22 @@ socks_connector_on_ready
 {
 	(void) pfd;
 
-	if (!self->on_data (self) || !socks_try_flush_write_buffer (self))
-	{
-		// We've failed this target, let's try to move on
-		// FIXME: we need to cancel all events
-		socks_connector_step (self);
-	}
-	// If we successfully establish the connection, then the FD is reset to -1
-	else if (self->socket_fd != -1)
+	if (self->on_data (self) && socks_try_flush_write_buffer (self))
 	{
 		poller_fd_set (&self->socket_event,
 			self->write_buffer.len ? (POLLIN | POLLOUT) : POLLIN);
 	}
+	else if (self->done)
+	{
+		int fd = self->socket_fd;
+		poller_fd_reset (&self->socket_event);
+		self->socket_fd = -1;
+		set_blocking (fd, true);
+		self->on_connected (self->user_data, fd);
+	}
+	else
+		// We've failed this target, let's try to move on
+		socks_connector_step (self);
 }
 
 static void