1 /*
2  * Copyright 2018 Google Inc.
3  * Author: Soheil Hassas Yeganeh (soheil@google.com)
4  *
5  * Simple example on how to use TCP_INQ and TCP_CM_INQ.
6  *
7  * License (GPLv2):
8  *
9  * This program is free software; you can redistribute it and/or modify it
10  * under the terms and conditions of the GNU General Public License,
11  * version 2, as published by the Free Software Foundation.
12  *
13  * This program is distributed in the hope it will be useful, but WITHOUT
14  * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
15  * FITNESS FOR A PARTICULAR PURPOSE. * See the GNU General Public License for
16  * more details.
17  */
18 #define _GNU_SOURCE
19 
20 #include <error.h>
21 #include <netinet/in.h>
22 #include <netinet/tcp.h>
23 #include <pthread.h>
24 #include <stdio.h>
25 #include <errno.h>
26 #include <stdlib.h>
27 #include <string.h>
28 #include <sys/socket.h>
29 #include <unistd.h>
30 
31 #ifndef TCP_INQ
32 #define TCP_INQ 36
33 #endif
34 
35 #ifndef TCP_CM_INQ
36 #define TCP_CM_INQ TCP_INQ
37 #endif
38 
39 #define BUF_SIZE 8192
40 #define CMSG_SIZE 32
41 
42 static int family = AF_INET6;
43 static socklen_t addr_len = sizeof(struct sockaddr_in6);
44 static int port = 4974;
45 
46 static void setup_loopback_addr(int family, struct sockaddr_storage *sockaddr)
47 {
48 	struct sockaddr_in6 *addr6 = (void *) sockaddr;
49 	struct sockaddr_in *addr4 = (void *) sockaddr;
50 
51 	switch (family) {
52 	case PF_INET:
53 		memset(addr4, 0, sizeof(*addr4));
54 		addr4->sin_family = AF_INET;
55 		addr4->sin_addr.s_addr = htonl(INADDR_LOOPBACK);
56 		addr4->sin_port = htons(port);
57 		break;
58 	case PF_INET6:
59 		memset(addr6, 0, sizeof(*addr6));
60 		addr6->sin6_family = AF_INET6;
61 		addr6->sin6_addr = in6addr_loopback;
62 		addr6->sin6_port = htons(port);
63 		break;
64 	default:
65 		error(1, 0, "illegal family");
66 	}
67 }
68 
69 void *start_server(void *arg)
70 {
71 	int server_fd = (int)(unsigned long)arg;
72 	struct sockaddr_in addr;
73 	socklen_t addrlen = sizeof(addr);
74 	char *buf;
75 	int fd;
76 	int r;
77 
78 	buf = malloc(BUF_SIZE);
79 
80 	for (;;) {
81 		fd = accept(server_fd, (struct sockaddr *)&addr, &addrlen);
82 		if (fd == -1) {
83 			perror("accept");
84 			break;
85 		}
86 		do {
87 			r = send(fd, buf, BUF_SIZE, 0);
88 		} while (r < 0 && errno == EINTR);
89 		if (r < 0)
90 			perror("send");
91 		if (r != BUF_SIZE)
92 			fprintf(stderr, "can only send %d bytes\n", r);
93 		/* TCP_INQ can overestimate in-queue by one byte if we send
94 		 * the FIN packet. Sleep for 1 second, so that the client
95 		 * likely invoked recvmsg().
96 		 */
97 		sleep(1);
98 		close(fd);
99 	}
100 
101 	free(buf);
102 	close(server_fd);
103 	pthread_exit(0);
104 }
105 
106 int main(int argc, char *argv[])
107 {
108 	struct sockaddr_storage listen_addr, addr;
109 	int c, one = 1, inq = -1;
110 	pthread_t server_thread;
111 	char cmsgbuf[CMSG_SIZE];
112 	struct iovec iov[1];
113 	struct cmsghdr *cm;
114 	struct msghdr msg;
115 	int server_fd, fd;
116 	char *buf;
117 
118 	while ((c = getopt(argc, argv, "46p:")) != -1) {
119 		switch (c) {
120 		case '4':
121 			family = PF_INET;
122 			addr_len = sizeof(struct sockaddr_in);
123 			break;
124 		case '6':
125 			family = PF_INET6;
126 			addr_len = sizeof(struct sockaddr_in6);
127 			break;
128 		case 'p':
129 			port = atoi(optarg);
130 			break;
131 		}
132 	}
133 
134 	server_fd = socket(family, SOCK_STREAM, 0);
135 	if (server_fd < 0)
136 		error(1, errno, "server socket");
137 	setup_loopback_addr(family, &listen_addr);
138 	if (setsockopt(server_fd, SOL_SOCKET, SO_REUSEADDR,
139 		       &one, sizeof(one)) != 0)
140 		error(1, errno, "setsockopt(SO_REUSEADDR)");
141 	if (bind(server_fd, (const struct sockaddr *)&listen_addr,
142 		 addr_len) == -1)
143 		error(1, errno, "bind");
144 	if (listen(server_fd, 128) == -1)
145 		error(1, errno, "listen");
146 	if (pthread_create(&server_thread, NULL, start_server,
147 			   (void *)(unsigned long)server_fd) != 0)
148 		error(1, errno, "pthread_create");
149 
150 	fd = socket(family, SOCK_STREAM, 0);
151 	if (fd < 0)
152 		error(1, errno, "client socket");
153 	setup_loopback_addr(family, &addr);
154 	if (connect(fd, (const struct sockaddr *)&addr, addr_len) == -1)
155 		error(1, errno, "connect");
156 	if (setsockopt(fd, SOL_TCP, TCP_INQ, &one, sizeof(one)) != 0)
157 		error(1, errno, "setsockopt(TCP_INQ)");
158 
159 	msg.msg_name = NULL;
160 	msg.msg_namelen = 0;
161 	msg.msg_iov = iov;
162 	msg.msg_iovlen = 1;
163 	msg.msg_control = cmsgbuf;
164 	msg.msg_controllen = sizeof(cmsgbuf);
165 	msg.msg_flags = 0;
166 
167 	buf = malloc(BUF_SIZE);
168 	iov[0].iov_base = buf;
169 	iov[0].iov_len = BUF_SIZE / 2;
170 
171 	if (recvmsg(fd, &msg, 0) != iov[0].iov_len)
172 		error(1, errno, "recvmsg");
173 	if (msg.msg_flags & MSG_CTRUNC)
174 		error(1, 0, "control message is truncated");
175 
176 	for (cm = CMSG_FIRSTHDR(&msg); cm; cm = CMSG_NXTHDR(&msg, cm))
177 		if (cm->cmsg_level == SOL_TCP && cm->cmsg_type == TCP_CM_INQ)
178 			inq = *((int *) CMSG_DATA(cm));
179 
180 	if (inq != BUF_SIZE - iov[0].iov_len) {
181 		fprintf(stderr, "unexpected inq: %d\n", inq);
182 		exit(1);
183 	}
184 
185 	printf("PASSED\n");
186 	free(buf);
187 	close(fd);
188 	return 0;
189 }
190