From 31699e830debae70b495dcec47c7fcae5331cff0 Mon Sep 17 00:00:00 2001
From: Ykkrosh <philip@wildfiregames.com>
Date: Tue, 6 Jul 2010 19:54:17 +0000
Subject: [PATCH] Handle disconnections better.

Remove local sessions (just use ENet for everything instead) because
they add far too much complexity.
Fix memory leaks.

This was SVN commit r7706.
---
 .../data/mods/public/gui/common/network.js    |  19 +++
 .../mods/public/gui/gamesetup/gamesetup.js    |   5 +-
 .../mods/public/gui/gamesetup/gamesetup.xml   |   1 +
 .../mods/public/gui/gamesetup/gamesetup_mp.js |  16 +-
 .../public/gui/gamesetup/gamesetup_mp.xml     |   8 +-
 source/graphics/MapReader.cpp                 |   4 +-
 source/gui/scripting/ScriptFunctions.cpp      |   2 +-
 source/main.cpp                               |   8 +
 source/network/NetClient.cpp                  |  28 ++--
 source/network/NetClient.h                    |  14 +-
 source/network/NetHost.h                      |  12 ++
 source/network/NetMessage.cpp                 |  19 ++-
 source/network/NetServer.cpp                  | 131 +++++++---------
 source/network/NetServer.h                    |  44 +++---
 source/network/NetSession.cpp                 | 142 ++++--------------
 source/network/NetSession.h                   | 141 ++++++-----------
 source/network/fsm.cpp                        |   2 +-
 source/network/fsm.h                          |   3 +-
 source/network/tests/test_Net.h               |   5 +-
 source/ps/GameSetup/GameSetup.cpp             |   2 +-
 20 files changed, 262 insertions(+), 344 deletions(-)
 create mode 100644 binaries/data/mods/public/gui/common/network.js

diff --git a/binaries/data/mods/public/gui/common/network.js b/binaries/data/mods/public/gui/common/network.js
new file mode 100644
index 0000000000..a46dde84fa
--- /dev/null
+++ b/binaries/data/mods/public/gui/common/network.js
@@ -0,0 +1,19 @@
+function getDisconnectReason(id)
+{
+	// Must be kept in sync with source/network/NetHost.h
+	switch (id)
+	{
+	case 0: return "Unknown reason";
+	case 1: return "Unexpected shutdown";
+	case 2: return "Incorrect network protocol version";
+	case 3: return "Game has already started";
+	default: return "[Invalid value "+id+"]";
+	}
+}
+
+function reportDisconnect(reason)
+{
+	messageBox(400, 200,
+		"Lost connection to the server.\n\nReason: " + getDisconnectReason(reason) + ".",
+		"Disconnected", 2);
+}
diff --git a/binaries/data/mods/public/gui/gamesetup/gamesetup.js b/binaries/data/mods/public/gui/gamesetup/gamesetup.js
index 15986d95c2..20ae46590b 100644
--- a/binaries/data/mods/public/gui/gamesetup/gamesetup.js
+++ b/binaries/data/mods/public/gui/gamesetup/gamesetup.js
@@ -89,8 +89,9 @@ function handleNetMessage(message)
 		case "disconnected":
 			Engine.DisconnectNetworkGame();
 			Engine.PopGuiPage();
-			messageBox(400, 200, "Connection to the server has been lost.", "Disconnected", 2);
+			reportDisconnect(message.reason);
 			break;
+
 		default:
 			error("Unrecognised netstatus type "+message.status);
 			break;
@@ -118,7 +119,7 @@ function handleNetMessage(message)
 		break;
 
 	case "start":
-		Engine.PushGuiPage("page_loading.xml", { "attribs": g_GameAttributes });
+		Engine.SwitchGuiPage("page_loading.xml", { "attribs": g_GameAttributes });
 		break;
 
 	case "chat":
diff --git a/binaries/data/mods/public/gui/gamesetup/gamesetup.xml b/binaries/data/mods/public/gui/gamesetup/gamesetup.xml
index 04fc080dd0..37ed42eca0 100644
--- a/binaries/data/mods/public/gui/gamesetup/gamesetup.xml
+++ b/binaries/data/mods/public/gui/gamesetup/gamesetup.xml
@@ -2,6 +2,7 @@
 
 <objects>
 
+	<script file="gui/common/network.js"/>
 	<script file="gui/common/functions_global_object.js"/>
 	<script file="gui/gamesetup/gamesetup.js"/>
 
diff --git a/binaries/data/mods/public/gui/gamesetup/gamesetup_mp.js b/binaries/data/mods/public/gui/gamesetup/gamesetup_mp.js
index fe88a00cef..daea45e749 100644
--- a/binaries/data/mods/public/gui/gamesetup/gamesetup_mp.js
+++ b/binaries/data/mods/public/gui/gamesetup/gamesetup_mp.js
@@ -5,6 +5,13 @@ function init()
 {
 }
 
+function cancelSetup()
+{
+	if (g_IsConnecting)
+		Engine.DisconnectNetworkGame();
+	Engine.PopGuiPage();	
+}
+
 function startConnectionStatus(type)
 {
 	g_GameType = type;
@@ -33,10 +40,17 @@ function onTick()
 			case "connected":
 				getGUIObjectByName("connectionStatus").caption = "Registering with server...";
 				break;
+
 			case "authenticated":
 				Engine.PopGuiPage();
 				Engine.PushGuiPage("page_gamesetup.xml", { "type": g_GameType });
-				return; // don't process any more messages
+				return; // don't process any more messages - leave them for the game GUI loop
+
+			case "disconnected":
+				cancelSetup();
+				reportDisconnect(message.reason);
+				return;
+
 			default:
 				error("Unrecognised netstatus type "+message.status);
 				break;
diff --git a/binaries/data/mods/public/gui/gamesetup/gamesetup_mp.xml b/binaries/data/mods/public/gui/gamesetup/gamesetup_mp.xml
index 179ba2f494..cd59a5ca4b 100644
--- a/binaries/data/mods/public/gui/gamesetup/gamesetup_mp.xml
+++ b/binaries/data/mods/public/gui/gamesetup/gamesetup_mp.xml
@@ -2,6 +2,8 @@
 
 <objects>
 
+	<script file="gui/common/network.js"/>
+	<script file="gui/common/functions_global_object.js"/>
 	<script file="gui/gamesetup/gamesetup_mp.js"/>
 
 	<!-- Add a translucent black background to fade out the menu page -->
@@ -17,9 +19,9 @@
 			Multiplayer
 		</object>
 	
-		<object type="button" style="wheatExit">
+		<object type="button" style="wheatExit" tooltip_style="snToolTip">
 			<action on="Press"><![CDATA[
-				Engine.PopGuiPage();
+				cancelSetup();
 			]]></action>
 		</object>
 
@@ -145,7 +147,7 @@
 			<object type="button" size="100%-110 200 100%-10 230" style="wheatButton">
 				Cancel
 				<action on="Press"><![CDATA[
-					error("not implemented yet");
+					cancelSetup();
 				]]></action>
 			</object>
 
diff --git a/source/graphics/MapReader.cpp b/source/graphics/MapReader.cpp
index d110e4d6ff..7ecb83ada8 100644
--- a/source/graphics/MapReader.cpp
+++ b/source/graphics/MapReader.cpp
@@ -448,7 +448,9 @@ void CXMLReader::ReadTerrain(XMBElement parent)
 	m_MapReader.m_PatchesPerSide = patches;
 
 	// Load the texture
-	CTextureEntry* texentry = g_TexMan.FindTexture(texture);
+	CTextureEntry* texentry = NULL;
+	if (CTextureManager::IsInitialised())
+		texentry = g_TexMan.FindTexture(texture);
 
 	m_MapReader.pTerrain->Initialize(patches, NULL);
 
diff --git a/source/gui/scripting/ScriptFunctions.cpp b/source/gui/scripting/ScriptFunctions.cpp
index d5c836f24a..0c865d01aa 100644
--- a/source/gui/scripting/ScriptFunctions.cpp
+++ b/source/gui/scripting/ScriptFunctions.cpp
@@ -203,7 +203,7 @@ void StartNetworkHost(void* UNUSED(cbdata), std::wstring playerName)
 	g_Game = new CGame();
 	g_NetClient = new CNetClient(g_Game);
 	g_NetClient->SetUserName(playerName);
-	g_NetClient->SetupLocalConnection(*g_NetServer);
+	g_NetClient->SetupConnection("127.0.0.1");
 }
 
 void StartNetworkJoin(void* UNUSED(cbdata), std::wstring playerName, std::string serverAddress)
diff --git a/source/main.cpp b/source/main.cpp
index f2b8c2691b..88a8ab30bd 100644
--- a/source/main.cpp
+++ b/source/main.cpp
@@ -330,6 +330,14 @@ static void Frame()
 	}
 	PROFILE_END( "game logic" );
 
+	// Immediately flush any messages produced by simulation code
+	PROFILE_START("network flush");
+	if (g_NetServer)
+		g_NetServer->Flush();
+	if (g_NetClient)
+		g_NetClient->Flush();
+	PROFILE_END("network flush");
+
 	PROFILE_START( "update console" );
 	g_Console->Update(TimeSinceLastFrame);
 	PROFILE_END( "update console" );
diff --git a/source/network/NetClient.cpp b/source/network/NetClient.cpp
index 80f6195ed1..6bfb1bccc5 100644
--- a/source/network/NetClient.cpp
+++ b/source/network/NetClient.cpp
@@ -70,8 +70,6 @@ CNetClient::CNetClient(CGame* game) :
 	AddTransition(NCS_INGAME, (uint)NMT_SYNC_ERROR, NCS_INGAME, (void*)&OnInGame, context);
 	AddTransition(NCS_INGAME, (uint)NMT_END_COMMAND_BATCH, NCS_INGAME, (void*)&OnInGame, context);
 
-	// TODO: add chat
-
 	// Set first state
 	SetFirstState(NCS_UNCONNECTED);
 }
@@ -90,18 +88,12 @@ void CNetClient::SetUserName(const CStrW& username)
 
 bool CNetClient::SetupConnection(const CStr& server)
 {
-	CNetClientSessionRemote* session = new CNetClientSessionRemote(*this);
+	CNetClientSession* session = new CNetClientSession(*this);
 	bool ok = session->Connect(PS_DEFAULT_PORT, server);
 	SetAndOwnSession(session);
 	return ok;
 }
 
-void CNetClient::SetupLocalConnection(CNetServer& server)
-{
-	CNetClientSessionLocal* session = new CNetClientSessionLocal(*this, server);
-	SetAndOwnSession(session);
-}
-
 void CNetClient::SetAndOwnSession(CNetClientSession* session)
 {
 	delete m_Session;
@@ -114,6 +106,12 @@ void CNetClient::Poll()
 		m_Session->Poll();
 }
 
+void CNetClient::Flush()
+{
+	if (m_Session)
+		m_Session->Flush();
+}
+
 CScriptValRooted CNetClient::GuiPoll()
 {
 	if (m_GuiMessageQueue.empty())
@@ -171,6 +169,9 @@ void CNetClient::PostPlayerAssignmentsToScript()
 
 bool CNetClient::SendMessage(const CNetMessage* message)
 {
+	if (!m_Session)
+		return false;
+
 	return m_Session->SendMessage(message);
 }
 
@@ -179,11 +180,18 @@ void CNetClient::HandleConnect()
 	Update((uint)NMT_CONNECT_COMPLETE, NULL);
 }
 
-void CNetClient::HandleDisconnect()
+void CNetClient::HandleDisconnect(u32 reason)
 {
 	CScriptValRooted msg;
 	GetScriptInterface().Eval("({'type':'netstatus','status':'disconnected'})", msg);
+	GetScriptInterface().SetProperty(msg.get(), "reason", (int)reason, false);
 	PushGuiMessage(msg);
+
+	SAFE_DELETE(m_Session);
+
+	// Update the state immediately to UNCONNECTED (don't bother with FSM transitions since
+	// we'd need one for every single state, and we don't need to use per-state actions)
+	SetCurrState(NCS_UNCONNECTED);
 }
 
 void CNetClient::SendChatMessage(const std::wstring& text)
diff --git a/source/network/NetClient.h b/source/network/NetClient.h
index 7fb5d8add0..c330b3d341 100644
--- a/source/network/NetClient.h
+++ b/source/network/NetClient.h
@@ -78,12 +78,6 @@ public:
 	 */
 	bool SetupConnection(const CStr& server);
 
-	/**
-	 * Set up a connection to the local server on the current machine.
-	 * @param server object to connect to
-	 */
-	void SetupLocalConnection(CNetServer& server);
-
 	/**
 	 * Poll the connection for messages from the server and process them, and send
 	 * any queued messages.
@@ -91,6 +85,12 @@ public:
 	 */
 	void Poll();
 
+	/**
+	 * Flush any queued outgoing network messages.
+	 * This should be called soon after sending a group of messages that may be batched together.
+	 */
+	void Flush();
+
 	/**
 	 * Retrieves the next queued GUI message, and removes it from the queue.
 	 * The returned value is in the GetScriptInterface() JS context.
@@ -140,7 +140,7 @@ public:
 	/**
 	 * Call when the network connection has been lost.
 	 */
-	void HandleDisconnect();
+	void HandleDisconnect(u32 reason);
 
 	/**
 	 * Call when a message has been received from the network.
diff --git a/source/network/NetHost.h b/source/network/NetHost.h
index 41a0e0c52a..c5806c9187 100644
--- a/source/network/NetHost.h
+++ b/source/network/NetHost.h
@@ -38,6 +38,18 @@ struct PlayerAssignment
 
 typedef std::map<CStr, PlayerAssignment> PlayerAssignmentMap; // map from GUID -> assignment
 
+/**
+ * Reasons sent by server to clients in disconnection messages.
+ * Must be kept in sync with binaries/data/mods/public/gui/common/network.js
+ */
+enum NetDisconnectReason
+{
+	NDR_UNKNOWN = 0,
+	NDR_UNEXPECTED_SHUTDOWN,
+	NDR_INCORRECT_PROTOCOL_VERSION,
+	NDR_SERVER_ALREADY_IN_GAME
+};
+
 class CNetHost
 {
 public:
diff --git a/source/network/NetMessage.cpp b/source/network/NetMessage.cpp
index 9cd7ff0f30..876bcb099b 100644
--- a/source/network/NetMessage.cpp
+++ b/source/network/NetMessage.cpp
@@ -218,9 +218,20 @@ CNetMessage* CNetMessageFactory::CloneMessage( const CNetMessage* message, Scrip
 
 	size_t len = message->GetSerializedLength();
 	u8* buffer = new u8[len];
-	u8* newbuf = message->Serialize(buffer);
-	if (!newbuf)
+
+	u8* bufend = message->Serialize(buffer);
+
+	if (!bufend)
+	{
+		delete[] buffer;
 		return NULL;
-	debug_assert(newbuf == buffer+len);
-	return CreateMessage(buffer, len, scriptInterface);
+	}
+
+	debug_assert(bufend == buffer+len);
+
+	CNetMessage* ret = CreateMessage(buffer, len, scriptInterface);
+
+	delete[] buffer;
+
+	return ret;
 }
diff --git a/source/network/NetServer.cpp b/source/network/NetServer.cpp
index 99752d5509..8ab3d4243d 100644
--- a/source/network/NetServer.cpp
+++ b/source/network/NetServer.cpp
@@ -53,7 +53,6 @@ CNetServer::CNetServer() :
 
 	m_ServerTurnManager = NULL;
 
-	m_Port = PS_DEFAULT_PORT;
 	m_ServerName = DEFAULT_SERVER_NAME;
 	m_WelcomeMessage = DEFAULT_WELCOME_MESSAGE;
 }
@@ -64,7 +63,7 @@ CNetServer::~CNetServer()
 
 	for (size_t i = 0; i < m_Sessions.size(); ++i)
 	{
-		m_Sessions[i]->Disconnect();
+		m_Sessions[i]->DisconnectNow(NDR_UNEXPECTED_SHUTDOWN);
 		delete m_Sessions[i];
 	}
 
@@ -73,6 +72,8 @@ CNetServer::~CNetServer()
 		enet_host_destroy(m_Host);
 	}
 
+	delete m_ServerTurnManager;
+
 	m_GameAttributes = CScriptValRooted(); // clear root before deleting its context
 	delete m_ScriptInterface;
 }
@@ -85,7 +86,7 @@ bool CNetServer::SetupConnection()
 	// Bind to default host
 	ENetAddress addr;
 	addr.host = ENET_HOST_ANY;
-	addr.port = m_Port;
+	addr.port = PS_DEFAULT_PORT;
 
 	// Create ENet server
 	m_Host = enet_host_create(&addr, MAX_CLIENTS, 0, 0);
@@ -139,20 +140,6 @@ void CNetServer::Poll()
 {
 	debug_assert(m_Host);
 
-	for (size_t i = 0; i < m_LocalMessageQueue.size(); ++i)
-	{
-		CNetMessage* msg = m_LocalMessageQueue[i].second;
-		CNetServerSession* session = m_LocalMessageQueue[i].first;
-
-		LOGMESSAGE(L"Net server: Received local message %hs of size %lu from %hs", msg->ToString().c_str(), (unsigned long)msg->GetSerializedLength(), DebugName(session).c_str());
-
-		bool ok = HandleMessageReceive(msg, session);
-		debug_assert(ok); // TODO
-
-		delete msg;
-	}
-	m_LocalMessageQueue.clear();
-
 	// Poll host for events
 	ENetEvent event;
 	while (enet_host_service(m_Host, &event, 0) > 0)
@@ -161,49 +148,48 @@ void CNetServer::Poll()
 		{
 		case ENET_EVENT_TYPE_CONNECT:
 		{
-			// If this is a new client to our server, save the peer reference
-			if (std::find(m_Peers.begin(), m_Peers.end(), event.peer) == m_Peers.end())
-				m_Peers.push_back(event.peer);
-
 			// Report the client address
 			char hostname[256] = "(error)";
 			enet_address_get_host_ip(&event.peer->address, hostname, ARRAY_SIZE(hostname));
 			LOGMESSAGE(L"Net server: Received connection from %hs:%u", hostname, event.peer->address.port);
 
-			// Set up a session object for this peer
-
-			CNetServerSession* session = new CNetServerSessionRemote(*this, event.peer);
-
-			SetupSession(session);
-
-			if (!HandleConnect(session))
+			if (m_State != SERVER_STATE_PREGAME)
 			{
-				delete session;
+				enet_peer_disconnect(event.peer, NDR_SERVER_ALREADY_IN_GAME);
 				break;
 			}
 
+			// Set up a session object for this peer
+
+			CNetServerSession* session = new CNetServerSession(*this, event.peer);
+
+			m_Sessions.push_back(session);
+
+			SetupSession(session);
+
 			debug_assert(event.peer->data == NULL);
 			event.peer->data = session;
 
+			HandleConnect(session);
+
 			break;
 		}
 
 		case ENET_EVENT_TYPE_DISCONNECT:
 		{
-			// Delete from our peer list
-			m_Peers.erase(remove(m_Peers.begin(), m_Peers.end(), event.peer), m_Peers.end());
-
 			// If there is an active session with this peer, then reset and delete it
 
 			CNetServerSession* session = static_cast<CNetServerSession*>(event.peer->data);
 			if (session)
 			{
-				LOGMESSAGE(L"Net server: %hs disconnected", DebugName(session).c_str());
+				LOGMESSAGE(L"Net server: Disconnected %hs", DebugName(session).c_str());
+
+				// Remove the session first, so we won't send player-update messages to it
+				// when updating the FSM
+				m_Sessions.erase(remove(m_Sessions.begin(), m_Sessions.end(), session), m_Sessions.end());
 
 				session->Update((uint)NMT_CONNECTION_LOST, NULL);
 
-				HandleDisconnect(session);
-
 				delete session;
 				event.peer->data = NULL;
 			}
@@ -224,8 +210,7 @@ void CNetServer::Poll()
 				{
 					LOGMESSAGE(L"Net server: Received message %hs of size %lu from %hs", msg->ToString().c_str(), (unsigned long)msg->GetSerializedLength(), DebugName(session).c_str());
 
-					bool ok = HandleMessageReceive(msg, session);
-					debug_assert(ok); // TODO
+					HandleMessageReceive(msg, session);
 
 					delete msg;
 				}
@@ -240,30 +225,19 @@ void CNetServer::Poll()
 	}
 }
 
-void CNetServer::AddLocalClientSession(CNetClientSessionLocal& clientSession)
+void CNetServer::Flush()
 {
-	LOGMESSAGE(L"Net server: Received local connection");
-	CNetServerSessionLocal* session = new CNetServerSessionLocal(*this, clientSession);
-	clientSession.SetServerSession(session);
-	SetupSession(session);
-	HandleConnect(session);
+	debug_assert(m_Host);
+
+	enet_host_flush(m_Host);
 }
 
-void CNetServer::SendLocalMessage(CNetClientSessionLocal& clientSession, const CNetMessage* message)
-{
-	CNetMessage* clonedMessage = CNetMessageFactory::CloneMessage(message, GetScriptInterface());
-	if (!clonedMessage)
-		return;
-	m_LocalMessageQueue.push_back(std::make_pair(clientSession.GetServerSession(), clonedMessage));
-}
-
-bool CNetServer::HandleMessageReceive(const CNetMessage* message, CNetServerSession* session)
+void CNetServer::HandleMessageReceive(const CNetMessage* message, CNetServerSession* session)
 {
 	// Update FSM
 	bool ok = session->Update(message->GetType(), (void*)message);
 	if (!ok)
 		LOGERROR(L"Net server: Error running FSM update (type=%d state=%d)", (int)message->GetType(), (int)session->GetCurrState());
-	return ok;
 }
 
 void CNetServer::SetupSession(CNetServerSession* session)
@@ -271,9 +245,13 @@ void CNetServer::SetupSession(CNetServerSession* session)
 	void* context = session;
 
 	// Set up transitions for session
-	session->AddTransition(NSS_HANDSHAKE, (uint)NMT_CONNECTION_LOST, NSS_UNCONNECTED, (void*)&OnDisconnect, context);
+
+	session->AddTransition(NSS_UNCONNECTED, (uint)NMT_CONNECTION_LOST, NSS_UNCONNECTED);
+
+	session->AddTransition(NSS_HANDSHAKE, (uint)NMT_CONNECTION_LOST, NSS_UNCONNECTED);
 	session->AddTransition(NSS_HANDSHAKE, (uint)NMT_CLIENT_HANDSHAKE, NSS_AUTHENTICATE, (void*)&OnClientHandshake, context);
 
+	session->AddTransition(NSS_AUTHENTICATE, (uint)NMT_CONNECTION_LOST, NSS_UNCONNECTED);
 	session->AddTransition(NSS_AUTHENTICATE, (uint)NMT_AUTHENTICATE, NSS_PREGAME, (void*)&OnAuthenticate, context);
 
 	session->AddTransition(NSS_PREGAME, (uint)NMT_CONNECTION_LOST, NSS_UNCONNECTED, (void*)&OnDisconnect, context);
@@ -292,9 +270,6 @@ void CNetServer::SetupSession(CNetServerSession* session)
 
 bool CNetServer::HandleConnect(CNetServerSession* session)
 {
-	m_Sessions.push_back(session);
-
-	// Player joined the game, start authentication
 	CSrvHandshakeMessage handshake;
 	handshake.m_Magic = PS_PROTOCOL_MAGIC;
 	handshake.m_ProtocolVersion = PS_PROTOCOL_VERSION;
@@ -302,13 +277,6 @@ bool CNetServer::HandleConnect(CNetServerSession* session)
 	return session->SendMessage(&handshake);
 }
 
-bool CNetServer::HandleDisconnect(CNetServerSession* session)
-{
-	m_Sessions.erase(remove(m_Sessions.begin(), m_Sessions.end(), session), m_Sessions.end());
-
-	return true;
-}
-
 void CNetServer::OnUserJoin(CNetServerSession* session)
 {
 	AddPlayer(session->GetGUID(), session->GetUserName());
@@ -411,21 +379,25 @@ bool CNetServer::OnClientHandshake(void* context, CFsmEvent* event)
 	CNetServerSession* session = (CNetServerSession*)context;
 	CNetServer& server = session->GetServer();
 
+	if (server.m_State != SERVER_STATE_PREGAME)
+	{
+		session->Disconnect(NDR_SERVER_ALREADY_IN_GAME);
+		return false;
+	}
+
 	CCliHandshakeMessage* message = (CCliHandshakeMessage*)event->GetParamRef();
 	if (message->m_ProtocolVersion != PS_PROTOCOL_VERSION)
 	{
-		// TODO: probably should report some error message (either locally or to the client)
-		session->Disconnect();
-	}
-	else
-	{
-		CSrvHandshakeResponseMessage handshakeResponse;
-		handshakeResponse.m_UseProtocolVersion = PS_PROTOCOL_VERSION;
-		handshakeResponse.m_Message = server.m_WelcomeMessage;
-		handshakeResponse.m_Flags = 0;
-		session->SendMessage(&handshakeResponse);
+		session->Disconnect(NDR_INCORRECT_PROTOCOL_VERSION);
+		return false;
 	}
 
+	CSrvHandshakeResponseMessage handshakeResponse;
+	handshakeResponse.m_UseProtocolVersion = PS_PROTOCOL_VERSION;
+	handshakeResponse.m_Message = server.m_WelcomeMessage;
+	handshakeResponse.m_Flags = 0;
+	session->SendMessage(&handshakeResponse);
+
 	return true;
 }
 
@@ -436,6 +408,12 @@ bool CNetServer::OnAuthenticate(void* context, CFsmEvent* event)
 	CNetServerSession* session = (CNetServerSession*)context;
 	CNetServer& server = session->GetServer();
 
+	if (server.m_State != SERVER_STATE_PREGAME)
+	{
+		session->Disconnect(NDR_SERVER_ALREADY_IN_GAME);
+		return false;
+	}
+
 	CAuthenticateMessage* message = (CAuthenticateMessage*)event->GetParamRef();
 
 	// TODO: check server password etc?
@@ -527,9 +505,7 @@ bool CNetServer::OnDisconnect(void* context, CFsmEvent* event)
 	CNetServerSession* session = (CNetServerSession*)context;
 	CNetServer& server = session->GetServer();
 
-	// If the user had authenticated, we need to handle their leaving
-	if (session->GetCurrState() == NSS_PREGAME || session->GetCurrState() == NSS_INGAME)
-		server.OnUserLeave(session);
+	server.OnUserLeave(session);
 
 	return true;
 }
@@ -603,6 +579,7 @@ CStrW CNetServer::DeduplicatePlayerName(const CStrW& original)
 {
 	CStrW name = original;
 
+	// Try names "Foo", "Foo (2)", "Foo (3)", etc
 	size_t id = 2;
 	while (true)
 	{
diff --git a/source/network/NetServer.h b/source/network/NetServer.h
index 18866b746a..0ef73aee09 100644
--- a/source/network/NetServer.h
+++ b/source/network/NetServer.h
@@ -24,7 +24,6 @@
 
 #include <vector>
 
-class CNetClientSessionLocal;
 class CNetServerSession;
 class CNetServerTurnManager;
 class CFsmEvent;
@@ -57,12 +56,26 @@ enum NetServerState
 /**
  * Server session representation of client state
  */
-enum
+enum NetServerSessionState
 {
+	// The client has disconnected or been disconnected
 	NSS_UNCONNECTED,
+
+	// The client has just connected and we're waiting for its handshake message,
+	// to agree on the protocol version
 	NSS_HANDSHAKE,
+
+	// The client has handshook and we're waiting for its authentication message,
+	// to find its name and check its password etc
 	NSS_AUTHENTICATE,
+
+	// The client has fully joined, and is in the pregame setup stage
+	// or is loading the game.
+	// Server must be in SERVER_STATE_PREGAME or SERVER_STATE_LOADING.
 	NSS_PREGAME,
+
+	// The client is running the game.
+	// Server must be in SERVER_STATE_LOADING or SERVER_STATE_INGAME.
 	NSS_INGAME
 };
 
@@ -73,9 +86,6 @@ enum
  *
  * TODO: ideally the ENet server would run in a separate thread so it can receive
  * and forward messages with minimal latency. But that's not supported now.
- *
- * TODO: we need to be much more careful at handling client states, to cope with
- * e.g. people being in the middle of connecting or authenticating when we start the game.
  */
 class CNetServer
 {
@@ -102,28 +112,23 @@ public:
 	 */
 	virtual void Poll();
 
+	/**
+	 * Flush any queued outgoing network messages.
+	 * This should be called soon after sending a group of messages that may be batched together.
+	 */
+	void Flush();
+
 	/**
 	 * Send a message to the given network peer.
 	 */
 	bool SendMessage(ENetPeer* peer, const CNetMessage* message);
 
-	/**
-	 * Send a message to the given local client.
-	 */
-	void SendLocalMessage(CNetClientSessionLocal& clientSession, const CNetMessage* message);
-
 	/**
 	 * Send a message to all clients who have completed the full connection process
 	 * (i.e. are in the pre-game or in-game states).
 	 */
 	bool Broadcast(const CNetMessage* message);
 
-	/**
-	 * Register a local client with this server (equivalent to a remote client connecting
-	 * over the network).
-	 */
-	void AddLocalClientSession(CNetClientSessionLocal& clientSession);
-
 	/**
 	 * Call from the GUI to update the player assignments.
 	 * The given GUID will be (re)assigned to the given player ID.
@@ -176,7 +181,6 @@ private:
 
 	void SetupSession(CNetServerSession* session);
 	bool HandleConnect(CNetServerSession* session);
-	bool HandleDisconnect(CNetServerSession* session);
 
 	void OnUserJoin(CNetServerSession* session);
 	void OnUserLeave(CNetServerSession* session);
@@ -192,7 +196,7 @@ private:
 
 	void ConstructPlayerAssignmentMessage(CPlayerAssignmentMessage& message);
 
-	bool HandleMessageReceive(const CNetMessage* message, CNetServerSession* session);
+	void HandleMessageReceive(const CNetMessage* message, CNetServerSession* session);
 
 	/**
 	 * Internal script context for (de)serializing script messages.
@@ -202,18 +206,14 @@ private:
 	ScriptInterface* m_ScriptInterface;
 
 	ENetHost* m_Host;
-	std::vector<ENetPeer*> m_Peers;
 	std::vector<CNetServerSession*> m_Sessions;
 
 	CNetStatsTable* m_Stats;
 
-	std::vector<std::pair<CNetServerSession*, CNetMessage*> > m_LocalMessageQueue;
-
 	NetServerState m_State;
 
 	CStrW m_ServerName;
 	CStrW m_WelcomeMessage;
-	int m_Port;
 
 	u32 m_NextHostID;
 
diff --git a/source/network/NetSession.cpp b/source/network/NetSession.cpp
index c59835e008..117787e3c9 100644
--- a/source/network/NetSession.cpp
+++ b/source/network/NetSession.cpp
@@ -28,34 +28,19 @@
 
 static const int CHANNEL_COUNT = 1;
 
-
 CNetClientSession::CNetClientSession(CNetClient& client) :
-	m_Client(client)
+	m_Client(client), m_Host(NULL), m_Server(NULL), m_Stats(NULL)
 {
 }
 
 CNetClientSession::~CNetClientSession()
-{
-}
-
-
-// TODO: the whole disconnection process is a tangled mess - need
-// to sort out exactly what happens at what times, and the interactions
-// between ENet and session and client and GUI
-
-CNetClientSessionRemote::CNetClientSessionRemote(CNetClient& client) :
-	CNetClientSession(client), m_Host(NULL), m_Server(NULL), m_Stats(NULL)
-{
-}
-
-CNetClientSessionRemote::~CNetClientSessionRemote()
 {
 	delete m_Stats;
 
 	if (m_Host && m_Server)
 	{
-		// Disconnect without waiting for confirmation
-		enet_peer_disconnect_now(m_Server, 0);
+		// Disconnect immediately (we can't wait for acks)
+		enet_peer_disconnect_now(m_Server, NDR_UNEXPECTED_SHUTDOWN);
 		enet_host_destroy(m_Host);
 
 		m_Host = NULL;
@@ -63,7 +48,7 @@ CNetClientSessionRemote::~CNetClientSessionRemote()
 	}
 }
 
-bool CNetClientSessionRemote::Connect(u16 port, const CStr& server)
+bool CNetClientSession::Connect(u16 port, const CStr& server)
 {
 	debug_assert(!m_Host);
 	debug_assert(!m_Server);
@@ -94,12 +79,12 @@ bool CNetClientSessionRemote::Connect(u16 port, const CStr& server)
 	return true;
 }
 
-void CNetClientSessionRemote::Disconnect()
+void CNetClientSession::Disconnect(u32 reason)
 {
 	debug_assert(m_Host && m_Server);
 
 	// TODO: ought to do reliable async disconnects, probably
-	enet_peer_disconnect_now(m_Server, 0);
+	enet_peer_disconnect_now(m_Server, reason);
 	enet_host_destroy(m_Host);
 
 	m_Host = NULL;
@@ -108,7 +93,7 @@ void CNetClientSessionRemote::Disconnect()
 	SAFE_DELETE(m_Stats);
 }
 
-void CNetClientSessionRemote::Poll()
+void CNetClientSession::Poll()
 {
 	debug_assert(m_Host && m_Server);
 
@@ -126,7 +111,7 @@ void CNetClientSessionRemote::Poll()
 			enet_address_get_host_ip(&event.peer->address, hostname, ARRAY_SIZE(hostname));
 			LOGMESSAGE(L"Net client: Connected to %hs:%u", hostname, event.peer->address.port);
 
-			GetClient().HandleConnect();
+			m_Client.HandleConnect();
 
 			break;
 		}
@@ -135,20 +120,19 @@ void CNetClientSessionRemote::Poll()
 		{
 			debug_assert(event.peer == m_Server);
 
-			GetClient().HandleDisconnect();
-
-			break;
+			LOGMESSAGE(L"Net client: Disconnected");
+			m_Client.HandleDisconnect(event.data);
+			return;
 		}
 
 		case ENET_EVENT_TYPE_RECEIVE:
 		{
-			CNetMessage* msg = CNetMessageFactory::CreateMessage(event.packet->data, event.packet->dataLength, GetClient().GetScriptInterface());
+			CNetMessage* msg = CNetMessageFactory::CreateMessage(event.packet->data, event.packet->dataLength, m_Client.GetScriptInterface());
 			if (msg)
 			{
 				LOGMESSAGE(L"Net client: Received message %hs of size %lu from server", msg->ToString().c_str(), (unsigned long)msg->GetSerializedLength());
 
-				bool ok = GetClient().HandleMessage(msg);
-				debug_assert(ok); // TODO
+				m_Client.HandleMessage(msg);
 
 				delete msg;
 			}
@@ -162,7 +146,14 @@ void CNetClientSessionRemote::Poll()
 
 }
 
-bool CNetClientSessionRemote::SendMessage(const CNetMessage* message)
+void CNetClientSession::Flush()
+{
+	debug_assert(m_Host && m_Server);
+
+	enet_host_flush(m_Host);
+}
+
+bool CNetClientSession::SendMessage(const CNetMessage* message)
 {
 	debug_assert(m_Host && m_Server);
 
@@ -171,95 +162,24 @@ bool CNetClientSessionRemote::SendMessage(const CNetMessage* message)
 
 
 
-CNetClientSessionLocal::CNetClientSessionLocal(CNetClient& client, CNetServer& server) :
-	CNetClientSession(client), m_Server(server), m_ServerSession(NULL)
-{
-	server.AddLocalClientSession(*this);
-	client.HandleConnect();
-}
-
-void CNetClientSessionLocal::Poll()
-{
-	for (size_t i = 0; i < m_LocalMessageQueue.size(); ++i)
-	{
-		CNetMessage* msg = m_LocalMessageQueue[i];
-
-		LOGMESSAGE(L"Net client: Received local message %hs of size %lu from server", msg->ToString().c_str(), (unsigned long)msg->GetSerializedLength());
-
-		bool ok = GetClient().HandleMessage(msg);
-		debug_assert(ok); // TODO
-
-		delete msg;
-	}
-	m_LocalMessageQueue.clear();
-}
-
-void CNetClientSessionLocal::Disconnect()
-{
-	GetClient().HandleDisconnect();
-}
-
-bool CNetClientSessionLocal::SendMessage(const CNetMessage* message)
-{
-	LOGMESSAGE(L"Net client: Sending local message %hs to server", message->ToString().c_str());
-	m_Server.SendLocalMessage(*this, message);
-	return true;
-}
-
-void CNetClientSessionLocal::AddLocalMessage(const CNetMessage* message)
-{
-	// Clone into the client's script context
-	CNetMessage* clonedMessage = CNetMessageFactory::CloneMessage(message, GetClient().GetScriptInterface());
-	if (!clonedMessage)
-		return;
-	m_LocalMessageQueue.push_back(clonedMessage);
-}
-
-
-
-CNetServerSession::CNetServerSession(CNetServer& server) :
-	m_Server(server)
+CNetServerSession::CNetServerSession(CNetServer& server, ENetPeer* peer) :
+	m_Server(server), m_Peer(peer)
 {
 }
 
-CNetServerSession::~CNetServerSession()
+void CNetServerSession::Disconnect(u32 reason)
 {
+	Update((uint)NMT_CONNECTION_LOST, NULL);
+
+	enet_peer_disconnect(m_Peer, reason);
 }
 
-
-
-CNetServerSessionRemote::CNetServerSessionRemote(CNetServer& server, ENetPeer* peer) :
-	CNetServerSession(server), m_Peer(peer)
+void CNetServerSession::DisconnectNow(u32 reason)
 {
-
+	enet_peer_disconnect_now(m_Peer, reason);
 }
 
-void CNetServerSessionRemote::Disconnect()
+bool CNetServerSession::SendMessage(const CNetMessage* message)
 {
-	// TODO: ought to do reliable async disconnects, probably
-	enet_peer_disconnect_now(m_Peer, 0);
-}
-
-bool CNetServerSessionRemote::SendMessage(const CNetMessage* message)
-{
-	return GetServer().SendMessage(m_Peer, message);
-}
-
-
-
-CNetServerSessionLocal::CNetServerSessionLocal(CNetServer& server, CNetClientSessionLocal& clientSession) :
-	CNetServerSession(server), m_ClientSession(clientSession)
-{
-}
-
-void CNetServerSessionLocal::Disconnect()
-{
-	m_ClientSession.Disconnect();
-}
-
-bool CNetServerSessionLocal::SendMessage(const CNetMessage* message)
-{
-	LOGMESSAGE(L"Net server: Sending local message %hs to %p", message->ToString().c_str(), &m_ClientSession);
-	m_ClientSession.AddLocalMessage(message);
-	return true;
+	return m_Server.SendMessage(m_Peer, message);
 }
diff --git a/source/network/NetSession.h b/source/network/NetSession.h
index 2cbbf5ce9b..dbc6010b90 100644
--- a/source/network/NetSession.h
+++ b/source/network/NetSession.h
@@ -26,8 +26,6 @@
 class CNetClient;
 class CNetServer;
 
-class CNetServerSessionLocal; // forward declaration, needed because of circular references
-
 class CNetStatsTable;
 
 /**
@@ -37,10 +35,6 @@ class CNetStatsTable;
  * Each session has two classes: CNetClientSession runs on the client,
  * and CNetServerSession runs on the server.
  * A client runs one session at once; a server typically runs many.
- *
- * There are two variants of each session: Remote (the normal ENet-based
- * network session) and Local (a shortcut when the client and server are
- * running inside the same process and don't need the network to communicate).
  */
 
 /**
@@ -53,70 +47,38 @@ class CNetClientSession
 
 public:
 	CNetClientSession(CNetClient& client);
-	virtual ~CNetClientSession();
-
-	virtual void Poll() = 0;
-	virtual void Disconnect() = 0;
-	virtual bool SendMessage(const CNetMessage* message) = 0;
-
-	CNetClient& GetClient() { return m_Client; }
-
-private:
-	CNetClient& m_Client;
-};
-
-/**
- * ENet-based implementation of CNetClientSession.
- */
-class CNetClientSessionRemote : public CNetClientSession
-{
-	NONCOPYABLE(CNetClientSessionRemote);
-
-public:
-	CNetClientSessionRemote(CNetClient& client);
-	~CNetClientSessionRemote();
+	~CNetClientSession();
 
 	bool Connect(u16 port, const CStr& server);
 
-	virtual void Poll();
-	virtual void Disconnect();
-	virtual bool SendMessage(const CNetMessage* message);
+	/**
+	 * Process queued incoming messages.
+	 */
+	void Poll();
 
-	ENetPacket* CreatePacket(const CNetMessage* message);
+	/**
+	 * Flush queued outgoing network messages.
+	 */
+	void Flush();
+
+	/**
+	 * Disconnect from the server.
+	 * Sends a disconnection notification to the server.
+	 */
+	void Disconnect(u32 reason);
+
+	/**
+	 * Send a message to the server.
+	 */
+	bool SendMessage(const CNetMessage* message);
 
 private:
+	CNetClient& m_Client;
 	ENetHost* m_Host;
 	ENetPeer* m_Server;
 	CNetStatsTable* m_Stats;
 };
 
-/**
- * Local implementation of CNetClientSession, for use with servers
- * running in the same process.
- */
-class CNetClientSessionLocal : public CNetClientSession
-{
-	NONCOPYABLE(CNetClientSessionLocal);
-
-public:
-	CNetClientSessionLocal(CNetClient& client, CNetServer& server);
-
-	void SetServerSession(CNetServerSessionLocal* session) { m_ServerSession = session; }
-	CNetServerSessionLocal* GetServerSession() { return m_ServerSession; }
-
-	virtual void Poll();
-	virtual void Disconnect();
-	virtual bool SendMessage(const CNetMessage* message);
-
-	void AddLocalMessage(const CNetMessage* message);
-
-private:
-	CNetServer& m_Server;
-	CNetServerSessionLocal* m_ServerSession;
-
-	std::vector<CNetMessage*> m_LocalMessageQueue;
-};
-
 
 /**
  * The server's end of a network session.
@@ -128,8 +90,7 @@ class CNetServerSession : public CFsm
 	NONCOPYABLE(CNetServerSession);
 
 public:
-	CNetServerSession(CNetServer& server);
-	virtual ~CNetServerSession();
+	CNetServerSession(CNetServer& server, ENetPeer* peer);
 
 	CNetServer& GetServer() { return m_Server; }
 
@@ -142,50 +103,34 @@ public:
 	u32 GetHostID() const { return m_HostID; }
 	void SetHostID(u32 id) { m_HostID = id; }
 
-	virtual void Disconnect() = 0;
-	virtual bool SendMessage(const CNetMessage* message) = 0;
+	/**
+	 * Sends a disconnection notification to the client,
+	 * and sends a NMT_CONNECTION_LOST message to the session FSM.
+	 * The server will receive a disconnection notification after a while.
+	 * The server will not receive any further messages sent via this session.
+	 */
+	void Disconnect(u32 reason);
+
+	/**
+	 * Sends an unreliable disconnection notification to the client.
+	 * The server will not receive any disconnection notification.
+	 * The server will not receive any further messages sent via this session.
+	 */
+	void DisconnectNow(u32 reason);
+
+	/**
+	 * Send a message to the client.
+	 */
+	bool SendMessage(const CNetMessage* message);
 
 private:
 	CNetServer& m_Server;
 
+	ENetPeer* m_Peer;
+
 	CStr m_GUID;
 	CStrW m_UserName;
 	u32 m_HostID;
 };
 
-/**
- * ENet-based implementation of CNetServerSession.
- */
-class CNetServerSessionRemote : public CNetServerSession
-{
-	NONCOPYABLE(CNetServerSessionRemote);
-
-public:
-	CNetServerSessionRemote(CNetServer& server, ENetPeer* peer);
-
-	virtual void Disconnect();
-	virtual bool SendMessage(const CNetMessage* message);
-
-private:
-	ENetPeer* m_Peer;
-};
-
-/**
- * Local implementation of CNetServerSession, for use with clients
- * running in the same process.
- */
-class CNetServerSessionLocal : public CNetServerSession
-{
-	NONCOPYABLE(CNetServerSessionLocal);
-
-public:
-	CNetServerSessionLocal(CNetServer& server, CNetClientSessionLocal& clientSession);
-
-	virtual void Disconnect();
-	virtual bool SendMessage(const CNetMessage* message);
-
-private:
-	CNetClientSessionLocal& m_ClientSession;
-};
-
 #endif	// NETSESSION_H
diff --git a/source/network/fsm.cpp b/source/network/fsm.cpp
index d216f5cfc1..a245ac14b5 100644
--- a/source/network/fsm.cpp
+++ b/source/network/fsm.cpp
@@ -141,7 +141,7 @@ bool CFsmTransition::ApplyConditions( void ) const
 
 //-----------------------------------------------------------------------------
 // Name: RunActions()
-// Desc: Execur actions for the transition
+// Desc: Execute actions for the transition
 // Note: If there are no actions, assume true
 //-----------------------------------------------------------------------------
 bool CFsmTransition::RunActions( void ) const
diff --git a/source/network/fsm.h b/source/network/fsm.h
index 45168bc1e9..a656ee7e68 100644
--- a/source/network/fsm.h
+++ b/source/network/fsm.h
@@ -150,6 +150,7 @@ public:
 										 unsigned int eventType ) const;
 	CFsmTransition*	GetEventTransition	( unsigned int eventType ) const;
 	void			SetFirstState		( unsigned int firstState );
+	void			SetCurrState		( unsigned int state );
 	unsigned int	GetCurrState		( void ) const		{ return m_CurrState; }
 	const StateSet&	GetStates			( void ) const		{ return m_States; }
 	const EventMap&	GetEvents			( void ) const		{ return m_Events; }
@@ -159,8 +160,6 @@ public:
 	bool			IsValidEvent		( unsigned int eventType ) const;
 	virtual bool	IsDone				( void ) const;
 
-protected:
-	void			SetCurrState		( unsigned int state );
 private:
 	bool			IsFirstTime			( void ) const;
 
diff --git a/source/network/tests/test_Net.h b/source/network/tests/test_Net.h
index e3b00350cc..b9ffbd79e1 100644
--- a/source/network/tests/test_Net.h
+++ b/source/network/tests/test_Net.h
@@ -64,8 +64,7 @@ public:
 	void connect(CNetServer& server, const std::vector<CNetClient*>& clients)
 	{
 		TS_ASSERT(server.SetupConnection());
-		clients[0]->SetupLocalConnection(server);
-		for (size_t j = 1; j < clients.size(); ++j)
+		for (size_t j = 0; j < clients.size(); ++j)
 			TS_ASSERT(clients[j]->SetupConnection("127.0.0.1"));
 
 		for (size_t i = 0; ; ++i)
@@ -138,7 +137,7 @@ public:
 		CNetServer server;
 
 		CScriptValRooted attrs;
-		server.GetScriptInterface().Eval("({map:'Latium',thing:'example'})", attrs);
+		server.GetScriptInterface().Eval("({map:'_default',thing:'example'})", attrs);
 		server.UpdateGameAttributes(attrs);
 
 		CNetClient client1(&client1Game);
diff --git a/source/ps/GameSetup/GameSetup.cpp b/source/ps/GameSetup/GameSetup.cpp
index 1fd93db893..1774138ff2 100644
--- a/source/ps/GameSetup/GameSetup.cpp
+++ b/source/ps/GameSetup/GameSetup.cpp
@@ -886,7 +886,7 @@ static bool Autostart(const CmdLineArgs& args)
 
 		g_NetClient = new CNetClient(g_Game);
 		// TODO: player name, etc
-		g_NetClient->SetupLocalConnection(*g_NetServer);
+		g_NetClient->SetupConnection("127.0.0.1");
 	}
 	else if (args.Has("autostart-client"))
 	{