xref: /mOS-networking-stack/samples/nat/nat.c (revision 602e2e46)
1 
2 #define _LARGEFILE64_SOURCE
3 #include <stdio.h>
4 #include <stdlib.h>
5 #include <unistd.h>
6 #include <stdint.h>
7 #include <sys/types.h>
8 #include <sys/stat.h>
9 #include <sys/socket.h>
10 #include <netinet/in.h>
11 #include <arpa/inet.h>
12 #include <fcntl.h>
13 #include <dirent.h>
14 #include <string.h>
15 #include <time.h>
16 #include <pthread.h>
17 #include <signal.h>
18 #include <linux/if_ether.h>
19 #include <linux/tcp.h>
20 
21 #include <sys/queue.h>
22 
23 #include <mos_api.h>
24 #include <mtcp_util.h>
25 
26 #include "cpu.h"
27 #include "http_parsing.h"
28 #include "debug.h"
29 #include "applib.h"
30 
31 /* update_curpacket()'s byte offsets */
32 #define OFFSET_DST_IP		16
33 #define OFFSET_SRC_IP		12
34 #define OFFSET_DST_PORT		 2
35 #define OFFSET_SRC_PORT		 0
36 
37 /* default configure file path */
38 #define MOS_CONFIG_FILE		"config/mos.conf"
39 
40 #define MAX_CORES 16
41 
42 enum {
43 	SRC,
44 	DST,
45 };
46 
47 struct port {
48 	uint16_t port;
49 	TAILQ_ENTRY(port) link;
50 };
51 
52 static int g_core_limit  = 1;   /* number of CPU cores used, WHY GLOBAL? */
53 static in_addr_t g_NATIP = 0;   /* NAT IP address */
54 static TAILQ_HEAD(, port) g_free_addrs;
55 static pthread_mutex_t g_addrlock;
56 
57 /*----------------------------------------------------------------------------*/
58 static void
assign_port(mctx_t mctx,int sock)59 assign_port(mctx_t mctx, int sock)
60 {
61 	struct port *w;
62 	struct sockaddr_in addr[2];
63 	socklen_t len = sizeof(struct sockaddr_in) * 2;
64 
65 	/* remove a NAT mapping for this connection */
66 	if (mtcp_getpeername(mctx, sock, (struct sockaddr *)&addr, &len,
67 						 MOS_SIDE_BOTH) < 0) {
68 		TRACE_ERROR("mtcp_getpeer() failed for sock=%d\n", sock);
69 		return;
70 	}
71 	/* assign a port number */
72 	pthread_mutex_lock(&g_addrlock);
73 	TAILQ_FOREACH(w, &g_free_addrs, link) {
74 		/* Don't compute RSS CPU mapping value if # cpus = 1 */
75 		if (g_core_limit == 1 || GetRSSCPUCore(g_NATIP, addr[MOS_SIDE_SVR].sin_addr.s_addr,
76 				  w->port, addr[MOS_SIDE_SVR].sin_port, g_core_limit)
77 			 == mctx->cpu)
78 			break;
79 	}
80 	if (w) {
81 		TAILQ_REMOVE(&g_free_addrs, w, link);
82 		mtcp_set_uctx(mctx, sock, w);
83 	} else {
84 		/* we're running out of available ports */
85 		/* FIXME: Handle this */
86 		TRACE_ERROR("No suitable port found! (CPU %d)\n", mctx->cpu);
87 	}
88 	pthread_mutex_unlock(&g_addrlock);
89 }
90 /*----------------------------------------------------------------------------*/
91 static int
set_addr(mctx_t mctx,int sock,uint32_t ip,uint16_t port,int part)92 set_addr(mctx_t mctx, int sock, uint32_t ip, uint16_t port, int part)
93 {
94 	int off_ip, off_port;
95 
96 	if (part == SRC) {
97 		off_ip = OFFSET_SRC_IP;
98 		off_port = OFFSET_SRC_PORT;
99 	} else /* if (part == DST) */ {
100 		off_ip = OFFSET_DST_IP;
101 		off_port = OFFSET_DST_PORT;
102 	}
103 
104 	if (mtcp_setlastpkt(mctx, sock, 0, off_ip,
105 				  (uint8_t *)&ip, sizeof(in_addr_t),
106 				  MOS_IP_HDR | MOS_OVERWRITE) < 0) {
107 		TRACE_ERROR("mtcp_setlastpkt() failed\n");
108 		return -1;
109 	}
110 	if (mtcp_setlastpkt(mctx, sock, 0, off_port,
111 				  (uint8_t *)&port, sizeof(in_port_t),
112 				  MOS_TCP_HDR | MOS_OVERWRITE
113 				  | MOS_UPDATE_IP_CHKSUM | MOS_UPDATE_TCP_CHKSUM) < 0) {
114 		TRACE_ERROR("mtcp_setlastpkt() failed\n");
115 		return -1;
116 	}
117 
118 	return 0;
119 }
120 /*----------------------------------------------------------------------------*/
121 static void
translate_addr(mctx_t mctx,int sock,int side,uint64_t events,filter_arg_t * arg)122 translate_addr(mctx_t mctx, int sock, int side, uint64_t events,
123 		filter_arg_t *arg)
124 {
125 	struct port *w;
126 
127 	if (!(w = mtcp_get_uctx(mctx, sock)))
128 		assign_port(mctx, sock);
129 
130 	/* Translate the addresses */
131 	if (side == MOS_SIDE_CLI) {
132 		/* CLI (LAN) ==> SVR (WAN) : SNAT */
133 		if (!(w = mtcp_get_uctx(mctx, sock)))
134 			return;
135 
136 		if (set_addr(mctx, sock, g_NATIP, w->port, SRC) < 0 &&
137 			mtcp_setlastpkt(mctx, sock, side, 0, NULL, 0, MOS_DROP) < 0)
138 			exit(EXIT_FAILURE);
139 	} else /* if (side == MOS_SIDE_SVR) */ {
140 		/* SVR (WAN) ==> CLI (LAN) : DNAT */
141 		struct sockaddr_in addr[2];
142 		socklen_t len = sizeof(struct sockaddr_in) * 2;
143 
144 		if (mtcp_getpeername(mctx, sock, (struct sockaddr *)&addr, &len,
145 					MOS_SIDE_BOTH) < 0) {
146 			TRACE_ERROR("mtcp_getpeer() failed sock=%d side=%d\n", sock, side);
147 			return;
148 		}
149 
150 		if (set_addr(mctx, sock, addr[MOS_SIDE_CLI].sin_addr.s_addr,
151 					 addr[MOS_SIDE_CLI].sin_port, DST) < 0 &&
152 			mtcp_setlastpkt(mctx, sock, side, 0, NULL, 0, MOS_DROP) < 0)
153 			exit(EXIT_FAILURE);
154 	}
155 }
156 /*----------------------------------------------------------------------------*/
157 static void
release_port(mctx_t mctx,int sock,int side,uint64_t events,filter_arg_t * arg)158 release_port(mctx_t mctx, int sock, int side, uint64_t events,
159 		filter_arg_t *arg)
160 {
161 	/* release the port number */
162 	struct port *w;
163 
164 	if (!(w = mtcp_get_uctx(mctx, sock)))
165 		return;
166 
167 	/* assign a port number */
168 	pthread_mutex_lock(&g_addrlock);
169 	TAILQ_INSERT_TAIL(&g_free_addrs, w, link);
170 	mtcp_set_uctx(mctx, sock, NULL);
171 	pthread_mutex_unlock(&g_addrlock);
172 }
173 /*----------------------------------------------------------------------------*/
174 static void
init_monitor(mctx_t mctx)175 init_monitor(mctx_t mctx)
176 {
177 	int lsock = mtcp_socket(mctx, AF_INET, MOS_SOCK_MONITOR_STREAM, 0);
178 	if (lsock < 0) {
179 		TRACE_ERROR("Failed to create monitor raw socket!\n");
180 		return;
181 	}
182 
183 	if (mtcp_register_callback(mctx, lsock, MOS_ON_PKT_IN, MOS_HK_SND,
184 							   translate_addr))
185 		exit(EXIT_FAILURE);
186 	if (mtcp_register_callback(mctx, lsock, MOS_ON_CONN_END, MOS_HK_SND,
187 							   release_port))
188 		exit(EXIT_FAILURE);
189 }
190 /*----------------------------------------------------------------------------*/
191 int
main(int argc,char ** argv)192 main(int argc, char **argv)
193 {
194 	int i, opt;
195 	char *fname = MOS_CONFIG_FILE;  /* path to the default mos config file */
196 	struct mtcp_conf mcfg;          /* mOS configuration */
197 	mctx_t mctx_list[MAX_CORES];                     /* mOS context */
198 
199 	/* get the total # of cpu cores */
200 	g_core_limit = GetNumCPUs();
201 
202 	/* Parse command line arguments */
203 	while ((opt = getopt(argc, argv, "c:f:i:")) != -1) {
204 		switch (opt) {
205 		case 'f':
206 			fname = optarg;
207 			break;
208 		case 'c':
209 			if (atoi(optarg) > g_core_limit) {
210 				printf("Available number of CPU cores is %d\n", g_core_limit);
211 				return -1;
212 			}
213 			g_core_limit = atoi(optarg);
214 			break;
215 		case 'i':
216 			g_NATIP = inet_addr(optarg);
217 			break;
218 		default:
219 			printf("Usage: %s [-f mos_config_file] [-c #_of_cpu] [-i ip_address]\n", argv[0]);
220 			return 0;
221 		}
222 	}
223 
224 	/* NAT IP address checking */
225 	if (!g_NATIP) {
226 		fprintf(stderr, "You have to specify IP address of NAT with '-i' option\n");
227 		exit(EXIT_FAILURE);
228 	}
229 
230 	/* parse mos configuration file */
231 	if (mtcp_init(fname)) {
232 		fprintf(stderr, "Failed to initialize mtcp.\n");
233 		exit(EXIT_FAILURE);
234 	}
235 
236 	/* set the core limit */
237 	mtcp_getconf(&mcfg);
238 	mcfg.num_cores = g_core_limit;
239 	mtcp_setconf(&mcfg);
240 
241 	/* Initialize global data structure */
242 	pthread_mutex_init(&g_addrlock, NULL);
243 	TAILQ_INIT(&g_free_addrs);
244 	for (i = 1025; i < 65535; i++) {
245 		struct port *p = malloc(sizeof(struct port));
246 		if (!p)
247 			exit(EXIT_FAILURE);
248 		p->port = htons(i);
249 		TAILQ_INSERT_TAIL(&g_free_addrs, p, link);
250 	}
251 
252 	for (i = 0; i < g_core_limit; i++) {
253 		/* Run mOS for each CPU core */
254 		if (!(mctx_list[i] = mtcp_create_context(i))) {
255 			fprintf(stderr, "Failed to craete mtcp context.\n");
256 			return -1;
257 		}
258 
259 		/* init monitor */
260 		init_monitor(mctx_list[i]);
261 	}
262 
263 	/* wait until mOS finishes */
264 	for (i = 0; i < g_core_limit; i++)
265 		mtcp_app_join(mctx_list[i]);
266 
267 	mtcp_destroy();
268 	return 0;
269 }
270 /*----------------------------------------------------------------------------*/
271