3 files changed, 67 insertions, 8 deletions

diff --git a/src/inspircd.cpp b/src/inspircd.cpp
index b9b9eb33e..0db94f5d4 100644
--- a/src/inspircd.cpp
+++ b/src/inspircd.cpp
@@ -2340,11 +2340,15 @@ int InspIRCd(char** argv, int argc)
 		for (std::vector<InspSocket*>::iterator a = module_sockets.begin(); a < module_sockets.end(); a++)
 		{
 			InspSocket* s = (InspSocket*)*a;
-			if (!s->Poll())
+			// Polling a listening socket class may result in the size of module_sockets increasing.
+			// This is still safe to do, however if the size of module_sockets is decreased, e.g.
+			// by a close or error, we cannot continue to use this iterator and must bail out asap.
+			if ((s) && (!s->Poll()))
 			{
+				log(DEBUG,"Socket poll returned false, close and bail");
 				s->Close();
-				delete s;
 				module_sockets.erase(a);
+				delete s;
 				break;
 			}
 		}
diff --git a/src/modules/m_spanningtree.cpp b/src/modules/m_spanningtree.cpp
index 95d4d4398..4cc4b50d6 100644
--- a/src/modules/m_spanningtree.cpp
+++ b/src/modules/m_spanningtree.cpp
@@ -39,6 +39,12 @@ class TreeSocket : public InspSocket
 		Srv->Log(DEBUG,"Create new");
 		myhost = host;
 	}
+
+	TreeSocket(int newfd)
+		: InspSocket(newfd)
+	{
+		// Associate with an existing file descriptor (accepted from incoming connection)
+	}
 	
         virtual bool OnConnected()
 	{
@@ -88,9 +94,17 @@ class TreeSocket : public InspSocket
 		Srv->SendToModeMask("o",WM_AND,"*** CLOSED ***");
 	}
 	
-	virtual int OnIncomingConnection()
+	virtual int OnIncomingConnection(int newsock, char* ip)
 	{
 		Srv->SendToModeMask("o",WM_AND,"*** INCOMING ***");
+		// use the (int) constructor to associate an incoming
+		// connection with a class, without actually creating
+		// a connection or binding from scratch.
+		TreeSocket* s = new TreeSocket(newsock);
+		Srv->AddSocket(s);
+		char message[1024];
+		sprintf(message,"Added new socket to list with fd %d from %s",newsock,ip);
+		Srv->Log(DEBUG,message);
 		return true;
 	}
 };
@@ -113,6 +127,8 @@ class ModuleSpanningTree : public Module
 		Srv = new Server;
 		Srv->AddCommand("CONNECTTEST",handle_connecttest,'o',1,"m_spanningtree.so");
 		Srv->Log(DEBUG,"ModCreate");
+		TreeSocket* listeningsock = new TreeSocket("127.0.0.1",11111,true,10);
+		Srv->AddSocket(listeningsock);
 	}
 	
 	virtual void OnUserJoin(userrec* user, chanrec* channel)
diff --git a/src/socket.cpp b/src/socket.cpp
index 48a58db0b..db95b852e 100644
--- a/src/socket.cpp
+++ b/src/socket.cpp
@@ -48,9 +48,41 @@ InspSocket::InspSocket()
 	this->state = I_DISCONNECTED;
 }
 
+InspSocket::InspSocket(int newfd)
+{
+	this->fd = newfd;
+	this->state = I_CONNECTED;
+}
+
 InspSocket::InspSocket(std::string host, int port, bool listening, unsigned long maxtime)
 {
 	if (listening) {
+		if ((this->fd = OpenTCPSocket()) == ERROR)
+		{
+			this->fd = -1;
+			this->state = I_ERROR;
+			this->OnError(I_ERR_SOCKET);
+			log(DEBUG,"OpenTCPSocket() error");
+                        return;
+		}
+		else
+		{
+			if (BindSocket(this->fd,this->client,this->server,port,(char*)host.c_str()) == ERROR)
+			{
+				this->Close();
+				this->fd = -1;
+				this->state = I_ERROR;
+				this->OnError(I_ERR_BIND);
+				log(DEBUG,"BindSocket() error %s",strerror(errno));
+				return;
+			}
+			else
+			{
+				this->state = I_LISTENING;
+				log(DEBUG,"New socket now in I_LISTENING state");
+				return;
+			}
+		}			
 	} else {
 		char* ip;
 		this->host = host;
@@ -149,10 +181,13 @@ int InspSocket::Write(std::string data)
 
 bool InspSocket::Poll()
 {
-	if (time(NULL) > timeout_end)
+	if ((time(NULL) > timeout_end) && (this->state == I_CONNECTING))
 	{
+		// for non-listening sockets, the timeout can occur
+		// which causes termination of the connection after
+		// the given number of seconds without a successful
+		// connection.
 		this->OnTimeout();
-		this->Close();
 		this->OnError(I_ERR_TIMEOUT);
         	timeout = true;
 		this->state = I_ERROR;
@@ -164,6 +199,8 @@ bool InspSocket::Poll()
 
         if (ret > 0)
 	{
+		int incoming = -1;
+		
 		switch (this->state)
 		{
 			case I_CONNECTING:
@@ -171,7 +208,10 @@ bool InspSocket::Poll()
 				return this->OnConnected();
 			break;
 			case I_LISTENING:
-				this->OnIncomingConnection();
+				length = sizeof (client);
+				incoming = accept (this->fd, (sockaddr*)&client,&length);
+				this->OnIncomingConnection(incoming,inet_ntoa(client.sin_addr));
+				return true;
 			break;
 			case I_CONNECTED:
 				return this->OnDataReady();
@@ -180,7 +220,6 @@ bool InspSocket::Poll()
 			break;
 		}
 	}
-
 	return true;
 }
 
@@ -193,7 +232,7 @@ void InspSocket::SetState(InspSocketState s)
 bool InspSocket::OnConnected() { return true; }
 void InspSocket::OnError(InspSocketError e) { return; }
 int InspSocket::OnDisconnect() { return 0; }
-int InspSocket::OnIncomingConnection() { return 0; }
+int InspSocket::OnIncomingConnection(int newfd, char* ip) { return 0; }
 bool InspSocket::OnDataReady() { return true; }
 void InspSocket::OnTimeout() { return; }
 void InspSocket::OnClose() { return; }