Add key set and key set ID classes for static keys.

2026-02-16 04:39:54 +01:00 · 2019-04-16 16:56:50 +01:00
parent 7acbe56197
commit 0f5f440f1c
23 changed files with 313 additions and 184 deletions
--- a/bramble-core/src/main/java/org/briarproject/bramble/db/Database.java
+++ b/bramble-core/src/main/java/org/briarproject/bramble/db/Database.java
@@ -5,6 +5,7 @@ import org.briarproject.bramble.api.contact.ContactId;
 import org.briarproject.bramble.api.crypto.SecretKey;
 import org.briarproject.bramble.api.db.DataTooNewException;
 import org.briarproject.bramble.api.db.DataTooOldException;
+import org.briarproject.bramble.api.db.DatabaseComponent;
 import org.briarproject.bramble.api.db.DbException;
 import org.briarproject.bramble.api.db.MessageDeletedException;
 import org.briarproject.bramble.api.db.Metadata;
@@ -23,8 +24,8 @@ import org.briarproject.bramble.api.sync.Message;
 import org.briarproject.bramble.api.sync.MessageId;
 import org.briarproject.bramble.api.sync.MessageStatus;
 import org.briarproject.bramble.api.sync.validation.MessageState;
-import org.briarproject.bramble.api.transport.KeySet;
-import org.briarproject.bramble.api.transport.KeySetId;
+import org.briarproject.bramble.api.transport.TransportKeySet;
+import org.briarproject.bramble.api.transport.TransportKeySetId;
 import org.briarproject.bramble.api.transport.TransportKeys;

 import java.util.Collection;
@@ -33,11 +34,14 @@ import java.util.Map;
 import javax.annotation.Nullable;

 /**
- * A low-level interface to the database (DatabaseComponent provides a
- * high-level interface). Most operations take a transaction argument, which is
- * obtained by calling {@link #startTransaction()}. Every transaction must be
- * terminated by calling either {@link #abortTransaction(Object) abortTransaction(T)} or
- * {@link #commitTransaction(Object) commitTransaction(T)}, even if an exception is thrown.
+ * A low-level interface to the database ({@link DatabaseComponent} provides a
+ * high-level interface).
+ * <p>
+ * Most operations take a transaction argument, which is obtained by calling
+ * {@link #startTransaction()}. Every transaction must be terminated by calling
+ * either {@link #abortTransaction(Object) abortTransaction(T)} or
+ * {@link #commitTransaction(Object) commitTransaction(T)}, even if an
+ * exception is thrown.
 */
@NotNullByDefault
 interface Database<T> {
@@ -131,7 +135,7 @@ interface Database<T> {
 	 * Stores the given transport keys for the given contact and returns a
 	 * key set ID.
 	 */
-	KeySetId addTransportKeys(T txn, ContactId c, TransportKeys k)
+	TransportKeySetId addTransportKeys(T txn, ContactId c, TransportKeys k)
 			throws DbException;

 	/**
@@ -489,13 +493,13 @@ interface Database<T> {
 	 * <p/>
 	 * Read-only.
 	 */
-	Collection<KeySet> getTransportKeys(T txn, TransportId t)
+	Collection<TransportKeySet> getTransportKeys(T txn, TransportId t)
 			throws DbException;

 	/**
 	 * Increments the outgoing stream counter for the given transport keys.
 	 */
-	void incrementStreamCounter(T txn, TransportId t, KeySetId k)
+	void incrementStreamCounter(T txn, TransportId t, TransportKeySetId k)
 			throws DbException;

 	/**
@@ -589,7 +593,7 @@ interface Database<T> {
 	/**
 	 * Removes the given transport keys from the database.
 	 */
-	void removeTransportKeys(T txn, TransportId t, KeySetId k)
+	void removeTransportKeys(T txn, TransportId t, TransportKeySetId k)
 			throws DbException;

 	/**
@@ -637,13 +641,13 @@ interface Database<T> {
 	 * Sets the reordering window for the given key set and transport in the
 	 * given time period.
 	 */
-	void setReorderingWindow(T txn, KeySetId k, TransportId t,
+	void setReorderingWindow(T txn, TransportKeySetId k, TransportId t,
 			long timePeriod, long base, byte[] bitmap) throws DbException;

 	/**
 	 * Marks the given transport keys as usable for outgoing streams.
 	 */
-	void setTransportKeysActive(T txn, TransportId t, KeySetId k)
+	void setTransportKeysActive(T txn, TransportId t, TransportKeySetId k)
 			throws DbException;

 	/**
@@ -657,5 +661,5 @@ interface Database<T> {
 	/**
 	 * Updates the given transport keys following key rotation.
 	 */
-	void updateTransportKeys(T txn, KeySet ks) throws DbException;
+	void updateTransportKeys(T txn, TransportKeySet ks) throws DbException;
 }
--- a/bramble-core/src/main/java/org/briarproject/bramble/db/DatabaseComponentImpl.java
+++ b/bramble-core/src/main/java/org/briarproject/bramble/db/DatabaseComponentImpl.java
@@ -59,8 +59,8 @@ import org.briarproject.bramble.api.sync.event.MessageToRequestEvent;
 import org.briarproject.bramble.api.sync.event.MessagesAckedEvent;
 import org.briarproject.bramble.api.sync.event.MessagesSentEvent;
 import org.briarproject.bramble.api.sync.validation.MessageState;
-import org.briarproject.bramble.api.transport.KeySet;
-import org.briarproject.bramble.api.transport.KeySetId;
+import org.briarproject.bramble.api.transport.TransportKeySet;
+import org.briarproject.bramble.api.transport.TransportKeySetId;
 import org.briarproject.bramble.api.transport.TransportKeys;

 import java.util.ArrayList;
@@ -291,8 +291,8 @@ class DatabaseComponentImpl<T> implements DatabaseComponent {
 	}

 	@Override
-	public KeySetId addTransportKeys(Transaction transaction, ContactId c,
-			TransportKeys k) throws DbException {
+	public TransportKeySetId addTransportKeys(Transaction transaction,
+			ContactId c, TransportKeys k) throws DbException {
 		if (transaction.isReadOnly()) throw new IllegalArgumentException();
 		T txn = unbox(transaction);
 		if (!db.containsContact(txn, c))
@@ -663,7 +663,7 @@ class DatabaseComponentImpl<T> implements DatabaseComponent {
 	}

 	@Override
-	public Collection<KeySet> getTransportKeys(Transaction transaction,
+	public Collection<TransportKeySet> getTransportKeys(Transaction transaction,
 			TransportId t) throws DbException {
 		T txn = unbox(transaction);
 		if (!db.containsTransport(txn, t))
@@ -673,7 +673,7 @@ class DatabaseComponentImpl<T> implements DatabaseComponent {

 	@Override
 	public void incrementStreamCounter(Transaction transaction, TransportId t,
-			KeySetId k) throws DbException {
+			TransportKeySetId k) throws DbException {
 		if (transaction.isReadOnly()) throw new IllegalArgumentException();
 		T txn = unbox(transaction);
 		if (!db.containsTransport(txn, t))
@@ -856,7 +856,7 @@ class DatabaseComponentImpl<T> implements DatabaseComponent {

 	@Override
 	public void removeTransportKeys(Transaction transaction,
-			TransportId t, KeySetId k) throws DbException {
+			TransportId t, TransportKeySetId k) throws DbException {
 		if (transaction.isReadOnly()) throw new IllegalArgumentException();
 		T txn = unbox(transaction);
 		if (!db.containsTransport(txn, t))
@@ -955,9 +955,9 @@ class DatabaseComponentImpl<T> implements DatabaseComponent {
 	}

 	@Override
-	public void setReorderingWindow(Transaction transaction, KeySetId k,
-			TransportId t, long timePeriod, long base, byte[] bitmap)
-			throws DbException {
+	public void setReorderingWindow(Transaction transaction,
+			TransportKeySetId k, TransportId t, long timePeriod, long base,
+			byte[] bitmap) throws DbException {
 		if (transaction.isReadOnly()) throw new IllegalArgumentException();
 		T txn = unbox(transaction);
 		if (!db.containsTransport(txn, t))
@@ -967,7 +967,7 @@ class DatabaseComponentImpl<T> implements DatabaseComponent {

 	@Override
 	public void setTransportKeysActive(Transaction transaction, TransportId t,
-			KeySetId k) throws DbException {
+			TransportKeySetId k) throws DbException {
 		if (transaction.isReadOnly()) throw new IllegalArgumentException();
 		T txn = unbox(transaction);
 		if (!db.containsTransport(txn, t))
@@ -977,11 +977,11 @@ class DatabaseComponentImpl<T> implements DatabaseComponent {

 	@Override
 	public void updateTransportKeys(Transaction transaction,
-			Collection<KeySet> keys) throws DbException {
+			Collection<TransportKeySet> keys) throws DbException {
 		if (transaction.isReadOnly()) throw new IllegalArgumentException();
 		T txn = unbox(transaction);
-		for (KeySet ks : keys) {
-			TransportId t = ks.getTransportKeys().getTransportId();
+		for (TransportKeySet ks : keys) {
+			TransportId t = ks.getKeys().getTransportId();
 			if (db.containsTransport(txn, t))
 				db.updateTransportKeys(txn, ks);
 		}
--- a/bramble-core/src/main/java/org/briarproject/bramble/db/JdbcDatabase.java
+++ b/bramble-core/src/main/java/org/briarproject/bramble/db/JdbcDatabase.java
@@ -27,9 +27,9 @@ import org.briarproject.bramble.api.sync.MessageStatus;
 import org.briarproject.bramble.api.sync.validation.MessageState;
 import org.briarproject.bramble.api.system.Clock;
 import org.briarproject.bramble.api.transport.IncomingKeys;
-import org.briarproject.bramble.api.transport.KeySet;
-import org.briarproject.bramble.api.transport.KeySetId;
 import org.briarproject.bramble.api.transport.OutgoingKeys;
+import org.briarproject.bramble.api.transport.TransportKeySet;
+import org.briarproject.bramble.api.transport.TransportKeySetId;
 import org.briarproject.bramble.api.transport.TransportKeys;

 import java.sql.Connection;
@@ -958,7 +958,7 @@ abstract class JdbcDatabase implements Database<Connection> {
 	}

 	@Override
-	public KeySetId addTransportKeys(Connection txn, ContactId c,
+	public TransportKeySetId addTransportKeys(Connection txn, ContactId c,
 			TransportKeys k) throws DbException {
 		PreparedStatement ps = null;
 		ResultSet rs = null;
@@ -985,7 +985,7 @@ abstract class JdbcDatabase implements Database<Connection> {
 			ps = txn.prepareStatement(sql);
 			rs = ps.executeQuery();
 			if (!rs.next()) throw new DbStateException();
-			KeySetId keySetId = new KeySetId(rs.getInt(1));
+			TransportKeySetId keySetId = new TransportKeySetId(rs.getInt(1));
 			if (rs.next()) throw new DbStateException();
 			rs.close();
 			ps.close();
@@ -2191,8 +2191,8 @@ abstract class JdbcDatabase implements Database<Connection> {
 	}

 	@Override
-	public Collection<KeySet> getTransportKeys(Connection txn, TransportId t)
-			throws DbException {
+	public Collection<TransportKeySet> getTransportKeys(Connection txn,
+			TransportId t) throws DbException {
 		PreparedStatement ps = null;
 		ResultSet rs = null;
 		try {
@@ -2225,11 +2225,12 @@ abstract class JdbcDatabase implements Database<Connection> {
 			ps = txn.prepareStatement(sql);
 			ps.setString(1, t.getString());
 			rs = ps.executeQuery();
-			Collection<KeySet> keys = new ArrayList<>();
+			Collection<TransportKeySet> keys = new ArrayList<>();
 			for (int i = 0; rs.next(); i++) {
 				// There should be three times as many incoming keys
 				if (inKeys.size() < (i + 1) * 3) throw new DbStateException();
-				KeySetId keySetId = new KeySetId(rs.getInt(1));
+				TransportKeySetId keySetId =
+						new TransportKeySetId(rs.getInt(1));
 				ContactId contactId = new ContactId(rs.getInt(2));
 				long timePeriod = rs.getLong(3);
 				SecretKey tagKey = new SecretKey(rs.getBytes(4));
@@ -2243,7 +2244,8 @@ abstract class JdbcDatabase implements Database<Connection> {
 				IncomingKeys inNext = inKeys.get(i * 3 + 2);
 				TransportKeys transportKeys = new TransportKeys(t, inPrev,
 						inCurr, inNext, outCurr);
-				keys.add(new KeySet(keySetId, contactId, transportKeys));
+				keys.add(new TransportKeySet(keySetId, contactId,
+						transportKeys));
 			}
 			rs.close();
 			ps.close();
@@ -2257,7 +2259,7 @@ abstract class JdbcDatabase implements Database<Connection> {

 	@Override
 	public void incrementStreamCounter(Connection txn, TransportId t,
-			KeySetId k) throws DbException {
+			TransportKeySetId k) throws DbException {
 		PreparedStatement ps = null;
 		try {
 			String sql = "UPDATE outgoingKeys SET stream = stream + 1"
@@ -2730,8 +2732,8 @@ abstract class JdbcDatabase implements Database<Connection> {
 	}

 	@Override
-	public void removeTransportKeys(Connection txn, TransportId t, KeySetId k)
-			throws DbException {
+	public void removeTransportKeys(Connection txn, TransportId t,
+			TransportKeySetId k) throws DbException {
 		PreparedStatement ps = null;
 		try {
 			// Delete any existing outgoing keys - this will also remove any
@@ -2934,8 +2936,9 @@ abstract class JdbcDatabase implements Database<Connection> {
 	}

 	@Override
-	public void setReorderingWindow(Connection txn, KeySetId k, TransportId t,
-			long timePeriod, long base, byte[] bitmap) throws DbException {
+	public void setReorderingWindow(Connection txn, TransportKeySetId k,
+			TransportId t, long timePeriod, long base, byte[] bitmap)
+			throws DbException {
 		PreparedStatement ps = null;
 		try {
 			String sql = "UPDATE incomingKeys SET base = ?, bitmap = ?"
@@ -2958,7 +2961,7 @@ abstract class JdbcDatabase implements Database<Connection> {

 	@Override
 	public void setTransportKeysActive(Connection txn, TransportId t,
-			KeySetId k) throws DbException {
+			TransportKeySetId k) throws DbException {
 		PreparedStatement ps = null;
 		try {
 			String sql = "UPDATE outgoingKeys SET active = true"
@@ -3013,7 +3016,7 @@ abstract class JdbcDatabase implements Database<Connection> {
 	}

 	@Override
-	public void updateTransportKeys(Connection txn, KeySet ks)
+	public void updateTransportKeys(Connection txn, TransportKeySet ks)
 			throws DbException {
 		PreparedStatement ps = null;
 		try {
@@ -3022,7 +3025,7 @@ abstract class JdbcDatabase implements Database<Connection> {
 					+ " tagKey = ?, headerKey = ?, stream = ?"
 					+ " WHERE transportId = ? AND keySetId = ?";
 			ps = txn.prepareStatement(sql);
-			TransportKeys k = ks.getTransportKeys();
+			TransportKeys k = ks.getKeys();
 			OutgoingKeys outCurr = k.getCurrentOutgoingKeys();
 			ps.setLong(1, outCurr.getTimePeriod());
 			ps.setBytes(2, outCurr.getTagKey().getBytes());
--- a/bramble-core/src/main/java/org/briarproject/bramble/transport/KeyManagerImpl.java
+++ b/bramble-core/src/main/java/org/briarproject/bramble/transport/KeyManagerImpl.java
@@ -19,8 +19,8 @@ import org.briarproject.bramble.api.plugin.TransportId;
 import org.briarproject.bramble.api.plugin.duplex.DuplexPluginFactory;
 import org.briarproject.bramble.api.plugin.simplex.SimplexPluginFactory;
 import org.briarproject.bramble.api.transport.KeyManager;
-import org.briarproject.bramble.api.transport.KeySetId;
 import org.briarproject.bramble.api.transport.StreamContext;
+import org.briarproject.bramble.api.transport.TransportKeySetId;

 import java.util.HashMap;
 import java.util.Map;
@@ -95,10 +95,10 @@ class KeyManagerImpl implements KeyManager, Service, EventListener {
 	}

 	@Override
-	public Map<TransportId, KeySetId> addContact(Transaction txn, ContactId c,
-			SecretKey rootKey, long timestamp, boolean alice, boolean active)
-			throws DbException {
-		Map<TransportId, KeySetId> ids = new HashMap<>();
+	public Map<TransportId, TransportKeySetId> addContact(Transaction txn,
+			ContactId c, SecretKey rootKey, long timestamp, boolean alice,
+			boolean active) throws DbException {
+		Map<TransportId, TransportKeySetId> ids = new HashMap<>();
 		for (Entry<TransportId, TransportKeyManager> e : managers.entrySet()) {
 			TransportId t = e.getKey();
 			TransportKeyManager m = e.getValue();
@@ -108,9 +108,9 @@ class KeyManagerImpl implements KeyManager, Service, EventListener {
 	}

 	@Override
-	public void activateKeys(Transaction txn, Map<TransportId, KeySetId> keys)
-			throws DbException {
-		for (Entry<TransportId, KeySetId> e : keys.entrySet()) {
+	public void activateKeys(Transaction txn, Map<TransportId,
+			TransportKeySetId> keys) throws DbException {
+		for (Entry<TransportId, TransportKeySetId> e : keys.entrySet()) {
 			TransportId t = e.getKey();
 			TransportKeyManager m = managers.get(t);
 			if (m == null) {
--- a/bramble-core/src/main/java/org/briarproject/bramble/transport/MutableKeySet.java
+++ b/bramble-core/src/main/java/org/briarproject/bramble/transport/MutableKeySet.java
@@ -1,22 +1,22 @@
 package org.briarproject.bramble.transport;

 import org.briarproject.bramble.api.contact.ContactId;
-import org.briarproject.bramble.api.transport.KeySetId;
+import org.briarproject.bramble.api.transport.TransportKeySetId;

 class MutableKeySet {

-	private final KeySetId keySetId;
+	private final TransportKeySetId keySetId;
 	private final ContactId contactId;
 	private final MutableTransportKeys transportKeys;

-	MutableKeySet(KeySetId keySetId, ContactId contactId,
+	MutableKeySet(TransportKeySetId keySetId, ContactId contactId,
 			MutableTransportKeys transportKeys) {
 		this.keySetId = keySetId;
 		this.contactId = contactId;
 		this.transportKeys = transportKeys;
 	}

-	KeySetId getKeySetId() {
+	TransportKeySetId getKeySetId() {
 		return keySetId;
 	}

--- a/bramble-core/src/main/java/org/briarproject/bramble/transport/TransportKeyManager.java
+++ b/bramble-core/src/main/java/org/briarproject/bramble/transport/TransportKeyManager.java
@@ -5,8 +5,8 @@ import org.briarproject.bramble.api.crypto.SecretKey;
 import org.briarproject.bramble.api.db.DbException;
 import org.briarproject.bramble.api.db.Transaction;
 import org.briarproject.bramble.api.nullsafety.NotNullByDefault;
-import org.briarproject.bramble.api.transport.KeySetId;
 import org.briarproject.bramble.api.transport.StreamContext;
+import org.briarproject.bramble.api.transport.TransportKeySetId;

 import javax.annotation.Nullable;

@@ -15,10 +15,11 @@ interface TransportKeyManager {

 	void start(Transaction txn) throws DbException;

-	KeySetId addContact(Transaction txn, ContactId c, SecretKey rootKey,
-			long timestamp, boolean alice, boolean active) throws DbException;
+	TransportKeySetId addContact(Transaction txn, ContactId c,
+			SecretKey rootKey, long timestamp, boolean alice, boolean active)
+			throws DbException;

-	void activateKeys(Transaction txn, KeySetId k) throws DbException;
+	void activateKeys(Transaction txn, TransportKeySetId k) throws DbException;

 	void removeContact(ContactId c);

--- a/bramble-core/src/main/java/org/briarproject/bramble/transport/TransportKeyManagerImpl.java
+++ b/bramble-core/src/main/java/org/briarproject/bramble/transport/TransportKeyManagerImpl.java
@@ -11,9 +11,9 @@ import org.briarproject.bramble.api.nullsafety.NotNullByDefault;
 import org.briarproject.bramble.api.plugin.TransportId;
 import org.briarproject.bramble.api.system.Clock;
 import org.briarproject.bramble.api.system.Scheduler;
-import org.briarproject.bramble.api.transport.KeySet;
-import org.briarproject.bramble.api.transport.KeySetId;
 import org.briarproject.bramble.api.transport.StreamContext;
+import org.briarproject.bramble.api.transport.TransportKeySet;
+import org.briarproject.bramble.api.transport.TransportKeySetId;
 import org.briarproject.bramble.api.transport.TransportKeys;
 import org.briarproject.bramble.transport.ReorderingWindow.Change;

@@ -56,7 +56,7 @@ class TransportKeyManagerImpl implements TransportKeyManager {
 	private final ReentrantLock lock = new ReentrantLock();

 	// The following are locking: lock
-	private final Map<KeySetId, MutableKeySet> keys = new HashMap<>();
+	private final Map<TransportKeySetId, MutableKeySet> keys = new HashMap<>();
 	private final Map<Bytes, TagContext> inContexts = new HashMap<>();
 	private final Map<ContactId, MutableKeySet> outContexts = new HashMap<>();

@@ -80,7 +80,8 @@ class TransportKeyManagerImpl implements TransportKeyManager {
 		lock.lock();
 		try {
 			// Load the transport keys from the DB
-			Collection<KeySet> loaded = db.getTransportKeys(txn, transportId);
+			Collection<TransportKeySet> loaded =
+					db.getTransportKeys(txn, transportId);
 			// Rotate the keys to the current time period
 			RotationResult rotationResult = rotateKeys(loaded, now);
 			// Initialise mutable state for all contacts
@@ -95,14 +96,16 @@ class TransportKeyManagerImpl implements TransportKeyManager {
 		scheduleKeyRotation(now);
 	}

-	private RotationResult rotateKeys(Collection<KeySet> keys, long now) {
+	private RotationResult rotateKeys(Collection<TransportKeySet> keys,
+			long now) {
 		RotationResult rotationResult = new RotationResult();
 		long timePeriod = now / timePeriodLength;
-		for (KeySet ks : keys) {
-			TransportKeys k = ks.getTransportKeys();
-			TransportKeys k1 =
-					transportCrypto.rotateTransportKeys(k, timePeriod);
-			KeySet ks1 = new KeySet(ks.getKeySetId(), ks.getContactId(), k1);
+		for (TransportKeySet ks : keys) {
+			TransportKeys k = ks.getKeys();
+			TransportKeys k1 = transportCrypto.rotateTransportKeys(k,
+					timePeriod);
+			TransportKeySet ks1 = new TransportKeySet(ks.getKeySetId(),
+					ks.getContactId(), k1);
 			if (k1.getTimePeriod() > k.getTimePeriod())
 				rotationResult.rotated.add(ks1);
 			rotationResult.current.add(ks1);
@@ -111,15 +114,15 @@ class TransportKeyManagerImpl implements TransportKeyManager {
 	}

 	// Locking: lock
-	private void addKeys(Collection<KeySet> keys) {
-		for (KeySet ks : keys) {
+	private void addKeys(Collection<TransportKeySet> keys) {
+		for (TransportKeySet ks : keys) {
 			addKeys(ks.getKeySetId(), ks.getContactId(),
-					new MutableTransportKeys(ks.getTransportKeys()));
+					new MutableTransportKeys(ks.getKeys()));
 		}
 	}

 	// Locking: lock
-	private void addKeys(KeySetId keySetId, ContactId contactId,
+	private void addKeys(TransportKeySetId keySetId, ContactId contactId,
 			MutableTransportKeys m) {
 		MutableKeySet ks = new MutableKeySet(keySetId, contactId, m);
 		keys.put(keySetId, ks);
@@ -130,7 +133,7 @@ class TransportKeyManagerImpl implements TransportKeyManager {
 	}

 	// Locking: lock
-	private void encodeTags(KeySetId keySetId, ContactId contactId,
+	private void encodeTags(TransportKeySetId keySetId, ContactId contactId,
 			MutableIncomingKeys inKeys) {
 		for (long streamNumber : inKeys.getWindow().getUnseen()) {
 			TagContext tagCtx =
@@ -170,8 +173,9 @@ class TransportKeyManagerImpl implements TransportKeyManager {
 	}

 	@Override
-	public KeySetId addContact(Transaction txn, ContactId c, SecretKey rootKey,
-			long timestamp, boolean alice, boolean active) throws DbException {
+	public TransportKeySetId addContact(Transaction txn, ContactId c,
+			SecretKey rootKey, long timestamp, boolean alice, boolean active)
+			throws DbException {
 		lock.lock();
 		try {
 			// Work out what time period the timestamp belongs to
@@ -183,7 +187,7 @@ class TransportKeyManagerImpl implements TransportKeyManager {
 			timePeriod = clock.currentTimeMillis() / timePeriodLength;
 			k = transportCrypto.rotateTransportKeys(k, timePeriod);
 			// Write the keys back to the DB
-			KeySetId keySetId = db.addTransportKeys(txn, c, k);
+			TransportKeySetId keySetId = db.addTransportKeys(txn, c, k);
 			// Initialise mutable state for the contact
 			addKeys(keySetId, c, new MutableTransportKeys(k));
 			return keySetId;
@@ -193,7 +197,8 @@ class TransportKeyManagerImpl implements TransportKeyManager {
 	}

 	@Override
-	public void activateKeys(Transaction txn, KeySetId k) throws DbException {
+	public void activateKeys(Transaction txn, TransportKeySetId k)
+			throws DbException {
 		lock.lock();
 		try {
 			MutableKeySet ks = keys.get(k);
@@ -323,10 +328,10 @@ class TransportKeyManagerImpl implements TransportKeyManager {
 		lock.lock();
 		try {
 			// Rotate the keys to the current time period
-			Collection<KeySet> snapshot = new ArrayList<>(keys.size());
+			Collection<TransportKeySet> snapshot = new ArrayList<>(keys.size());
 			for (MutableKeySet ks : keys.values()) {
-				snapshot.add(new KeySet(ks.getKeySetId(), ks.getContactId(),
-						ks.getTransportKeys().snapshot()));
+				snapshot.add(new TransportKeySet(ks.getKeySetId(),
+						ks.getContactId(), ks.getTransportKeys().snapshot()));
 			}
 			RotationResult rotationResult = rotateKeys(snapshot, now);
 			// Rebuild the mutable state for all contacts
@@ -346,12 +351,12 @@ class TransportKeyManagerImpl implements TransportKeyManager {

 	private static class TagContext {

-		private final KeySetId keySetId;
+		private final TransportKeySetId keySetId;
 		private final ContactId contactId;
 		private final MutableIncomingKeys inKeys;
 		private final long streamNumber;

-		private TagContext(KeySetId keySetId, ContactId contactId,
+		private TagContext(TransportKeySetId keySetId, ContactId contactId,
 				MutableIncomingKeys inKeys, long streamNumber) {
 			this.keySetId = keySetId;
 			this.contactId = contactId;
@@ -362,7 +367,7 @@ class TransportKeyManagerImpl implements TransportKeyManager {

 	private static class RotationResult {

-		private final Collection<KeySet> current = new ArrayList<>();
-		private final Collection<KeySet> rotated = new ArrayList<>();
+		private final Collection<TransportKeySet> current = new ArrayList<>();
+		private final Collection<TransportKeySet> rotated = new ArrayList<>();
 	}
 }