irC

commit 7ce7e02b536104b9399ff01b355a7f931e8527b2
parent 59a2d87ee14a416b3151c096051d736a1e8f3305
Author: aabacchus <ben@bvnf.space>
Date:   Mon, 21 Feb 2022 01:00:55 +0000

nc improvements: don't assume text data

Diffstat:
M
nc.c
 | 
58
++++++++++++++++++++++++++++++++++------------------------

1 file changed, 34 insertions(+), 24 deletions(-)
diff --git a/nc.c b/nc.c
@@ -19,11 +19,11 @@
 #define BUF_LEN 1024
 
 void usage(const char *);
-int sock_write(int, char *);
+int sock_write(int, char *, size_t);
 int sock_read(int);
 
 int main(int argc, char *argv[]) {
-    int s, sfd;
+    int s, sfd = -1;
     struct addrinfo hints, *result, *rp;
     if (argc < 3) {
         usage(argv[0]);
@@ -42,7 +42,6 @@ int main(int argc, char *argv[]) {
         return 1;
     }
 
-    sfd = socket(result->ai_family, result->ai_socktype, result->ai_protocol);
     for (rp = result; rp != NULL; rp = rp->ai_next) {
         sfd = socket(rp->ai_family, rp->ai_socktype, rp->ai_protocol);
         if (sfd == -1)
@@ -59,10 +58,10 @@ int main(int argc, char *argv[]) {
 
     freeaddrinfo(result);
 
-    if (rp == NULL)
+    if (rp == NULL) {
+        fprintf(stderr, "connection refused\n");
         return 1;
-
-    size_t linelen = 0;
+    }
 
     int pollret;
     struct pollfd fds[2];
@@ -71,33 +70,45 @@ int main(int argc, char *argv[]) {
     fds[1].fd = sfd;
     fds[1].events = POLLIN | POLLWRBAND;
 
+    char *line = NULL;
+    size_t linelen = 0;
+
     while (1) {
         pollret = poll(fds, 2, -1);
         if (pollret <= 0) {
-            if (pollret == 0 || errno == EAGAIN)
+            if (errno == EAGAIN)
                 continue;
+            /*if (pollret == 0)
+                fprintf(stderr, "timed out\n");*/
             perror("poll");
+            free(line);
             return 1;
         }
-        if (fds[1].revents & POLLIN)
-            if (sock_read(sfd) != 0)
-                return 1;
         if (fds[0].revents & POLLIN) {
-            char *line = NULL;
-            if (getline(&line, &linelen, stdin) == -1) {
-                free(line);
+            ssize_t n;
+            if ((n = getline(&line, &linelen, stdin)) == -1) {
                 break;
             }
-            if (sock_write(sfd, line) != 0) {
+            if (sock_write(sfd, line, (size_t)n) != 0) {
                 free(line);
                 return 1;
             }
-            free(line);
+        }
+        if (fds[1].revents & POLLIN) {
+            int rec = sock_read(sfd);
+            if (rec < 0) {
+                free(line);
+                return 1;
+            } else if (rec == 0) {
+                break;
+            }
         }
         fds[0].revents = 0;
         fds[1].revents = 0;
     }
 
+    free(line);
+
     if (close(sfd) == -1) {
         perror("close");
         return 1;
@@ -110,9 +121,8 @@ void usage(const char *argv0){
     exit(1);
 }
 
-int sock_write(int sfd, char *msg){
+int sock_write(int sfd, char *msg, size_t len){
     ssize_t ret;
-    int len = strlen(msg);
     ret = send(sfd, msg, len, 0);
     if (ret != (ssize_t) len) {
         perror("sock_write");
@@ -122,16 +132,16 @@ int sock_write(int sfd, char *msg){
 }
 
 int sock_read(int sfd){
-    char buff[BUF_LEN];
-    memset(&buff, 0, sizeof buff);
-    ssize_t ret;
+    char buff[BUF_LEN] = {'\0'};
 
-    ret = recv(sfd, buff, BUF_LEN, 0);
+    ssize_t ret = recv(sfd, buff, BUF_LEN, 0);
 
     if (ret < 0) {
         perror("sock_read");
-        return 1;
+        return -1;
     }
-    printf("%s", buff);
-    return 0;
+    if (ret == 0)
+        return 0;
+    write(1, buff, ret);
+    return 1;
 }