Update cidr-trie

cidr-trie.c

@@ -221,3 +221,40 @@ int trie128_lookup(trie_t *trie, uint8_t *ip)
 	}
 	return TRIE_NO_VALUE;
 }
+
+cidr_trie_t *cidr_trie_new()
+{
+	cidr_trie_t *cidr_trie = calloc(1, sizeof(*cidr_trie));
+	cidr_trie->cidr4_trie = NULL;
+	cidr_trie->cidr6_trie = NULL;
+	return cidr_trie;
+}
+
+void cidr_trie_insert_str(cidr_trie_t *cidr_trie, const char *ipstr, int value)
+{
+	char *p;
+	p = strchr(ipstr, ':');
+	if (p) {
+		if (cidr_trie->cidr6_trie == NULL)
+			cidr_trie->cidr6_trie = trie_new();
+		trie128_insert_str(cidr_trie->cidr6_trie, ipstr, value);
+	} else {
+		if (cidr_trie->cidr4_trie == NULL)
+			cidr_trie->cidr4_trie = trie_new();
+		trie32_insert_str(cidr_trie->cidr4_trie, ipstr, value);
+	}
+}
+
+int cidr4_trie_lookup(cidr_trie_t *cidr_trie, uint32_t ip)
+{
+	if (cidr_trie->cidr4_trie == NULL)
+		return TRIE_NO_VALUE;
+	return trie32_lookup(cidr_trie->cidr4_trie, ip);
+}
+
+int cidr6_trie_lookup(cidr_trie_t *cidr_trie, uint8_t *ip)
+{
+	if (cidr_trie->cidr6_trie == NULL)
+		return TRIE_NO_VALUE;
+	return trie128_lookup(cidr_trie->cidr6_trie, ip);
+}

cidr-trie.h

@@ -52,4 +52,14 @@ void trie128_insert(trie_t *trie, struct cidr6_s *cidr6, int value);
 int trie128_insert_str(trie_t *trie, const char *ipstr, int value);
 int trie128_lookup(trie_t *trie, uint8_t *ip);
 
+typedef struct {
+	trie_t *cidr4_trie;
+	trie_t *cidr6_trie;
+} cidr_trie_t;
+
+cidr_trie_t *cidr_trie_new();
+void cidr_trie_insert_str(cidr_trie_t *cidr_trie, const char *ipstr, int value);
+int cidr4_trie_lookup(cidr_trie_t *cidr_trie, uint32_t ip);
+int cidr6_trie_lookup(cidr_trie_t *cidr_trie, uint8_t *ip);
+
 #endif

graftcp.c

@@ -44,12 +44,12 @@ char *DEFAULT_LOCAL_PIPE_PAHT    = "/tmp/graftcplocal.fifo";
 bool DEFAULT_IGNORE_LOCAL        = true;
 int LOCAL_PIPE_FD;
 
-trie_t *BLACKLIST_IP     = NULL;
-trie_t *WHITELACKLIST_IP = NULL;
+cidr_trie_t *BLACKLIST_IP     = NULL;
+cidr_trie_t *WHITELACKLIST_IP = NULL;
 
 static int exit_code = 0;
 
-static void load_ip_file(char *path, trie_t **trie)
+static void load_ip_file(char *path, cidr_trie_t **trie)
 {
 	FILE *f;
 	char *line = NULL;
@@ -61,14 +61,12 @@ static void load_ip_file(char *path, trie_t **trie)
 		perror("fopen");
 		exit(1);
 	}
-	if (*trie == NULL)
-		*trie = trie_new();
 	while ((read = getline(&line, &len, f)) != -1) {
 		/* 7 is the shortest ip: (x.x.x.x) */
 		if (read < 7)
 			continue;
 		line[read - 1] = '\0';
-		trie32_insert_str(*trie, line, 1);
+		cidr_trie_insert_str(*trie, line, 1);
 		line = NULL;
 	}
 	fclose(f);
@@ -84,14 +82,27 @@ static void load_whiteip_file(char *path)
 	load_ip_file(path, &WHITELACKLIST_IP);
 }
 
-static bool is_ignore(uint32_t ip)
+static bool ip4_is_ignore(uint32_t ip)
 {
 	if (BLACKLIST_IP) {
-		if (trie32_lookup(BLACKLIST_IP, ntohl(ip)))
+		if (cidr4_trie_lookup(BLACKLIST_IP, ntohl(ip)))
 			return true;
 	}
 	if (WHITELACKLIST_IP) {
-		if (!trie32_lookup(WHITELACKLIST_IP, ntohl(ip)))
+		if (!cidr4_trie_lookup(WHITELACKLIST_IP, ntohl(ip)))
+			return true;
+	}
+	return false;
+}
+
+static bool ip6_is_ignore(uint8_t *ip)
+{
+	if (BLACKLIST_IP) {
+		if (cidr6_trie_lookup(BLACKLIST_IP, ip))
+			return true;
+	}
+	if (WHITELACKLIST_IP) {
+		if (!cidr6_trie_lookup(WHITELACKLIST_IP, ip))
 			return true;
 	}
 	return false;
@@ -198,14 +209,15 @@ void connect_pre_handle(struct proc_info *pinfp)
 		dest_ip_port = SOCKPORT(dest_sa);
 		dest_ip_addr.s_addr = SOCKADDR(dest_sa);
 		dest_ip_addr_str = inet_ntoa(dest_ip_addr);
-		if (is_ignore(dest_ip_addr.s_addr))
+		if (ip4_is_ignore(dest_ip_addr.s_addr))
 			return;
 	} else if (dest_sa.sin_family == AF_INET6) { /* IPv6 */
 		getdata(pinfp->pid, addr, (char *)&dest_sa6, sizeof(dest_sa6));
 		dest_ip_port = SOCKPORT6(dest_sa6);
+		if (ip6_is_ignore(dest_sa6.sin6_addr.s6_addr))
+			return;
 		inet_ntop(AF_INET6, &dest_sa6.sin6_addr, dest_str, INET6_ADDRSTRLEN);
 		dest_ip_addr_str = dest_str;
-		// TODO: is_ignore128()
 	} else {
 		return;
 	}
@@ -565,9 +577,9 @@ int client_main(int argc, char **argv)
 		load_whiteip_file(conf.whiteip_file_path);
 	if (*conf.ignore_local) {
 		if (BLACKLIST_IP == NULL)
-			BLACKLIST_IP = trie_new();
-		trie32_insert_str(BLACKLIST_IP, conf.local_addr, 1);
-		trie32_insert_str(BLACKLIST_IP, LOCAL_DEFAULT_ADDR, 1);
+			BLACKLIST_IP = cidr_trie_new();
+		cidr_trie_insert_str(BLACKLIST_IP, conf.local_addr, 1);
+		cidr_trie_insert_str(BLACKLIST_IP, LOCAL_DEFAULT_ADDR, 1);
 	}
 	PROXY_SA.sin_family = AF_INET;
 	PROXY_SA.sin_port = htons(*conf.local_port);