client: Support read timeouts on receive & close.

* web/socket/client.scm (close-websocket): Add timeout parameter.  Call
read-close-frame with timeout.
(websocket-receive): Add timeout parameter.  Call read-data-frame with
timeout.  Don't close the websocket on timeout.
* web/socket/frame.scm (port-readable?, port-writable?, frame-reader):
New helper procedures.
(read-close-frame): New procedure.
(read-data-frame): Add timeout parameter.  Use frame-reader to support a
composed read timeout.

2 files changed, 82 insertions, 22 deletions

diff --git a/web/socket/client.scm b/web/socket/client.scm
index c867427..4510493 100644
--- a/web/socket/client.scm
+++ b/web/socket/client.scm
@@ -183,19 +183,15 @@ true, verify HTTPS server certificates."
           ws)
         (error "not a websocket uri" uri))))
 
-(define (close-websocket ws)
-  "Close the WebSocket connection for the client WS."
+(define* (close-websocket ws #:key (timeout 1))
+  "Close the WebSocket connection for the client WS, waiting up to
+TIMEOUT seconds (default 1 second) for a graceful shutdown."
   (let ((socket (websocket-socket ws)))
     (set-websocket-state! ws 'closing)
     (write-frame (make-close-frame (make-bytevector 0)) socket)
     ;; Per section 5.5.1 , wait for the server to close the connection
     ;; for a reasonable amount of time.
-    (let loop ()
-      (match (select #() (vector socket) #() 1) ; 1 second timeout
-        ((#() #(socket) #()) ; there is output to read
-         (unless (port-eof? socket)
-           (read-frame socket) ; throw it away
-           (loop)))))
+    (read-close-frame socket #:timeout timeout)
     (close-port socket)
     (close-port (websocket-entropy-port ws))
     (set-websocket-state! ws 'closed)
@@ -213,15 +209,18 @@ connected to."
   (write-frame (make-text-frame data (generate-masking-key ws))
                (websocket-socket ws)))
 
-(define (websocket-receive ws)
+(define* (websocket-receive ws #:key timeout)
   "Read data from the server that WS is connected to.  Returns a string
 if text data was received, a bytevector if binary data was received,
-or #f if the WebSocket connection was closed."
-  (let ((frame (read-data-frame (websocket-socket ws))))
+or #f if the WebSocket connection was closed or TIMEOUT seconds
+elapsed without receiving a data frame."
+  (let* ((socket (websocket-socket ws))
+         (frame (read-data-frame socket #:timeout timeout)))
     (cond ((not frame)
-           (close-port (websocket-socket ws))
-           (close-port (websocket-entropy-port ws))
-           (set-websocket-state! ws 'closed)
+           (when (port-closed? socket)
+             ;; EOF - clean up the websocket.
+             (close-port (websocket-entropy-port ws))
+             (set-websocket-state! ws 'closed))
            #f)
           ((binary-frame? frame)
            (frame-data frame))
diff --git a/web/socket/frame.scm b/web/socket/frame.scm
index b5c05e5..9319c18 100644
--- a/web/socket/frame.scm
+++ b/web/socket/frame.scm
@@ -29,6 +29,7 @@
   #:use-module (rnrs io ports)
   #:use-module (ice-9 format)
   #:use-module (ice-9 match)
+  #:use-module (ice-9 ports internal)
   #:use-module (ice-9 rdelim)
   #:use-module (srfi srfi-1)
   #:use-module (srfi srfi-9)
@@ -66,6 +67,7 @@
 
             read-frame
             read-data-frame
+            read-close-frame
             write-frame))
 
 ;;;
@@ -324,20 +326,65 @@ MASKING-KEY."
          (data (read-data type masking-key length)))
     (make-frame final? type masking-key data)))
 
-(define* (read-data-frame port #:key echo-close?)
-  "Read frames from PORT until a complete data frame is received,
-handling any control frames transparently.  ECHO-CLOSE? should be #t
-if PORT represents the server end of a WebSocket connection."
+(define* (port-readable? port #:optional timeout)
+  "Wait for PORT to become readable or TIMEOUT seconds to elapse.
+Returns #t if the port is readable, or #f on timeout."
+  ;; TODO: Integrate with suspendable-ports & fibers.
+  (match (select (vector (port-read-wait-fd port)) #() #() timeout)
+    ((#(port) #() #()) #t)
+    (else #f)))
+
+(define* (port-writable? port #:optional timeout)
+  "Wait for PORT to become writable, or TIMEOUT seconds to elapse.
+Returns #t if the port is writable, or #f on timeout."
+  ;; TODO: Integrate with suspendable-ports & fibers.
+  (match (select #() (vector (port-write-wait-fd port)) #() timeout)
+    ((#() #(port) #()) #t)
+    (else #f)))
+
+(define* (frame-reader port #:key timeout)
+  "Return a generator which will read a frame from PORT on each call,
+until TIMEOUT seconds have elapsed.  This is used to implement timeouts
+over composed read operations."
+  (define timeout-at-real-time
+    (and timeout
+         (exact->inexact
+          (floor (+ (get-internal-real-time)
+                    (* internal-time-units-per-second
+                       timeout))))))
+  (define (timeout-remaining)
+    (and timeout
+         (/ (- timeout-at-real-time
+               (get-internal-real-time))
+            internal-time-units-per-second)))
+  (lambda ()
+    (cond ((port-closed? port)
+           #f)
+          ((not (port-readable? port (timeout-remaining)))
+           #f)
+          ((port-eof? port)
+           (eof-object))
+          (else
+           (read-frame port)))))
+
+(define* (read-data-frame port #:key echo-close? timeout)
+  "Read frames from PORT until a complete data frame is received or
+TIMEOUT seconds have elapsed, handling any control frames transparently.
+ECHO-CLOSE? should be #t if PORT represents the server end of a
+WebSocket connection."
+  (define next-frame
+    (frame-reader port #:timeout timeout))
   (let loop ((fragments '())
              (type #f))
-    (let ((frame (and (not (port-closed? port))
-                      (not (port-eof? port))
-                      (read-frame port))))
+    (let ((frame (next-frame)))
       (cond
        ;; EOF - port is closed.
-       ((not frame)
+       ((eof-object? frame)
         (close-port port)
         #f)
+       ;; Read timeout or already closed.
+       ((not frame)
+        #f)
        ;; Per section 5.4, control frames may appear interspersed
        ;; along with a fragmented message.
        ((close-frame? frame)
@@ -367,6 +414,20 @@ if PORT represents the server end of a WebSocket connection."
        (else
         (websocket-error "unexpected frame: " frame))))))
 
+(define* (read-close-frame port #:key timeout)
+  "Read frames from PORT until a close frame is received or TIMEOUT
+seconds have elapsed, ignoring control and data frames."
+  (define next-frame
+    (frame-reader port #:timeout timeout))
+  (let loop ()
+    (let ((frame (next-frame)))
+      (cond ((not (frame? frame))
+             #f)
+            ((close-frame? frame)
+             #t)
+            (else
+             (loop))))))
+
 
 ;;;
 ;;; Frame writer