diff --git a/bitchat/Protocols/Packets.swift b/bitchat/Protocols/Packets.swift index 8914d4b78..4cd727501 100644 --- a/bitchat/Protocols/Packets.swift +++ b/bitchat/Protocols/Packets.swift @@ -114,20 +114,19 @@ struct PrivateMessagePacket { } func encode() -> Data? { + guard let messageIDData = messageID.data(using: .utf8), messageIDData.count <= 255 else { return nil } + guard let contentData = content.data(using: .utf8), contentData.count <= 65535 else { return nil } + var data = Data() - data.reserveCapacity(2 + min(messageID.count, 255) + 2 + min(content.count, 255)) + let estimatedMessageLength = 1 + 1 + messageIDData.count + let estimatedContentLength = contentData.count >= 255 ? 1 + 3 + contentData.count : 1 + 2 + contentData.count + data.reserveCapacity(estimatedMessageLength + estimatedContentLength) - // TLV for messageID - guard let messageIDData = messageID.data(using: .utf8), messageIDData.count <= 255 else { return nil } data.append(TLVType.messageID.rawValue) data.append(UInt8(messageIDData.count)) data.append(messageIDData) - // TLV for content - guard let contentData = content.data(using: .utf8), contentData.count <= 255 else { return nil } - data.append(TLVType.content.rawValue) - data.append(UInt8(contentData.count)) - data.append(contentData) + appendContentTLV(value: contentData, into: &data) return data } @@ -141,17 +140,20 @@ struct PrivateMessagePacket { guard let type = TLVType(rawValue: data[offset]) else { return nil } offset += 1 - let length = Int(data[offset]) - offset += 1 - - guard offset + length <= data.count else { return nil } - let value = data[offset.. Int? { + guard offset < data.count else { return nil } + let firstByte = Int(data[offset]) + offset += 1 + + if firstByte < 255 { + return firstByte + } + + guard offset + 2 <= data.count else { return nil } + let high = Int(data[offset]) + let low = Int(data[offset + 1]) + offset += 2 + + let length = (high << 8) | low + return length <= 65535 ? length : nil + } } diff --git a/bitchatTests/Protocol/PrivateMessagePacketTests.swift b/bitchatTests/Protocol/PrivateMessagePacketTests.swift new file mode 100644 index 000000000..96d37f2e1 --- /dev/null +++ b/bitchatTests/Protocol/PrivateMessagePacketTests.swift @@ -0,0 +1,40 @@ +import XCTest +@testable import bitchat + +final class PrivateMessagePacketTests: XCTestCase { + func testEncodeDecodeSupportsLargeContent() throws { + let longContent = String(repeating: "A", count: 1024) + let packet = PrivateMessagePacket(messageID: "msg-123", content: longContent) + + let encoded = try XCTUnwrap(packet.encode()) + let decoded = try XCTUnwrap(PrivateMessagePacket.decode(from: encoded)) + + XCTAssertEqual(decoded.messageID, "msg-123") + XCTAssertEqual(decoded.content, longContent) + } + + func testEncodeDecodeContentExactly255Bytes() throws { + let content = String(repeating: "B", count: 255) + let packet = PrivateMessagePacket(messageID: "msg-255", content: content) + + let encoded = try XCTUnwrap(packet.encode()) + let decoded = try XCTUnwrap(PrivateMessagePacket.decode(from: encoded)) + + XCTAssertEqual(decoded.content.count, 255) + XCTAssertEqual(decoded.content, content) + } + + func testEncodeRejectsOversizedContent() { + let oversizedContent = String(repeating: "C", count: 70_000) + let packet = PrivateMessagePacket(messageID: "msg-oversize", content: oversizedContent) + + XCTAssertNil(packet.encode()) + } + + func testEncodeRejectsOversizedMessageID() { + let longID = String(repeating: "x", count: 256) + let packet = PrivateMessagePacket(messageID: longID, content: "ok") + + XCTAssertNil(packet.encode()) + } +}