better logging on failed decrypt

// FREEBIE

AxolotlKit/Classes/CipherMessage/WhisperMessage.m

@@ -1,9 +1,5 @@
 //
-//  WhisperMessage.m
-//  AxolotlKit
-//
-//  Created by Frederic Jacobs on 23/07/14.
-//  Copyright (c) 2014 Frederic Jacobs. All rights reserved.
+//  Copyright (c) 2017 Open Whisper Systems. All rights reserved.
 //
 
 #import "AxolotlExceptions.h"
@@ -98,8 +94,21 @@
                                  serialized:data];
     
     if (![theirMac isEqualToData:ourMac]) {
+        DDLogError(@"%@ Bad Mac! Their Mac: %@ Our Mac: %@", self.tag, theirMac, ourMac);
         @throw [NSException exceptionWithName:InvalidMessageException reason:@"Bad Mac!" userInfo:@{}];
     }
 }
 
+#pragma mark - Logging
+
++ (NSString *)tag
+{
+    return [NSString stringWithFormat:@"[%@]", self.class];
+}
+
+- (NSString *)tag
+{
+    return self.class.tag;
+}
+
 @end

AxolotlKit/Classes/SessionCipher.m

@@ -32,12 +32,13 @@ static dispatch_queue_t _sessionCipherDispatchQueue;
 
 @interface SessionCipher ()
 
-@property NSString* recipientId;
-@property int deviceId;
+@property (nonatomic, readonly) NSString *recipientId;
+@property (nonatomic, readonly) int deviceId;
 
-@property (nonatomic, retain) id<SessionStore> sessionStore;
-@property (nonatomic, retain) SessionBuilder   *sessionBuilder;
-@property (nonatomic, retain) id<PreKeyStore>  prekeyStore;
+@property (nonatomic, readonly) id<IdentityKeyStore> identityKeyStore;
+@property (nonatomic, readonly) id<SessionStore> sessionStore;
+@property (nonatomic, readonly) SessionBuilder *sessionBuilder;
+@property (nonatomic, readonly) id<PreKeyStore> prekeyStore;
 
 @end
 
@@ -63,16 +64,16 @@ static dispatch_queue_t _sessionCipherDispatchQueue;
     self = [super init];
 
     if (self){
-        self.recipientId       = recipientId;
-        self.deviceId          = deviceId;
-        self.sessionStore      = sessionStore;
-        self.sessionBuilder    = [[SessionBuilder alloc] initWithSessionStore:sessionStore
-                                                              preKeyStore:preKeyStore
-                                                        signedPreKeyStore:signedPreKeyStore
-                                                         identityKeyStore:identityKeyStore
-                                                              recipientId:recipientId
-                                                                 deviceId:deviceId];
-        
+        _recipientId = recipientId;
+        _deviceId = deviceId;
+        _sessionStore = sessionStore;
+        _identityKeyStore = identityKeyStore;
+        _sessionBuilder = [[SessionBuilder alloc] initWithSessionStore:sessionStore
+                                                           preKeyStore:preKeyStore
+                                                     signedPreKeyStore:signedPreKeyStore
+                                                      identityKeyStore:identityKeyStore
+                                                           recipientId:recipientId
+                                                              deviceId:deviceId];
     }
     
     return self;
@@ -106,38 +107,53 @@ static dispatch_queue_t _sessionCipherDispatchQueue;
 - (id<CipherMessage>)encryptMessage:(NSData*)paddedMessage{
     [self assertOnSessionCipherDispatchQueue];
     SessionRecord *sessionRecord = [self.sessionStore loadSession:self.recipientId deviceId:self.deviceId];
-    SessionState  *session       = sessionRecord.sessionState;
-    ChainKey *chainKey           = session.senderChainKey;
+    SessionState *sessionState = sessionRecord.sessionState;
+    ChainKey *chainKey = sessionState.senderChainKey;
     MessageKeys *messageKeys     = chainKey.messageKeys;
-    NSData *senderRatchetKey     = session.senderRatchetKey;
-    int previousCounter          = session.previousCounter;
-    int sessionVersion           = session.version;
-    
+    NSData *senderRatchetKey = sessionState.senderRatchetKey;
+    int previousCounter = sessionState.previousCounter;
+    int sessionVersion = sessionState.version;
+
+    if (![self.identityKeyStore isTrustedIdentityKey:sessionState.remoteIdentityKey
+                                         recipientId:self.recipientId
+                                           direction:TSMessageDirectionOutgoing]) {
+        DDLogWarn(
+            @"%@ Previously known identity key for while encrypting for recipient: %@", self.tag, self.recipientId);
+        @throw [NSException exceptionWithName:UntrustedIdentityKeyException
+                                       reason:@"There is a previously known identity key."
+                                     userInfo:@{}];
+    }
+
+    if ([self.identityKeyStore saveRemoteIdentity:sessionState.remoteIdentityKey recipientId:self.recipientId]) {
+        [sessionRecord removePreviousSessionStates];
+    }
 
     NSData *ciphertextBody = [AES_CBC encryptCBCMode:paddedMessage withKey:messageKeys.cipherKey withIV:messageKeys.iv];
 
-    id<CipherMessage> cipherMessage = [[WhisperMessage alloc] initWithVersion:sessionVersion
-                                                                       macKey:messageKeys.macKey
-                                                             senderRatchetKey:senderRatchetKey.prependKeyType
-                                                                      counter:chainKey.index
-                                                              previousCounter:previousCounter
-                                                                   cipherText:ciphertextBody
-                                                            senderIdentityKey:session.localIdentityKey.prependKeyType
-                                                          receiverIdentityKey:session.remoteIdentityKey.prependKeyType];
-    
-    if ([session hasUnacknowledgedPreKeyMessage]){
-        PendingPreKey *items = [session unacknowledgedPreKeyMessageItems];
-        int localRegistrationId = [session localRegistrationId];
-        
-        cipherMessage = [[PreKeyWhisperMessage alloc] initWithWhisperMessage:cipherMessage
-                                                              registrationId:localRegistrationId
-                                                                    prekeyId:items.preKeyId
-                                                              signedPrekeyId:items.signedPreKeyId
-                                                                     baseKey:items.baseKey.prependKeyType
-                                                                 identityKey:session.localIdentityKey.prependKeyType];
+    id<CipherMessage> cipherMessage =
+        [[WhisperMessage alloc] initWithVersion:sessionVersion
+                                         macKey:messageKeys.macKey
+                               senderRatchetKey:senderRatchetKey.prependKeyType
+                                        counter:chainKey.index
+                                previousCounter:previousCounter
+                                     cipherText:ciphertextBody
+                              senderIdentityKey:sessionState.localIdentityKey.prependKeyType
+                            receiverIdentityKey:sessionState.remoteIdentityKey.prependKeyType];
+
+    if ([sessionState hasUnacknowledgedPreKeyMessage]) {
+        PendingPreKey *items = [sessionState unacknowledgedPreKeyMessageItems];
+        int localRegistrationId = [sessionState localRegistrationId];
+
+        cipherMessage =
+            [[PreKeyWhisperMessage alloc] initWithWhisperMessage:cipherMessage
+                                                  registrationId:localRegistrationId
+                                                        prekeyId:items.preKeyId
+                                                  signedPrekeyId:items.signedPreKeyId
+                                                         baseKey:items.baseKey.prependKeyType
+                                                     identityKey:sessionState.localIdentityKey.prependKeyType];
     }
-    
-    [session setSenderChainKey:[chainKey nextChainKey]];
+
+    [sessionState setSenderChainKey:[chainKey nextChainKey]];
     [self.sessionStore storeSession:self.recipientId deviceId:self.deviceId session:sessionRecord];
     
     return cipherMessage;
@@ -176,7 +192,22 @@ static dispatch_queue_t _sessionCipherDispatchQueue;
     
     SessionRecord  *sessionRecord  = [self.sessionStore loadSession:self.recipientId deviceId:self.deviceId];
     NSData         *plaintext      = [self decryptWithSessionRecord:sessionRecord whisperMessage:message];
-    
+
+    if (![self.identityKeyStore isTrustedIdentityKey:sessionRecord.sessionState.remoteIdentityKey
+                                         recipientId:self.recipientId
+                                           direction:TSMessageDirectionIncoming]) {
+        DDLogWarn(
+            @"%@ Previously known identity key for while decrypting from recipient: %@", self.tag, self.recipientId);
+        @throw [NSException exceptionWithName:UntrustedIdentityKeyException
+                                       reason:@"There is a previously known identity key."
+                                     userInfo:@{}];
+    }
+
+    if ([self.identityKeyStore saveRemoteIdentity:sessionRecord.sessionState.remoteIdentityKey
+                                      recipientId:self.recipientId]) {
+        [sessionRecord removePreviousSessionStates];
+    }
+
     [self.sessionStore storeSession:self.recipientId deviceId:self.deviceId session:sessionRecord];
     
     return plaintext;