From c14996aa9371f8b8917cbd2b076fabc4f1cf222b Mon Sep 17 00:00:00 2001
From: litong12 <litong12@xiaomi.com>
Date: Fri, 17 Dec 2021 17:19:01 +0800
Subject: [PATCH] examples/rpmsgsocket: Fix rpmsg example's bug

Signed-off-by: litong12 <litong12@xiaomi.com>
---
 examples/rpmsgsocket/rpsock_client.c | 32 +++++++++++++++++++++--
 examples/rpmsgsocket/rpsock_server.c | 39 +++++++++++++++++++++++-----
 2 files changed, 62 insertions(+), 9 deletions(-)

diff --git a/examples/rpmsgsocket/rpsock_client.c b/examples/rpmsgsocket/rpsock_client.c
index a4af16010..092283a2f 100644
--- a/examples/rpmsgsocket/rpsock_client.c
+++ b/examples/rpmsgsocket/rpsock_client.c
@@ -279,7 +279,21 @@ static int rpsock_stream_client(int argc, char *argv[])
 
   printf("client: Connecting to %s,%s...\n", myaddr.rp_cpu, myaddr.rp_name);
   ret = connect(sockfd, (struct sockaddr *)&myaddr, sizeof(myaddr));
-  if (ret < 0)
+  if (ret < 0 && errno == EINPROGRESS)
+    {
+      struct pollfd pfd;
+      memset(&pfd, 0, sizeof(struct pollfd));
+      pfd.fd = sockfd;
+      pfd.events = POLLOUT;
+
+      ret = poll(&pfd, 1, -1);
+      if (ret < 0)
+        {
+          printf("client: poll failure: %d\n", errno);
+          goto errout_with_socket;
+        }
+    }
+  else if (ret < 0)
     {
       printf("client: connect failure: %d\n", errno);
       goto errout_with_socket;
@@ -459,7 +473,21 @@ static int rpsock_dgram_client(int argc, char *argv[])
 
   printf("client: Connecting to %s,%s...\n", myaddr.rp_cpu, myaddr.rp_name);
   ret = connect(sockfd, (struct sockaddr *)&myaddr, sizeof(myaddr));
-  if (ret < 0)
+  if (ret < 0 && errno == EINPROGRESS)
+    {
+      struct pollfd pfd;
+      memset(&pfd, 0, sizeof(struct pollfd));
+      pfd.fd = sockfd;
+      pfd.events = POLLOUT;
+
+      ret = poll(&pfd, 1, -1);
+      if (ret < 0)
+        {
+          printf("[client] poll failure: %d\n", errno);
+          goto errout_with_socket;
+        }
+    }
+  else if (ret < 0)
     {
       printf("client: connect failure: %d\n", errno);
       goto errout_with_socket;
diff --git a/examples/rpmsgsocket/rpsock_server.c b/examples/rpmsgsocket/rpsock_server.c
index da8194bb3..c4f3a6876 100644
--- a/examples/rpmsgsocket/rpsock_server.c
+++ b/examples/rpmsgsocket/rpsock_server.c
@@ -43,7 +43,6 @@ struct rpsock_arg_s
 {
   int  fd;
   bool nonblock;
-  bool skippoll;
 };
 
 /****************************************************************************
@@ -62,7 +61,7 @@ static void *rpsock_thread(pthread_addr_t pvarg)
       char *tmp;
       int snd;
 
-      if (args->nonblock && !args->skippoll)
+      if (args->nonblock)
         {
           memset(&pfd, 0, sizeof(struct pollfd));
           pfd.fd = args->fd;
@@ -76,8 +75,6 @@ static void *rpsock_thread(pthread_addr_t pvarg)
             }
         }
 
-      args->skippoll = false;
-
       ret = recv(args->fd, buf, sizeof(buf), 0);
       if (ret == 0 || (ret < 0 && errno == ECONNRESET))
         {
@@ -89,6 +86,20 @@ static void *rpsock_thread(pthread_addr_t pvarg)
           usleep(10);
           continue;
         }
+      else if (ret < 0 && errno == EINPROGRESS)
+        {
+          memset(&pfd, 0, sizeof(struct pollfd));
+          pfd.fd = args->fd;
+          pfd.events = POLLOUT;
+          ret = poll(&pfd, 1, -1);
+          if (ret < 0)
+            {
+              printf("server: poll failure: %d\n", errno);
+              break;
+            }
+
+          continue;
+        }
       else if (ret < 0)
         {
           printf("server recv data failed ret %d, errno %d\n", ret, errno);
@@ -238,7 +249,6 @@ static int rpsock_stream_server(int argc, char *argv[])
 
       args->fd       = new;
       args->nonblock = nonblock;
-      args->skippoll = false;
 
       pthread_create(&thread, NULL, rpsock_thread,
                      (pthread_addr_t)args);
@@ -303,7 +313,23 @@ static int rpsock_dgram_server(int argc, char *argv[])
   printf("server: bind cpu %s, name %s ...\n",
           myaddr.rp_cpu, myaddr.rp_name);
   ret = bind(fd, (struct sockaddr *)&myaddr, sizeof(myaddr));
-  if (ret < 0)
+  ret = connect(fd, (struct sockaddr *)&myaddr, sizeof(myaddr));
+  if (ret < 0 && errno == EINPROGRESS)
+    {
+      struct pollfd pfd;
+      memset(&pfd, 0, sizeof(struct pollfd));
+      pfd.fd = fd;
+      pfd.events = POLLOUT;
+
+      ret = poll(&pfd, 1, -1);
+      if (ret < 0)
+        {
+          printf("server: poll failure: %d\n", errno);
+          close(fd);
+          return ret;
+        }
+    }
+  else if (ret < 0)
     {
       printf("server: bind failure: %d\n", errno);
       close(fd);
@@ -315,7 +341,6 @@ static int rpsock_dgram_server(int argc, char *argv[])
 
   args->fd       = fd;
   args->nonblock = nonblock;
-  args->skippoll = true;
 
   rpsock_thread(args);