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