diff --git a/Rosetta/Core/Network/Protocol/Packets/Packet.swift b/Rosetta/Core/Network/Protocol/Packets/Packet.swift index 4e70867..e2df592 100644 --- a/Rosetta/Core/Network/Protocol/Packets/Packet.swift +++ b/Rosetta/Core/Network/Protocol/Packets/Packet.swift @@ -86,6 +86,39 @@ struct MessageAttachment: Codable, Equatable { var preview: String = "" var blob: String = "" var type: AttachmentType = .image + var transportTag: String = "" + var transportServer: String = "" + + /// Effective download tag: prefers per-attachment `transportTag`, falls back to preview-embedded tag. + /// Backward compat: old messages (before multi-transport) have empty `transportTag` but tag in preview. + var effectiveDownloadTag: String { + if !transportTag.isEmpty { return transportTag } + return AttachmentPreviewCodec.downloadTag(from: preview) + } + + // Explicit memberwise init (preserves all existing call sites). + init(id: String = "", preview: String = "", blob: String = "", + type: AttachmentType = .image, + transportTag: String = "", transportServer: String = "") { + self.id = id + self.preview = preview + self.blob = blob + self.type = type + self.transportTag = transportTag + self.transportServer = transportServer + } + + // Custom Decodable — `decodeIfPresent` for backward compat with old DB JSON + // that doesn't have transportTag/transportServer fields. + init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + id = try container.decodeIfPresent(String.self, forKey: .id) ?? "" + preview = try container.decodeIfPresent(String.self, forKey: .preview) ?? "" + blob = try container.decodeIfPresent(String.self, forKey: .blob) ?? "" + type = try container.decodeIfPresent(AttachmentType.self, forKey: .type) ?? .image + transportTag = try container.decodeIfPresent(String.self, forKey: .transportTag) ?? "" + transportServer = try container.decodeIfPresent(String.self, forKey: .transportServer) ?? "" + } } // MARK: - Reply Message Data (desktop parity) diff --git a/Rosetta/Core/Network/Protocol/Packets/PacketMessage.swift b/Rosetta/Core/Network/Protocol/Packets/PacketMessage.swift index f420019..8412fed 100644 --- a/Rosetta/Core/Network/Protocol/Packets/PacketMessage.swift +++ b/Rosetta/Core/Network/Protocol/Packets/PacketMessage.swift @@ -30,6 +30,8 @@ struct PacketMessage: Packet { stream.writeString(attachment.preview) stream.writeString(attachment.blob) stream.writeInt8(attachment.type.rawValue) + stream.writeString(attachment.transportTag) + stream.writeString(attachment.transportServer) } stream.writeString(aesChachaKey) } @@ -43,14 +45,22 @@ struct PacketMessage: Packet { privateKey = stream.readString() messageId = stream.readString() - let attachmentCount = stream.readInt8() + let attachmentCount = max(stream.readInt8(), 0) var list: [MessageAttachment] = [] for _ in 0.. Data { - guard bytes.isEmpty == false else { - return Data() - } - var data = Data(count: bytes.count) - data.replaceSubrange(0..> 3 + guard byteCount > 0 else { return Data() } + return Data(bytes[0..> 3 - let shift = 7 - (writePointer & 7) - bytes[byteIndex] = bytes[byteIndex] | (bit << shift) - writePointer += 1 + writeBits(value & 1, count: 1) } func readBit() -> Int { - guard readPointer < bytes.count * 8 else { - return 0 - } + guard remainingBits() >= 1 else { return 0 } let byteIndex = readPointer >> 3 let shift = 7 - (readPointer & 7) - let bit = (bytes[byteIndex] >> shift) & 1 + let bit = Int((bytes[byteIndex] >> shift) & 1) readPointer += 1 - return Int(bit) + return bit } // MARK: - Bool @@ -65,156 +62,266 @@ final class PacketBitStream: NSObject { readBit() == 1 } - // MARK: - Int8 (9 bits: 1 sign + 8 data) + // MARK: - Byte + + func writeByte(_ value: UInt8) { + writeUInt8(Int(value)) + } + + func readByte() -> UInt8 { + UInt8(truncatingIfNeeded: readUInt8()) + } + + // MARK: - UInt8 / Int8 (8 bits) + + func writeUInt8(_ value: Int) { + let v = UInt8(truncatingIfNeeded: value) + + // Fast path: byte-aligned + if (writePointer & 7) == 0 { + ensureCapacityForUpcomingBits(8) + bytes[writePointer >> 3] = v + writePointer += 8 + return + } + + writeBits(Int(v), count: 8) + } + + func readUInt8() -> Int { + guard remainingBits() >= 8 else { return 0 } + + // Fast path: byte-aligned + if (readPointer & 7) == 0 { + let value = Int(bytes[readPointer >> 3]) + readPointer += 8 + return value + } + + return readBits(8) + } func writeInt8(_ value: Int) { - let negationBit: UInt8 = value < 0 ? 1 : 0 - let int8Value = UInt8(abs(value) & 0xFF) - ensureCapacityForUpcomingBits(9) - - let byteIndex = writePointer >> 3 - let signShift = 7 - (writePointer & 7) - bytes[byteIndex] = bytes[byteIndex] | (negationBit << signShift) - writePointer += 1 - - for i in 0..<8 { - let bit = (int8Value >> (7 - i)) & 1 - let idx = writePointer >> 3 - let shift = 7 - (writePointer & 7) - bytes[idx] = bytes[idx] | (bit << shift) - writePointer += 1 - } + writeUInt8(value) } func readInt8() -> Int { - guard readPointer + 9 <= bytes.count * 8 else { - readPointer = bytes.count * 8 - return 0 - } - - var value = 0 - let signShift = 7 - (readPointer & 7) - let negationBit = Int((bytes[readPointer >> 3] >> signShift) & 1) - readPointer += 1 - - for i in 0..<8 { - let shift = 7 - (readPointer & 7) - let bit = Int((bytes[readPointer >> 3] >> shift) & 1) - value = value | (bit << (7 - i)) - readPointer += 1 - } - - return negationBit == 1 ? -value : value + Int(Int8(truncatingIfNeeded: readUInt8())) } - // MARK: - Int16 (2 × Int8) + // MARK: - UInt16 / Int16 (16 bits) + + func writeUInt16(_ value: Int) { + let v = value & 0xFFFF + writeUInt8((v >> 8) & 0xFF) + writeUInt8(v & 0xFF) + } + + func readUInt16() -> Int { + let hi = readUInt8() + let lo = readUInt8() + return (hi << 8) | lo + } func writeInt16(_ value: Int) { - writeInt8(value >> 8) - writeInt8(value & 0xFF) + writeUInt16(value) } func readInt16() -> Int { - let high = readInt8() << 8 - return high | readInt8() + Int(Int16(truncatingIfNeeded: readUInt16())) } - // MARK: - Int32 (2 × Int16) + // MARK: - UInt32 / Int32 (32 bits) + + func writeUInt32(_ value: Int) { + writeUInt8((value >> 24) & 0xFF) + writeUInt8((value >> 16) & 0xFF) + writeUInt8((value >> 8) & 0xFF) + writeUInt8(value & 0xFF) + } + + func readUInt32() -> Int { + let b1 = readUInt8() + let b2 = readUInt8() + let b3 = readUInt8() + let b4 = readUInt8() + return (b1 << 24) | (b2 << 16) | (b3 << 8) | b4 + } func writeInt32(_ value: Int) { - writeInt16(value >> 16) - writeInt16(value & 0xFFFF) + writeUInt32(value) } func readInt32() -> Int { - let high = readInt16() << 16 - return high | readInt16() + Int(Int32(truncatingIfNeeded: readUInt32())) } - // MARK: - Int64 (2 × Int32) + // MARK: - UInt64 / Int64 (64 bits) - func writeInt64(_ value: Int64) { - let high = Int((value >> 32) & 0xFFFFFFFF) - let low = Int(value & 0xFFFFFFFF) - writeInt32(high) - writeInt32(low) + func writeUInt64(_ value: Int64) { + writeUInt8(Int((value >> 56) & 0xFF)) + writeUInt8(Int((value >> 48) & 0xFF)) + writeUInt8(Int((value >> 40) & 0xFF)) + writeUInt8(Int((value >> 32) & 0xFF)) + writeUInt8(Int((value >> 24) & 0xFF)) + writeUInt8(Int((value >> 16) & 0xFF)) + writeUInt8(Int((value >> 8) & 0xFF)) + writeUInt8(Int(value & 0xFF)) } - func readInt64() -> Int64 { - let high = Int64(readInt32()) - let low = Int64(readInt32()) & 0xFFFFFFFF + func readUInt64() -> Int64 { + let high = Int64(readUInt32()) & 0xFFFFFFFF + let low = Int64(readUInt32()) & 0xFFFFFFFF return (high << 32) | low } - // MARK: - String (Int32 length + UTF-16 code units) + func writeInt64(_ value: Int64) { + writeUInt64(value) + } + + func readInt64() -> Int64 { + readUInt64() + } + + // MARK: - Float32 + + func writeFloat32(_ value: Float) { + writeInt32(Int(Int32(bitPattern: value.bitPattern))) + } + + func readFloat32() -> Float { + Float(bitPattern: UInt32(bitPattern: Int32(truncatingIfNeeded: readInt32()))) + } + + // MARK: - String (UInt32 length + UInt16 chars) func writeString(_ value: String) { let utf16Units = Array(value.utf16) - let requiredBits = 36 + utf16Units.count * 18 - ensureCapacityForUpcomingBits(requiredBits) - writeInt32(utf16Units.count) + let length = utf16Units.count + writeUInt32(length) + + guard length > 0 else { return } + + ensureCapacityForUpcomingBits(length * 16) for codeUnit in utf16Units { - writeInt16(Int(codeUnit)) + writeUInt16(Int(codeUnit)) } } func readString() -> String { - let length = readInt32() - let bitsAvailable = bytes.count * 8 - readPointer - let bytesAvailable = max(bitsAvailable / 8, 0) - if length < 0 || (length * 2) > bytesAvailable { - return "" - } + let length = readUInt32() + guard length > 0 else { return "" } + + let requiredBits = length * 16 + guard requiredBits <= remainingBits() else { return "" } var codeUnits = [UInt16]() codeUnits.reserveCapacity(length) for _ in 0.. 0 else { return } + + ensureCapacityForUpcomingBits(length * 8) + + // Fast path: byte-aligned + if (writePointer & 7) == 0 { + let byteIndex = writePointer >> 3 + for (i, byte) in value.enumerated() { + bytes[byteIndex + i] = byte + } + writePointer += length << 3 + return + } + for byte in value { - writeInt8(Int(byte)) + writeUInt8(Int(byte)) } } func readBytes() -> Data { - let length = readInt32() - guard length >= 0 else { - return Data() - } + let length = readUInt32() + guard length > 0 else { return Data() } - let bitsAvailable = bytes.count * 8 - readPointer - let bytesAvailable = max(bitsAvailable / 9, 0) - guard length <= bytesAvailable else { - return Data() + let requiredBits = length * 8 + guard requiredBits <= remainingBits() else { return Data() } + + // Fast path: byte-aligned + if (readPointer & 7) == 0 { + let byteIndex = readPointer >> 3 + let result = Data(bytes[byteIndex..<(byteIndex + length)]) + readPointer += length << 3 + return result } var result = Data(capacity: length) for _ in 0.. Int { + writePointer - readPointer + } + + private func writeBits(_ value: Int, count: Int) { + guard count > 0 else { return } + ensureCapacityForUpcomingBits(count) + + for i in stride(from: count - 1, through: 0, by: -1) { + let bit = UInt8((value >> i) & 1) + let byteIndex = writePointer >> 3 + let shift = 7 - (writePointer & 7) + if bit == 1 { + bytes[byteIndex] |= (1 << shift) + } else { + bytes[byteIndex] &= ~(1 << shift) + } + writePointer += 1 + } + } + + private func readBits(_ count: Int) -> Int { + guard count > 0, remainingBits() >= count else { return 0 } + + var value = 0 + for _ in 0..> 3 + let shift = 7 - (readPointer & 7) + let bit = Int((bytes[byteIndex] >> shift) & 1) + value = (value << 1) | bit + readPointer += 1 + } + return value + } + private func ensureCapacityForUpcomingBits(_ bitCount: Int) { guard bitCount > 0 else { return } let lastBitIndex = writePointer + bitCount - 1 ensureCapacity(lastBitIndex >> 3) } - private func ensureCapacity(_ index: Int) { - if bytes.count <= index { - bytes.append(contentsOf: repeatElement(0, count: index - bytes.count + 1)) + private func ensureCapacity(_ byteIndex: Int) { + let requiredSize = byteIndex + 1 + guard requiredSize > bytes.count else { return } + + var newSize = bytes.isEmpty ? 32 : bytes.count + while newSize < requiredSize { + newSize <<= 1 } + + bytes.append(contentsOf: repeatElement(0 as UInt8, count: newSize - bytes.count)) } } diff --git a/Rosetta/Core/Network/TransportManager.swift b/Rosetta/Core/Network/TransportManager.swift index 9ab8bf9..1229ca8 100644 --- a/Rosetta/Core/Network/TransportManager.swift +++ b/Rosetta/Core/Network/TransportManager.swift @@ -71,16 +71,16 @@ final class TransportManager: @unchecked Sendable { // MARK: - Upload /// Uploads file content to the transport server. - /// Desktop parity: `TransportProvider.tsx` `uploadFile()`. + /// Desktop parity: `usePrepareAttachment.ts` `uploadFile()`. /// /// - Parameters: /// - id: Unique file identifier (used as filename in multipart). /// - content: Raw file content to upload. - /// - Returns: Server-assigned tag for later download. + /// - Returns: `(tag, server)` — server-assigned tag and the transport server URL used. /// Android parity: retry with exponential backoff (1s, 2s, 4s) on upload failure. private static let maxUploadRetries = 3 - func uploadFile(id: String, content: Data) async throws -> String { + func uploadFile(id: String, content: Data) async throws -> (tag: String, server: String) { guard let serverUrl = await MainActor.run(body: { transportServer }) else { throw TransportError.noTransportServer } @@ -125,8 +125,8 @@ final class TransportManager: @unchecked Sendable { throw TransportError.missingTag } - Self.logger.info("Upload complete: id=\(id), tag=\(tag)") - return tag + Self.logger.info("Upload complete: id=\(id), tag=\(tag), server=\(serverUrl)") + return (tag: tag, server: serverUrl) } catch { lastError = error if attempt < Self.maxUploadRetries - 1 { @@ -141,14 +141,22 @@ final class TransportManager: @unchecked Sendable { // MARK: - Download - /// Downloads file content from the transport server. - /// Desktop parity: `TransportProvider.tsx` `downloadFile()`. + /// Downloads file content from a transport server. + /// Desktop parity: `useAttachment.ts` `downloadFile(id, tag, server)`. /// - /// - Parameter tag: Server-assigned file tag from upload response. + /// - Parameters: + /// - tag: Server-assigned file tag from upload response. + /// - server: Per-attachment transport server URL. Falls back to global transport if empty/nil. /// - Returns: Raw file content. - func downloadFile(tag: String) async throws -> Data { - guard let serverUrl = await MainActor.run(body: { transportServer }) else { - throw TransportError.noTransportServer + func downloadFile(tag: String, server: String? = nil) async throws -> Data { + let serverUrl: String + if let explicit = server, !explicit.isEmpty { + serverUrl = explicit + } else { + guard let global = await MainActor.run(body: { transportServer }) else { + throw TransportError.noTransportServer + } + serverUrl = global } guard let url = URL(string: "\(serverUrl)/d/\(tag)") else { diff --git a/Rosetta/Core/Services/ParityTestSupportInterfaces.swift b/Rosetta/Core/Services/ParityTestSupportInterfaces.swift index a4bd9c6..972a690 100644 --- a/Rosetta/Core/Services/ParityTestSupportInterfaces.swift +++ b/Rosetta/Core/Services/ParityTestSupportInterfaces.swift @@ -1,8 +1,8 @@ import Foundation protocol AttachmentFlowTransporting { - func uploadFile(id: String, content: Data) async throws -> String - func downloadFile(tag: String) async throws -> Data + func uploadFile(id: String, content: Data) async throws -> (tag: String, server: String) + func downloadFile(tag: String, server: String?) async throws -> Data } protocol PacketFlowSending { @@ -19,12 +19,12 @@ protocol SearchResultDispatching { } struct LiveAttachmentFlowTransport: AttachmentFlowTransporting { - func uploadFile(id: String, content: Data) async throws -> String { + func uploadFile(id: String, content: Data) async throws -> (tag: String, server: String) { try await TransportManager.shared.uploadFile(id: id, content: content) } - func downloadFile(tag: String) async throws -> Data { - try await TransportManager.shared.downloadFile(tag: tag) + func downloadFile(tag: String, server: String?) async throws -> Data { + try await TransportManager.shared.downloadFile(tag: tag, server: server) } } diff --git a/Rosetta/Core/Services/SessionManager.swift b/Rosetta/Core/Services/SessionManager.swift index eda8e00..b2cb1d2 100644 --- a/Rosetta/Core/Services/SessionManager.swift +++ b/Rosetta/Core/Services/SessionManager.swift @@ -409,9 +409,9 @@ final class SessionManager { DialogRepository.shared.updateDialogFromMessages(opponentKey: packet.toPublicKey) // Upload encrypted blob to transport server in background (desktop: uploadFile) - let tag: String + let upload: (tag: String, server: String) do { - tag = try await attachmentFlowTransport.uploadFile( + upload = try await attachmentFlowTransport.uploadFile( id: attachmentId, content: Data(encryptedBlob.utf8) ) @@ -424,13 +424,15 @@ final class SessionManager { } // Update preview with CDN tag (tag::blurhash) - let preview = "\(tag)::\(blurhash)" + let preview = "\(upload.tag)::\(blurhash)" packet.attachments = [ MessageAttachment( id: attachmentId, preview: preview, blob: "", // Desktop parity: blob cleared after upload - type: .avatar + type: .avatar, + transportTag: upload.tag, + transportServer: upload.server ), ] @@ -442,14 +444,14 @@ final class SessionManager { DialogRepository.shared.updateDeliveryStatus(messageId: messageId, opponentKey: toPublicKey, status: .delivered) // Send to server for multi-device sync (unlike text Saved Messages) packetFlowSender.sendPacket(packet) - Self.logger.info("📤 Avatar synced to Saved Messages (multi-device) tag=\(tag)") + Self.logger.info("📤 Avatar synced to Saved Messages (multi-device) tag=\(upload.tag)") return } packetFlowSender.sendPacket(packet) registerOutgoingRetry(for: packet) MessageRepository.shared.persistNow() - Self.logger.info("📤 Avatar sent to \(toPublicKey.prefix(12))… tag=\(tag)") + Self.logger.info("📤 Avatar sent to \(toPublicKey.prefix(12))… tag=\(upload.tag)") } /// Sends a message with image/file attachments. @@ -662,23 +664,27 @@ final class SessionManager { // ── Phase 2: Upload in background, then send packet ── let flowTransport = attachmentFlowTransport let messageAttachments: [MessageAttachment] = try await withThrowingTaskGroup( - of: (Int, String).self + of: (Int, String, String).self ) { group in for (index, item) in encryptedAttachments.enumerated() { group.addTask { - let tag = try await flowTransport.uploadFile( + let result = try await flowTransport.uploadFile( id: item.original.id, content: item.encryptedData ) - return (index, tag) + return (index, result.tag, result.server) } } - var tags = [Int: String]() - for try await (index, tag) in group { tags[index] = tag } + var uploads = [Int: (tag: String, server: String)]() + for try await (index, tag, server) in group { uploads[index] = (tag, server) } return encryptedAttachments.enumerated().map { index, item in - let tag = tags[index] ?? "" - let preview = AttachmentPreviewCodec.compose(downloadTag: tag, payload: item.preview) - Self.logger.info("📤 Attachment uploaded: type=\(String(describing: item.original.type)), tag=\(tag)") - return MessageAttachment(id: item.original.id, preview: preview, blob: "", type: item.original.type) + let upload = uploads[index] ?? (tag: "", server: "") + let preview = AttachmentPreviewCodec.compose(downloadTag: upload.tag, payload: item.preview) + Self.logger.info("📤 Attachment uploaded: type=\(String(describing: item.original.type)), tag=\(upload.tag), server=\(upload.server)") + return MessageAttachment( + id: item.original.id, preview: preview, blob: "", + type: item.original.type, + transportTag: upload.tag, transportServer: upload.server + ) } } @@ -789,7 +795,7 @@ final class SessionManager { Self.logger.debug("📤 Forward re-upload: \(originalId) → \(newAttId) (\(jpegData.count) bytes JPEG, \(encryptedBlob.count) encrypted)") #endif - let tag = try await attachmentFlowTransport.uploadFile( + let upload = try await attachmentFlowTransport.uploadFile( id: newAttId, content: Data(encryptedBlob.utf8) ) @@ -799,7 +805,7 @@ final class SessionManager { .flatMap { $0.attachments } .first(where: { $0.id == originalId })?.preview ?? "" let blurhash = AttachmentPreviewCodec.blurHash(from: originalPreview) - let newPreview = AttachmentPreviewCodec.compose(downloadTag: tag, payload: blurhash) + let newPreview = AttachmentPreviewCodec.compose(downloadTag: upload.tag, payload: blurhash) attachmentIdMap[originalId] = (newAttId, newPreview) // Cache locally under new ID for ForwardedImagePreviewCell @@ -808,7 +814,7 @@ final class SessionManager { } #if DEBUG - Self.logger.debug("📤 Forward re-upload OK: \(newAttId) tag=\(tag) preview=\(newPreview)") + Self.logger.debug("📤 Forward re-upload OK: \(newAttId) tag=\(upload.tag) preview=\(newPreview)") #endif } } @@ -831,7 +837,7 @@ final class SessionManager { Self.logger.debug("📤 Forward file re-upload: \(originalId) → \(newAttId) (\(fileInfo.data.count) bytes, \(fileInfo.fileName))") #endif - let tag = try await attachmentFlowTransport.uploadFile( + let upload = try await attachmentFlowTransport.uploadFile( id: newAttId, content: Data(encryptedBlob.utf8) ) @@ -846,11 +852,11 @@ final class SessionManager { fallbackFileSize: fileInfo.data.count ) let fileMeta = "\(filePreview.fileSize)::\(filePreview.fileName)" - let newPreview = AttachmentPreviewCodec.compose(downloadTag: tag, payload: fileMeta) + let newPreview = AttachmentPreviewCodec.compose(downloadTag: upload.tag, payload: fileMeta) attachmentIdMap[originalId] = (newAttId, newPreview) #if DEBUG - Self.logger.debug("📤 Forward file re-upload OK: \(newAttId) tag=\(tag) preview=\(newPreview)") + Self.logger.debug("📤 Forward file re-upload OK: \(newAttId) tag=\(upload.tag) preview=\(newPreview)") #endif } } diff --git a/Rosetta/Features/Chats/ChatDetail/MessageAvatarView.swift b/Rosetta/Features/Chats/ChatDetail/MessageAvatarView.swift index ef90498..e938d51 100644 --- a/Rosetta/Features/Chats/ChatDetail/MessageAvatarView.swift +++ b/Rosetta/Features/Chats/ChatDetail/MessageAvatarView.swift @@ -233,7 +233,7 @@ struct MessageAvatarView: View { private func downloadAvatar() { guard !isDownloading else { return } - let tag = extractTag(from: attachment.preview) + let tag = attachment.effectiveDownloadTag guard !tag.isEmpty else { downloadError = true return @@ -247,9 +247,10 @@ struct MessageAvatarView: View { isDownloading = true downloadError = false + let server = attachment.transportServer Task { do { - let encryptedData = try await TransportManager.shared.downloadFile(tag: tag) + let encryptedData = try await TransportManager.shared.downloadFile(tag: tag, server: server) let encryptedString = String(decoding: encryptedData, as: UTF8.self) let passwords = MessageCrypto.attachmentPasswordCandidates(from: storedPassword) diff --git a/Rosetta/Features/Chats/ChatDetail/MessageFileView.swift b/Rosetta/Features/Chats/ChatDetail/MessageFileView.swift index fff3000..d7ddb8f 100644 --- a/Rosetta/Features/Chats/ChatDetail/MessageFileView.swift +++ b/Rosetta/Features/Chats/ChatDetail/MessageFileView.swift @@ -98,7 +98,10 @@ struct MessageFileView: View { private var fileName: String { fileMetadata.name } private var fileSize: Int { fileMetadata.size } - private var fileTag: String { fileMetadata.tag } + private var fileTag: String { + let effective = attachment.effectiveDownloadTag + return effective.isEmpty ? fileMetadata.tag : effective + } private var formattedFileSize: String { let bytes = fileSize @@ -140,7 +143,7 @@ struct MessageFileView: View { Task { do { - let encryptedData = try await TransportManager.shared.downloadFile(tag: fileTag) + let encryptedData = try await TransportManager.shared.downloadFile(tag: fileTag, server: attachment.transportServer) let encryptedString = String(decoding: encryptedData, as: UTF8.self) let passwords = MessageCrypto.attachmentPasswordCandidates(from: storedPassword) diff --git a/Rosetta/Features/Chats/ChatDetail/MessageImageView.swift b/Rosetta/Features/Chats/ChatDetail/MessageImageView.swift index 178fe8b..09c2e5a 100644 --- a/Rosetta/Features/Chats/ChatDetail/MessageImageView.swift +++ b/Rosetta/Features/Chats/ChatDetail/MessageImageView.swift @@ -251,7 +251,7 @@ struct MessageImageView: View { private func downloadImage() { guard !isDownloading, image == nil else { return } - let tag = extractTag(from: attachment.preview) + let tag = attachment.effectiveDownloadTag guard !tag.isEmpty else { downloadError = true return @@ -265,9 +265,10 @@ struct MessageImageView: View { isDownloading = true downloadError = false + let server = attachment.transportServer Task { do { - let encryptedData = try await TransportManager.shared.downloadFile(tag: tag) + let encryptedData = try await TransportManager.shared.downloadFile(tag: tag, server: server) let encryptedString = String(decoding: encryptedData, as: UTF8.self) let passwords = MessageCrypto.attachmentPasswordCandidates(from: storedPassword) diff --git a/Rosetta/Features/Chats/ChatDetail/NativeMessageCell.swift b/Rosetta/Features/Chats/ChatDetail/NativeMessageCell.swift index bb7edad..4f75b04 100644 --- a/Rosetta/Features/Chats/ChatDetail/NativeMessageCell.swift +++ b/Rosetta/Features/Chats/ChatDetail/NativeMessageCell.swift @@ -901,10 +901,11 @@ final class NativeMessageCell: UICollectionViewCell, UIContextMenuInteractionDel /// Downloads avatar from CDN, decrypts, caches to disk, and returns the image. /// Shared logic with `MessageAvatarView.downloadAvatar()`. private static func downloadAndCacheAvatar( - tag: String, attachmentId: String, storedPassword: String, senderKey: String + tag: String, attachmentId: String, storedPassword: String, senderKey: String, + server: String = "" ) async -> UIImage? { do { - let encryptedData = try await TransportManager.shared.downloadFile(tag: tag) + let encryptedData = try await TransportManager.shared.downloadFile(tag: tag, server: server) let encryptedString = String(decoding: encryptedData, as: UTF8.self) let passwords = MessageCrypto.attachmentPasswordCandidates(from: storedPassword) @@ -1109,7 +1110,7 @@ final class NativeMessageCell: UICollectionViewCell, UIContextMenuInteractionDel // Already downloaded? if AttachmentCache.shared.cachedImage(forAttachmentId: id) != nil { return } // Download from CDN - let tag = AttachmentPreviewCodec.downloadTag(from: avatarAtt.preview) + let tag = avatarAtt.effectiveDownloadTag guard !tag.isEmpty else { return } guard let password = message.attachmentPassword, !password.isEmpty else { return } @@ -1118,10 +1119,12 @@ final class NativeMessageCell: UICollectionViewCell, UIContextMenuInteractionDel let messageId = message.id let senderKey = message.fromPublicKey + let server = avatarAtt.transportServer Task.detached(priority: .userInitiated) { let downloaded = await Self.downloadAndCacheAvatar( tag: tag, attachmentId: id, - storedPassword: password, senderKey: senderKey + storedPassword: password, senderKey: senderKey, + server: server ) await MainActor.run { [weak self] in guard let self, self.message?.id == messageId else { return } @@ -1569,7 +1572,7 @@ final class NativeMessageCell: UICollectionViewCell, UIContextMenuInteractionDel private func downloadPhotoAttachment(attachment: MessageAttachment, message: ChatMessage) { if photoDownloadTasks[attachment.id] != nil { return } - let tag = Self.extractTag(from: attachment.preview) + let tag = attachment.effectiveDownloadTag guard !tag.isEmpty, let storedPassword = message.attachmentPassword, !storedPassword.isEmpty else { @@ -1593,9 +1596,10 @@ final class NativeMessageCell: UICollectionViewCell, UIContextMenuInteractionDel photoTileDownloadArrows[tileIndex].isHidden = true } + let server = attachment.transportServer photoDownloadTasks[attachmentId] = Task { [weak self] in do { - let encryptedData = try await TransportManager.shared.downloadFile(tag: tag) + let encryptedData = try await TransportManager.shared.downloadFile(tag: tag, server: server) let encryptedString = String(decoding: encryptedData, as: UTF8.self) let passwords = MessageCrypto.attachmentPasswordCandidates(from: storedPassword) let image = Self.decryptAndParseImage(encryptedString: encryptedString, passwords: passwords) diff --git a/RosettaTests/AttachmentParityTests.swift b/RosettaTests/AttachmentParityTests.swift index 4bf7afb..2d516e7 100644 --- a/RosettaTests/AttachmentParityTests.swift +++ b/RosettaTests/AttachmentParityTests.swift @@ -160,12 +160,13 @@ private final class MockAttachmentFlowTransport: AttachmentFlowTransporting { var tagsById: [String: String] = [:] private(set) var uploadedIds: [String] = [] - func uploadFile(id: String, content: Data) async throws -> String { + func uploadFile(id: String, content: Data) async throws -> (tag: String, server: String) { uploadedIds.append(id) - return tagsById[id] ?? UUID().uuidString.lowercased() + let tag = tagsById[id] ?? UUID().uuidString.lowercased() + return (tag: tag, server: "https://mock-transport.test") } - func downloadFile(tag: String) async throws -> Data { + func downloadFile(tag: String, server: String?) async throws -> Data { Data() } }