Use fstat() to check if there is a socket or not

recv/send function can only be applied to a socket. We can use fstat() function to check if the fd is a socket or not. If so, we would like to use recv/send function to replace read/write. Related to:#30 Signed-off-by: Ke Zhao <[email protected]>
enarx-archive · Apr 4, 2019 · 950ddf3 · 950ddf3
1 parent 72f2c71
commit 950ddf3
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 2 deletions.
diff --git a/bin/main.c b/bin/main.c
@@ -25,6 +25,7 @@
 #include "exe.h"
 #include "non.h"
 
+#include <sys/stat.h>
 #include <sys/types.h>
 #include <sys/socket.h>
 #include <netdb.h>
@@ -183,12 +184,22 @@ on_conn(options_t *opts, int con, int in, int out, const struct addrinfo *ai)
   while (poll(pfds, 2, -1) >= 0) {
     char buffer[64 * 1024] = {};
     ssize_t ret;
+    struct stat st;
 
     for (int i = 0; i < 2; i++) {
       if (!pfds[i].revents)
         continue;
 
-      ret = read(pfds[i].fd, buffer, sizeof(buffer));
+      if (fstat(pfds[i].fd, &st)) {
+        fprintf(stderr, "Error in fstat. Errno is %d\n", errno);
+        return -1;
+      }
+
+      if (S_ISSOCK(st.st_mode))
+        ret = recv(pfds[i].fd, buffer, sizeof(buffer), 0);
+      else
+        ret = read(pfds[i].fd, buffer, sizeof(buffer));
+
       if (ret <= 0) {
         if (pfds[i].revents != POLLHUP &&
             (errno == EAGAIN || errno == EWOULDBLOCK))

diff --git a/bin/non.c b/bin/non.c
@@ -21,6 +21,8 @@
 
 #include "non.h"
 
+#include <sys/stat.h>
+#include <stdio.h>
 #include <stdint.h>
 #include <unistd.h>
 #include <errno.h>
@@ -97,8 +99,18 @@ non_write(int fd, void *buf, size_t len)
   struct pollfd pfd = { .fd = fd, .events = POLLOUT };
   uint8_t *b = buf;
   ssize_t ret;
+  struct stat st;
+
+  if (fstat(fd, &st)) {
+    fprintf(stderr, "Error in fstat. Errno is %d\n", errno);
+    return -1;
+  }
+
+  if (S_ISSOCK(st.st_mode))
+    ret = send(fd, buf, len, 0);
+  else
+    ret = write(fd, buf, len);
 
-  ret = write(fd, buf, len);
   if (ret < 0) {
     if (errno != EAGAIN)
       return ret;