Gemini/GeminiProtocol/GeminiProtocol.swift

//
//  GeminiProtocol.swift
//  Gemini
//
//  Created by Shadowfacts on 7/12/20.
//

import Network

class GeminiProtocol: NWProtocolFramerImplementation {
    static let definition = NWProtocolFramer.Definition(implementation: GeminiProtocol.self)
    
    static let label = "Gemini"
    
    private var tempStatusCode: GeminiResponseHeader.StatusCode?
    private var tempMeta: String?
    private var lastAttemptedMetaLength: Int?
    private var lastFoundCR = false
    
    required init(framer: NWProtocolFramer.Instance) {
    }
    
    func start(framer: NWProtocolFramer.Instance) -> NWProtocolFramer.StartResult {
        return .ready
    }
    
    func wakeup(framer: NWProtocolFramer.Instance) {
    }
    
    func stop(framer: NWProtocolFramer.Instance) -> Bool {
        return true
    }
    
    func cleanup(framer: NWProtocolFramer.Instance) {
    }
    
    func handleInput(framer: NWProtocolFramer.Instance) -> Int {
        if tempStatusCode == nil {
            _ = framer.parseInput(minimumIncompleteLength: 3, maximumLength: 3) { (buffer, isComplete) -> Int in
                guard let buffer = buffer,
                      buffer.count == 3 else { return 0 }
                self.tempStatusCode = GeminiResponseHeader.StatusCode(buffer)
                return 3
            }
        }
        guard let statusCode = tempStatusCode else {
            return 3
        }
        
        if tempMeta == nil {
            let min: Int
            // if we previously tried to get the meta but failed (because the <CR><LF> was not found,
            // the minimum amount we need before trying to parse is at least 1 or 2 (depending on whether we found the <CR>) bytes more
            if let lastAttemptedMetaLength = lastAttemptedMetaLength {
                min = lastAttemptedMetaLength + (lastFoundCR ? 1 : 2)
            } else {
                // Minimum length is 2 bytes, spec does not say meta string is required
                min = 2
            }
            _ = framer.parseInput(minimumIncompleteLength: min, maximumLength: 1024 + 2) { (buffer, isComplete) -> Int in
                guard let buffer = buffer,
                      buffer.count >= 2 else { return 0 }
                self.lastAttemptedMetaLength = buffer.count
                
                let lastPossibleCRIndex = buffer.index(before: buffer.index(before: buffer.endIndex))
                var index = buffer.startIndex
                var found = false
                while index <= lastPossibleCRIndex {
                    // <CR><LF>
                    if buffer[index] == 13 && buffer[buffer.index(after: index)] == 10 {
                        found = true
                        break
                    }
                    index = buffer.index(after: index)
                }
                
                if !found {
                    if buffer[index] == 13 {
                        // if we found <CR>, but not <LF>, save that info so that next time we only wait for 1 more byte instead of 2
                        self.lastFoundCR = true
                    }
                    if buffer.count < 1026 {
                        return 0
                    } else {
                        fatalError("Didn't find <CR><LF> in buffer. Meta string was longer than 1024 bytes")
                    }
                }
                
                self.tempMeta = String(bytes: buffer[..<index], encoding: .utf8)
                return buffer.startIndex.distance(to: index) + 2
            }
        }
        guard let meta = tempMeta else {
            if let attempted = self.lastAttemptedMetaLength {
                return attempted + (lastFoundCR ? 1 : 2)
            } else {
                return 2
            }
        }
        
        let header = GeminiResponseHeader(status: statusCode, meta: meta)
        
        let message = NWProtocolFramer.Message(geminiResponseHeader: header)
        // Deliver all the input (the response body) to the client without copying.
        _ = framer.deliverInputNoCopy(length: statusCode.isSuccess ? .max : 0, message: message, isComplete: true)
        // Just in case, set the framer to pass-through input so it never invokes this method again.
        // todo: this should work according to an apple engineer, but the request seems to hang forever on a real device
        // sometimes works fine when stepping through w/ debugger => race condition?
//        framer.passThroughInput()
        return 0
    }
    
    func handleOutput(framer: NWProtocolFramer.Instance, message: NWProtocolFramer.Message, messageLength: Int, isComplete: Bool) {
        guard let request = message.geminiRequest else { fatalError("GeminiProtocol can't send message that doesn't have an associated GeminiRequest") }
        framer.writeOutput(data: request.data)
    }
    
}

fileprivate extension GeminiResponseHeader.StatusCode {
    init?(_ buffer: UnsafeMutableRawBufferPointer) {
        guard let str = String(bytes: buffer[...buffer.index(after: buffer.startIndex)], encoding: .utf8),
              let value = Int(str, radix: 10) else { return nil }
        self.init(rawValue: value)
    }
}