1 files changed, 85 insertions, 53 deletions

diff --git a/docs/examples/websocket.c b/docs/examples/websocket.c
index 758ee48e0..0df12c84b 100644
--- a/docs/examples/websocket.c
+++ b/docs/examples/websocket.c
@@ -37,51 +37,77 @@
 
 #include <curl/curl.h>
 
-static int ping(CURL *curl, const char *send_payload)
+static CURLcode ping(CURL *curl, const char *send_payload)
 {
-  size_t sent;
-  CURLcode result =
-    curl_ws_send(curl, send_payload, strlen(send_payload), &sent, 0,
-                 CURLWS_PING);
-  return (int)result;
+  CURLcode res = CURLE_OK;
+  const char *buf = send_payload;
+  size_t sent, blen = strlen(send_payload);
+
+  while(blen) {
+    res = curl_ws_send(curl, buf, blen, &sent, 0, CURLWS_PING);
+    if(!res) {
+      buf += sent; /* deduct what was sent */
+      blen -= sent;
+    }
+    else if(res == CURLE_AGAIN) {  /* blocked on sending */
+      fprintf(stderr, "ws: sent PING blocked, waiting a second\n");
+      sleep(1);  /* either select() on socket or max timeout would
+                    be good here. */
+    }
+    else /* real error sending */
+      break;
+  }
+  if(!res)
+    fprintf(stderr, "ws: sent PING with payload\n");
+  return res;
 }
 
-static int recv_pong(CURL *curl, const char *expected_payload)
+static CURLcode recv_pong(CURL *curl, const char *expected_payload)
 {
   size_t rlen;
   const struct curl_ws_frame *meta;
   char buffer[256];
-  CURLcode result = curl_ws_recv(curl, buffer, sizeof(buffer), &rlen, &meta);
-  if(!result) {
+  CURLcode res;
+
+retry:
+  res = curl_ws_recv(curl, buffer, sizeof(buffer), &rlen, &meta);
+  if(!res) {
+    /* on small PING content, this example assumes the complete
+     * PONG content arrives in one go. Larger frames will arrive
+     * in chunks, however. */
     if(meta->flags & CURLWS_PONG) {
       int same = 0;
-      fprintf(stderr, "ws: got PONG back\n");
       if(rlen == strlen(expected_payload)) {
-        if(!memcmp(expected_payload, buffer, rlen)) {
-          fprintf(stderr, "ws: got the same payload back\n");
+        if(!memcmp(expected_payload, buffer, rlen))
           same = 1;
-        }
       }
-      if(!same)
-        fprintf(stderr, "ws: did NOT get the same payload back\n");
+      fprintf(stderr, "ws: received PONG with %s payload back\n",
+              same ? "same" : "different");
+    }
+    else if(meta->flags & CURLWS_TEXT) {
+      fprintf(stderr, "ws: received TEXT frame '%.*s'\n", (int)rlen,
+              buffer);
+    }
+    else if(meta->flags & CURLWS_BINARY) {
+      fprintf(stderr, "ws: received BINARY frame of %u bytes\n", (int)rlen);
     }
     else {
-      fprintf(stderr, "recv_pong: got %u bytes rflags %x\n", (int)rlen,
+      /* some other frame arrived. */
+      fprintf(stderr, "ws: received frame of %u bytes rflags %x\n", (int)rlen,
               meta->flags);
+      goto retry;
     }
   }
-  fprintf(stderr, "ws: curl_ws_recv returned %u, received %u\n",
-          (unsigned int)result, (unsigned int)rlen);
-  return (int)result;
-}
-
-static CURLcode recv_any(CURL *curl)
-{
-  size_t rlen;
-  const struct curl_ws_frame *meta;
-  char buffer[256];
-
-  return curl_ws_recv(curl, buffer, sizeof(buffer), &rlen, &meta);
+  else if(res == CURLE_AGAIN) {  /* blocked on receiving */
+    fprintf(stderr, "ws: PONG not there yet, waiting a second\n");
+    sleep(1);  /* either select() on socket or max timeout would
+                  be good here. */
+    goto retry;
+  }
+  if(res)
+    fprintf(stderr, "ws: curl_ws_recv returned %u, received %u\n",
+            (unsigned int)res, (unsigned int)rlen);
+  return res;
 }
 
 /* close the connection */
@@ -91,45 +117,51 @@ static void websocket_close(CURL *curl)
   (void)curl_ws_send(curl, "", 0, &sent, 0, CURLWS_CLOSE);
 }
 
-static void websocket(CURL *curl)
+static CURLcode websocket(CURL *curl)
 {
+  CURLcode res;
   int i = 0;
   do {
-    recv_any(curl);
-    if(ping(curl, "foobar"))
-      return;
-    if(recv_pong(curl, "foobar")) {
-      return;
-    }
-    sleep(2);
+    res = ping(curl, "foobar");
+    if(res)
+      break;
+    res = recv_pong(curl, "foobar");
+    if(res)
+      break;
+    sleep(1);
   } while(i++ < 10);
   websocket_close(curl);
+  return res;
 }
 
-int main(void)
+int main(int argc, const char *argv[])
 {
   CURL *curl;
   CURLcode res;
 
   curl = curl_easy_init();
-  if(curl) {
+  if(!curl) {
+    return 1; /* memory failure */
+  }
+  if(argc == 2)
+    curl_easy_setopt(curl, CURLOPT_URL, argv[1]);
+  else
     curl_easy_setopt(curl, CURLOPT_URL, "wss://example.com");
 
-    curl_easy_setopt(curl, CURLOPT_CONNECT_ONLY, 2L); /* websocket style */
-
-    /* Perform the request, res gets the return code */
-    res = curl_easy_perform(curl);
-    /* Check for errors */
-    if(res != CURLE_OK)
-      fprintf(stderr, "curl_easy_perform() failed: %s\n",
-              curl_easy_strerror(res));
-    else {
-      /* connected and ready */
-      websocket(curl);
-    }
+  curl_easy_setopt(curl, CURLOPT_CONNECT_ONLY, 2L); /* websocket style */
 
-    /* always cleanup */
-    curl_easy_cleanup(curl);
+  /* Perform the request, res gets the return code */
+  res = curl_easy_perform(curl);
+  /* Check for errors */
+  if(res != CURLE_OK)
+    fprintf(stderr, "curl_easy_perform() failed: %s\n",
+            curl_easy_strerror(res));
+  else {
+    /* connected and ready */
+    res = websocket(curl);
   }
-  return 0;
+
+  /* always cleanup */
+  curl_easy_cleanup(curl);
+  return (int)res;
 }