SOCKS: finishing touches

Making sure that I handle all corner cases appropriately.

1 files changed, 48 insertions, 32 deletions

diff --git a/common.c b/common.c
index 99b3b8e..de00445 100644
--- a/common.c
+++ b/common.c
@@ -383,23 +383,20 @@ struct socks_connector
 
 	// Negotiation:
 
+	struct poller_timer timeout;        ///< Timeout timer
+
 	int socket_fd;                      ///< Current socket file descriptor
 	struct poller_fd socket_event;      ///< Socket can be read from/written to
 	struct str read_buffer;             ///< Read buffer
 	struct str write_buffer;            ///< Write buffer
 
-	struct poller_timer timeout;        ///< Timeout timer
-
-	bool done;                          ///< We're connected
+	bool done;                          ///< Tunnel succesfully established
 	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
+	size_t data_needed;                 ///< How much data "on_data" needs
 
 	/// Process incoming data if there's enough of it available
 	bool (*on_data) (struct socks_connector *, struct msg_unpacker *);
 
-	size_t data_needed;                 ///< How much data the callback needs
-
 	// Configuration:
 
 	const char *hostname;               ///< SOCKS server hostname
@@ -413,6 +410,11 @@ struct socks_connector
 
 	void *user_data;                    ///< User data for callbacks
 
+	// Additional results:
+
+	struct socks_addr bound_address;    ///< Bound address at the server
+	uint16_t bound_port;                ///< Bound port at the server
+
 	// You may destroy the connector object in these two main callbacks:
 
 	/// Connection has been successfully established
@@ -429,6 +431,8 @@ struct socks_connector
 	void (*on_error) (void *user_data, const char *error);
 };
 
+// I've tried to make the actual protocol handlers as simple as possible
+
 #define SOCKS_FAIL(...)                                                        \
 	BLOCK_START                                                                \
 		char *error = xstrdup_printf (__VA_ARGS__);                            \
@@ -713,17 +717,37 @@ socks_5_auth_start (struct socks_connector *self)
 static void socks_connector_start (struct socks_connector *self);
 
 static void
-socks_connector_fail (struct socks_connector *self)
+socks_connector_destroy_connector (struct socks_connector *self)
 {
+	if (self->connector)
+	{
+		connector_free (self->connector);
+		free (self->connector);
+		self->connector = NULL;
+	}
+}
+
+static void
+socks_connector_cancel_events (struct socks_connector *self)
+{
+	// Before calling the final callbacks, we should cancel events that
+	// could potentially fire; caller should destroy us immediately, though
 	poller_fd_reset (&self->socket_event);
+	poller_timer_reset (&self->timeout);
+}
+
+static void
+socks_connector_fail (struct socks_connector *self)
+{
+	socks_connector_cancel_events (self);
 	self->on_failure (self->user_data);
 }
 
 static bool
 socks_connector_step_iterators (struct socks_connector *self)
 {
-	// At the lowest level we iterate over all addresses for the SOCKS server;
-	// this is done automatically by the connector
+	// At the lowest level we iterate over all addresses for the SOCKS server
+	// and just try to connect; this is done automatically by the connector
 
 	// Then we iterate over available protocols
 	if (++self->protocol_iter != SOCKS_MAX)
@@ -747,13 +771,7 @@ socks_connector_step (struct socks_connector *self)
 		self->socket_fd = -1;
 	}
 
-	if (self->connector)
-	{
-		connector_free (self->connector);
-		free (self->connector);
-		self->connector = NULL;
-	}
-
+	socks_connector_destroy_connector (self);
 	if (socks_connector_step_iterators (self))
 		socks_connector_start (self);
 	else
@@ -766,6 +784,7 @@ socks_connector_on_timeout (struct socks_connector *self)
 	if (self->on_error)
 		self->on_error (self->user_data, "timeout");
 
+	socks_connector_destroy_connector (self);
 	socks_connector_fail (self);
 }
 
@@ -783,11 +802,9 @@ socks_connector_on_connected (void *user_data, int socket_fd)
 	str_reset (&self->read_buffer);
 	str_reset (&self->write_buffer);
 
-	if ((self->protocol_iter == SOCKS_5  && socks_5_auth_start (self))
-	 || (self->protocol_iter == SOCKS_4A && socks_4a_start     (self)))
-		return;
-
-	socks_connector_fail (self);
+	if (!(self->protocol_iter == SOCKS_5  && socks_5_auth_start (self))
+	 && !(self->protocol_iter == SOCKS_4A && socks_4a_start     (self)))
+		socks_connector_fail (self);
 }
 
 static void
@@ -838,9 +855,13 @@ socks_connector_start (struct socks_connector *self)
 	hard_assert (connector_add_target (connector,
 		self->hostname, self->service, NULL));
 
-	connector_step (connector);
 	poller_timer_set (&self->timeout, 60 * 1000);
+	connector_step (connector);
 	self->done = false;
+
+	self->bound_port = 0;
+	socks_addr_free (&self->bound_address);
+	memset (&self->bound_address, 0, sizeof self->bound_address);
 }
 
 // - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
@@ -923,8 +944,9 @@ socks_connector_on_ready
 	}
 	else if (self->done)
 	{
+		socks_connector_cancel_events (self);
+
 		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);
@@ -956,14 +978,8 @@ socks_connector_init (struct socks_connector *self, struct poller *poller)
 static void
 socks_connector_free (struct socks_connector *self)
 {
-	if (self->connector)
-	{
-		connector_free (self->connector);
-		free (self->connector);
-	}
-
-	poller_fd_reset (&self->socket_event);
-	poller_timer_reset (&self->timeout);
+	socks_connector_destroy_connector (self);
+	socks_connector_cancel_events (self);
 
 	if (self->socket_fd != -1)
 		xclose (self->socket_fd);