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 static void 58 assign_port(mctx_t mctx, int sock) 59 { 60 struct port *w; 61 struct sockaddr_in addr[2]; 62 socklen_t len = sizeof(addr) * 2; 63 64 /* remove a NAT mapping for this connection */ 65 if (mtcp_getpeername(mctx, sock, (struct sockaddr *)&addr, &len, 66 MOS_SIDE_CLI) < 0) { 67 TRACE_ERROR("mtcp_getpeer() failed for sock=%d\n", sock); 68 return; 69 } 70 /* assign a port number */ 71 pthread_mutex_lock(&g_addrlock); 72 TAILQ_FOREACH(w, &g_free_addrs, link) 73 if (GetRSSCPUCore(g_NATIP, addr[MOS_SIDE_SVR].sin_addr.s_addr, 74 w->port, addr[MOS_SIDE_SVR].sin_port, g_core_limit) 75 == mctx->cpu) 76 break; 77 if (w) { 78 TAILQ_REMOVE(&g_free_addrs, w, link); 79 mtcp_set_uctx(mctx, sock, w); 80 } else { 81 /* we're running out of available ports */ 82 /* FIXME: Handle this */ 83 TRACE_ERROR("No suitable port found! (CPU %d)\n", mctx->cpu); 84 } 85 pthread_mutex_unlock(&g_addrlock); 86 } 87 88 static int 89 set_addr(mctx_t mctx, int sock, uint32_t ip, uint16_t port, int part) 90 { 91 int off_ip, off_port; 92 93 if (part == SRC) { 94 off_ip = OFFSET_SRC_IP; 95 off_port = OFFSET_SRC_PORT; 96 } else /* if (part == DST) */ { 97 off_ip = OFFSET_DST_IP; 98 off_port = OFFSET_DST_PORT; 99 } 100 101 if (mtcp_setlastpkt(mctx, sock, 0, off_ip, 102 (uint8_t *)&ip, sizeof(in_addr_t), 103 MOS_IP_HDR | MOS_OVERWRITE) < 0) { 104 TRACE_ERROR("mtcp_setlastpkt() failed\n"); 105 return -1; 106 } 107 if (mtcp_setlastpkt(mctx, sock, 0, off_port, 108 (uint8_t *)&port, sizeof(in_port_t), 109 MOS_TCP_HDR | MOS_OVERWRITE 110 | MOS_UPDATE_IP_CHKSUM | MOS_UPDATE_TCP_CHKSUM) < 0) { 111 TRACE_ERROR("mtcp_setlastpkt() failed\n"); 112 return -1; 113 } 114 115 return 0; 116 } 117 118 static void 119 translate_addr(mctx_t mctx, int sock, int side, uint64_t events, 120 filter_arg_t *arg) 121 { 122 struct port *w; 123 124 if (!(w = mtcp_get_uctx(mctx, sock))) 125 assign_port(mctx, sock); 126 127 /* Translate the addresses */ 128 if (side == MOS_SIDE_CLI) { 129 /* CLI (LAN) ==> SVR (WAN) : SNAT */ 130 if (!(w = mtcp_get_uctx(mctx, sock))) 131 return; 132 133 if (set_addr(mctx, sock, g_NATIP, w->port, SRC) < 0 && 134 mtcp_setlastpkt(mctx, sock, side, 0, NULL, 0, MOS_DROP) < 0) 135 exit(EXIT_FAILURE); 136 } else /* if (side == MOS_SIDE_SVR) */ { 137 /* SVR (WAN) ==> CLI (LAN) : DNAT */ 138 struct sockaddr_in addr[2]; 139 socklen_t len = sizeof(addr) * 2; 140 141 if (mtcp_getpeername(mctx, sock, (struct sockaddr *)&addr, &len, 142 MOS_SIDE_CLI) < 0) { 143 TRACE_ERROR("mtcp_getpeer() failed sock=%d side=%d\n", sock, side); 144 return; 145 } 146 147 if (set_addr(mctx, sock, addr[MOS_SIDE_CLI].sin_addr.s_addr, 148 addr[MOS_SIDE_CLI].sin_port, DST) < 0 && 149 mtcp_setlastpkt(mctx, sock, side, 0, NULL, 0, MOS_DROP) < 0) 150 exit(EXIT_FAILURE); 151 } 152 } 153 154 static void 155 release_port(mctx_t mctx, int sock, int side, uint64_t events, 156 filter_arg_t *arg) 157 { 158 /* release the port number */ 159 struct port *w; 160 161 if (!(w = mtcp_get_uctx(mctx, sock))) 162 return; 163 164 /* assign a port number */ 165 pthread_mutex_lock(&g_addrlock); 166 TAILQ_INSERT_TAIL(&g_free_addrs, w, link); 167 mtcp_set_uctx(mctx, sock, NULL); 168 pthread_mutex_unlock(&g_addrlock); 169 } 170 /*----------------------------------------------------------------------------*/ 171 static void 172 init_monitor(mctx_t mctx) 173 { 174 int lsock = mtcp_socket(mctx, AF_INET, MOS_SOCK_MONITOR_STREAM, 0); 175 if (lsock < 0) { 176 TRACE_ERROR("Failed to create monitor raw socket!\n"); 177 return; 178 } 179 180 if (mtcp_register_callback(mctx, lsock, MOS_ON_PKT_IN, MOS_HK_SND, 181 translate_addr)) 182 exit(EXIT_FAILURE); 183 if (mtcp_register_callback(mctx, lsock, MOS_ON_CONN_END, MOS_HK_SND, 184 release_port)) 185 exit(EXIT_FAILURE); 186 } 187 /*----------------------------------------------------------------------------*/ 188 int 189 main(int argc, char **argv) 190 { 191 int i, opt; 192 char *fname = MOS_CONFIG_FILE; /* path to the default mos config file */ 193 struct mtcp_conf mcfg; /* mOS configuration */ 194 mctx_t mctx_list[MAX_CORES]; /* mOS context */ 195 196 /* get the total # of cpu cores */ 197 g_core_limit = GetNumCPUs(); 198 199 /* Parse command line arguments */ 200 while ((opt = getopt(argc, argv, "c:f:i:")) != -1) { 201 switch (opt) { 202 case 'f': 203 fname = optarg; 204 break; 205 case 'c': 206 if (atoi(optarg) > g_core_limit) { 207 printf("Available number of CPU cores is %d\n", g_core_limit); 208 return -1; 209 } 210 g_core_limit = atoi(optarg); 211 break; 212 case 'i': 213 g_NATIP = inet_addr(optarg); 214 break; 215 default: 216 printf("Usage: %s [-f mos_config_file] [-c #_of_cpu] [-i ip_address]\n", argv[0]); 217 return 0; 218 } 219 } 220 221 /* NAT IP address checking */ 222 if (!g_NATIP) { 223 fprintf(stderr, "You have to specify IP address of NAT with '-i' option\n"); 224 exit(EXIT_FAILURE); 225 } 226 227 /* parse mos configuration file */ 228 if (mtcp_init(fname)) { 229 fprintf(stderr, "Failed to initialize mtcp.\n"); 230 exit(EXIT_FAILURE); 231 } 232 233 /* set the core limit */ 234 mtcp_getconf(&mcfg); 235 mcfg.num_cores = g_core_limit; 236 mtcp_setconf(&mcfg); 237 238 /* Initialize global data structure */ 239 pthread_mutex_init(&g_addrlock, NULL); 240 TAILQ_INIT(&g_free_addrs); 241 for (i = 1025; i < 65535; i++) { 242 struct port *p = malloc(sizeof(struct port)); 243 if (!p) 244 exit(EXIT_FAILURE); 245 p->port = htons(i); 246 TAILQ_INSERT_TAIL(&g_free_addrs, p, link); 247 } 248 249 for (i = 0; i < g_core_limit; i++) { 250 /* Run mOS for each CPU core */ 251 if (!(mctx_list[i] = mtcp_create_context(i))) { 252 fprintf(stderr, "Failed to craete mtcp context.\n"); 253 return -1; 254 } 255 256 /* init monitor */ 257 init_monitor(mctx_list[i]); 258 } 259 260 /* wait until mOS finishes */ 261 for (i = 0; i < g_core_limit; i++) 262 mtcp_app_join(mctx_list[i]); 263 264 mtcp_destroy(); 265 return 0; 266 } 267 /*----------------------------------------------------------------------------*/ 268