@@ -1,7 +1,8 @@
 /*
  * Copyright (c) 2011, Jonathan Schleifer <js@webkeks.org>
+ * Copyright (c) 2012, Florian Zeitz <florob@babelmonkeys.de>
  *
  * https://webkeks.org/hg/objxmpp/
  *
  * Permission to use, copy, modify, and/or distribute this software for any
  * purpose with or without fee is hereby granted, provided that the above
@@ -22,10 +23,12 @@
 
 #ifdef HAVE_CONFIG_H
 # include "config.h"
 #endif
 
+#define XMPP_ROSTER_M
+
 #include <assert.h>
 
 #import "XMPPRoster.h"
 #import "XMPPRosterItem.h"
 #import "XMPPConnection.h"
@@ -33,16 +36,18 @@
 #import "XMPPJID.h"
 #import "XMPPMulticastDelegate.h"
 #import "namespaces.h"
 
 @implementation XMPPRoster
-- initWithConnection: (XMPPConnection*)conn
+- initWithConnection: (XMPPConnection*)connection_
 {
 	self = [super init];
 
 	@try {
 		rosterItems = [[OFMutableDictionary alloc] init];
+		connection = connection_;
+		[connection addDelegate: self];
 	} @catch (id e) {
 		[self release];
 		@throw e;
 	}
 
@@ -50,30 +55,14 @@
 }
 
 - (void)dealloc
 {
 	[rosterItems release];
-	[rosterID release];
 
 	[super dealloc];
 }
 
-- (void)XMPP_addRosterItem: (XMPPRosterItem*)rosterItem
-{
-	return [self XMPP_updateRosterItem: rosterItem];
-}
-
-- (void)XMPP_updateRosterItem: (XMPPRosterItem*)rosterItem
-{
-	[rosterItems setObject: rosterItem
-			forKey: [[rosterItem JID] bareJID]];
-}
-
-- (void)XMPP_deleteRosterItem: (XMPPRosterItem*)rosterItem
-{
-	[rosterItems removeObjectForKey: [[rosterItem JID] bareJID]];
-}
 
 - (OFDictionary*)rosterItems
 {
 	return [[rosterItems copy] autorelease];
 }
@@ -80,106 +69,53 @@
 
 - (void)requestRoster
 {
 	XMPPIQ *iq;
 
-	if (rosterID != nil)
-		assert(0);
-
-	rosterID = [[connection generateStanzaID] retain];
 	iq = [XMPPIQ IQWithType: @"get"
-			     ID: rosterID];
+			     ID: [connection generateStanzaID]];
 	[iq addChild: [OFXMLElement elementWithName: @"query"
 					  namespace: XMPP_NS_ROSTER]];
-	[connection sendStanza: iq];
+	[connection     sendIQ: iq
+	    withCallbackObject: self
+		      selector: @selector(XMPP_handleInitialRoster:)];
 }
 
-- (BOOL)handleIQ: (XMPPIQ*)iq
+- (BOOL)connection: (XMPPConnection*)connection_
+      didReceiveIQ: (XMPPIQ*)iq
 {
 	OFXMLElement *rosterElement;
 	OFXMLElement *element;
-	XMPPRosterItem *rosterItem = nil;
-	OFString *subscription;
-	OFEnumerator *enumerator;
-	BOOL isPush = ![[iq ID] isEqual: rosterID];
+	XMPPRosterItem *rosterItem;
 
 	rosterElement = [iq elementForName: @"query"
 				 namespace: XMPP_NS_ROSTER];
 
 	if (rosterElement == nil)
 		return NO;
 
-	if (isPush) {
-		if (![[iq type] isEqual: @"set"])
-			return NO;
-	} else {
-		if (![[iq type] isEqual: @"result"])
-			return NO;
-	}
-
-	enumerator = [[rosterElement children] objectEnumerator];
-	while ((element = [enumerator nextObject]) != nil) {
-		OFMutableArray *groups = [OFMutableArray array];
-		OFEnumerator *groupEnumerator;
-		OFXMLElement *groupElement;
-
-		if (![[element name] isEqual: @"item"] ||
-		    ![[element namespace] isEqual: XMPP_NS_ROSTER])
-			continue;
-
-		rosterItem = [XMPPRosterItem rosterItem];
-		[rosterItem setJID: [XMPPJID JIDWithString:
-		    [[element attributeForName: @"jid"] stringValue]]];
-		[rosterItem setName:
-		    [[element attributeForName: @"name"] stringValue]];
-
-		subscription = [[element attributeForName:
-		    @"subscription"] stringValue];
-
-		if (![subscription isEqual: @"none"] &&
-		    ![subscription isEqual: @"to"] &&
-		    ![subscription isEqual: @"from"] &&
-		    ![subscription isEqual: @"both"] &&
-		    (![subscription isEqual: @"remove"] || !isPush))
-			subscription = @"none";
-
-		[rosterItem setSubscription: subscription];
-
-		groupEnumerator = [[element
-		    elementsForName: @"group"
-			  namespace: XMPP_NS_ROSTER] objectEnumerator];
-		while ((groupElement = [groupEnumerator nextObject]) != nil)
-			[groups addObject: [groupElement stringValue]];
-
-		if ([groups count] > 0)
-			[rosterItem setGroups: groups];
-
-		if ([subscription isEqual: @"remove"])
+	if (![[iq type] isEqual: @"set"])
+		return NO;
+
+	element = [rosterElement elementForName: @"item"
+				      namespace: XMPP_NS_ROSTER];
+
+	if (element != nil) {
+		rosterItem = [self XMPP_rosterItemWithXMLElement: element];
+
+		if ([[rosterItem subscription] isEqual: @"remove"])
 			[self XMPP_deleteRosterItem: rosterItem];
 		else
 			[self XMPP_addRosterItem: rosterItem];
 
-		if (isPush) {
-			SEL sel = @selector(connection:didReceiveRosterItem:);
-
-			[[connection XMPP_delegates]
-			    broadcastSelector: sel
-				forConnection: connection
-				   withObject: rosterItem];
-		}
-	}
-
-	if (isPush) {
-		[connection sendStanza: [iq resultIQ]];
-	} else {
-		[[connection XMPP_delegates]
-		    broadcastSelector: @selector(connectionDidReceiveRoster:)
-			forConnection: connection];
-
-		[rosterID release];
-		rosterID = nil;
-	}
+		if ([delegate respondsToSelector:
+		    @selector(roster:didReceiveRosterItem:)])
+			[delegate         roster: self
+			    didReceiveRosterItem: rosterItem];
+	}
+
+	[connection_ sendStanza: [iq resultIQ]];
 
 	return YES;
 }
 
 - (void)addRosterItem: (XMPPRosterItem*)rosterItem
@@ -233,6 +169,107 @@
 	[query addChild: item];
 	[iq addChild: query];
 
 	[connection sendStanza: iq];
 }
+
+- (void)setDelegate: (id <XMPPRosterDelegate>)delegate_
+{
+	delegate = (id <XMPPRosterDelegate, OFObject>)delegate_;
+}
+
+- (id <XMPPRosterDelegate>)delegate
+{
+	return delegate;
+}
+
+- (void)XMPP_addRosterItem: (XMPPRosterItem*)rosterItem
+{
+	return [self XMPP_updateRosterItem: rosterItem];
+}
+
+- (void)XMPP_updateRosterItem: (XMPPRosterItem*)rosterItem
+{
+	[rosterItems setObject: rosterItem
+			forKey: [[rosterItem JID] bareJID]];
+}
+
+- (void)XMPP_deleteRosterItem: (XMPPRosterItem*)rosterItem
+{
+	[rosterItems removeObjectForKey: [[rosterItem JID] bareJID]];
+}
+
+- (XMPPRosterItem*)XMPP_rosterItemWithXMLElement: (OFXMLElement*)element
+{
+	OFString *subscription;
+	OFEnumerator *groupEnumerator;
+	OFXMLElement *groupElement;
+	OFMutableArray *groups = [OFMutableArray array];
+	XMPPRosterItem *rosterItem = [XMPPRosterItem rosterItem];
+	[rosterItem setJID: [XMPPJID JIDWithString:
+	    [[element attributeForName: @"jid"] stringValue]]];
+	[rosterItem setName:
+	    [[element attributeForName: @"name"] stringValue]];
+
+	subscription = [[element attributeForName:
+	    @"subscription"] stringValue];
+
+	if (![subscription isEqual: @"none"] &&
+	    ![subscription isEqual: @"to"] &&
+	    ![subscription isEqual: @"from"] &&
+	    ![subscription isEqual: @"both"] &&
+	    ![subscription isEqual: @"remove"])
+		subscription = @"none";
+
+	[rosterItem setSubscription: subscription];
+
+	groupEnumerator = [[element
+	    elementsForName: @"group"
+		  namespace: XMPP_NS_ROSTER] objectEnumerator];
+	while ((groupElement = [groupEnumerator nextObject]) != nil)
+		[groups addObject: [groupElement stringValue]];
+
+	if ([groups count] > 0)
+		[rosterItem setGroups: groups];
+
+	return rosterItem;
+}
+
+- (void)XMPP_handleInitialRoster: (XMPPIQ*)iq
+{
+	OFXMLElement *rosterElement;
+	OFEnumerator *enumerator;
+	OFXMLElement *element;
+	XMPPRosterItem *rosterItem = nil;
+
+	rosterElement = [iq elementForName: @"query"
+				 namespace: XMPP_NS_ROSTER];
+
+	enumerator = [[rosterElement children] objectEnumerator];
+	while ((element = [enumerator nextObject]) != nil) {
+		if (![[element name] isEqual: @"item"] ||
+		    ![[element namespace] isEqual: XMPP_NS_ROSTER])
+			continue;
+
+		rosterItem = [self XMPP_rosterItemWithXMLElement: element];
+
+		if ([[rosterItem subscription] isEqual: @"remove"])
+			[self XMPP_deleteRosterItem: rosterItem];
+		else
+			[self XMPP_addRosterItem: rosterItem];
+	}
+
+	if ([delegate respondsToSelector: @selector(rosterWasReceived:)])
+		[delegate rosterWasReceived: self];
+}
+@end
+
+@implementation OFObject (XMPPRosterDelegate)
+- (void)rosterWasReceived: (XMPPRoster*)roster
+{
+}
+
+-         (void)roster: (XMPPRoster*)roster
+  didReceiveRosterItem: (XMPPRosterItem*)rosterItem
+{
+}
 @end