Provide process function for key transport message

4 files changed, 84 insertions, 19 deletions

diff --git a/src/main/java/eu/siacs/conversations/crypto/axolotl/AxolotlService.java b/src/main/java/eu/siacs/conversations/crypto/axolotl/AxolotlService.java
index 4cb1c90c..7a92a1f7 100644
--- a/src/main/java/eu/siacs/conversations/crypto/axolotl/AxolotlService.java
+++ b/src/main/java/eu/siacs/conversations/crypto/axolotl/AxolotlService.java
@@ -150,8 +150,13 @@ public class AxolotlService {
 		@Override
 		public void put(AxolotlAddress address, XmppAxolotlSession value) {
 			super.put(address, value);
+			value.setNotFresh();
 			xmppConnectionService.syncRosterToDisk(account);
 		}
+
+		public void put(XmppAxolotlSession session) {
+			this.put(session.getRemoteAddress(), session);
+		}
 	}
 
 	private static enum FetchStatus {
@@ -641,24 +646,32 @@ public class AxolotlService {
 		return axolotlMessage;
 	}
 
-	public XmppAxolotlMessage.XmppAxolotlPlaintextMessage processReceiving(XmppAxolotlMessage message) {
-		XmppAxolotlMessage.XmppAxolotlPlaintextMessage plaintextMessage = null;
+	private XmppAxolotlSession recreateUncachedSession(AxolotlAddress address) {
+		IdentityKey identityKey = axolotlStore.loadSession(address).getSessionState().getRemoteIdentityKey();
+		return (identityKey != null)
+				? new XmppAxolotlSession(account, axolotlStore, address,
+						identityKey.getFingerprint().replaceAll("\\s", ""))
+				: null;
+	}
+
+	private XmppAxolotlSession getReceivingSession(XmppAxolotlMessage message) {
 		AxolotlAddress senderAddress = new AxolotlAddress(message.getFrom().toString(),
 				message.getSenderDeviceId());
-
-		boolean newSession = false;
 		XmppAxolotlSession session = sessions.get(senderAddress);
 		if (session == null) {
 			Log.d(Config.LOGTAG, AxolotlService.getLogprefix(account) + "Account: " + account.getJid() + " No axolotl session found while parsing received message " + message);
-			IdentityKey identityKey = axolotlStore.loadSession(senderAddress).getSessionState().getRemoteIdentityKey();
-			if (identityKey != null) {
-				session = new XmppAxolotlSession(account, axolotlStore, senderAddress, identityKey.getFingerprint().replaceAll("\\s", ""));
-			} else {
+			session = recreateUncachedSession(senderAddress);
+			if (session == null) {
 				session = new XmppAxolotlSession(account, axolotlStore, senderAddress);
 			}
-			newSession = true;
 		}
+		return session;
+	}
 
+	public XmppAxolotlMessage.XmppAxolotlPlaintextMessage processReceivingPayloadMessage(XmppAxolotlMessage message) {
+		XmppAxolotlMessage.XmppAxolotlPlaintextMessage plaintextMessage = null;
+
+		XmppAxolotlSession session = getReceivingSession(message);
 		try {
 			plaintextMessage = message.decrypt(session, getOwnDeviceId());
 			Integer preKeyId = session.getPreKeyId();
@@ -670,10 +683,23 @@ public class AxolotlService {
 			Log.w(Config.LOGTAG, getLogprefix(account) + "Failed to decrypt message: " + e.getMessage());
 		}
 
-		if (newSession && plaintextMessage != null) {
-			sessions.put(senderAddress, session);
+		if (session.isFresh() && plaintextMessage != null) {
+			sessions.put(session);
 		}
 
 		return plaintextMessage;
 	}
+
+	public XmppAxolotlMessage.XmppAxolotlKeyTransportMessage processReceivingKeyTransportMessage(XmppAxolotlMessage message) {
+		XmppAxolotlMessage.XmppAxolotlKeyTransportMessage keyTransportMessage = null;
+
+		XmppAxolotlSession session = getReceivingSession(message);
+		keyTransportMessage = message.getParameters(session, getOwnDeviceId());
+
+		if (session.isFresh() && keyTransportMessage != null) {
+			sessions.put(session);
+		}
+
+		return keyTransportMessage;
+	}
 }
diff --git a/src/main/java/eu/siacs/conversations/crypto/axolotl/XmppAxolotlMessage.java b/src/main/java/eu/siacs/conversations/crypto/axolotl/XmppAxolotlMessage.java
index fa6d895a..cf950d6d 100644
--- a/src/main/java/eu/siacs/conversations/crypto/axolotl/XmppAxolotlMessage.java
+++ b/src/main/java/eu/siacs/conversations/crypto/axolotl/XmppAxolotlMessage.java
@@ -64,6 +64,30 @@ public class XmppAxolotlMessage {
 		}
 	}
 
+	public static class XmppAxolotlKeyTransportMessage {
+		private final String fingerprint;
+		private final byte[] key;
+		private final byte[] iv;
+
+		public XmppAxolotlKeyTransportMessage(String fingerprint, byte[] key, byte[] iv) {
+			this.fingerprint = fingerprint;
+			this.key = key;
+			this.iv = iv;
+		}
+
+		public String getFingerprint() {
+			return fingerprint;
+		}
+
+		public byte[] getKey() {
+			return key;
+		}
+
+		public byte[] getIv() {
+			return iv;
+		}
+	}
+
 	private XmppAxolotlMessage(final Element axolotlMessage, final Jid from) throws IllegalArgumentException {
 		this.from = from;
 		Element header = axolotlMessage.findChild(HEADER);
@@ -188,11 +212,18 @@ public class XmppAxolotlMessage {
 		return encryptionElement;
 	}
 
-	public byte[] unpackKey(XmppAxolotlSession session, Integer sourceDeviceId) {
+	private byte[] unpackKey(XmppAxolotlSession session, Integer sourceDeviceId) {
 		byte[] encryptedKey = keys.get(sourceDeviceId);
 		return (encryptedKey != null) ? session.processReceiving(encryptedKey) : null;
 	}
 
+	public XmppAxolotlKeyTransportMessage getParameters(XmppAxolotlSession session, Integer sourceDeviceId) {
+		byte[] key = unpackKey(session, sourceDeviceId);
+		return (key != null)
+				? new XmppAxolotlKeyTransportMessage(session.getFingerprint(), key, getIV())
+				: null;
+	}
+
 	public XmppAxolotlPlaintextMessage decrypt(XmppAxolotlSession session, Integer sourceDeviceId) throws CryptoFailedException {
 		XmppAxolotlPlaintextMessage plaintextMessage = null;
 		byte[] key = unpackKey(session, sourceDeviceId);
diff --git a/src/main/java/eu/siacs/conversations/crypto/axolotl/XmppAxolotlSession.java b/src/main/java/eu/siacs/conversations/crypto/axolotl/XmppAxolotlSession.java
index 6ed73da6..46004a1a 100644
--- a/src/main/java/eu/siacs/conversations/crypto/axolotl/XmppAxolotlSession.java
+++ b/src/main/java/eu/siacs/conversations/crypto/axolotl/XmppAxolotlSession.java
@@ -23,16 +23,12 @@ import eu.siacs.conversations.entities.Account;
 
 public class XmppAxolotlSession {
 	private final SessionCipher cipher;
-	private Integer preKeyId = null;
 	private final SQLiteAxolotlStore sqLiteAxolotlStore;
-
-	public AxolotlAddress getRemoteAddress() {
-		return remoteAddress;
-	}
-
 	private final AxolotlAddress remoteAddress;
 	private final Account account;
 	private String fingerprint = null;
+	private Integer preKeyId = null;
+	private boolean fresh = true;
 
 	public XmppAxolotlSession(Account account, SQLiteAxolotlStore store, AxolotlAddress remoteAddress, String fingerprint) {
 		this(account, store, remoteAddress);
@@ -59,6 +55,18 @@ public class XmppAxolotlSession {
 		return fingerprint;
 	}
 
+	public AxolotlAddress getRemoteAddress() {
+		return remoteAddress;
+	}
+
+	public boolean isFresh() {
+		return fresh;
+	}
+
+	public void setNotFresh() {
+		this.fresh = false;
+	}
+
 	protected void setTrust(SQLiteAxolotlStore.Trust trust) {
 		sqLiteAxolotlStore.setFingerprintTrust(fingerprint, trust);
 	}
diff --git a/src/main/java/eu/siacs/conversations/parser/MessageParser.java b/src/main/java/eu/siacs/conversations/parser/MessageParser.java
index a53e3a50..5786487f 100644
--- a/src/main/java/eu/siacs/conversations/parser/MessageParser.java
+++ b/src/main/java/eu/siacs/conversations/parser/MessageParser.java
@@ -100,7 +100,7 @@ public class MessageParser extends AbstractParser implements
 		Message finishedMessage = null;
 		AxolotlService service = conversation.getAccount().getAxolotlService();
 		XmppAxolotlMessage xmppAxolotlMessage = XmppAxolotlMessage.fromElement(axolotlMessage, from.toBareJid());
-		XmppAxolotlMessage.XmppAxolotlPlaintextMessage plaintextMessage = service.processReceiving(xmppAxolotlMessage);
+		XmppAxolotlMessage.XmppAxolotlPlaintextMessage plaintextMessage = service.processReceivingPayloadMessage(xmppAxolotlMessage);
 		if(plaintextMessage != null) {
 			finishedMessage = new Message(conversation, plaintextMessage.getPlaintext(), Message.ENCRYPTION_AXOLOTL, status);
 			finishedMessage.setAxolotlFingerprint(plaintextMessage.getFingerprint());