--- a/wokkel/mam.py	Sun Jan 03 18:36:41 2016 +0100
+++ b/wokkel/mam.py	Sun Jan 03 18:36:41 2016 +0100
@@ -2,6 +2,7 @@
 # -*- test-case-name: wokkel.test.test_mam -*-
 #
 # SàT Wokkel extension for Message Archive Management (XEP-0313)
+# Copyright (C) 2015 Jérôme Poisson (goffi@goffi.org)
 # Copyright (C) 2015 Adien Cossa (souliane@mailoo.org)
 
 # This program is free software: you can redistribute it and/or modify
@@ -16,9 +17,6 @@
 
 # You should have received a copy of the GNU Affero General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
-#
-# Copyright (c) Adrien Cossa.
-# See LICENSE for details.
 
 """
 XMPP Message Archive Management protocol.
@@ -34,6 +32,7 @@
 from twisted.words.protocols.jabber.xmlstream import IQ, toResponse
 from twisted.words.xish import domish
 from twisted.words.protocols.jabber import jid
+from twisted.internet import defer
 from twisted.python import log
 
 from wokkel.subprotocols import IQHandlerMixin, XMPPHandler
@@ -41,7 +40,7 @@
 
 import rsm
 
-NS_MAM = 'urn:xmpp:mam:0'
+NS_MAM = 'urn:xmpp:mam:1'
 NS_FORWARD = 'urn:xmpp:forward:0'
 
 FIELDS_REQUEST = "/iq[@type='get']/query[@xmlns='%s']" % NS_MAM
@@ -84,14 +83,19 @@
 
     @ivar node: pubsub node id if querying a pubsub node, else None.
     @itype form: C{unicode}
+
+    @ivar query_id: id to use to track the query
+    @itype form: C{unicode}
     """
 
-    def __init__(self, form=None, rsm_=None, node=None):
+    def __init__(self, form=None, rsm_=None, node=None, query_id=None):
         if form is not None:
             assert form.formType == 'submit'
+            assert form.formNamespace == NS_MAM
         self.form = form
         self.rsm = rsm_
         self.node = node
+        self.query_id = query_id
 
     @classmethod
     def parse(cls, element):
@@ -110,8 +114,9 @@
             rsm_request = rsm.RSMRequest.parse(element)
         except rsm.RSMNotFoundError:
             rsm_request = None
-        node = element['node'] if element.hasAttribute('node') else None
-        return MAMQueryRequest(form, rsm_request, node)
+        node = element.getAttribute('node')
+        query_id = element.getAttribute('queryid')
+        return MAMQueryRequest(form, rsm_request, node, query_id)
 
     def toElement(self):
         """
@@ -122,6 +127,8 @@
         mam_elt = domish.Element((NS_MAM, 'query'))
         if self.node is not None:
             mam_elt['node'] = self.node
+        if self.query_id is not None:
+            mam_elt['queryid'] = self.query_id
         if self.form is not None:
             mam_elt.addChild(self.form.toElement())
         if self.rsm is not None:
@@ -158,9 +165,8 @@
     @type never: C{list}
     """
 
-    def __init__(self, default=None, always=None, never=None):
-        if default is not None:
-            assert default in ('always', 'never', 'roster')
+    def __init__(self, default, always=None, never=None):
+        assert default in ('always', 'never', 'roster')
         self.default = default
         if always is not None:
             assert isinstance(always, list)
@@ -174,27 +180,34 @@
         self.never = never
 
     @classmethod
-    def parse(cls, element):
+    def parse(cls, prefs_elt):
         """Parse the DOM representation of a MAM <prefs/> request.
 
-        @param element: MAM <prefs/> request element.
-        @type element: L{Element<twisted.words.xish.domish.Element>}
+        @param prefs_elt: MAM <prefs/> request element.
+        @type prefs_elt: L{Element<twisted.words.xish.domish.Element>}
 
         @return: MAMPrefs instance.
         @rtype: L{MAMPrefs}
         """
-        if element.uri != NS_MAM or element.name != 'prefs':
+        if prefs_elt.uri != NS_MAM or prefs_elt.name != 'prefs':
             raise MAMError('Element provided is not a MAM <prefs/> request')
-        default = element['default'] if element.hasAttribute('default') else None
+        try:
+            default = prefs_elt['default']
+        except KeyError:
+            # FIXME: return proper error here
+            raise MAMError('Element provided is not a valid MAM <prefs/> request')
+
         prefs = {}
         for attr in ('always', 'never'):
             prefs[attr] = []
             try:
-                pref = domish.generateElementsNamed(element.elements(), attr).next()
-                for jid_s in domish.generateElementsNamed(pref.elements(), 'jid'):
+                pref = prefs_elt.elements(NS_MAM, attr).next()
+            except StopIteration:
+                # FIXME: return proper error here
+                raise MAMError('Element provided is not a valid MAM <prefs/> request')
+            else:
+                for jid_s in pref.elements(NS_MAM, 'jid'):
                     prefs[attr].append(jid.JID(jid_s))
-            except StopIteration:
-                pass
         return MAMPrefs(default, **prefs)
 
     def toElement(self):
@@ -235,7 +248,7 @@
     This handler implements the protocol for sending out MAM requests.
     """
 
-    def queryArchive(self, service=None, form=None, rsm=None, node=None, sender=None):
+    def queryArchive(self, service=None, form=None, rsm=None, node=None, sender=None, query_id=None):
         """Query a user, MUC or pubsub archive.
 
         @param service: Entity offering the MAM service (None for user archives).
@@ -253,11 +266,14 @@
         @param sender: Optional sender address.
         @type sender: L{JID<twisted.words.protocols.jabber.jid.JID>}
 
+        @param query_id: Optional query id
+        @type query_id: C{unicode}
+
         @return: A deferred that fires upon receiving a response.
         @rtype: L{Deferred<twisted.internet.defer.Deferred>}
         """
         iq = IQ(self.xmlstream, 'set')
-        MAMQueryRequest(form, rsm, node).render(iq)
+        MAMQueryRequest(form, rsm, node, query_id).render(iq)
         if sender is not None:
             iq['from'] = unicode(sender)
         return iq.send(to=service.full() if service else None)
@@ -271,7 +287,7 @@
         @param sender: Optional sender address.
         @type sender: L{JID<twisted.words.protocols.jabber.jid.JID>}
 
-        @return: A deferred that fires upon receiving a response.
+        @return: data Form with the fields, or None if not found
         @rtype: L{Deferred<twisted.internet.defer.Deferred>}
         """
         # http://xmpp.org/extensions/xep-0313.html#query-form
@@ -279,7 +295,10 @@
         MAMQueryRequest().render(iq)
         if sender is not None:
             iq['from'] = unicode(sender)
-        return iq.send(to=service.full() if service else None)
+        d = iq.send(to=service.full() if service else None)
+        d.addCallback(lambda iq_result: iq_result.elements(NS_MAM, 'query').next())
+        d.addCallback(data_form.findForm, NS_MAM)
+        return d
 
     def queryPrefs(self, service=None, sender=None):
         """Retrieve the current user preferences.
@@ -407,7 +426,7 @@
         @type resource: L{object}
         """
         self.resource = resource
-        self.extra_filters = {}
+        self.extra_fields = {}
 
     def connectionInitialized(self):
         """
@@ -424,11 +443,10 @@
         """
         Add a new filter for querying MAM archive.
 
-        @param field: dictionary specifying the attributes to build a
-            wokkel.data_form.Field.
-        @type field: C{dict}
+        @param field: data form field of the filter
+        @type field: L{Form<wokkel.data_form.Field>}
         """
-        self.extra_filters[field.var] = field
+        self.extra_fields[field.var] = field
 
     def _onFieldsRequest(self, iq):
         """
@@ -438,7 +456,7 @@
         """
         response = toResponse(iq, 'result')
         query = response.addElement((NS_MAM, 'query'))
-        query.addChild(buildForm('form', extra=self.extra_filters).toElement())
+        query.addChild(buildForm('form', extra=self.extra_fields).toElement())
         self.xmlstream.send(response)
         iq.handled = True
 
@@ -446,12 +464,10 @@
         """
         Called when a message archive request has been received.
 
-        This immediately replies with a result response, followed by the
-        list of archived message and the finally the <fin/> message.
+        This replies with the list of archived message and the <iq> result
+        @return: A tuple with list of message data (id, element, data) and RSM element
+        @rtype: C{tuple}
         """
-        response = toResponse(iq, 'result')
-        self.xmlstream.send(response)
-
         mam_ = MAMQueryRequest.parse(iq.query)
         requestor = jid.JID(iq['from'])
 
@@ -459,7 +475,7 @@
         unsupported_fields = []
         if mam_.form:
             for key, field in mam_.form.fields.iteritems():
-                if key not in self._legacyFilters and key not in self.extra_filters:
+                if key not in self._legacyFilters and key not in self.extra_fields:
                     log.msg('Ignored unsupported MAM filter: %s' % field)
                     unsupported_fields.append(key)
         for key in unsupported_fields:
@@ -468,11 +484,13 @@
         def forward_message(id_, elt, date):
             msg = domish.Element((None, 'message'))
             msg['to'] = iq['from']
-            result = msg.addElement('result', NS_MAM)
-            if iq.hasAttribute('queryid'):
+            result = msg.addElement((NS_MAM, 'result'))
+            try:
                 result['queryid'] = iq.query['queryid']
+            except KeyError:
+                pass
             result['id'] = id_
-            forward = result.addElement('forwarded', NS_FORWARD)
+            forward = result.addElement((NS_FORWARD, 'forwarded'))
             forward.addChild(delay.Delay(date).toElement())
             forward.addChild(elt)
             self.xmlstream.send(msg)
@@ -481,16 +499,16 @@
             msg_data, rsm_elt = result
             for data in msg_data:
                 forward_message(*data)
-            msg = domish.Element((None, 'message'))
-            msg['to'] = iq['from']
-            fin = msg.addElement('fin', NS_MAM)
-            if iq.hasAttribute('queryid'):
-                fin['queryid'] = iq.query['queryid']
+
+            response = toResponse(iq, 'result')
+            fin = response.addElement((NS_MAM, 'fin'))
+
             if rsm_elt is not None:
                 fin.addChild(rsm_elt)
-            self.xmlstream.send(msg)
+            self.xmlstream.send(response)
 
-        self.resource.onArchiveRequest(mam_, requestor).addCallback(cb)
+        d = defer.maybeDeferred(self.resource.onArchiveRequest, mam_, requestor)
+        d.addCallback(cb)
         iq.handled = True
 
     def _onPrefsGetRequest(self, iq):
@@ -550,7 +568,7 @@
     return datetime.astimezone(tzutc()).strftime(stampFormat)
 
 
-def buildForm(formType='submit', start=None, end=None, with_jid=None, extra=None):
+def buildForm(formType='submit', start=None, end=None, with_jid=None, extra_fields=None):
     """Prepare a Data Form for MAM.
 
     @param formType: The type of the Data Form ('submit' or 'form').
@@ -565,32 +583,28 @@
     @param with_jid: JID against which to match messages.
     @type with_jid: L{JID<twisted.words.protocols.jabber.jid.JID>}
 
-    @param extra: list of extra fields that are not defined by the
-        specification. Each element of the list must be a dictionary
-        specifying the attributes to build a wokkel.data_form.Field.
+    @param extra_fields: list of extra data form fields that are not defined by the
+        specification.
     @type: C{list}
 
     @return: XEP-0004 Data Form object.
     @rtype: L{Form<wokkel.data_form.Form>}
     """
     form = data_form.Form(formType, formNamespace=NS_MAM)
-    filters = []
 
     if formType == 'form':
-        filters.extend(MAMService._legacyFilters.values())
+        for kwargs in MAMService._legacyFilters.values():
+            form.addField(data_form.Field(**kwargs))
     elif formType == 'submit':
         if start:
-            filters.append({'var': 'start', 'value': datetime2utc(start)})
+            form.addField(data_form.Field(var='start', value=datetime2utc(start)))
         if end:
-            filters.append({'var': 'end', 'value': datetime2utc(end)})
+            form.addField(data_form.Field(var='end', value=datetime2utc(end)))
         if with_jid:
-            # must specify the field type to overwrite default value in Field.__init__
-            filters.append({'fieldType': 'jid-single', 'var': 'with', 'value': with_jid.full()})
+            form.addField(data_form.Field(fieldType='jid-single', var='with', value=with_jid.full()))
 
-    if extra is not None:
-        filters.extend(extra)
-
-    for field in filters:
-        form.addField(data_form.Field(**field))
+    if extra_fields is not None:
+        for field in extra_fields:
+            form.addField(field)
 
     return form