From 2c023bc02e495d81cc1b7da0bd170b1eb7059b03 Mon Sep 17 00:00:00 2001
From: Shadowfacts <me@shadowfacts.net>
Date: Thu, 23 Jul 2020 09:18:59 -0400
Subject: [PATCH] Minor protocol fixes

---
 Gemini/BrowserWindowController.swift |  2 +-
 GeminiProtocol/GeminiDataTask.swift  | 15 ++++++--
 GeminiProtocol/GeminiProtocol.swift  | 54 ++++++++++++++++++----------
 3 files changed, 48 insertions(+), 23 deletions(-)

diff --git a/Gemini/BrowserWindowController.swift b/Gemini/BrowserWindowController.swift
index 9010993..c886bdf 100644
--- a/Gemini/BrowserWindowController.swift
+++ b/Gemini/BrowserWindowController.swift
@@ -29,7 +29,7 @@ class BrowserWindowController: NSWindowController {
         
         contentViewController = NSHostingController(rootView: ContentView(navigator: navigator))
         
-        urlUpdater = navigator.$currentURL.sink(receiveValue: self.currentURLChanged(_:))
+        urlUpdater = navigator.$currentURL.receive(on: DispatchQueue.main).sink(receiveValue: self.currentURLChanged(_:))
     }
     
     private func currentURLChanged(_ newValue: URL) {
diff --git a/GeminiProtocol/GeminiDataTask.swift b/GeminiProtocol/GeminiDataTask.swift
index eb7c544..1442252 100644
--- a/GeminiProtocol/GeminiDataTask.swift
+++ b/GeminiProtocol/GeminiDataTask.swift
@@ -21,13 +21,16 @@ public class GeminiDataTask {
     public init(request: GeminiRequest, completion: @escaping Completion) {
         self.request = request
         self.completion = completion
-        self.connection = NWConnection(to: .url(request.url), using: .gemini)
+        let port = request.url.port != nil ? UInt16(request.url.port!) : 1965
+        let endpoint: NWEndpoint = .hostPort(host: NWEndpoint.Host(request.url.host!), port: NWEndpoint.Port(rawValue: port)!)
+        self.connection = NWConnection(to: endpoint, using: .gemini)
         
         self.connection.stateUpdateHandler = { (newState) in
             switch newState {
             case .ready:
                 self.sendRequest()
             case let .failed(error):
+                self.state = .completed
                 self.connection.cancel()
                 self.completion(.failure(.connectionError(error)))
             default:
@@ -47,6 +50,7 @@ public class GeminiDataTask {
     
     public func resume() {
         guard state == .unstarted else { return }
+        state = .started
         connection.start(queue: GeminiDataTask.queue)
     }
     
@@ -59,8 +63,13 @@ public class GeminiDataTask {
     private func sendRequest() {
         let message = NWProtocolFramer.Message(geminiRequest: request)
         let context = NWConnection.ContentContext(identifier: "GeminiRequest", metadata: [message])
-        connection.send(content: nil, contentContext: context, isComplete: true, completion: .contentProcessed({ (_) in }))
-        state = .started
+        connection.send(content: nil, contentContext: context, isComplete: true, completion: .contentProcessed({ (error) in
+            if let error = error {
+                self.state = .completed
+                self.connection.cancel()
+                self.completion(.failure(.connectionError(error)))
+            }
+        }))
         receive()
     }
     
diff --git a/GeminiProtocol/GeminiProtocol.swift b/GeminiProtocol/GeminiProtocol.swift
index 32ae41a..2957a1f 100644
--- a/GeminiProtocol/GeminiProtocol.swift
+++ b/GeminiProtocol/GeminiProtocol.swift
@@ -12,6 +12,9 @@ class GeminiProtocol: NWProtocolFramerImplementation {
     
     static let label = "Gemini"
     
+    private var tempStatusCode: GeminiResponseHeader.StatusCode?
+    private var tempMeta: String?
+    
     required init(framer: NWProtocolFramer.Instance) {
     }
     
@@ -30,23 +33,25 @@ class GeminiProtocol: NWProtocolFramerImplementation {
     }
     
     func handleInput(framer: NWProtocolFramer.Instance) -> Int {
-        while true {
-            var tempStatusCode: GeminiResponseHeader.StatusCode?
-            let parsedStatusCodeAndSpace = framer.parseInput(minimumIncompleteLength: 3, maximumLength: 3) { (buffer, isComplete) -> Int in
+        if tempStatusCode == nil {
+            _ = framer.parseInput(minimumIncompleteLength: 3, maximumLength: 3) { (buffer, isComplete) -> Int in
                 guard let buffer = buffer,
                       buffer.count == 3 else { return 0 }
-                tempStatusCode = GeminiResponseHeader.StatusCode(buffer)
+                self.tempStatusCode = GeminiResponseHeader.StatusCode(buffer)
                 return 3
             }
-            guard parsedStatusCodeAndSpace, let statusCode = tempStatusCode else {
-                return 3
-            }
-            
-            var tempMeta: String?
+        }
+        guard let statusCode = tempStatusCode else {
+            return 3
+        }
+        
+        var attemptedMetaLength: Int?
+        if tempMeta == nil {
             // Minimum length is 2 bytes, spec does not say meta string is required
-            let parsedMeta = framer.parseInput(minimumIncompleteLength: 2, maximumLength: 1024 + 2) { (buffer, isComplete) -> Int in
+            _ = framer.parseInput(minimumIncompleteLength: 2, maximumLength: 1024 + 2) { (buffer, isComplete) -> Int in
                 guard let buffer = buffer,
                       buffer.count >= 2 else { return 0 }
+                attemptedMetaLength = buffer.count
                 
                 let lastPossibleCRIndex = buffer.index(before: buffer.index(before: buffer.endIndex))
                 var index = buffer.startIndex
@@ -60,19 +65,30 @@ class GeminiProtocol: NWProtocolFramerImplementation {
                     index = buffer.index(after: index)
                 }
                 
-                guard found else { fatalError("Didn't find <CR><LF> in buffer. Meta string was longer than 1024 bytes") }
+                if !found {
+                    if buffer.count < 1026 {
+                        return 0
+                    } else {
+                        fatalError("Didn't find <CR><LF> in buffer. Meta string was longer than 1024 bytes")
+                    }
+                }
                 
-                tempMeta = String(bytes: buffer[..<index], encoding: .utf8)
-                return buffer.startIndex.distance(to: index)
+                self.tempMeta = String(bytes: buffer[..<index], encoding: .utf8)
+                return buffer.startIndex.distance(to: index) + 2
             }
-            guard parsedMeta, let meta = tempMeta else {
-                // todo: what should actually be returned here? the minimum number of bytes necessary?
+        }
+        guard let meta = tempMeta else {
+            if let attempted = attemptedMetaLength {
+                return attempted + 1
+            } else {
                 return 2
             }
-            
-            let header = GeminiResponseHeader(status: statusCode, meta: meta)
-            
-            let message = NWProtocolFramer.Message(geminiResponseHeader: header)
+        }
+        
+        let header = GeminiResponseHeader(status: statusCode, meta: meta)
+        
+        let message = NWProtocolFramer.Message(geminiResponseHeader: header)
+        while true {
             if !framer.deliverInputNoCopy(length: .max, message: message, isComplete: true) {
                 return 0
             }