diff --git a/SignalServiceKit/src/TestUtils/Factories.swift b/SignalServiceKit/src/TestUtils/Factories.swift index e9255c7509..1c896c8032 100644 --- a/SignalServiceKit/src/TestUtils/Factories.swift +++ b/SignalServiceKit/src/TestUtils/Factories.swift @@ -142,7 +142,7 @@ public class OutgoingMessageFactory: NSObject, Factory { contactShare: contactShareBuilder(), linkPreview: linkPreviewBuilder(), messageSticker: messageStickerBuilder(), - perMessageExpirationDurationSeconds: 0) + perMessageExpirationDurationSeconds: perMessageExpirationDurationSecondsBuilder()) return item } @@ -219,6 +219,11 @@ public class OutgoingMessageFactory: NSObject, Factory { return nil } + @objc + public var perMessageExpirationDurationSecondsBuilder: () -> UInt32 = { + return 0 + } + // MARK: Delivery Receipts @objc @@ -266,7 +271,7 @@ public class IncomingMessageFactory: NSObject, Factory { messageSticker: messageStickerBuilder(), serverTimestamp: serverTimestampBuilder(), wasReceivedByUD: wasReceivedByUDBuilder(), - perMessageExpirationDurationSeconds: 0) + perMessageExpirationDurationSeconds: perMessageExpirationDurationSecondsBuilder()) item.anyInsert(transaction: transaction) @@ -349,6 +354,11 @@ public class IncomingMessageFactory: NSObject, Factory { public var wasReceivedByUDBuilder: () -> Bool = { return false } + + @objc + public var perMessageExpirationDurationSecondsBuilder: () -> UInt32 = { + return 0 + } } @objc diff --git a/SignalServiceKit/src/Util/PerMessageExpiration.swift b/SignalServiceKit/src/Util/PerMessageExpiration.swift index 61c1e8e269..fa5d00dd8a 100644 --- a/SignalServiceKit/src/Util/PerMessageExpiration.swift +++ b/SignalServiceKit/src/Util/PerMessageExpiration.swift @@ -75,32 +75,40 @@ public class PerMessageExpiration: NSObject { let perMessageExpiresAtMS = message.perMessageExpiresAt guard perMessageExpiresAtMS > nowMs else { + Logger.verbose("Expiring immediately.") // Message has expired; remove it immediately. completePerMessageExpiration(forMessage: message, transaction: transaction) return } + Logger.verbose("Scheduling expiration.") let delaySeconds: TimeInterval = Double(perMessageExpiresAtMS - nowMs) / 1000 DispatchQueue.global().asyncAfter(wallDeadline: .now() + delaySeconds) { - self.completePerMessageExpiration(forMessage: message) - } - } - - private class func completePerMessageExpiration(forMessage message: TSMessage) { - databaseStorage.write { (transaction) in - self.completePerMessageExpiration(forMessage: message, - transaction: transaction) + databaseStorage.write { (transaction) in + self.completePerMessageExpiration(forMessage: message, + transaction: transaction) + } } } private class func completePerMessageExpiration(forMessage message: TSMessage, transaction: SDSAnyWriteTransaction) { - guard message.perMessageExpirationHasExpired else { + Logger.verbose("Expiring message per schedule.") + guard let uniqueId = message.uniqueId else { + owsFailDebug("Missing uniqueId.") + return + } + guard let latestCopy = TSMessage.anyFetch(uniqueId: uniqueId, transaction: transaction) as? TSMessage else { + // Message has been deleted. + Logger.warn("Couldn't expire message; not in database.") + return + } + guard latestCopy.perMessageExpirationHasExpired else { // Already expired, no need to expire again. return } - message.updateWithHasPerMessageExpiredAndRemoveRenderableContent(with: transaction) + latestCopy.updateWithHasPerMessageExpiredAndRemoveRenderableContent(with: transaction) } // MARK: - Events diff --git a/SignalServiceKit/tests/Util/PerMessageExpirationTest.swift b/SignalServiceKit/tests/Util/PerMessageExpirationTest.swift new file mode 100644 index 0000000000..6049dc4a76 --- /dev/null +++ b/SignalServiceKit/tests/Util/PerMessageExpirationTest.swift @@ -0,0 +1,101 @@ +// +// Copyright (c) 2019 Open Whisper Systems. All rights reserved. +// + +import Foundation +import XCTest + +@testable import SignalServiceKit + +class PerMessageExpirationTest: SSKBaseTestSwift { + + override func setUp() { + super.setUp() + } + + override func tearDown() { + super.tearDown() + } + + // MARK: + + func test_expiration() { + + let messageCount = { () -> Int in + return self.databaseStorage.readReturningResult { transaction in + return TSInteraction.anyFetchAll(transaction: transaction).count + } + } + let latestCopy = { (message: TSMessage) -> TSMessage in + let uniqueId = message.uniqueId! + return self.databaseStorage.readReturningResult { transaction in + return TSMessage.anyFetch(uniqueId: uniqueId, transaction: transaction) as! TSMessage + } + } + + // Factory 1 builds messages without per-message expiration. + let factory1 = IncomingMessageFactory() + factory1.perMessageExpirationDurationSecondsBuilder = { + return 0 + } + + self.write { transaction in + _ = factory1.create(transaction: transaction) + } + + XCTAssertEqual(1, messageCount()) + + // Factory 2 builds messages with per-message expiration. + let expirationSeconds: UInt32 = 1 + let factory2 = IncomingMessageFactory() + factory2.perMessageExpirationDurationSecondsBuilder = { + return expirationSeconds + } + + var disappearingMessage: TSMessage? + self.write { transaction in + disappearingMessage = factory2.create(transaction: transaction) + } + + XCTAssertEqual(2, messageCount()) + + guard let message = disappearingMessage else { + XCTFail("Missing message.") + return + } + + XCTAssertTrue(message.hasPerMessageExpiration) + XCTAssertFalse(message.hasPerMessageExpirationStarted) + XCTAssertFalse(message.perMessageExpirationHasExpired) + + XCTAssertTrue(latestCopy(message).hasPerMessageExpiration) + XCTAssertFalse(latestCopy(message).hasPerMessageExpirationStarted) + XCTAssertFalse(latestCopy(message).perMessageExpirationHasExpired) + + self.write { transaction in + PerMessageExpiration.startPerMessageExpiration(forMessage: message, + transaction: transaction) + } + + XCTAssertTrue(message.hasPerMessageExpiration) + XCTAssertTrue(message.hasPerMessageExpirationStarted) + XCTAssertFalse(message.perMessageExpirationHasExpired) + + XCTAssertTrue(latestCopy(message).hasPerMessageExpiration) + XCTAssertTrue(latestCopy(message).hasPerMessageExpirationStarted) + XCTAssertFalse(latestCopy(message).perMessageExpirationHasExpired) + + XCTAssertEqual(2, messageCount()) + + Logger.verbose("Sleeping for \(expirationSeconds + 1)") + // Sleep for a little extra time to avoid races. + sleep(expirationSeconds + 1) + Logger.verbose("Slept for \(expirationSeconds + 1)") + + XCTAssertTrue(latestCopy(message).hasPerMessageExpiration) + XCTAssertTrue(latestCopy(message).hasPerMessageExpirationStarted) + XCTAssertTrue(latestCopy(message).perMessageExpirationHasExpired) + + XCTAssertEqual(2, messageCount()) + } +}