client: Handle control frames in websocket-receive.

* web/socket/frame.scm (read-data-frame): New procedure, extracted from
serve-client.
* web/socket/client.scm (websocket-receive): Call read-data-frame.
Clean up the websocket if the connection was closed.
* web/socket/server.scm (serve-client): Call read-data-frame.  Flush the
port after writing the handshake response.

3 files changed, 64 insertions, 42 deletions

diff --git a/web/socket/client.scm b/web/socket/client.scm
index 6561c26..c867427 100644
--- a/web/socket/client.scm
+++ b/web/socket/client.scm
@@ -214,9 +214,16 @@ connected to."
                (websocket-socket ws)))
 
 (define (websocket-receive ws)
-  "Read a response from the server that WS is connected to."
-  ;; TODO: Handle fragmented frames and control frames.
-  (let ((frame (read-frame (websocket-socket ws))))
-    (if (binary-frame? frame)
-        (frame-data frame)
-        (text-frame->string frame))))
+  "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))))
+    (cond ((not frame)
+           (close-port (websocket-socket ws))
+           (close-port (websocket-entropy-port ws))
+           (set-websocket-state! ws 'closed)
+           #f)
+          ((binary-frame? frame)
+           (frame-data frame))
+          (else
+           (text-frame->string frame)))))
diff --git a/web/socket/frame.scm b/web/socket/frame.scm
index f3709f6..b5c05e5 100644
--- a/web/socket/frame.scm
+++ b/web/socket/frame.scm
@@ -1,5 +1,6 @@
 ;;; guile-websocket --- WebSocket client/server
 ;;; Copyright © 2015 David Thompson <davet@gnu.org>
+;;; Copyright © 2023 Andrew Whatson <whatson@tailcall.au>
 ;;;
 ;;; This file is part of guile-websocket.
 ;;;
@@ -64,6 +65,7 @@
             text-frames->string
 
             read-frame
+            read-data-frame
             write-frame))
 
 ;;;
@@ -322,6 +324,49 @@ 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."
+  (let loop ((fragments '())
+             (type #f))
+    (let ((frame (and (not (port-closed? port))
+                      (not (port-eof? port))
+                      (read-frame port))))
+      (cond
+       ;; EOF - port is closed.
+       ((not frame)
+        (close-port port)
+        #f)
+       ;; Per section 5.4, control frames may appear interspersed
+       ;; along with a fragmented message.
+       ((close-frame? frame)
+        (when echo-close?
+          ;; Per section 5.5.1, echo the close frame back to the
+          ;; client before closing the socket.  The client may no
+          ;; longer be listening.
+          (false-if-exception
+           (write-frame (make-close-frame (frame-data frame)) port)))
+        (close-port port)
+        #f)
+       ((ping-frame? frame)
+        ;; Per section 5.5.3, a pong frame must include the exact
+        ;; same data as the ping frame.
+        (write-frame (make-pong-frame (frame-data frame)) port)
+        (loop fragments type))
+       ((pong-frame? frame) ; silently ignore pongs
+        (loop fragments type))
+       ((first-fragment-frame? frame) ; begin accumulating fragments
+        (loop (list frame) (frame-type frame)))
+       ((final-fragment-frame? frame) ; concatenate all fragments
+        (make-frame #t type #f (frame-concatenate (reverse fragments))))
+       ((fragment-frame? frame) ; add a fragment
+        (loop (cons frame fragments) type))
+       ((data-frame? frame) ; unfragmented data frame
+        frame)
+       (else
+        (websocket-error "unexpected frame: " frame))))))
+
 
 ;;;
 ;;; Frame writer
diff --git a/web/socket/server.scm b/web/socket/server.scm
index 2d1694d..3394d3b 100644
--- a/web/socket/server.scm
+++ b/web/socket/server.scm
@@ -1,6 +1,7 @@
 ;;; guile-websocket --- WebSocket client/server
 ;;; Copyright © 2015 David Thompson <davet@gnu.org>
 ;;; Copyright © 2021 Jan (janneke) Nieuwenhuizen <janneke@gnu.org>
+;;; Copyright © 2023 Andrew Whatson <whatson@tailcall.au>
 ;;;
 ;;; This file is part of guile-websocket.
 ;;;
@@ -87,48 +88,17 @@ called for each complete message that is received."
       (when response
         (write-frame response client-socket))))
 
-  (define (read-frame-maybe)
-    (and (not (port-eof? client-socket))
-         (read-frame client-socket)))
-
   ;; Perform the HTTP handshake and upgrade to WebSocket protocol.
   (let* ((request (read-handshake-request client-socket))
          (client-key (assoc-ref (request-headers request) 'sec-websocket-key))
          (response (make-handshake-response client-key)))
     (write-response response client-socket)
-    (let loop ((fragments '())
-               (type #f))
-      (let ((frame (read-frame-maybe)))
-        (cond
-         ;; EOF - port is closed.
-         ((not frame)
-          (close-port client-socket))
-         ;; Per section 5.4, control frames may appear interspersed
-         ;; along with a fragmented message.
-         ((close-frame? frame)
-          ;; Per section 5.5.1, echo the close frame back to the
-          ;; client before closing the socket.  The client may no
-          ;; longer be listening.
-          (false-if-exception
-           (write-frame (make-close-frame (frame-data frame)) client-socket))
-          (close-port client-socket))
-         ((ping-frame? frame)
-          ;; Per section 5.5.3, a pong frame must include the exact
-          ;; same data as the ping frame.
-          (write-frame (make-pong-frame (frame-data frame)) client-socket)
-          (loop fragments type))
-         ((pong-frame? frame) ; silently ignore pongs
-          (loop fragments type))
-         ((first-fragment-frame? frame) ; begin accumulating fragments
-          (loop (list frame) (frame-type frame)))
-         ((final-fragment-frame? frame) ; concatenate all fragments
-          (handle-data-frame type (frame-concatenate (reverse fragments)))
-          (loop '() #f))
-         ((fragment-frame? frame) ; add a fragment
-          (loop (cons frame fragments) type))
-         ((data-frame? frame) ; unfragmented data frame
+    (force-output client-socket)
+    (let loop ()
+      (let ((frame (read-data-frame client-socket #:echo-close? #t)))
+        (when frame
           (handle-data-frame (frame-type frame) (frame-data frame))
-          (loop '() #f)))))))
+          (loop))))))
 
 (define* (run-server handler #:optional (server-socket (make-server-socket)))
   "Run WebSocket server on SERVER-SOCKET.  HANDLER, a procedure that