|
21 | 21 | #include <sys/types.h> |
22 | 22 | #include <sys/wait.h> |
23 | 23 | #include <sys/random.h> |
| 24 | +#include <sys/ioctl.h> |
24 | 25 |
|
25 | 26 | #include <netdb.h> |
26 | 27 | #include <netinet/in.h> |
@@ -712,6 +713,157 @@ static void process_one_client(int fd, int pipefd) |
712 | 713 | close(fd); |
713 | 714 | } |
714 | 715 |
|
| 716 | +static void get_tcp_inq(struct msghdr *msgh, unsigned int *inqv) |
| 717 | +{ |
| 718 | + struct cmsghdr *cmsg; |
| 719 | + |
| 720 | + for (cmsg = CMSG_FIRSTHDR(msgh); cmsg ; cmsg = CMSG_NXTHDR(msgh, cmsg)) { |
| 721 | + if (cmsg->cmsg_level == IPPROTO_TCP && cmsg->cmsg_type == TCP_CM_INQ) { |
| 722 | + memcpy(inqv, CMSG_DATA(cmsg), sizeof(*inqv)); |
| 723 | + return; |
| 724 | + } |
| 725 | + } |
| 726 | + |
| 727 | + xerror("could not find TCP_CM_INQ cmsg type"); |
| 728 | +} |
| 729 | + |
| 730 | +static void process_one_client_inq(int fd, int unixfd) |
| 731 | +{ |
| 732 | + unsigned int tcp_inq; |
| 733 | + size_t expect_len; |
| 734 | + char msg_buf[4096]; |
| 735 | + char buf[4096]; |
| 736 | + char tmp[16]; |
| 737 | + struct iovec iov = { |
| 738 | + .iov_base = buf, |
| 739 | + .iov_len = 1, |
| 740 | + }; |
| 741 | + struct msghdr msg = { |
| 742 | + .msg_iov = &iov, |
| 743 | + .msg_iovlen = 1, |
| 744 | + .msg_control = msg_buf, |
| 745 | + .msg_controllen = sizeof(msg_buf), |
| 746 | + }; |
| 747 | + ssize_t ret, tot; |
| 748 | + int on = 1; |
| 749 | + |
| 750 | + if (-1 == setsockopt(fd, IPPROTO_TCP, TCP_INQ, &on, sizeof(on))) |
| 751 | + die_perror("setsockopt"); |
| 752 | + |
| 753 | + ret = write(unixfd, "xmit", 4); |
| 754 | + assert(ret == 4); |
| 755 | + |
| 756 | + ret = read(unixfd, &expect_len, sizeof(expect_len)); |
| 757 | + assert(ret == (ssize_t)sizeof(expect_len)); |
| 758 | + |
| 759 | + if (expect_len > sizeof(buf)) |
| 760 | + xerror("expect len %zu exceeds buffer size", expect_len); |
| 761 | + |
| 762 | + for (;;) { |
| 763 | + struct timespec req; |
| 764 | + unsigned int queued; |
| 765 | + |
| 766 | + ret = ioctl(fd, FIONREAD, &queued); |
| 767 | + if (ret < 0) |
| 768 | + die_perror("FIONREAD"); |
| 769 | + if (queued > expect_len) |
| 770 | + xerror("FIONREAD returned %u, but only %zu expected\n", |
| 771 | + queued, expect_len); |
| 772 | + if (queued == expect_len) |
| 773 | + break; |
| 774 | + |
| 775 | + req.tv_sec = 0; |
| 776 | + req.tv_nsec = 1000 * 1000ul; |
| 777 | + nanosleep(&req, NULL); |
| 778 | + } |
| 779 | + |
| 780 | + /* read one byte, expect cmsg to return expected - 1 */ |
| 781 | + ret = recvmsg(fd, &msg, 0); |
| 782 | + if (ret < 0) |
| 783 | + die_perror("recvmsg"); |
| 784 | + |
| 785 | + if (msg.msg_controllen == 0) |
| 786 | + xerror("msg_controllen is 0"); |
| 787 | + |
| 788 | + get_tcp_inq(&msg, &tcp_inq); |
| 789 | + |
| 790 | + assert((size_t)tcp_inq == (expect_len - 1)); |
| 791 | + |
| 792 | + iov.iov_len = sizeof(buf); |
| 793 | + ret = recvmsg(fd, &msg, 0); |
| 794 | + if (ret < 0) |
| 795 | + die_perror("recvmsg"); |
| 796 | + |
| 797 | + /* should have gotten exact remainder of all pending data */ |
| 798 | + assert(ret == (ssize_t)tcp_inq); |
| 799 | + |
| 800 | + /* should be 0, all drained */ |
| 801 | + get_tcp_inq(&msg, &tcp_inq); |
| 802 | + assert(tcp_inq == 0); |
| 803 | + |
| 804 | + /* request a large swath of data. */ |
| 805 | + ret = write(unixfd, "huge", 4); |
| 806 | + assert(ret == 4); |
| 807 | + |
| 808 | + ret = read(unixfd, &expect_len, sizeof(expect_len)); |
| 809 | + assert(ret == (ssize_t)sizeof(expect_len)); |
| 810 | + |
| 811 | + /* peer should send us a few mb of data */ |
| 812 | + if (expect_len <= sizeof(buf)) |
| 813 | + xerror("expect len %zu too small\n", expect_len); |
| 814 | + |
| 815 | + tot = 0; |
| 816 | + do { |
| 817 | + iov.iov_len = sizeof(buf); |
| 818 | + ret = recvmsg(fd, &msg, 0); |
| 819 | + if (ret < 0) |
| 820 | + die_perror("recvmsg"); |
| 821 | + |
| 822 | + tot += ret; |
| 823 | + |
| 824 | + get_tcp_inq(&msg, &tcp_inq); |
| 825 | + |
| 826 | + if (tcp_inq > expect_len - tot) |
| 827 | + xerror("inq %d, remaining %d total_len %d\n", |
| 828 | + tcp_inq, expect_len - tot, (int)expect_len); |
| 829 | + |
| 830 | + assert(tcp_inq <= expect_len - tot); |
| 831 | + } while ((size_t)tot < expect_len); |
| 832 | + |
| 833 | + ret = write(unixfd, "shut", 4); |
| 834 | + assert(ret == 4); |
| 835 | + |
| 836 | + /* wait for hangup. Should have received one more byte of data. */ |
| 837 | + ret = read(unixfd, tmp, sizeof(tmp)); |
| 838 | + assert(ret == 6); |
| 839 | + assert(strncmp(tmp, "closed", 6) == 0); |
| 840 | + |
| 841 | + sleep(1); |
| 842 | + |
| 843 | + iov.iov_len = 1; |
| 844 | + ret = recvmsg(fd, &msg, 0); |
| 845 | + if (ret < 0) |
| 846 | + die_perror("recvmsg"); |
| 847 | + assert(ret == 1); |
| 848 | + |
| 849 | + get_tcp_inq(&msg, &tcp_inq); |
| 850 | + |
| 851 | + /* tcp_inq should be 1 due to received fin. */ |
| 852 | + assert(tcp_inq == 1); |
| 853 | + |
| 854 | + iov.iov_len = 1; |
| 855 | + ret = recvmsg(fd, &msg, 0); |
| 856 | + if (ret < 0) |
| 857 | + die_perror("recvmsg"); |
| 858 | + |
| 859 | + /* expect EOF */ |
| 860 | + assert(ret == 0); |
| 861 | + get_tcp_inq(&msg, &tcp_inq); |
| 862 | + assert(tcp_inq == 1); |
| 863 | + |
| 864 | + close(fd); |
| 865 | +} |
| 866 | + |
715 | 867 | static int xaccept(int s) |
716 | 868 | { |
717 | 869 | int fd = accept(s, NULL, 0); |
@@ -744,7 +896,10 @@ static int server(int ipcfd) |
744 | 896 | alarm(15); |
745 | 897 | r = xaccept(fd); |
746 | 898 |
|
747 | | - process_one_client(r, ipcfd); |
| 899 | + if (inq) |
| 900 | + process_one_client_inq(r, ipcfd); |
| 901 | + else |
| 902 | + process_one_client(r, ipcfd); |
748 | 903 |
|
749 | 904 | return 0; |
750 | 905 | } |
|
0 commit comments