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