From 488596f5418ffbf874f33f0cfbccacb0c6ce3b03 Mon Sep 17 00:00:00 2001
From: aixiao <aixiao@aixiao.me>
Date: Sun, 16 Aug 2020 07:56:52 +0800
Subject: [PATCH] Add whether to enable whitelist configuration option.

---
 Makefile  |  2 +-
 README.md |  2 ++
 ais.c     | 23 ++++++++++++++---------
 ais.conf  |  3 ++-
 conf.c    |  3 +++
 conf.h    |  1 +
 info.sh   |  4 ----
 start.sh  |  4 ----
 stop.sh   |  4 ----
 9 files changed, 23 insertions(+), 23 deletions(-)

diff --git a/Makefile b/Makefile
index abdbfba..bb5340c 100644
--- a/Makefile
+++ b/Makefile
@@ -2,7 +2,7 @@ CROSS_COMPILE ?=
 CC := $(CROSS_COMPILE)gcc
 STRIP := $(CROSS_COMPILE)strip
 CFLAGS += -g -O2 -Wall
-LIBS = -static
+LIBS = 
 OBJ := ais
 
 all: conf.o ais.o
diff --git a/README.md b/README.md
index c94b667..4d5ed11 100644
--- a/README.md
+++ b/README.md
@@ -12,6 +12,8 @@
         
 # 配置文件
     global {
+        // 是否开启白名单（1开启，0关闭）
+        IP_RESTRICTION = 1;
         // 白名单IP段, 判断前两段IP空格隔开冒号结尾
         IP_SEGMENT= 115.60 115.61 115.62 223.88;
     }
diff --git a/ais.c b/ais.c
index 5d3f629..4546989 100644
--- a/ais.c
+++ b/ais.c
@@ -12,6 +12,7 @@
 #include <sys/wait.h>
 #include <netinet/in.h>
 #include <string.h>
+#include <pthread.h>
 #include "ais.h"
 #include "conf.h"
 
@@ -170,7 +171,6 @@ void extract_server_path(const char *header, char *output)
 
 int extract_host(const char *header)
 {
-
     char *_p = strstr(header, "CONNECT"); /* 在 CONNECT 方法中解析 隧道主机名称及端口号 */
     if (_p) {
         char *_p1 = strchr(_p, ' ');
@@ -258,7 +258,6 @@ void hand_mproxy_info_req(int sock, char *header)
                      </body></html>\n", info_buf);
 
     write(sock, response, strlen(response));
-
 }
 
 /* 获取运行的基本信息输出到指定的缓冲区 */
@@ -531,7 +530,8 @@ void sigchld_handler(int signal)
 // IP段白名单
 int whitelist(char *client_ip, char (*whitelist_ip)[32])
 {
-    for (int i = 1; i < WHITELIST_IP_NUM - 1; i++) {
+    int i;
+    for (i = 1; i < WHITELIST_IP_NUM - 1; i++) {
         if (strcmp(whitelist_ip[i], "\0") == 0) { //  如果字符串为空就跳出循环
             break;
         }
@@ -545,6 +545,7 @@ int whitelist(char *client_ip, char (*whitelist_ip)[32])
 
 void server_loop()
 {
+    int i;
     char ipstr[128];
     char client_ip[32];         // 客户端IP
     struct sockaddr_in client_addr;
@@ -554,24 +555,27 @@ void server_loop()
     read_conf("ais.conf", configure);
     printf("%s\n", configure->IP_SEGMENT);
 
-    char whitelist_ip[WHITELIST_IP_NUM][32] = { 0 };
+    char whitelist_ip[WHITELIST_IP_NUM][32] = {{ 0 }, { 0 }};
     split_string(configure->IP_SEGMENT, " ", whitelist_ip);
 
-    for (int i = 1; i <= WHITELIST_IP_NUM - 1; i++) {
+    for (i = 1; i <= WHITELIST_IP_NUM - 1; i++) {
         if (*whitelist_ip[i] != '\0')
             printf("%s\n", whitelist_ip[i]);
     }
-    //exit(0);
 
     while (1) {
         client_sock = accept(server_sock, (struct sockaddr *)&client_addr, &addrlen);
         if (client_sock > 0) {
             LOG("Client Ip %s Client Port %d\n", inet_ntop(AF_INET, &client_addr.sin_addr.s_addr, ipstr, sizeof(ipstr)), ntohs(client_addr.sin_port));
             strcpy(client_ip, inet_ntop(AF_INET, &client_addr.sin_addr.s_addr, ipstr, sizeof(ipstr))); // 复制客户端IP到client_ip
-            if (whitelist(client_ip, whitelist_ip) == 0) {
-                LOG("非法客户端, 拒绝连接\n");
-                continue;
+            
+            if (configure->IP_RESTRICTION == 1) {
+                if (whitelist(client_ip, whitelist_ip) == 0) {
+                    LOG("非法客户端, 拒绝连接\n");
+                    continue;
+                }
             }
+            
         }
 
         if (fork() == 0) {      // 创建子进程处理客户端连接请求
@@ -689,5 +693,6 @@ int _main(int argc, char *argv[])
     get_info(info_buf);
     LOG("%s\n", info_buf);
     start_server(daemon);
+
     return 0;
 }
diff --git a/ais.conf b/ais.conf
index 17e635f..3aeb990 100644
--- a/ais.conf
+++ b/ais.conf
@@ -1,3 +1,4 @@
 global {
-    IP_SEGMENT= 115.60 115.61 115.62 223.88;
+    IP_RESTRICTION = 1;
+    IP_SEGMENT= 223.104 115.60 115.61 115.62 223.88 223.89 106.33 117.136 61.158 171.10 171.9 61.158;
 }
diff --git a/conf.c b/conf.c
index aa16825..9d791be 100644
--- a/conf.c
+++ b/conf.c
@@ -74,6 +74,9 @@ static void parse_global_module(char *content, conf * p)
             memset(p->IP_SEGMENT, 0, val_begin_len);
             memcpy(p->IP_SEGMENT, val_begin, val_begin_len);
         }
+        if (strcasecmp(var, "IP_RESTRICTION") == 0) {
+            p->IP_RESTRICTION = atoi(val_begin);
+        }
         content = strchr(lineEnd + 1, '\n');
     }
 }
diff --git a/conf.h b/conf.h
index 7d53e50..02ca457 100644
--- a/conf.h
+++ b/conf.h
@@ -9,6 +9,7 @@
 
 // 配置文件结构
 typedef struct CONF {
+    int IP_RESTRICTION;
     char *IP_SEGMENT;
 } conf;
 
diff --git a/info.sh b/info.sh
index 90a5880..61e3410 100644
--- a/info.sh
+++ b/info.sh
@@ -1,8 +1,4 @@
 #!/bin/bash
-#
-# GET info
-# date 20200526
-#
 
 cat info.txt | grep "Client Ip" | awk '{print $7}' | uniq -c
 
diff --git a/start.sh b/start.sh
index 8c79d5d..5e44394 100644
--- a/start.sh
+++ b/start.sh
@@ -1,8 +1,4 @@
 #!/bin/bash
-#
-# Start AIS
-# date: 20200526
-#
 
 SHELL_FOLDER=$(cd "$(dirname "$0")"; pwd)       #脚本所在目录
 SHELL_FOLDER=$(dirname $(readlink -f "$0"))
diff --git a/stop.sh b/stop.sh
index ba99f89..d977c2d 100644
--- a/stop.sh
+++ b/stop.sh
@@ -1,8 +1,4 @@
 #!/bin/bash
-#
-# Stop AIS
-# date 20200526
-#
 
 killall ais