--- a/sat_pubsub/delegation.py	Mon May 04 18:40:47 2015 +0200
+++ b/sat_pubsub/delegation.py	Sun Aug 16 01:13:23 2015 +0200
@@ -27,6 +27,7 @@
 from wokkel import pubsub
 from wokkel import data_form
 from wokkel import disco, iwokkel
+from wokkel.iwokkel import IPubSubService
 from twisted.python import log
 from twisted.words.protocols.jabber import jid, error
 from twisted.words.protocols.jabber.xmlstream import toResponse
@@ -44,13 +45,56 @@
 class InvalidStanza(Exception):
     pass
 
+
+
 class DelegationsHandler(XMPPHandler):
     implements(iwokkel.IDisco)
+    _service_hacked = False
 
     def __init__(self):
         super(DelegationsHandler, self).__init__()
 
+    def _service_hack(self):
+        """Patch the PubSubService to track delegated stanzas"""
+        # XXX: we need to monkey patch to track origin of the stanza in PubSubRequest.
+        #      As PubSubRequest from sat.tmp.wokkel.pubsub use _request_class while
+        #      original wokkel.pubsub use directly pubsub.PubSubRequest, we need to
+        #      check which version is used before monkeypatching
+        for handler in self.parent.handlers:
+            if IPubSubService.providedBy(handler):
+                if hasattr(handler, '_request_class'):
+                    request_base_class = handler._request_class
+                else:
+                    request_base_class = pubsub.PubSubRequest
+
+                class PubSubRequestWithDelegation(request_base_class):
+                    """A PubSubReques which put an indicator if the stanza comme from delegation"""
+
+                    @classmethod
+                    def fromElement(cls, element):
+                        """Check if element comme from delegation, and set a delegated flags
+
+                        delegated flaf is either False, or it's a jid of the delegating server
+                        the delegated flag must be set on element before use
+                        """
+                        try:
+                            # __getattr__ is overriden in domish.Element, so we use __getattribute__
+                            delegated = element.__getattribute__('delegated')
+                        except AttributeError:
+                            delegated = False
+                        instance = cls.__base__.fromElement(element)
+                        instance.delegated = delegated
+                        return instance
+
+                if hasattr(handler, '_request_class'):
+                    handler._request_class = PubSubRequestWithDelegation
+                else:
+                    pubsub.PubSubRequest = PubSubRequestWithDelegation
+        DelegationsHandler._service_hacked = True
+
     def connectionInitialized(self):
+        if not self._service_hacked:
+            self._service_hack()
         self.xmlstream.addObserver(DELEGATION_ADV_XPATH, self.onAdvertise)
         self.xmlstream.addObserver(DELEGATION_FWD_XPATH, self._obsWrapper, 0, self.onForward)
         self._current_iqs = {} # dict of iq being handler by delegation
@@ -64,17 +108,21 @@
         """
         if isinstance(elt, domish.Element) and elt.name=='iq':
             try:
-                ori_iq, managed_entity = self._current_iqs.pop(elt.getAttribute('id'))
+                id_ = elt.getAttribute('id')
+                ori_iq, managed_entity = self._current_iqs[id_]
                 if jid.JID(elt['to']) != managed_entity:
-                    log.msg("IQ id conflict: the managed entity doesn't match")
+                    log.msg("IQ id conflict: the managed entity doesn't match (got {got} was expecting {expected})"
+                            .format(got=jid.JID(elt['to']), expected=managed_entity))
                     raise KeyError
             except KeyError:
                 # the iq is not a delegated one
                 self._xs_send(elt)
             else:
+                del self._current_iqs[id_]
                 iq_result_elt = toResponse(ori_iq, 'result')
                 fwd_elt = iq_result_elt.addElement('delegation', DELEGATION_NS).addElement('forwarded', FORWARDED_NS)
                 fwd_elt.addChild(elt)
+                elt.uri = elt.defaultUri = 'jabber:client'
                 self._xs_send(iq_result_elt)
         else:
             self._xs_send(elt)
@@ -132,7 +180,7 @@
         except StopIteration:
             raise error.StanzaError('not-acceptable')
 
-        managed_entity = jid.JID(fwd_iq.getAttribute('to') or fwd_iq['from'])
+        managed_entity = jid.JID(fwd_iq['from'])
 
         if managed_entity.host != iq['from']:
             log.msg((u"SECURITY WARNING: forwarded stanza doesn't come from the emitting server: {}"
@@ -140,6 +188,13 @@
             raise error.StanzaError('not-allowed')
 
         self._current_iqs[fwd_iq['id']] = (iq, managed_entity)
+        fwd_iq.delegated = True
+
+        # we need a recipient in pubsub request for PEP
+        # so we set "to" attribute if it doesn't exist
+        if not fwd_iq.hasAttribute('to'):
+            fwd_iq["to"] = jid.JID(fwd_iq["from"]).userhost()
+
         # we now inject the element in the stream
         self.xmlstream.dispatch(fwd_iq)