Musubi/Packages/MusubiKit/Sources/Intercom/RTC/PeerSession.swift

import Foundation

/// One phone: ICE-lite credentials, DTLS, SRTP and the two outgoing RTP streams.
/// Owned and called by `IntercomEngine` on its queue only.
final class PeerSession {
    struct Offer {
        let answer: String
        let session: PeerSession
    }

    enum Kind { case stun, dtls, rtcp, other }

    let id = UUID()
    let name: String
    let role: String
    let localUfrag: String
    let localPassword: String
    let remoteUfrag: String
    let created = Date()

    private(set) var lastConsent: Date?
    private(set) var connectedSince: Date?
    private var dtls: DTLSServer
    private var srtp: SRTPContext?
    private var audio: RTPStream?
    private var video: RTPStream?
    private var lastAudioTimestamp: UInt32 = 0
    private var lastVideoTimestamp: UInt32 = 0
    private var reportsDue = Date()

    var isMediaReady: Bool { srtp != nil }
    var dtlsState: DTLSServer.State { dtls.state }

    private init(name: String, role: String, remoteUfrag: String, certificate: DTLSCertificate, remoteFingerprint: String) {
        self.name = name
        self.role = role
        self.remoteUfrag = remoteUfrag
        localUfrag = Self.token(8)
        localPassword = Self.token(24)
        dtls = DTLSServer(certificate: certificate, expectedFingerprint: remoteFingerprint)
    }

    /// Accepts a browser offer and returns the answer; nil if the offer has nothing Musubi can send.
    static func accept(offerSDP: String, name: String, role: String, certificate: DTLSCertificate,
                       address: String, port: UInt16) throws -> Offer? {
        let offer = try SDPOffer(offerSDP)
        let session = PeerSession(name: name, role: role, remoteUfrag: offer.iceUfrag, certificate: certificate,
                                  remoteFingerprint: offer.fingerprint)
        var tracks: [SDPAnswer.Track] = []
        for media in offer.media {
            let ssrc = UInt32.random(in: 1...UInt32.max)
            var payloadType: Int?
            if media.port != 0 {
                if media.kind == "audio", session.audio == nil, let pt = offer.opusPayloadType(in: media) {
                    payloadType = pt
                    session.audio = RTPStream(ssrc: ssrc, payloadType: UInt8(pt))
                } else if media.kind == "video", session.video == nil, let pt = offer.h264PayloadType(in: media) {
                    payloadType = pt
                    session.video = RTPStream(ssrc: ssrc, payloadType: UInt8(pt))
                }
            }
            tracks.append(.init(mid: media.mid, kind: media.kind, payloadType: payloadType, ssrc: ssrc))
        }
        guard session.audio != nil || session.video != nil else { return nil }
        let answer = SDPAnswer(iceUfrag: session.localUfrag, icePwd: session.localPassword,
                               fingerprint: certificate.fingerprint, address: address, port: port, tracks: tracks)
        return Offer(answer: answer.text, session: session)
    }

    static func classify(_ packet: [UInt8]) -> Kind {
        if STUN.looksLikeSTUN(packet) { return .stun }
        if DTLSServer.looksLikeDTLS(packet) { return .dtls }
        if RTCPFeedback.isRTCP(packet) { return .rtcp }
        return .other
    }

    /// Answers a connectivity check if it carries this session's credentials.
    func answerSTUN(_ packet: [UInt8], message: STUN.Message, address: [UInt8], port: UInt16) -> [UInt8]? {
        guard message.type == STUN.bindingRequest, message.username == "\(localUfrag):\(remoteUfrag)",
              STUN.verify(packet, message: message, key: Array(localPassword.utf8)) else { return nil }
        lastConsent = Date()
        return STUN.bindingResponse(transactionID: message.transactionID, address: address, port: port,
                                    key: Array(localPassword.utf8))
    }

    /// Feeds DTLS; returns datagrams to send and whether media just became ready.
    func receiveDTLS(_ packet: [UInt8]) -> (replies: [[UInt8]], becameReady: Bool) {
        let replies = dtls.handle(packet)
        guard srtp == nil, let keys = dtls.keyMaterial else { return (replies, false) }
        // DTLS server = answerer: we send with the server keys.
        srtp = SRTPContext(masterKey: keys.serverKey, masterSalt: keys.serverSalt)
        connectedSince = Date()
        return (replies, true)
    }

    /// True if the browser asks for a key frame.
    func receiveRTCP(_ packet: [UInt8]) -> Bool {
        guard let srtpIn = inboundRTCP, let plain = srtpIn.unprotectRTCP(packet) else { return false }
        return RTCPFeedback.requestsKeyFrame(plain)
    }

    private lazy var inboundRTCP: SRTPContext? = {
        guard let keys = dtls.keyMaterial else { return nil }
        return SRTPContext(masterKey: keys.clientKey, masterSalt: keys.clientSalt)
    }()

    func audioPacket(opus: [UInt8], timestamp: UInt32) -> [UInt8]? {
        guard let srtp, audio != nil else { return nil }
        lastAudioTimestamp = timestamp
        return srtp.protectRTP(audio!.packet(timestamp: timestamp, marker: false, payload: opus))
    }

    func videoPackets(payloads: [[UInt8]], timestamp: UInt32) -> [[UInt8]] {
        guard let srtp, video != nil else { return [] }
        lastVideoTimestamp = timestamp
        return payloads.enumerated().map { index, payload in
            srtp.protectRTP(video!.packet(timestamp: timestamp, marker: index == payloads.count - 1, payload: payload))
        }
    }

    /// Sender reports roughly once per second.
    func senderReports(now: Date = Date()) -> [[UInt8]] {
        guard let srtp, now >= reportsDue else { return [] }
        reportsDue = now.addingTimeInterval(1)
        var out: [[UInt8]] = []
        if let audio {
            out.append(srtp.protectRTCP(RTCPFeedback.senderReport(ssrc: audio.ssrc, rtpTimestamp: lastAudioTimestamp,
                                                                    packets: audio.packetCount, octets: audio.octetCount, now: now)))
        }
        if let video {
            out.append(srtp.protectRTCP(RTCPFeedback.senderReport(ssrc: video.ssrc, rtpTimestamp: lastVideoTimestamp,
                                                                    packets: video.packetCount, octets: video.octetCount, now: now)))
        }
        return out
    }

    /// No consent for 30 s (RFC 7675), or never connected within 30 s, or DTLS ended.
    func isExpired(now: Date = Date()) -> Bool {
        switch dtls.state {
        case .closed, .failed: return true
        default: break
        }
        return now.timeIntervalSince(lastConsent ?? created) > 30
    }

    private static func token(_ length: Int) -> String {
        let alphabet = Array("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789")
        return String((0..<length).map { _ in alphabet.randomElement()! })
    }
}