Refactor axolotl message processing workflow

XmppAxolotlMessage is now entirely responsible for handling encryption
and decryption of messages, only leveraging XmppAxolotlSession as a
packing/unpacking primitive for payload keys.

Removed pseudo-dead session generation code step from prepareMessage
function, as sessions have been created by invoking the
TrustKeysActivity for a while now.

Added prepareKeyTransportMessage function, which creates a message with
no payload. The key that is packed into the header keyElements can then
be used for other purposes (e.g. encrypted file transfer).

3 files changed, 116 insertions, 152 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 4cc61b37..4cb1c90c 100644
--- a/src/main/java/eu/siacs/conversations/crypto/axolotl/AxolotlService.java
+++ b/src/main/java/eu/siacs/conversations/crypto/axolotl/AxolotlService.java
@@ -256,7 +256,7 @@ public class AxolotlService {
 			for (Integer deviceId : deviceIds) {
 				AxolotlAddress ownDeviceAddress = new AxolotlAddress(jid.toBareJid().toString(), deviceId);
 				if (sessions.get(ownDeviceAddress) == null) {
-					buildSessionFromPEP(null, ownDeviceAddress, false);
+					buildSessionFromPEP(ownDeviceAddress);
 				}
 			}
 		}
@@ -422,7 +422,7 @@ public class AxolotlService {
 		axolotlStore.setFingerprintTrust(fingerprint, trust);
 	}
 
-	private void buildSessionFromPEP(final Conversation conversation, final AxolotlAddress address, final boolean flushWaitingQueueAfterFetch) {
+	private void buildSessionFromPEP(final AxolotlAddress address) {
 		Log.i(Config.LOGTAG, AxolotlService.getLogprefix(account) + "Building new sesstion for " + address.getDeviceId());
 
 		try {
@@ -434,15 +434,6 @@ public class AxolotlService {
 					AxolotlAddress ownAddress = new AxolotlAddress(account.getJid().toBareJid().toString(), 0);
 					if (!fetchStatusMap.getAll(ownAddress).containsValue(FetchStatus.PENDING)
 							&& !fetchStatusMap.getAll(address).containsValue(FetchStatus.PENDING)) {
-						if (flushWaitingQueueAfterFetch && conversation != null) {
-							conversation.findUnsentMessagesWithEncryption(Message.ENCRYPTION_AXOLOTL,
-									new Conversation.OnMessageFound() {
-										@Override
-										public void onMessageFound(Message message) {
-											processSending(message, false);
-										}
-									});
-						}
 						mXmppConnectionService.keyStatusUpdated();
 					}
 				}
@@ -537,7 +528,7 @@ public class AxolotlService {
 		return addresses;
 	}
 
-	public boolean createSessionsIfNeeded(final Conversation conversation, final boolean flushWaitingQueueAfterFetch) {
+	public boolean createSessionsIfNeeded(final Conversation conversation) {
 		Log.i(Config.LOGTAG, AxolotlService.getLogprefix(account) + "Creating axolotl sessions if needed...");
 		boolean newSessions = false;
 		Set<AxolotlAddress> addresses = findDevicesWithoutSession(conversation);
@@ -546,7 +537,9 @@ public class AxolotlService {
 			FetchStatus status = fetchStatusMap.get(address);
 			if (status == null || status == FetchStatus.ERROR) {
 				fetchStatusMap.put(address, FetchStatus.PENDING);
-				this.buildSessionFromPEP(conversation, address, flushWaitingQueueAfterFetch);
+				this.buildSessionFromPEP(address);
+				newSessions = true;
+			} else if (status == FetchStatus.PENDING) {
 				newSessions = true;
 			} else {
 				Log.d(Config.LOGTAG, AxolotlService.getLogprefix(account) + "Already fetching bundle for " + address.toString());
@@ -565,40 +558,52 @@ public class AxolotlService {
 	}
 
 	@Nullable
-	public XmppAxolotlMessage encrypt(Message message) {
-		final String content;
-		if (message.hasFileOnRemoteHost()) {
-			content = message.getFileParams().url.toString();
-		} else {
-			content = message.getBody();
-		}
-		final XmppAxolotlMessage axolotlMessage;
-		try {
-			axolotlMessage = new XmppAxolotlMessage(message.getContact().getJid().toBareJid(),
-					getOwnDeviceId(), content);
-		} catch (CryptoFailedException e) {
-			Log.w(Config.LOGTAG, getLogprefix(account) + "Failed to encrypt message: " + e.getMessage());
-			return null;
-		}
+	private XmppAxolotlMessage buildHeader(Contact contact) {
+		final XmppAxolotlMessage axolotlMessage = new XmppAxolotlMessage(
+				contact.getJid().toBareJid(), getOwnDeviceId());
 
-		if (findSessionsforContact(message.getContact()).isEmpty()) {
+		Set<XmppAxolotlSession> contactSessions = findSessionsforContact(contact);
+		Set<XmppAxolotlSession> ownSessions = findOwnSessions();
+		if (contactSessions.isEmpty()) {
 			return null;
 		}
 		Log.d(Config.LOGTAG, AxolotlService.getLogprefix(account) + "Building axolotl foreign keyElements...");
-		for (XmppAxolotlSession session : findSessionsforContact(message.getContact())) {
+		for (XmppAxolotlSession session : contactSessions) {
 			Log.v(Config.LOGTAG, AxolotlService.getLogprefix(account) + session.getRemoteAddress().toString());
-			axolotlMessage.addKeyElement(session.processSending(axolotlMessage.getInnerKey()));
+			axolotlMessage.addDevice(session);
 		}
 		Log.d(Config.LOGTAG, AxolotlService.getLogprefix(account) + "Building axolotl own keyElements...");
-		for (XmppAxolotlSession session : findOwnSessions()) {
+		for (XmppAxolotlSession session : ownSessions) {
 			Log.v(Config.LOGTAG, AxolotlService.getLogprefix(account) + session.getRemoteAddress().toString());
-			axolotlMessage.addKeyElement(session.processSending(axolotlMessage.getInnerKey()));
+			axolotlMessage.addDevice(session);
+		}
+
+		return axolotlMessage;
+	}
+
+	@Nullable
+	public XmppAxolotlMessage encrypt(Message message) {
+		XmppAxolotlMessage axolotlMessage = buildHeader(message.getContact());
+
+		if (axolotlMessage != null) {
+			final String content;
+			if (message.hasFileOnRemoteHost()) {
+				content = message.getFileParams().url.toString();
+			} else {
+				content = message.getBody();
+			}
+			try {
+				axolotlMessage.encrypt(content);
+			} catch (CryptoFailedException e) {
+				Log.w(Config.LOGTAG, getLogprefix(account) + "Failed to encrypt message: " + e.getMessage());
+				return null;
+			}
 		}
 
 		return axolotlMessage;
 	}
 
-	private void processSending(final Message message, final boolean delay) {
+	public void preparePayloadMessage(final Message message, final boolean delay) {
 		executor.execute(new Runnable() {
 			@Override
 			public void run() {
@@ -615,13 +620,14 @@ public class AxolotlService {
 		});
 	}
 
-	public void prepareMessage(final Message message, final boolean delay) {
-		if (!messageCache.containsKey(message.getUuid())) {
-			boolean newSessions = createSessionsIfNeeded(message.getConversation(), true);
-			if (!newSessions) {
-				this.processSending(message, delay);
+	public void prepareKeyTransportMessage(final Contact contact, final OnMessageCreatedCallback onMessageCreatedCallback) {
+		executor.execute(new Runnable() {
+			@Override
+			public void run() {
+				XmppAxolotlMessage axolotlMessage = buildHeader(contact);
+				onMessageCreatedCallback.run(axolotlMessage);
 			}
-		}
+		});
 	}
 
 	public XmppAxolotlMessage fetchAxolotlMessageFromCache(Message message) {
@@ -653,26 +659,15 @@ public class AxolotlService {
 			newSession = true;
 		}
 
-		for (XmppAxolotlMessage.XmppAxolotlKeyElement keyElement : message.getKeyElements()) {
-			if (keyElement.getRecipientDeviceId() == getOwnDeviceId()) {
-				Log.d(Config.LOGTAG, AxolotlService.getLogprefix(account) + "Found axolotl keyElement matching own device ID, processing...");
-				byte[] payloadKey = session.processReceiving(keyElement);
-				if (payloadKey != null) {
-					Log.d(Config.LOGTAG, AxolotlService.getLogprefix(account) + "Got payload key from axolotl keyElement. Decrypting message...");
-					try {
-						plaintextMessage = message.decrypt(session, payloadKey, session.getFingerprint());
-					} catch (CryptoFailedException e) {
-						Log.w(Config.LOGTAG, getLogprefix(account) + "Failed to decrypt message: " + e.getMessage());
-						break;
-					}
-				}
-				Integer preKeyId = session.getPreKeyId();
-				if (preKeyId != null) {
-					publishBundlesIfNeeded();
-					session.resetPreKeyId();
-				}
-				break;
+		try {
+			plaintextMessage = message.decrypt(session, getOwnDeviceId());
+			Integer preKeyId = session.getPreKeyId();
+			if (preKeyId != null) {
+				publishBundlesIfNeeded();
+				session.resetPreKeyId();
 			}
+		} catch (CryptoFailedException e) {
+			Log.w(Config.LOGTAG, getLogprefix(account) + "Failed to decrypt message: " + e.getMessage());
 		}
 
 		if (newSession && plaintextMessage != null) {
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 4e954a72..fa6d895a 100644
--- a/src/main/java/eu/siacs/conversations/crypto/axolotl/XmppAxolotlMessage.java
+++ b/src/main/java/eu/siacs/conversations/crypto/axolotl/XmppAxolotlMessage.java
@@ -1,6 +1,5 @@
 package eu.siacs.conversations.crypto.axolotl;
 
-import android.support.annotation.Nullable;
 import android.util.Base64;
 import android.util.Log;
 
@@ -9,8 +8,9 @@ import java.security.InvalidKeyException;
 import java.security.NoSuchAlgorithmException;
 import java.security.NoSuchProviderException;
 import java.security.SecureRandom;
-import java.util.HashSet;
-import java.util.Set;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
 
 import javax.crypto.BadPaddingException;
 import javax.crypto.Cipher;
@@ -26,9 +26,11 @@ import eu.siacs.conversations.xml.Element;
 import eu.siacs.conversations.xmpp.jid.Jid;
 
 public class XmppAxolotlMessage {
-	public static final String TAGNAME = "encrypted";
+	public static final String CONTAINERTAG = "encrypted";
 	public static final String HEADER = "header";
 	public static final String SOURCEID = "sid";
+	public static final String KEYTAG = "key";
+	public static final String REMOTEID = "rid";
 	public static final String IVTAG = "iv";
 	public static final String PAYLOAD = "payload";
 
@@ -39,54 +41,15 @@ public class XmppAxolotlMessage {
 	private byte[] innerKey;
 	private byte[] ciphertext = null;
 	private byte[] iv = null;
-	private final Set<XmppAxolotlKeyElement> keyElements;
+	private final Map<Integer, byte[]> keys;
 	private final Jid from;
 	private final int sourceDeviceId;
 
-	public static class XmppAxolotlKeyElement {
-		public static final String TAGNAME = "key";
-		public static final String REMOTEID = "rid";
-
-		private final int recipientDeviceId;
-		private final byte[] content;
-
-		public XmppAxolotlKeyElement(int deviceId, byte[] content) {
-			this.recipientDeviceId = deviceId;
-			this.content = content;
-		}
-
-		public XmppAxolotlKeyElement(Element keyElement) {
-			if (TAGNAME.equals(keyElement.getName())) {
-				this.recipientDeviceId = Integer.parseInt(keyElement.getAttribute(REMOTEID));
-				this.content = Base64.decode(keyElement.getContent(), Base64.DEFAULT);
-			} else {
-				throw new IllegalArgumentException("Argument not a <" + TAGNAME + "> Element!");
-			}
-		}
-
-		public int getRecipientDeviceId() {
-			return recipientDeviceId;
-		}
-
-		public byte[] getContents() {
-			return content;
-		}
-
-		public Element toXml() {
-			Element keyElement = new Element(TAGNAME);
-			keyElement.setAttribute(REMOTEID, getRecipientDeviceId());
-			keyElement.setContent(Base64.encodeToString(getContents(), Base64.DEFAULT));
-			return keyElement;
-		}
-	}
-
 	public static class XmppAxolotlPlaintextMessage {
-		private final XmppAxolotlSession session;
 		private final String plaintext;
 		private final String fingerprint;
 
-		public XmppAxolotlPlaintextMessage(XmppAxolotlSession session, String plaintext, String fingerprint) {
-			this.session = session;
+		public XmppAxolotlPlaintextMessage(String plaintext, String fingerprint) {
 			this.plaintext = plaintext;
 			this.fingerprint = fingerprint;
 		}
@@ -95,24 +58,28 @@ public class XmppAxolotlMessage {
 			return plaintext;
 		}
 
-		public XmppAxolotlSession getSession() {
-			return session;
-		}
 
 		public String getFingerprint() {
 			return fingerprint;
 		}
 	}
 
-	public XmppAxolotlMessage(Jid from, Element axolotlMessage) throws IllegalArgumentException {
+	private XmppAxolotlMessage(final Element axolotlMessage, final Jid from) throws IllegalArgumentException {
 		this.from = from;
 		Element header = axolotlMessage.findChild(HEADER);
 		this.sourceDeviceId = Integer.parseInt(header.getAttribute(SOURCEID));
-		this.keyElements = new HashSet<>();
-		for (Element keyElement : header.getChildren()) {
+		List<Element> keyElements = header.getChildren();
+		this.keys = new HashMap<>(keyElements.size());
+		for (Element keyElement : keyElements) {
 			switch (keyElement.getName()) {
-				case XmppAxolotlKeyElement.TAGNAME:
-					keyElements.add(new XmppAxolotlKeyElement(keyElement));
+				case KEYTAG:
+					try {
+						Integer recipientId = Integer.parseInt(keyElement.getAttribute(REMOTEID));
+						byte[] key = Base64.decode(keyElement.getContent(), Base64.DEFAULT);
+						this.keys.put(recipientId, key);
+					} catch (NumberFormatException e) {
+						throw new IllegalArgumentException(e);
+					}
 					break;
 				case IVTAG:
 					if (this.iv != null) {
@@ -134,14 +101,13 @@ public class XmppAxolotlMessage {
 	public XmppAxolotlMessage(Jid from, int sourceDeviceId) {
 		this.from = from;
 		this.sourceDeviceId = sourceDeviceId;
-		this.keyElements = new HashSet<>();
+		this.keys = new HashMap<>();
 		this.iv = generateIv();
 		this.innerKey = generateKey();
 	}
 
-	public XmppAxolotlMessage(Jid from, int sourceDeviceId, String plaintext) throws CryptoFailedException {
-		this(from, sourceDeviceId);
-		this.encrypt(plaintext);
+	public static XmppAxolotlMessage fromElement(Element element, Jid from) {
+		return new XmppAxolotlMessage(element, from);
 	}
 
 	private static byte[] generateKey() {
@@ -162,7 +128,7 @@ public class XmppAxolotlMessage {
 		return iv;
 	}
 
-	private void encrypt(String plaintext) throws CryptoFailedException {
+	public void encrypt(String plaintext) throws CryptoFailedException {
 		try {
 			SecretKey secretKey = new SecretKeySpec(innerKey, KEYTYPE);
 			IvParameterSpec ivSpec = new IvParameterSpec(iv);
@@ -189,13 +155,10 @@ public class XmppAxolotlMessage {
 		return ciphertext;
 	}
 
-	public Set<XmppAxolotlKeyElement> getKeyElements() {
-		return keyElements;
-	}
-
-	public void addKeyElement(@Nullable XmppAxolotlKeyElement keyElement) {
-		if (keyElement != null) {
-			keyElements.add(keyElement);
+	public void addDevice(XmppAxolotlSession session) {
+		byte[] key = session.processSending(innerKey);
+		if (key != null) {
+			keys.put(session.getRemoteAddress().getDeviceId(), key);
 		}
 	}
 
@@ -207,12 +170,15 @@ public class XmppAxolotlMessage {
 		return this.iv;
 	}
 
-	public Element toXml() {
-		Element encryptionElement = new Element(TAGNAME, AxolotlService.PEP_PREFIX);
+	public Element toElement() {
+		Element encryptionElement = new Element(CONTAINERTAG, AxolotlService.PEP_PREFIX);
 		Element headerElement = encryptionElement.addChild(HEADER);
 		headerElement.setAttribute(SOURCEID, sourceDeviceId);
-		for (XmppAxolotlKeyElement header : keyElements) {
-			headerElement.addChild(header.toXml());
+		for (Map.Entry<Integer, byte[]> keyEntry : keys.entrySet()) {
+			Element keyElement = new Element(KEYTAG);
+			keyElement.setAttribute(REMOTEID, keyEntry.getKey());
+			keyElement.setContent(Base64.encodeToString(keyEntry.getValue(), Base64.DEFAULT));
+			headerElement.addChild(keyElement);
 		}
 		headerElement.addChild(IVTAG).setContent(Base64.encodeToString(iv, Base64.DEFAULT));
 		if (ciphertext != null) {
@@ -222,24 +188,30 @@ public class XmppAxolotlMessage {
 		return encryptionElement;
 	}
 
+	public byte[] unpackKey(XmppAxolotlSession session, Integer sourceDeviceId) {
+		byte[] encryptedKey = keys.get(sourceDeviceId);
+		return (encryptedKey != null) ? session.processReceiving(encryptedKey) : null;
+	}
 
-	public XmppAxolotlPlaintextMessage decrypt(XmppAxolotlSession session, byte[] key, String fingerprint) throws CryptoFailedException {
+	public XmppAxolotlPlaintextMessage decrypt(XmppAxolotlSession session, Integer sourceDeviceId) throws CryptoFailedException {
 		XmppAxolotlPlaintextMessage plaintextMessage = null;
-		try {
-
-			Cipher cipher = Cipher.getInstance(CIPHERMODE, PROVIDER);
-			SecretKeySpec keySpec = new SecretKeySpec(key, KEYTYPE);
-			IvParameterSpec ivSpec = new IvParameterSpec(iv);
-
-			cipher.init(Cipher.DECRYPT_MODE, keySpec, ivSpec);
-
-			String plaintext = new String(cipher.doFinal(ciphertext));
-			plaintextMessage = new XmppAxolotlPlaintextMessage(session, plaintext, fingerprint);
-
-		} catch (NoSuchAlgorithmException | NoSuchPaddingException | InvalidKeyException
-				| InvalidAlgorithmParameterException | IllegalBlockSizeException
-				| BadPaddingException | NoSuchProviderException e) {
-			throw new CryptoFailedException(e);
+		byte[] key = unpackKey(session, sourceDeviceId);
+		if (key != null) {
+			try {
+				Cipher cipher = Cipher.getInstance(CIPHERMODE, PROVIDER);
+				SecretKeySpec keySpec = new SecretKeySpec(key, KEYTYPE);
+				IvParameterSpec ivSpec = new IvParameterSpec(iv);
+
+				cipher.init(Cipher.DECRYPT_MODE, keySpec, ivSpec);
+
+				String plaintext = new String(cipher.doFinal(ciphertext));
+				plaintextMessage = new XmppAxolotlPlaintextMessage(plaintext, session.getFingerprint());
+
+			} catch (NoSuchAlgorithmException | NoSuchPaddingException | InvalidKeyException
+					| InvalidAlgorithmParameterException | IllegalBlockSizeException
+					| BadPaddingException | NoSuchProviderException e) {
+				throw new CryptoFailedException(e);
+			}
 		}
 		return plaintextMessage;
 	}
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 d60e7715..6ed73da6 100644
--- a/src/main/java/eu/siacs/conversations/crypto/axolotl/XmppAxolotlSession.java
+++ b/src/main/java/eu/siacs/conversations/crypto/axolotl/XmppAxolotlSession.java
@@ -69,7 +69,7 @@ public class XmppAxolotlSession {
 	}
 
 	@Nullable
-	public byte[] processReceiving(XmppAxolotlMessage.XmppAxolotlKeyElement incomingHeader) {
+	public byte[] processReceiving(byte[] encryptedKey) {
 		byte[] plaintext = null;
 		SQLiteAxolotlStore.Trust trust = getTrust();
 		switch (trust) {
@@ -79,7 +79,7 @@ public class XmppAxolotlSession {
 			case TRUSTED:
 				try {
 					try {
-						PreKeyWhisperMessage message = new PreKeyWhisperMessage(incomingHeader.getContents());
+						PreKeyWhisperMessage message = new PreKeyWhisperMessage(encryptedKey);
 						Log.i(Config.LOGTAG, AxolotlService.getLogprefix(account) + "PreKeyWhisperMessage received, new session ID:" + message.getSignedPreKeyId() + "/" + message.getPreKeyId());
 						String fingerprint = message.getIdentityKey().getFingerprint().replaceAll("\\s", "");
 						if (this.fingerprint != null && !this.fingerprint.equals(fingerprint)) {
@@ -93,7 +93,7 @@ public class XmppAxolotlSession {
 						}
 					} catch (InvalidMessageException | InvalidVersionException e) {
 						Log.i(Config.LOGTAG, AxolotlService.getLogprefix(account) + "WhisperMessage received");
-						WhisperMessage message = new WhisperMessage(incomingHeader.getContents());
+						WhisperMessage message = new WhisperMessage(encryptedKey);
 						plaintext = cipher.decrypt(message);
 					} catch (InvalidKeyException | InvalidKeyIdException | UntrustedIdentityException e) {
 						Log.w(Config.LOGTAG, AxolotlService.getLogprefix(account) + "Error decrypting axolotl header, " + e.getClass().getName() + ": " + e.getMessage());
@@ -117,14 +117,11 @@ public class XmppAxolotlSession {
 	}
 
 	@Nullable
-	public XmppAxolotlMessage.XmppAxolotlKeyElement processSending(@NonNull byte[] outgoingMessage) {
+	public byte[] processSending(@NonNull byte[] outgoingMessage) {
 		SQLiteAxolotlStore.Trust trust = getTrust();
 		if (trust == SQLiteAxolotlStore.Trust.TRUSTED) {
 			CiphertextMessage ciphertextMessage = cipher.encrypt(outgoingMessage);
-			XmppAxolotlMessage.XmppAxolotlKeyElement header =
-					new XmppAxolotlMessage.XmppAxolotlKeyElement(remoteAddress.getDeviceId(),
-							ciphertextMessage.serialize());
-			return header;
+			return ciphertextMessage.serialize();
 		} else {
 			return null;
 		}