From 950ddf335bc43c7c45b2877a9fe8859bfa4c231c Mon Sep 17 00:00:00 2001
From: Ke Zhao <kzhao@redhat.com>
Date: Tue, 26 Mar 2019 14:49:15 -0400
Subject: [PATCH] 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 <kzhao@redhat.com>
---
 bin/main.c | 13 ++++++++++++-
 bin/non.c  | 14 +++++++++++++-
 2 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/bin/main.c b/bin/main.c
index ffd9655..a0ff2db 100644
--- 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
index c20db0c..a86c8ed 100644
--- 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;