persist hashes and origin signatures for PDUs

3 files changed, 31 insertions, 5 deletions

diff --git a/synapse/storage/__init__.py b/synapse/storage/__init__.py
index 6dadeb8cce..bfeab7d1e8 100644
--- a/synapse/storage/__init__.py
+++ b/synapse/storage/__init__.py
@@ -40,6 +40,8 @@ from .stream import StreamStore
 from .pdu import StatePduStore, PduStore, PdusTable
 from .transactions import TransactionStore
 from .keys import KeyStore
+from .signatures import SignatureStore
+
 
 import json
 import logging
@@ -59,6 +61,7 @@ SCHEMAS = [
     "room_aliases",
     "keys",
     "redactions",
+    "signatures",
 ]
 
 
@@ -76,7 +79,7 @@ class _RollbackButIsFineException(Exception):
 class DataStore(RoomMemberStore, RoomStore,
                 RegistrationStore, StreamStore, ProfileStore, FeedbackStore,
                 PresenceStore, PduStore, StatePduStore, TransactionStore,
-                DirectoryStore, KeyStore):
+                DirectoryStore, KeyStore, SignatureStore):
 
     def __init__(self, hs):
         super(DataStore, self).__init__(hs)
@@ -144,6 +147,8 @@ class DataStore(RoomMemberStore, RoomStore,
     def _persist_event_pdu_txn(self, txn, pdu):
         cols = dict(pdu.__dict__)
         unrec_keys = dict(pdu.unrecognized_keys)
+        del cols["hashes"]
+        del cols["signatures"]
         del cols["content"]
         del cols["prev_pdus"]
         cols["content_json"] = json.dumps(pdu.content)
@@ -157,6 +162,20 @@ class DataStore(RoomMemberStore, RoomStore,
 
         logger.debug("Persisting: %s", repr(cols))
 
+        for hash_alg, hash_base64 in pdu.hashes.items():
+            hash_bytes = decode_base64(hash_base64)
+            self._store_pdu_hash_txn(
+                txn, pdu.pdu_id, pdu.origin, hash_alg, hash_bytes,
+            )
+
+        signatures = pdu.sigatures.get(pdu.orgin, {})
+
+        for key_id, signature_base64 in signatures:
+            signature_bytes = decode_base64(signature_base64)
+            self.store_pdu_origin_signatures_txn(
+                txn, pdu.pdu_id, pdu.origin, key_id, signature_bytes,
+            )
+
         if pdu.is_state:
             self._persist_state_txn(txn, pdu.prev_pdus, cols)
         else:
diff --git a/synapse/storage/pdu.py b/synapse/storage/pdu.py
index d70467dcd6..9d624429b7 100644
--- a/synapse/storage/pdu.py
+++ b/synapse/storage/pdu.py
@@ -64,6 +64,11 @@ class PduStore(SQLBaseStore):
                 for r in PduEdgesTable.decode_results(txn.fetchall())
             ]
 
+            hashes = self._get_pdu_hashes_txn(txn, pdu_id, origin)
+            signatures = self._get_pdu_origin_signatures_txn(
+                txn, pdu_id, origin
+            )
+
             query = (
                 "SELECT %(fields)s FROM %(pdus)s as p "
                 "LEFT JOIN %(state)s as s "
@@ -80,7 +85,9 @@ class PduStore(SQLBaseStore):
 
             row = txn.fetchone()
             if row:
-                results.append(PduTuple(PduEntry(*row), edges))
+                results.append(PduTuple(
+                    PduEntry(*row), edges, hashes, signatures
+                ))
 
         return results
 
@@ -908,7 +915,7 @@ This does not include a prev_pdus key.
 
 PduTuple = namedtuple(
     "PduTuple",
-    ("pdu_entry", "prev_pdu_list")
+    ("pdu_entry", "prev_pdu_list", "hashes", "signatures")
 )
 """ This is a tuple of a `PduEntry` and a list of `PduIdTuple` that represent
 the `prev_pdus` key of a PDU.
diff --git a/synapse/storage/schema/signatures.sql b/synapse/storage/schema/signatures.sql
index ba3bbb5471..86ee0f2377 100644
--- a/synapse/storage/schema/signatures.sql
+++ b/synapse/storage/schema/signatures.sql
@@ -28,9 +28,9 @@ CREATE TABLE IF NOT EXISTS pdu_origin_signatures (
   origin TEXT,
   key_id TEXT,
   signature BLOB,
-  CONSTRAINT uniqueness UNIQUE (pdu_id, origin, algorithm)
+  CONSTRAINT uniqueness UNIQUE (pdu_id, origin, key_id)
 );
 
 CREATE INDEX IF NOT EXISTS pdu_origin_signatures_id ON pdu_origin_signatures (
-    pdu_id, origin,
+    pdu_id, origin
 );