4 files changed, 65 insertions, 39 deletions
diff --git a/synapse/events/__init__.py b/synapse/events/__init__.py
index 4252e5ab5c..bf07951027 100644
--- a/synapse/events/__init__.py
+++ b/synapse/events/__init__.py
@@ -18,7 +18,7 @@ from synapse.util.frozenutils import freeze, unfreeze
 
 class _EventInternalMetadata(object):
     def __init__(self, internal_metadata_dict):
-        self.__dict__ = internal_metadata_dict
+        self.__dict__ = dict(internal_metadata_dict)
 
     def get_dict(self):
         return dict(self.__dict__)
diff --git a/synapse/events/builder.py b/synapse/events/builder.py
index a9b1b99a10..9d45bdb892 100644
--- a/synapse/events/builder.py
+++ b/synapse/events/builder.py
@@ -23,14 +23,15 @@ import copy
 
 
 class EventBuilder(EventBase):
-    def __init__(self, key_values={}):
+    def __init__(self, key_values={}, internal_metadata_dict={}):
         signatures = copy.deepcopy(key_values.pop("signatures", {}))
         unsigned = copy.deepcopy(key_values.pop("unsigned", {}))
 
         super(EventBuilder, self).__init__(
             key_values,
             signatures=signatures,
-            unsigned=unsigned
+            unsigned=unsigned,
+            internal_metadata_dict=internal_metadata_dict,
         )
 
     def build(self):
diff --git a/synapse/events/snapshot.py b/synapse/events/snapshot.py
index 6bbba8d6ba..7e98bdef28 100644
--- a/synapse/events/snapshot.py
+++ b/synapse/events/snapshot.py
@@ -20,3 +20,4 @@ class EventContext(object):
         self.current_state = current_state
         self.auth_events = auth_events
         self.state_group = None
+        self.rejected = False
diff --git a/synapse/events/utils.py b/synapse/events/utils.py
index e391aca4cc..1aa952150e 100644
--- a/synapse/events/utils.py
+++ b/synapse/events/utils.py
@@ -45,12 +45,14 @@ def prune_event(event):
         "membership",
     ]
 
+    event_dict = event.get_dict()
+
     new_content = {}
 
     def add_fields(*fields):
         for field in fields:
             if field in event.content:
-                new_content[field] = event.content[field]
+                new_content[field] = event_dict["content"][field]
 
     if event_type == EventTypes.Member:
         add_fields("membership")
@@ -75,7 +77,7 @@ def prune_event(event):
 
     allowed_fields = {
         k: v
-        for k, v in event.get_dict().items()
+        for k, v in event_dict.items()
         if k in allowed_keys
     }
 
@@ -86,10 +88,53 @@ def prune_event(event):
     if "age_ts" in event.unsigned:
         allowed_fields["unsigned"]["age_ts"] = event.unsigned["age_ts"]
 
-    return type(event)(allowed_fields)
+    return type(event)(
+        allowed_fields,
+        internal_metadata_dict=event.internal_metadata.get_dict()
+    )
+
+
+def format_event_raw(d):
+    return d
 
 
-def serialize_event(e, time_now_ms, client_event=True):
+def format_event_for_client_v1(d):
+    d["user_id"] = d.pop("sender", None)
+
+    move_keys = ("age", "redacted_because", "replaces_state", "prev_content")
+    for key in move_keys:
+        if key in d["unsigned"]:
+            d[key] = d["unsigned"][key]
+
+    drop_keys = (
+        "auth_events", "prev_events", "hashes", "signatures", "depth",
+        "unsigned", "origin", "prev_state"
+    )
+    for key in drop_keys:
+        d.pop(key, None)
+    return d
+
+
+def format_event_for_client_v2(d):
+    drop_keys = (
+        "auth_events", "prev_events", "hashes", "signatures", "depth",
+        "origin", "prev_state",
+    )
+    for key in drop_keys:
+        d.pop(key, None)
+    return d
+
+
+def format_event_for_client_v2_without_event_id(d):
+    d = format_event_for_client_v2(d)
+    d.pop("room_id", None)
+    d.pop("event_id", None)
+    return d
+
+
+def serialize_event(e, time_now_ms, as_client_event=True,
+                    event_format=format_event_for_client_v1,
+                    token_id=None):
     # FIXME(erikj): To handle the case of presence events and the like
     if not isinstance(e, EventBase):
         return e
@@ -99,43 +144,22 @@ def serialize_event(e, time_now_ms, client_event=True):
     # Should this strip out None's?
     d = {k: v for k, v in e.get_dict().items()}
 
-    if not client_event:
-        # set the age and keep all other keys
-        if "age_ts" in d["unsigned"]:
-            d["unsigned"]["age"] = time_now_ms - d["unsigned"]["age_ts"]
-        return d
-
     if "age_ts" in d["unsigned"]:
-        d["age"] = time_now_ms - d["unsigned"]["age_ts"]
+        d["unsigned"]["age"] = time_now_ms - d["unsigned"]["age_ts"]
         del d["unsigned"]["age_ts"]
 
-    d["user_id"] = d.pop("sender", None)
-
     if "redacted_because" in e.unsigned:
-        d["redacted_because"] = serialize_event(
+        d["unsigned"]["redacted_because"] = serialize_event(
             e.unsigned["redacted_because"], time_now_ms
         )
 
-        del d["unsigned"]["redacted_because"]
+    if token_id is not None:
+        if token_id == getattr(e.internal_metadata, "token_id", None):
+            txn_id = getattr(e.internal_metadata, "txn_id", None)
+            if txn_id is not None:
+                d["unsigned"]["transaction_id"] = txn_id
 
-    if "redacted_by" in e.unsigned:
-        d["redacted_by"] = e.unsigned["redacted_by"]
-        del d["unsigned"]["redacted_by"]
-
-    if "replaces_state" in e.unsigned:
-        d["replaces_state"] = e.unsigned["replaces_state"]
-        del d["unsigned"]["replaces_state"]
-
-    if "prev_content" in e.unsigned:
-        d["prev_content"] = e.unsigned["prev_content"]
-        del d["unsigned"]["prev_content"]
-
-    del d["auth_events"]
-    del d["prev_events"]
-    del d["hashes"]
-    del d["signatures"]
-    d.pop("depth", None)
-    d.pop("unsigned", None)
-    d.pop("origin", None)
-
-    return d
+    if as_client_event:
+        return event_format(d)
+    else:
+        return d