1 files changed, 13 insertions, 14 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))