Fix redactions. Fix 'age' key

6 files changed, 83 insertions, 30 deletions

diff --git a/synapse/events/__init__.py b/synapse/events/__init__.py
index 7103b937af..98d7f0e324 100644
--- a/synapse/events/__init__.py
+++ b/synapse/events/__init__.py
@@ -97,6 +97,7 @@ class EventBase(object):
     origin_server_ts = _event_dict_property("origin_server_ts")
     prev_events = _event_dict_property("prev_events")
     prev_state = _event_dict_property("prev_state")
+    redacts = _event_dict_property("redacts")
     room_id = _event_dict_property("room_id")
     sender = _event_dict_property("sender")
     state_key = _event_dict_property("state_key")
diff --git a/synapse/events/builder.py b/synapse/events/builder.py
index 642264e9f3..9579b1fe8b 100644
--- a/synapse/events/builder.py
+++ b/synapse/events/builder.py
@@ -19,11 +19,18 @@ from synapse.types import EventID
 
 from synapse.util.stringutils import random_string
 
+import copy
+
 
 class EventBuilder(EventBase):
     def __init__(self, key_values={}):
+        signatures = copy.deepcopy(key_values.pop("signatures", {}))
+        unsigned = copy.deepcopy(key_values.pop("unsigned", {}))
+
         super(EventBuilder, self).__init__(
             key_values,
+            signatures=signatures,
+            unsigned=unsigned
         )
 
     def update_event_key(self, key, value):
@@ -61,9 +68,9 @@ class EventBuilderFactory(object):
         key_values.setdefault("origin", self.hostname)
         key_values.setdefault("origin_server_ts", time_now)
 
-        if "unsigned" in key_values:
-            age = key_values["unsigned"].pop("age", 0)
-            key_values["unsigned"].setdefault("age_ts", time_now - age)
+        key_values.setdefault("unsigned", {})
+        age = key_values["unsigned"].pop("age", 0)
+        key_values["unsigned"].setdefault("age_ts", time_now - age)
 
         key_values["signatures"] = {}
 
diff --git a/synapse/events/utils.py b/synapse/events/utils.py
index f5e135e3d0..6d9c9352e2 100644
--- a/synapse/events/utils.py
+++ b/synapse/events/utils.py
@@ -80,6 +80,11 @@ def prune_event(event):
 
     allowed_fields["content"] = new_content
 
+    allowed_fields["unsigned"] = {}
+
+    if "age_ts" in event.unsigned:
+        allowed_fields["unsigned"]["age_ts"] = event.unsigned["age_ts"]
+
     return type(event)(allowed_fields)
 
 
@@ -97,4 +102,20 @@ def serialize_event(hs, e):
 
     d["user_id"] = d.pop("sender", None)
 
+    if "redacted_because" in e.unsigned:
+        d["redacted_because"] = serialize_event(
+            hs, e.unsigned["redacted_because"]
+        )
+
+        del d["unsigned"]["redacted_because"]
+
+    if "redacted_by" in e.unsigned:
+        d["redacted_by"] = e.unsigned["redacted_by"]
+        del d["unsigned"]["redacted_by"]
+
+    del d["auth_events"]
+    del d["prev_events"]
+    del d["hashes"]
+    del d["signatures"]
+
     return d
diff --git a/synapse/state.py b/synapse/state.py
index 7fdf596006..5bfa73fb46 100644
--- a/synapse/state.py
+++ b/synapse/state.py
@@ -155,6 +155,12 @@ class StateHandler(object):
             else:
                 context.auth_events = {}
 
+            if event.is_state():
+                key = (event.type, event.state_key)
+                if key in context.current_state:
+                    replaces = context.current_state[key]
+                    event.unsigned["replaces_state"] = replaces.event_id
+
             defer.returnValue([])
 
         if event.is_state():
@@ -177,6 +183,12 @@ class StateHandler(object):
             prev_state
         )
 
+        if event.is_state():
+            key = (event.type, event.state_key)
+            if key in context.current_state:
+                replaces = context.current_state[key]
+                event.unsigned["replaces_state"] = replaces.event_id
+
         if hasattr(event, "auth_events") and event.auth_events:
             auth_ids = zip(*event.auth_events)[0]
             context.auth_events = {
diff --git a/synapse/storage/_base.py b/synapse/storage/_base.py
index 12239fa074..ffc26d4a61 100644
--- a/synapse/storage/_base.py
+++ b/synapse/storage/_base.py
@@ -444,38 +444,50 @@ class SQLBaseStore(object):
     def _get_events_txn(self, txn, event_ids):
         events = []
         for e_id in event_ids:
-            js = self._simple_select_one_onecol_txn(
-                txn,
-                table="event_json",
-                keyvalues={"event_id": e_id},
-                retcol="json",
-                allow_none=True,
-            )
+            ev = self._get_event_txn(txn, e_id)
 
-            if not js:
-                # FIXME (erikj): What should we actually do here?
-                continue
+            if ev:
+                events.append(ev)
 
-            d = json.loads(js)
+        return events
 
-            ev = FrozenEvent(d)
+    def _get_event_txn(self, txn, event_id, check_redacted=True):
+        sql = (
+            "SELECT json, r.event_id FROM event_json as e "
+            "LEFT JOIN redactions as r ON e.event_id = r.redacts "
+            "WHERE e.event_id = ? "
+            "LIMIT 1 "
+        )
 
-            if hasattr(ev, "redacted") and ev.redacted:
-                # Get the redaction event.
-                select_event_sql = "SELECT * FROM events WHERE event_id = ?"
-                txn.execute(select_event_sql, (ev.redacted,))
+        txn.execute(sql, (event_id,))
 
-                del_evs = self._parse_events_txn(
-                    txn, self.cursor_to_dict(txn)
-                )
+        res = txn.fetchone()
 
-                if del_evs:
-                    ev = prune_event(ev)
-                    ev.redacted_because = del_evs[0]
+        if not res:
+            return None
 
-            events.append(ev)
+        js, redacted = res
 
-        return events
+        d = json.loads(js)
+
+        ev = FrozenEvent(d)
+
+        if check_redacted and redacted:
+            ev = prune_event(ev)
+
+            ev.unsigned["redacted_by"] = redacted
+            # Get the redaction event.
+
+            because = self._get_event_txn(
+                txn,
+                redacted,
+                check_redacted=False
+            )
+
+            if because:
+                ev.unsigned["redacted_because"] = because
+
+        return ev
 
     def _parse_events(self, rows):
         return self.runInteraction(
diff --git a/tests/storage/test_redaction.py b/tests/storage/test_redaction.py
index e8671ae3a2..d81f7add1c 100644
--- a/tests/storage/test_redaction.py
+++ b/tests/storage/test_redaction.py
@@ -149,7 +149,7 @@ class RedactionTestCase(unittest.TestCase):
             event,
         )
 
-        self.assertFalse(hasattr(event, "redacted_because"))
+        self.assertFalse("redacted_because" in event.unsigned)
 
         # Redact event
         reason = "Because I said so"
@@ -179,7 +179,7 @@ class RedactionTestCase(unittest.TestCase):
             event,
         )
 
-        self.assertTrue(hasattr(event, "redacted_because"))
+        self.assertTrue("redacted_because" in event.unsigned)
 
         self.assertObjectHasAttributes(
             {
@@ -187,7 +187,7 @@ class RedactionTestCase(unittest.TestCase):
                 "user_id": self.u_alice.to_string(),
                 "content": {"reason": reason},
             },
-            event.redacted_because,
+            event.unsigned["redacted_because"],
         )
 
     @defer.inlineCallbacks