diff --git a/UA2F/scripts/test.py b/UA2F/scripts/test.py index 019655cf9..6e936649e 100644 --- a/UA2F/scripts/test.py +++ b/UA2F/scripts/test.py @@ -50,10 +50,12 @@ def start_ua2f(u: str): # iptables 设置函数 def setup_iptables(): os.system(f"sudo iptables -A OUTPUT -p tcp --dport {PORT} -j NFQUEUE --queue-num 10010") + # os.system(f"sudo ip6tables -A OUTPUT -p tcp --dport {PORT} -j NFQUEUE --queue-num 10010") def cleanup_iptables(): os.system(f"sudo iptables -D OUTPUT -p tcp --dport {PORT} -j NFQUEUE --queue-num 10010") + # os.system(f"sudo ip6tables -D OUTPUT -p tcp --dport {PORT} -j NFQUEUE --queue-num 10010") if __name__ == "__main__": @@ -81,13 +83,21 @@ if __name__ == "__main__": time.sleep(3) - for i in tqdm(range(10000)): + for i in tqdm(range(2000)): nxt = ua.random - response = requests.get(f"http://localhost:{PORT}", headers={ + response = requests.get(f"http://127.0.0.1:{PORT}", headers={ "User-Agent": nxt }) assert response.ok assert response.text == str(len(nxt)) + # for i in tqdm(range(2000)): + # nxt = ua.random + # response = requests.get(f"http://[::1]:{PORT}", headers={ + # "User-Agent": nxt + # }) + # assert response.ok + # assert response.text == str(len(nxt)) + # clean cleanup_iptables() diff --git a/UA2F/src/cache.c b/UA2F/src/cache.c index 329d5011f..a97af3cf3 100644 --- a/UA2F/src/cache.c +++ b/UA2F/src/cache.c @@ -63,18 +63,10 @@ bool cache_contains(struct addr_port target) { pthread_rwlock_unlock(&cacheLock); if (s != NULL) { - bool ret; pthread_rwlock_wrlock(&cacheLock); - if (difftime(time(NULL), s->last_time) > check_interval * 2) { - HASH_DEL(not_http_dst_cache, s); - free(s); - ret = false; - } else { - s->last_time = time(NULL); - ret = true; - } + s->last_time = time(NULL); pthread_rwlock_unlock(&cacheLock); - return ret; + return true; } return false; diff --git a/UA2F/src/handler.c b/UA2F/src/handler.c index 436f60eaf..fe477725d 100644 --- a/UA2F/src/handler.c +++ b/UA2F/src/handler.c @@ -9,6 +9,7 @@ #include "config.h" #endif +#include #include #include #include @@ -114,7 +115,7 @@ end: } static void add_to_cache(const struct nf_packet *pkt) { - struct addr_port target = { + const struct addr_port target = { .addr = pkt->orig.dst, .port = pkt->orig.dst_port, }; @@ -171,6 +172,64 @@ bool should_ignore(const struct nf_packet *pkt) { return retval; } +enum { + IP_UNK = 0, +}; + +static bool ipv4_set_transport_header(struct pkt_buff *pkt_buff) { + struct iphdr *ip_hdr = nfq_ip_get_hdr(pkt_buff); + if (ip_hdr == NULL) { + return false; + } + + if (nfq_ip_set_transport_header(pkt_buff, ip_hdr) < 0) { + syslog(LOG_ERR, "Failed to set ipv4 transport header"); + return false; + } + return true; +} + +static bool ipv6_set_transport_header(struct pkt_buff *pkt_buff) { + struct ip6_hdr *ip_hdr = nfq_ip6_get_hdr(pkt_buff); + if (ip_hdr == NULL) { + return false; + } + + if (nfq_ip6_set_transport_header(pkt_buff, ip_hdr, IPPROTO_TCP) < 0) { + syslog(LOG_ERR, "Failed to set ipv6 transport header"); + return false; + } + return true; +} + +static int set_transport_header(struct pkt_buff* pkt_buff, const int ip_type) { + if (ip_type == IPV4) { + if (ipv4_set_transport_header(pkt_buff)) { + count_ipv4_packet(); + return IPV4; + } + return IP_UNK; + } + if (ip_type == IPV6) { + if (ipv6_set_transport_header(pkt_buff)) { + count_ipv6_packet(); + return IPV6; + } + return IP_UNK; + } + + // unknown ip type + if (ipv4_set_transport_header(pkt_buff)) { + count_ipv4_packet(); + return IPV4; + } + if (ipv6_set_transport_header(pkt_buff)) { + count_ipv6_packet(); + return IPV6; + } + return IP_UNK; +} + void handle_packet(const struct nf_queue *queue, const struct nf_packet *pkt) { if (use_conntrack) { if (!pkt->has_conntrack) { @@ -191,36 +250,16 @@ void handle_packet(const struct nf_queue *queue, const struct nf_packet *pkt) { } struct pkt_buff *pkt_buff = pktb_alloc(AF_INET, pkt->payload, pkt->payload_len, 0); - ASSERT(pkt_buff != NULL); + assert(pkt_buff != NULL); int type; if (use_conntrack) { type = pkt->orig.ip_version; + set_transport_header(pkt_buff, type); } else { - const __auto_type ip_hdr = nfq_ip_get_hdr(pkt_buff); - if (ip_hdr == NULL) { - type = IPV6; - } else { - type = IPV4; - } - } - - if (type == IPV4) { - count_ipv4_packet(); - } else { - count_ipv6_packet(); - } - - if (type == IPV4) { - const __auto_type ip_hdr = nfq_ip_get_hdr(pkt_buff); - if (nfq_ip_set_transport_header(pkt_buff, ip_hdr) < 0) { - syslog(LOG_ERR, "Failed to set ipv4 transport header"); - goto end; - } - } else { - const __auto_type ip_hdr = nfq_ip6_get_hdr(pkt_buff); - if (nfq_ip6_set_transport_header(pkt_buff, ip_hdr, IPPROTO_TCP) < 0) { - syslog(LOG_ERR, "Failed to set ipv6 transport header"); + type = set_transport_header(pkt_buff, IP_UNK); + if (type == IP_UNK) { + syslog(LOG_ERR, "Failed to set transport header"); goto end; } } @@ -234,9 +273,13 @@ void handle_packet(const struct nf_queue *queue, const struct nf_packet *pkt) { } const __auto_type tcp_payload = nfq_tcp_get_payload(tcp_hdr, pkt_buff); - const __auto_type tcp_payload_len = nfq_tcp_get_payload_len(tcp_hdr, pkt_buff); + if (tcp_payload == NULL) { + send_verdict(queue, pkt, get_next_mark(pkt, false), NULL); + goto end; + } - if (tcp_payload == NULL || tcp_payload_len < USER_AGENT_MATCH_LENGTH) { + const __auto_type tcp_payload_len = nfq_tcp_get_payload_len(tcp_hdr, pkt_buff); + if (tcp_payload_len < USER_AGENT_MATCH_LENGTH) { send_verdict(queue, pkt, get_next_mark(pkt, false), NULL); goto end; } diff --git a/UA2F/src/ua2f.c b/UA2F/src/ua2f.c index eed647f18..e3d59b7f6 100644 --- a/UA2F/src/ua2f.c +++ b/UA2F/src/ua2f.c @@ -8,10 +8,12 @@ #include "config.h" #endif +#include #include #include #include #include +#include #pragma clang diagnostic push #pragma ide diagnostic ignored "EndlessLoop" @@ -28,11 +30,9 @@ int parse_packet(const struct nf_queue *queue, struct nf_buffer *buf) { while (!should_exit) { const __auto_type status = nfqueue_next(buf, packet); - switch (status) { - case IO_READY: + if (status == IO_READY) { handle_packet(queue, packet); - break; - default: + } else { return status; } } @@ -42,12 +42,10 @@ int parse_packet(const struct nf_queue *queue, struct nf_buffer *buf) { int read_buffer(struct nf_queue *queue, struct nf_buffer *buf) { const __auto_type buf_status = nfqueue_receive(queue, buf, 0); - switch (buf_status) { - case IO_READY: + if (buf_status == IO_READY) { return parse_packet(queue, buf); - default: - return buf_status; } + return buf_status; } bool retry_without_conntrack(struct nf_queue *queue) { @@ -109,6 +107,8 @@ int main(const int argc, char *argv[]) { syslog(LOG_ERR, "Failed to open nfqueue"); return EXIT_FAILURE; } + assert(queue->queue_num == QUEUE_NUM); + assert(queue->nl_socket != NULL); main_loop(queue); @@ -117,4 +117,4 @@ int main(const int argc, char *argv[]) { return EXIT_SUCCESS; } -#pragma clang diagnostic pop +#pragma clang diagnostic pop \ No newline at end of file diff --git a/UA2F/test/cache_test.cc b/UA2F/test/cache_test.cc index bfca95a34..bd1f7dc52 100644 --- a/UA2F/test/cache_test.cc +++ b/UA2F/test/cache_test.cc @@ -6,7 +6,7 @@ extern "C" { class CacheTest : public ::testing::Test { protected: - struct addr_port test_addr{}; + addr_port test_addr{}; void SetUp() override { test_addr.addr.ip4 = 12345; @@ -39,11 +39,11 @@ TEST_F(CacheTest, AddAndRemoveFromCache) { cache_add(test_addr); EXPECT_TRUE(cache_contains(test_addr)); sleep(5); - EXPECT_FALSE(cache_contains(test_addr)); + EXPECT_TRUE(cache_contains(test_addr)); } TEST_F(CacheTest, CacheDoesNotContainNonexistentEntry) { - struct addr_port nonexistent_addr{}; + addr_port nonexistent_addr{}; nonexistent_addr.addr.ip4 = 54321; EXPECT_FALSE(cache_contains(nonexistent_addr)); } \ No newline at end of file