7 files changed, 249 insertions, 76 deletions
diff --git a/Swiften/Client/Client.cpp b/Swiften/Client/Client.cpp
index 763c83e..2406b0f 100644
--- a/Swiften/Client/Client.cpp
+++ b/Swiften/Client/Client.cpp
@@ -81,9 +81,10 @@ void Client::handleConnectorFinished(boost::shared_ptr<Connection> connection, C
 
 		session_ = ClientSession::create(jid_, sessionStream_);
 		session_->onInitialized.connect(boost::bind(boost::ref(onConnected)));
+		session_->onStanzaAcked.connect(boost::bind(&Client::handleStanzaAcked, this, _1));
 		session_->onFinished.connect(boost::bind(&Client::handleSessionFinished, this, _1));
 		session_->onNeedCredentials.connect(boost::bind(&Client::handleNeedCredentials, this));
-		session_->onElementReceived.connect(boost::bind(&Client::handleElement, this, _1));
+		session_->onStanzaReceived.connect(boost::bind(&Client::handleStanza, this, _1));
 		session_->start();
 	}
 }
@@ -115,7 +116,7 @@ void Client::send(boost::shared_ptr<Stanza> stanza) {
 		std::cerr << "Warning: Client: Trying to send a stanza while disconnected." << std::endl;
 		return;
 	}
-	session_->sendElement(stanza);
+	session_->sendStanza(stanza);
 }
 
 void Client::sendIQ(boost::shared_ptr<IQ> iq) {
@@ -134,20 +135,20 @@ String Client::getNewIQID() {
 	return idGenerator_.generateID();
 }
 
-void Client::handleElement(boost::shared_ptr<Element> element) {
-	boost::shared_ptr<Message> message = boost::dynamic_pointer_cast<Message>(element);
+void Client::handleStanza(boost::shared_ptr<Stanza> stanza) {
+	boost::shared_ptr<Message> message = boost::dynamic_pointer_cast<Message>(stanza);
 	if (message) {
 		onMessageReceived(message);
 		return;
 	}
 
-	boost::shared_ptr<Presence> presence = boost::dynamic_pointer_cast<Presence>(element);
+	boost::shared_ptr<Presence> presence = boost::dynamic_pointer_cast<Presence>(stanza);
 	if (presence) {
 		onPresenceReceived(presence);
 		return;
 	}
 
-	boost::shared_ptr<IQ> iq = boost::dynamic_pointer_cast<IQ>(element);
+	boost::shared_ptr<IQ> iq = boost::dynamic_pointer_cast<IQ>(stanza);
 	if (iq) {
 		onIQReceived(iq);
 		return;
@@ -222,6 +223,13 @@ void Client::handleNeedCredentials() {
 	session_->sendCredentials(password_);
 }
 
+bool Client::getStreamManagementEnabled() const {
+	if (session_) {
+		return session_->getStreamManagementEnabled();
+	}
+	return false;
+}
+
 void Client::handleDataRead(const String& data) {
 	onDataRead(data);
 }
@@ -230,4 +238,8 @@ void Client::handleDataWritten(const String& data) {
 	onDataWritten(data);
 }
 
+void Client::handleStanzaAcked(boost::shared_ptr<Stanza> stanza) {
+	onStanzaAcked(stanza);
+}
+
 }
diff --git a/Swiften/Client/Client.h b/Swiften/Client/Client.h
index 313f856..10e7c38 100644
--- a/Swiften/Client/Client.h
+++ b/Swiften/Client/Client.h
@@ -44,6 +44,8 @@ namespace Swift {
 			
 			bool isAvailable();
 
+			bool getStreamManagementEnabled() const;
+
 			virtual void sendIQ(boost::shared_ptr<IQ>);
 			virtual void sendMessage(boost::shared_ptr<Message>);
 			virtual void sendPresence(boost::shared_ptr<Presence>);
@@ -53,16 +55,18 @@ namespace Swift {
 			boost::signal<void ()> onConnected;
 			boost::signal<void (const String&)> onDataRead;
 			boost::signal<void (const String&)> onDataWritten;
+			boost::signal<void (boost::shared_ptr<Stanza>)> onStanzaAcked;
 
 		private:
 			void handleConnectorFinished(boost::shared_ptr<Connection>, Connector::ref);
 			void send(boost::shared_ptr<Stanza>);
 			virtual String getNewIQID();
-			void handleElement(boost::shared_ptr<Element>);
+			void handleStanza(boost::shared_ptr<Stanza>);
 			void handleSessionFinished(boost::shared_ptr<Error>);
 			void handleNeedCredentials();
 			void handleDataRead(const String&);
 			void handleDataWritten(const String&);
+			void handleStanzaAcked(boost::shared_ptr<Stanza>);
 
 			void closeConnection();
 
diff --git a/Swiften/Client/ClientSession.cpp b/Swiften/Client/ClientSession.cpp
index 720bc6d..551d835 100644
--- a/Swiften/Client/ClientSession.cpp
+++ b/Swiften/Client/ClientSession.cpp
@@ -24,7 +24,12 @@
 #include "Swiften/Elements/Compressed.h"
 #include "Swiften/Elements/CompressFailure.h"
 #include "Swiften/Elements/CompressRequest.h"
+#include "Swiften/Elements/EnableStreamManagement.h"
+#include "Swiften/Elements/StreamManagementEnabled.h"
+#include "Swiften/Elements/StreamManagementFailed.h"
 #include "Swiften/Elements/StartSession.h"
+#include "Swiften/Elements/StanzaAck.h"
+#include "Swiften/Elements/StanzaAckRequest.h"
 #include "Swiften/Elements/IQ.h"
 #include "Swiften/Elements/ResourceBind.h"
 #include "Swiften/SASL/PLAINClientAuthenticator.h"
@@ -42,6 +47,7 @@ ClientSession::ClientSession(
 			stream(stream),
 			allowPLAINOverNonTLS(false),
 			needSessionStart(false),
+			needResourceBind(false),
 			authenticator(NULL) {
 }
 
@@ -65,8 +71,11 @@ void ClientSession::sendStreamHeader() {
 	stream->writeHeader(header);
 }
 
-void ClientSession::sendElement(boost::shared_ptr<Element> element) {
-	stream->writeElement(element);
+void ClientSession::sendStanza(boost::shared_ptr<Stanza> stanza) {
+	stream->writeElement(stanza);
+	if (stanzaAckRequester_) {
+		stanzaAckRequester_->handleStanzaSent(stanza);
+	}
 }
 
 void ClientSession::handleStreamStart(const ProtocolHeader&) {
@@ -75,8 +84,77 @@ void ClientSession::handleStreamStart(const ProtocolHeader&) {
 }
 
 void ClientSession::handleElement(boost::shared_ptr<Element> element) {
-	if (getState() == Initialized) {
-		onElementReceived(element);
+	if (boost::shared_ptr<Stanza> stanza = boost::dynamic_pointer_cast<Stanza>(element)) {
+		if (stanzaAckResponder_) {
+			stanzaAckResponder_->handleStanzaReceived();
+		}
+		if (getState() == Initialized) {
+			onStanzaReceived(stanza);
+		}
+		else if (boost::shared_ptr<IQ> iq = boost::dynamic_pointer_cast<IQ>(element)) {
+			if (state == BindingResource) {
+				boost::shared_ptr<ResourceBind> resourceBind(iq->getPayload<ResourceBind>());
+				if (iq->getType() == IQ::Error && iq->getID() == "session-bind") {
+					finishSession(Error::ResourceBindError);
+				}
+				else if (!resourceBind) {
+					finishSession(Error::UnexpectedElementError);
+				}
+				else if (iq->getType() == IQ::Result) {
+					localJID = resourceBind->getJID();
+					if (!localJID.isValid()) {
+						finishSession(Error::ResourceBindError);
+					}
+					needResourceBind = false;
+					continueSessionInitialization();
+				}
+				else {
+					finishSession(Error::UnexpectedElementError);
+				}
+			}
+			else if (state == StartingSession) {
+				if (iq->getType() == IQ::Result) {
+					needSessionStart = false;
+					continueSessionInitialization();
+				}
+				else if (iq->getType() == IQ::Error) {
+					finishSession(Error::SessionStartError);
+				}
+				else {
+					finishSession(Error::UnexpectedElementError);
+				}
+			}
+			else {
+				finishSession(Error::UnexpectedElementError);
+			}
+		}
+	}
+	else if (boost::dynamic_pointer_cast<StanzaAckRequest>(element)) {
+		if (stanzaAckResponder_) {
+			stanzaAckResponder_->handleAckRequestReceived();
+		}
+	}
+	else if (boost::shared_ptr<StanzaAck> ack = boost::dynamic_pointer_cast<StanzaAck>(element)) {
+		if (stanzaAckRequester_) {
+			if (ack->isValid()) {
+				stanzaAckRequester_->handleAckReceived(ack->getHandledStanzasCount());
+			}
+			else {
+				std::cerr << "Warning: Got invalid ack from server" << std::endl;
+			}
+		}
+		else {
+			std::cerr << "Warning: Ignoring ack" << std::endl;
+		}
+	}
+	else if (getState() == Initialized) {
+		boost::shared_ptr<Stanza> stanza = boost::dynamic_pointer_cast<Stanza>(element);
+		if (stanza) {
+			if (stanzaAckResponder_) {
+				stanzaAckResponder_->handleStanzaReceived();
+			}
+			onStanzaReceived(stanza);
+		}
 	}
 	else if (StreamFeatures* streamFeatures = dynamic_cast<StreamFeatures*>(element.get())) {
 		if (!checkState(Negotiating)) {
@@ -133,25 +211,14 @@ void ClientSession::handleElement(boost::shared_ptr<Element> element) {
 		else {
 			// Start the session
 			stream->setWhitespacePingEnabled(true);
-
-			if (streamFeatures->hasSession()) {
-				needSessionStart = true;
-			}
-
-			if (streamFeatures->hasResourceBind()) {
-				state = BindingResource;
-				boost::shared_ptr<ResourceBind> resourceBind(new ResourceBind());
-				if (!localJID.getResource().isEmpty()) {
-					resourceBind->setResource(localJID.getResource());
-				}
-				stream->writeElement(IQ::createRequest(IQ::Set, JID(), "session-bind", resourceBind));
-			}
-			else if (needSessionStart) {
-				sendSessionStart();
+			needSessionStart = streamFeatures->hasSession();
+			needResourceBind = streamFeatures->hasResourceBind();
+			if (streamFeatures->hasStreamManagement()) {
+				state = EnablingSessionManagement;
+				stream->writeElement(boost::shared_ptr<EnableStreamManagement>(new EnableStreamManagement()));
 			}
 			else {
-				state = Initialized;
-				onInitialized();
+				continueSessionInitialization();
 			}
 		}
 	}
@@ -165,6 +232,17 @@ void ClientSession::handleElement(boost::shared_ptr<Element> element) {
 	else if (boost::dynamic_pointer_cast<CompressFailure>(element)) {
 		finishSession(Error::CompressionFailedError);
 	}
+	else if (boost::dynamic_pointer_cast<StreamManagementEnabled>(element)) {
+		stanzaAckRequester_ = boost::shared_ptr<StanzaAckRequester>(new StanzaAckRequester());
+		stanzaAckRequester_->onRequestAck.connect(boost::bind(&ClientSession::requestAck, this));
+		stanzaAckRequester_->onStanzaAcked.connect(boost::bind(&ClientSession::handleStanzaAcked, this, _1));
+		stanzaAckResponder_ = boost::shared_ptr<StanzaAckResponder>(new StanzaAckResponder());
+		stanzaAckResponder_->onAck.connect(boost::bind(&ClientSession::ack, this, _1));
+		continueSessionInitialization();
+	}
+	else if (boost::dynamic_pointer_cast<StreamManagementFailed>(element)) {
+		continueSessionInitialization();
+	}
 	else if (AuthChallenge* challenge = dynamic_cast<AuthChallenge*>(element.get())) {
 		checkState(Authenticating);
 		assert(authenticator);
@@ -201,47 +279,6 @@ void ClientSession::handleElement(boost::shared_ptr<Element> element) {
 	else if (dynamic_cast<StartTLSFailure*>(element.get())) {
 		finishSession(Error::TLSError);
 	}
-	else if (IQ* iq = dynamic_cast<IQ*>(element.get())) {
-		if (state == BindingResource) {
-			boost::shared_ptr<ResourceBind> resourceBind(iq->getPayload<ResourceBind>());
-			if (iq->getType() == IQ::Error && iq->getID() == "session-bind") {
-				finishSession(Error::ResourceBindError);
-			}
-			else if (!resourceBind) {
-				finishSession(Error::UnexpectedElementError);
-			}
-			else if (iq->getType() == IQ::Result) {
-				localJID = resourceBind->getJID();
-				if (!localJID.isValid()) {
-					finishSession(Error::ResourceBindError);
-				}
-				if (needSessionStart) {
-					sendSessionStart();
-				}
-				else {
-					state = Initialized;
-				}
-			}
-			else {
-				finishSession(Error::UnexpectedElementError);
-			}
-		}
-		else if (state == StartingSession) {
-			if (iq->getType() == IQ::Result) {
-				state = Initialized;
-				onInitialized();
-			}
-			else if (iq->getType() == IQ::Error) {
-				finishSession(Error::SessionStartError);
-			}
-			else {
-				finishSession(Error::UnexpectedElementError);
-			}
-		}
-		else {
-			finishSession(Error::UnexpectedElementError);
-		}
-	}
 	else {
 		// FIXME Not correct?
 		state = Initialized;
@@ -249,9 +286,23 @@ void ClientSession::handleElement(boost::shared_ptr<Element> element) {
 	}
 }
 
-void ClientSession::sendSessionStart() {
-	state = StartingSession;
-	stream->writeElement(IQ::createRequest(IQ::Set, JID(), "session-start", boost::shared_ptr<StartSession>(new StartSession())));
+void ClientSession::continueSessionInitialization() {
+	if (needResourceBind) {
+		state = BindingResource;
+		boost::shared_ptr<ResourceBind> resourceBind(new ResourceBind());
+		if (!localJID.getResource().isEmpty()) {
+			resourceBind->setResource(localJID.getResource());
+		}
+		sendStanza(IQ::createRequest(IQ::Set, JID(), "session-bind", resourceBind));
+	}
+	else if (needSessionStart) {
+		state = StartingSession;
+		sendStanza(IQ::createRequest(IQ::Set, JID(), "session-start", boost::shared_ptr<StartSession>(new StartSession())));
+	}
+	else {
+		state = Initialized;
+		onInitialized();
+	}
 }
 
 bool ClientSession::checkState(State state) {
@@ -298,4 +349,16 @@ void ClientSession::finishSession(boost::shared_ptr<Swift::Error> error) {
 }
 
 
+void ClientSession::requestAck() {
+	stream->writeElement(boost::shared_ptr<StanzaAckRequest>(new StanzaAckRequest()));
+}
+
+void ClientSession::handleStanzaAcked(boost::shared_ptr<Stanza> stanza) {
+	onStanzaAcked(stanza);
+}
+
+void ClientSession::ack(unsigned int handledStanzasCount) {
+	stream->writeElement(boost::shared_ptr<StanzaAck>(new StanzaAck(handledStanzasCount)));
+}
+
 }
diff --git a/Swiften/Client/ClientSession.h b/Swiften/Client/ClientSession.h
index 936fd18..c15508a 100644
--- a/Swiften/Client/ClientSession.h
+++ b/Swiften/Client/ClientSession.h
@@ -16,6 +16,8 @@
 #include "Swiften/Base/String.h"
 #include "Swiften/JID/JID.h"
 #include "Swiften/Elements/Element.h"
+#include "Swiften/StreamManagement/StanzaAckRequester.h"
+#include "Swiften/StreamManagement/StanzaAckResponder.h"
 
 namespace Swift {
 	class ClientAuthenticator;
@@ -31,6 +33,7 @@ namespace Swift {
 				Encrypting,
 				WaitingForCredentials,
 				Authenticating,
+				EnablingSessionManagement,
 				BindingResource,
 				StartingSession,
 				Initialized,
@@ -65,17 +68,22 @@ namespace Swift {
 				allowPLAINOverNonTLS = b;
 			}
 
+			bool getStreamManagementEnabled() const {
+				return stanzaAckRequester_;
+			}
+
 			void start();
 			void finish();
 
 			void sendCredentials(const String& password);
-			void sendElement(boost::shared_ptr<Element> element);
+			void sendStanza(boost::shared_ptr<Stanza>);
 
 		public:
 			boost::signal<void ()> onNeedCredentials;
 			boost::signal<void ()> onInitialized;
 			boost::signal<void (boost::shared_ptr<Swift::Error>)> onFinished;
-			boost::signal<void (boost::shared_ptr<Element>)> onElementReceived;
+			boost::signal<void (boost::shared_ptr<Stanza>)> onStanzaReceived;
+			boost::signal<void (boost::shared_ptr<Stanza>)> onStanzaAcked;
 		
 		private:
 			ClientSession(
@@ -90,7 +98,6 @@ namespace Swift {
 			}
 
 			void sendStreamHeader();
-			void sendSessionStart();
 
 			void handleElement(boost::shared_ptr<Element>);
 			void handleStreamStart(const ProtocolHeader&);
@@ -99,6 +106,11 @@ namespace Swift {
 			void handleTLSEncrypted();
 
 			bool checkState(State);
+			void continueSessionInitialization();
+
+			void requestAck();
+			void handleStanzaAcked(boost::shared_ptr<Stanza> stanza);
+			void ack(unsigned int handledStanzasCount);
 
 		private:
 			JID localJID;
@@ -106,6 +118,9 @@ namespace Swift {
 			boost::shared_ptr<SessionStream> stream;
 			bool allowPLAINOverNonTLS;
 			bool needSessionStart;
+			bool needResourceBind;
 			ClientAuthenticator* authenticator;
+			boost::shared_ptr<StanzaAckRequester> stanzaAckRequester_;
+			boost::shared_ptr<StanzaAckResponder> stanzaAckResponder_;
 	};
 }
diff --git a/Swiften/Client/UnitTest/ClientSessionTest.cpp b/Swiften/Client/UnitTest/ClientSessionTest.cpp
index 180eab8..e27f130 100644
--- a/Swiften/Client/UnitTest/ClientSessionTest.cpp
+++ b/Swiften/Client/UnitTest/ClientSessionTest.cpp
@@ -19,6 +19,11 @@
 #include "Swiften/Elements/AuthRequest.h"
 #include "Swiften/Elements/AuthSuccess.h"
 #include "Swiften/Elements/AuthFailure.h"
+#include "Swiften/Elements/StreamManagementEnabled.h"
+#include "Swiften/Elements/StreamManagementFailed.h"
+#include "Swiften/Elements/EnableStreamManagement.h"
+#include "Swiften/Elements/IQ.h"
+#include "Swiften/Elements/ResourceBind.h"
 
 using namespace Swift;
 
@@ -31,6 +36,8 @@ class ClientSessionTest : public CppUnit::TestFixture {
 		CPPUNIT_TEST(testAuthenticate);
 		CPPUNIT_TEST(testAuthenticate_Unauthorized);
 		CPPUNIT_TEST(testAuthenticate_NoValidAuthMechanisms);
+		CPPUNIT_TEST(testStreamManagement);
+		CPPUNIT_TEST(testStreamManagement_Failed);
 		/*
 		CPPUNIT_TEST(testResourceBind);
 		CPPUNIT_TEST(testResourceBind_ChangeResource);
@@ -163,6 +170,46 @@ class ClientSessionTest : public CppUnit::TestFixture {
 			CPPUNIT_ASSERT(sessionFinishedError);
 		}
 
+		void testStreamManagement() {
+			boost::shared_ptr<ClientSession> session(createSession());
+			session->start();
+			server->receiveStreamStart();
+			server->sendStreamStart();
+			server->sendStreamFeaturesWithPLAINAuthentication();
+			session->sendCredentials("mypass");
+			server->receiveAuthRequest("PLAIN");
+			server->sendAuthSuccess();
+			server->receiveStreamStart();
+			server->sendStreamStart();
+			server->sendStreamFeaturesWithBindAndStreamManagement();
+			server->receiveStreamManagementEnable();
+			server->sendStreamManagementEnabled();
+
+			CPPUNIT_ASSERT(session->getStreamManagementEnabled());
+			server->receiveBind();
+			// TODO: Test if the requesters & responders do their work
+		}
+
+		void testStreamManagement_Failed() {
+			boost::shared_ptr<ClientSession> session(createSession());
+			session->start();
+			server->receiveStreamStart();
+			server->sendStreamStart();
+			server->sendStreamFeaturesWithPLAINAuthentication();
+			session->sendCredentials("mypass");
+			server->receiveAuthRequest("PLAIN");
+			server->sendAuthSuccess();
+			server->receiveStreamStart();
+			server->sendStreamStart();
+			server->sendStreamFeaturesWithBindAndStreamManagement();
+			server->receiveStreamManagementEnable();
+			server->sendStreamManagementFailed();
+
+			CPPUNIT_ASSERT(!session->getStreamManagementEnabled());
+			server->receiveBind();
+		}
+
+
 	private:
 		boost::shared_ptr<ClientSession> createSession() {
 			boost::shared_ptr<ClientSession> session = ClientSession::create(JID("me@foo.com"), server);
@@ -277,6 +324,13 @@ class ClientSessionTest : public CppUnit::TestFixture {
 					onElementReceived(streamFeatures);
 				}
 
+				void sendStreamFeaturesWithBindAndStreamManagement() {
+					boost::shared_ptr<StreamFeatures> streamFeatures(new StreamFeatures());
+					streamFeatures->setHasResourceBind();
+					streamFeatures->setHasStreamManagement();
+					onElementReceived(streamFeatures);
+				}
+
 				void sendAuthSuccess() {
 					onElementReceived(boost::shared_ptr<AuthSuccess>(new AuthSuccess()));
 				}
@@ -285,6 +339,14 @@ class ClientSessionTest : public CppUnit::TestFixture {
 					onElementReceived(boost::shared_ptr<AuthFailure>(new AuthFailure()));
 				}
 
+				void sendStreamManagementEnabled() {
+					onElementReceived(boost::shared_ptr<StreamManagementEnabled>(new StreamManagementEnabled()));
+				}
+
+				void sendStreamManagementFailed() {
+					onElementReceived(boost::shared_ptr<StreamManagementFailed>(new StreamManagementFailed()));
+				}
+
 				void receiveStreamStart() {
 					Event event = popEvent();
 					CPPUNIT_ASSERT(event.header);
@@ -304,6 +366,20 @@ class ClientSessionTest : public CppUnit::TestFixture {
 					CPPUNIT_ASSERT_EQUAL(mech, request->getMechanism());
 				}
 
+				void receiveStreamManagementEnable() {
+					Event event = popEvent();
+					CPPUNIT_ASSERT(event.element);
+					CPPUNIT_ASSERT(boost::dynamic_pointer_cast<EnableStreamManagement>(event.element));
+				}
+
+				void receiveBind() {
+					Event event = popEvent();
+					CPPUNIT_ASSERT(event.element);
+					boost::shared_ptr<IQ> iq = boost::dynamic_pointer_cast<IQ>(event.element);
+					CPPUNIT_ASSERT(iq);
+					CPPUNIT_ASSERT(iq->getPayload<ResourceBind>());
+				}
+
 				Event popEvent() {
 					CPPUNIT_ASSERT(receivedEvents.size() > 0);
 					Event event = receivedEvents.front();
diff --git a/Swiften/Elements/StanzaAck.h b/Swiften/Elements/StanzaAck.h
index eaf4e26..53b62b4 100644
--- a/Swiften/Elements/StanzaAck.h
+++ b/Swiften/Elements/StanzaAck.h
@@ -13,6 +13,7 @@ namespace Swift {
 	class StanzaAck : public Element, public Shared<StanzaAck> {
 		public:
 			StanzaAck() : valid(false), handledStanzasCount(0) {}
+			StanzaAck(unsigned int handledStanzasCount) : valid(true), handledStanzasCount(handledStanzasCount) {}
 
 			unsigned int getHandledStanzasCount() const {
 				return handledStanzasCount;
diff --git a/Swiften/StreamManagement/StanzaAckResponder.cpp b/Swiften/StreamManagement/StanzaAckResponder.cpp
index 05ab5c4..5b71f91 100644
--- a/Swiften/StreamManagement/StanzaAckResponder.cpp
+++ b/Swiften/StreamManagement/StanzaAckResponder.cpp
@@ -20,5 +20,7 @@ void StanzaAckResponder::handleStanzaReceived() {
 }
 
 void StanzaAckResponder::handleAckRequestReceived() {
-	onAck(handledStanzasCount);}
+	onAck(handledStanzasCount);
+}
+
 }