API tools faq
paste
Socolin

OpenBSD nc patch with kqueue

Socolin
Jun 14th, 2014
286
0
Never
Add comment
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Diff 8.22 KB | None | 0 0
raw download clone embed print report
  1. Index: netcat.c
  2. ===================================================================
  3. RCS file: /cvs/src/usr.bin/nc/netcat.c,v
  4. retrieving revision 1.121
  5. diff -u -p -r1.121 netcat.c
  6. --- netcat.c    10 Jun 2014 16:35:42 -0000      1.121
  7. +++ netcat.c    14 Jun 2014 17:46:31 -0000
  8. @@ -36,6 +36,7 @@
  9.  #include <sys/time.h>
  10.  #include <sys/uio.h>
  11.  #include <sys/un.h>
  12. +#include <sys/event.h>
  13.  
  14.  #include <netinet/in.h>
  15.  #include <netinet/in_systm.h>
  16. @@ -734,58 +735,124 @@ readwrite(int nfd)
  17.  {
  18.         struct pollfd pfd[2];
  19.         unsigned char buf[16 * 1024];
  20. +       struct kevent keventlist[4];
  21. +       quad_t write_space_network;
  22. +       quad_t write_space_stdout;
  23. +       quad_t available;
  24.         int n, wfd = fileno(stdin);
  25.         int lfd = fileno(stdout);
  26.         int plen;
  27. +       int kq;
  28. +       int evcount;
  29. +       int evid;
  30.  
  31.         plen = sizeof(buf);
  32.  
  33. -       /* Setup Network FD */
  34. -       pfd[0].fd = nfd;
  35. -       pfd[0].events = POLLIN;
  36. +       /* Set up kqueue*/
  37. +       kq = kqueue();
  38. +       if (kq == -1) {
  39. +               close(nfd);
  40. +               err(1, "Kqueue Error");
  41. +       }
  42. +
  43. +       /* Set up STDOUT FD. */
  44. +       EV_SET(&keventlist[0], nfd, EVFILT_WRITE, EV_ADD, 0, 0, NULL);
  45. +
  46. +       /* Set up Network FD */
  47. +       EV_SET(&keventlist[1], lfd, EVFILT_WRITE, EV_ADD, 0, 0, NULL);
  48. +       EV_SET(&keventlist[2], nfd, EVFILT_READ, EV_ADD, 0, 0, NULL);
  49.  
  50.         /* Set up STDIN FD. */
  51. -       pfd[1].fd = wfd;
  52. -       pfd[1].events = POLLIN;
  53. +       EV_SET(&keventlist[3], wfd, EVFILT_READ, EV_ADD, 0, 0, NULL);
  54.  
  55. -       while (pfd[0].fd != -1) {
  56. +       /* Registering events */
  57. +       evcount = kevent(kq, keventlist, 4, NULL, 0, NULL);
  58. +       if (evcount == -1) {
  59. +               close(nfd);
  60. +               err(1, "Kevent Error");
  61. +       }
  62. +
  63. +       while (evcount != -1) {
  64.                 if (iflag)
  65.                         sleep(iflag);
  66. -
  67. -               if ((n = poll(pfd, 2 - dflag, timeout)) < 0) {
  68. +               evcount = kevent(kq, NULL, 0, keventlist, 4, NULL);
  69. +               if (evcount == -1) {
  70.                         close(nfd);
  71. -                       err(1, "Polling Error");
  72. +                       err(1, "Kevent Error");
  73.                 }
  74. +               for (evid = 0; evid < evcount; evid++) {
  75. +                       struct kevent event = keventlist[evid];
  76. +                       if (event.flags & EV_ERROR) {
  77. +                               close(nfd);
  78. +                               err(1, "Kevent Error");
  79. +                       }
  80.  
  81. -               if (n == 0)
  82. -                       return;
  83. -
  84. -               if (pfd[0].revents & POLLIN) {
  85. -                       if ((n = read(nfd, buf, plen)) < 0)
  86. -                               return;
  87. -                       else if (n == 0) {
  88. -                               shutdown(nfd, SHUT_RD);
  89. -                               pfd[0].fd = -1;
  90. -                               pfd[0].events = 0;
  91. -                       } else {
  92. -                               if (tflag)
  93. -                                       atelnet(nfd, buf, n);
  94. -                               if (atomicio(vwrite, lfd, buf, n) != n)
  95. +                       if (event.ident == nfd) {
  96. +                               if (event.flags & EV_EOF) {
  97. +                                       shutdown(nfd, SHUT_RD);
  98.                                         return;
  99. +                               } else {
  100. +                                       if (event.filter == EVFILT_WRITE) {
  101. +                                               write_space_network = event.data;
  102. +                                       }
  103. +                                       else if (event.filter == EVFILT_READ) {
  104. +                                               available = MIN(event.data, write_space_stdout);
  105. +                                               available = MIN(available, plen);
  106. +                                               if (available > 0) {
  107. +                                                       n = read(nfd, buf, available);
  108. +                                                       if (n != available) {
  109. +                                                               return;
  110. +                                                       } else {
  111. +                                                               n = write(lfd, buf, available);
  112. +                                                               if (n != available) {
  113. +                                                                       return;
  114. +                                                               }
  115. +                                                       }
  116. +                                                       write_space_stdout -= available;
  117. +                                                       /* Break to ignore EVFILT_WRITE that could
  118. +                                                        * be in current keventlist, after this event
  119. +                                                        * and which is now invalid, because we just
  120. +                                                        * write in the fd.
  121. +                                                        */
  122. +                                                       break;
  123. +                                               }
  124. +                                       }
  125. +                               }
  126.                         }
  127. -               }
  128. -
  129. -               if (!dflag && pfd[1].revents & POLLIN) {
  130. -                       if ((n = read(wfd, buf, plen)) < 0)
  131. -                               return;
  132. -                       else if (n == 0) {
  133. -                               if (Nflag)
  134. -                                       shutdown(nfd, SHUT_WR);
  135. -                               pfd[1].fd = -1;
  136. -                               pfd[1].events = 0;
  137. -                       } else {
  138. -                               if (atomicio(vwrite, nfd, buf, n) != n)
  139. +                       else if (event.ident == lfd) {
  140. +                               if (event.filter == EVFILT_WRITE) {
  141. +                                       write_space_stdout = event.data;
  142. +                               }
  143. +                       }
  144. +                       else if (event.ident == wfd) {
  145. +                               if (event.flags & EV_EOF) {
  146. +                                       if (Nflag)
  147. +                                               shutdown(nfd, SHUT_WR);
  148.                                         return;
  149. +                               } else {
  150. +                                       if (event.filter == EVFILT_READ) {
  151. +                                               available = MIN(event.data, write_space_network);
  152. +                                               available = MIN(available, plen);
  153. +                                               if (available > 0) {
  154. +                                                       n = read(wfd, buf, available);
  155. +                                                       if (n != available) {
  156. +                                                               return;
  157. +                                                       } else {
  158. +                                                               n = write(nfd, buf, available);
  159. +                                                               if (n != available) {
  160. +                                                                       return;
  161. +                                                               }
  162. +                                                       }
  163. +                                                       write_space_network -= available;
  164. +                                                       /* Break to ignore EVFILT_WRITE that could
  165. +                                                        * be in current keventlist, after this event
  166. +                                                        * and which is now invalid, because we just
  167. +                                                        * write in the fd.
  168. +                                                        */
  169. +                                                       break;
  170. +                                               }
  171. +                                       }
  172. +                               }
  173.                         }
  174.                 }
  175.         }
Advertisement
Add Comment
Please, Sign In to add comment
Public Pastes
We use cookies for various purposes including analytics. By continuing to use Pastebin, you agree to our use of cookies as described in the Cookies Policy.  OK, I Understand
Not a member of Pastebin yet?
Sign Up, it unlocks many cool features!