diff --git a/src/include/nftset.h b/src/include/nftset.h index 699fac4..18de0da 100644 --- a/src/include/nftset.h +++ b/src/include/nftset.h @@ -1,6 +1,6 @@ /************************************************************************* * - * Copyright (C) 2018-2022Ruilin Peng (Nick) . + * Copyright (C) 2018-2022 Ruilin Peng (Nick) . * * smartdns is free software: you can redistribute it and/or modify * it under the terms of the GNU General Public License as published by diff --git a/src/lib/nftset.c b/src/lib/nftset.c index 7ac87b2..5b07d79 100644 --- a/src/lib/nftset.c +++ b/src/lib/nftset.c @@ -89,6 +89,7 @@ static int _nftset_addattr(struct nlmsghdr *n, int maxlen, __u16 type, const voi void *rta_data = RTA_DATA(attr); memcpy(rta_data, data, alen); + memset((uint8_t *)rta_data + alen, 0, RTA_ALIGN(len) - len); n->nlmsg_len = newlen; @@ -195,20 +196,29 @@ static int _nftset_socket_init(void) return 0; } -static int _nftset_socket_send(void *msg, int msg_len) +static int _nftset_socket_request(void *msg, int msg_len, void *ret_msg, int ret_msg_len) { - char recvbuff[1024]; int ret = -1; struct pollfd pfds; int do_recv = 0; int last_errno = 0; + int len = 0; if (_nftset_socket_init() != 0) { return -1; } + /* clear pending error message*/ for (;;) { - int len = send(nftset_fd, msg, msg_len, 0); + uint8_t buff[1024]; + ret = recv(nftset_fd, buff, sizeof(buff), MSG_DONTWAIT); + if (ret < 0) { + break; + } + } + + for (;;) { + len = send(nftset_fd, msg, msg_len, 0); if (len == msg_len) { break; } @@ -224,7 +234,7 @@ static int _nftset_socket_send(void *msg, int msg_len) return -1; } - if (dns_conf_nftset_debug_enable == 0) { + if (ret_msg == NULL || ret_msg_len <= 0) { return 0; } @@ -240,9 +250,10 @@ static int _nftset_socket_send(void *msg, int msg_len) return -1; } - memset(recvbuff, 0, sizeof(recvbuff)); + memset(ret_msg, 0, ret_msg_len); + len = 0; for (;;) { - ret = recv(nftset_fd, recvbuff, 1024, 0); + ret = recv(nftset_fd, ret_msg + len, ret_msg_len - len, 0); if (ret < 0) { if (errno == EAGAIN && do_recv == 1) { break; @@ -251,12 +262,14 @@ static int _nftset_socket_send(void *msg, int msg_len) if (errno == EAGAIN && last_errno != 0) { errno = last_errno; } + return -1; } do_recv = 1; + len += ret; - struct nlmsghdr *nlh = (struct nlmsghdr *)recvbuff; + struct nlmsghdr *nlh = (struct nlmsghdr *)ret_msg; if (nlh->nlmsg_type == NLMSG_ERROR) { struct nlmsgerr *err = (struct nlmsgerr *)NLMSG_DATA(nlh); if (err->error != 0) { @@ -264,13 +277,94 @@ static int _nftset_socket_send(void *msg, int msg_len) last_errno = errno; return -1; } + + continue; + } + + if (nlh->nlmsg_type & (NFNL_SUBSYS_NFTABLES << 8)) { + if (nlh->nlmsg_type & NLMSG_DONE) { + break; + } + } + + errno = ENOTSUP; + return -1; + } + + return 0; +} + +static int _nftset_socket_send(void *msg, int msg_len) +{ + char recvbuff[1024]; + + if (dns_conf_nftset_debug_enable == 0) { + return _nftset_socket_request(msg, msg_len, NULL, 0); + } + + return _nftset_socket_request(msg, msg_len, recvbuff, sizeof(recvbuff)); +} + +static int _nftset_get_nftset(int nffamily, const char *table_name, const char *setname, void *buf, void **nextbuf) +{ + struct nlmsgreq *req = (struct nlmsgreq *)buf; + memset(buf, 0, sizeof(struct nlmsgreq)); + + req->h.nlmsg_len = NLMSG_LENGTH(sizeof(struct nfgenmsg)); + req->h.nlmsg_flags = NLM_F_REQUEST | NLM_F_ACK; + req->h.nlmsg_type = NFNL_SUBSYS_NFTABLES << 8 | NFT_MSG_GETSET; + req->h.nlmsg_seq = time(NULL); + + req->m.nfgen_family = nffamily; + req->m.res_id = NFNL_SUBSYS_NFTABLES; + req->m.version = 0; + + struct nlmsghdr *n = &req->h; + + _nftset_addattr_string(n, PAYLOAD_MAX, NFTA_SET_ELEM_LIST_SET, setname); + _nftset_addattr_string(n, PAYLOAD_MAX, NFTA_SET_ELEM_LIST_TABLE, table_name); + + if (nextbuf) { + *nextbuf = (uint8_t *)buf + req->h.nlmsg_len; + } + + return 0; +} + +static int _nftset_get_flags(int nffamily, const char *tablename, const char *setname, uint32_t *flags) +{ + uint8_t buf[PAYLOAD_MAX]; + uint8_t result[PAYLOAD_MAX]; + void *next = buf; + int buffer_len = 0; + + if (flags == NULL) { + return -1; + } + + _nftset_get_nftset(nffamily, tablename, setname, next, &next); + buffer_len = (uint8_t *)next - buf; + int ret = _nftset_socket_request(buf, buffer_len, result, sizeof(result)); + if (ret < 0) { + return -1; + } + + struct nlmsghdr *nlh = (struct nlmsghdr *)result; + struct nfgenmsg *nfmsg = (struct nfgenmsg *)NLMSG_DATA(nlh); + struct nfattr *nfa = (struct nfattr *)NFM_NFA(nfmsg); + *flags = 0; + for (; NFA_OK(nfa, nlh->nlmsg_len); nfa = NFA_NEXT(nfa, nlh->nlmsg_len)) { + if (nfa->nfa_type == NFTA_SET_FLAGS) { + *flags = ntohl(*(uint32_t *)NFA_DATA(nfa)); + break; } } return 0; } -static int _nftset_del_element(int nffamily, const char *table_name, const char *setname, const void *data, int data_len, void *buf, +static int _nftset_del_element(int nffamily, const char *table_name, const char *setname, const void *data, + int data_len, const void *data_interval, int data_interval_len, void *buf, void **nextbuf) { struct nlmsgreq *req = (struct nlmsgreq *)buf; @@ -298,6 +392,18 @@ static int _nftset_del_element(int nffamily, const char *table_name, const char _nftset_addattr(n, PAYLOAD_MAX, NFTA_DATA_VALUE, data, data_len); _nftset_addattr_nest_end(n, nest_elem_key); _nftset_addattr_nest_end(n, nest_elem); + + /* interval attribute */ + if (data_interval && data_interval_len > 0) { + struct rtattr *nest_interval_end = _nftset_addattr_nest(n, PAYLOAD_MAX, NLA_F_NESTED | NFTA_LIST_ELEM); + _nftset_addattr_uint32(n, PAYLOAD_MAX, NFTA_SET_ELEM_FLAGS, htonl(NFT_SET_ELEM_INTERVAL_END)); + struct rtattr *nest_elem_interval_key = _nftset_addattr_nest(n, PAYLOAD_MAX, NLA_F_NESTED | NFTA_SET_ELEM_KEY); + + _nftset_addattr(n, PAYLOAD_MAX, NFTA_DATA_VALUE, data_interval, data_interval_len); + _nftset_addattr_nest_end(n, nest_elem_interval_key); + _nftset_addattr_nest_end(n, nest_interval_end); + } + _nftset_addattr_nest_end(n, nest_list); if (nextbuf) { @@ -307,8 +413,9 @@ static int _nftset_del_element(int nffamily, const char *table_name, const char return 0; } -static int _nftset_add_element(int nffamily, const char *table_name, const char *setname, const void *data, int data_len, - unsigned long timeout, void *buf, void **nextbuf) +static int _nftset_add_element(int nffamily, const char *table_name, const char *setname, const void *data, + int data_len, const void *data_interval, int data_interval_len, unsigned long timeout, + void *buf, void **nextbuf) { struct nlmsgreq *req = (struct nlmsgreq *)buf; memset(buf, 0, sizeof(struct nlmsgreq)); @@ -329,18 +436,28 @@ static int _nftset_add_element(int nffamily, const char *table_name, const char _nftset_addattr_string(n, PAYLOAD_MAX, NFTA_SET_ELEM_LIST_TABLE, table_name); _nftset_addattr_string(n, PAYLOAD_MAX, NFTA_SET_ELEM_LIST_SET, setname); struct rtattr *nest_list = _nftset_addattr_nest(n, PAYLOAD_MAX, NLA_F_NESTED | NFTA_SET_ELEM_LIST_ELEMENTS); - struct rtattr *nest_elem = _nftset_addattr_nest(n, PAYLOAD_MAX, NLA_F_NESTED); + struct rtattr *nest_elem = _nftset_addattr_nest(n, PAYLOAD_MAX, NLA_F_NESTED | NFTA_LIST_ELEM); struct rtattr *nest_elem_key = _nftset_addattr_nest(n, PAYLOAD_MAX, NLA_F_NESTED | NFTA_SET_ELEM_KEY); _nftset_addattr(n, PAYLOAD_MAX, NFTA_DATA_VALUE, data, data_len); _nftset_addattr_nest_end(n, nest_elem_key); - if (timeout > 0) { uint64_t timeout_value = htobe64(timeout * 1000); _nftset_addattr(n, PAYLOAD_MAX, NFTA_SET_ELEM_TIMEOUT, &timeout_value, sizeof(timeout_value)); } - _nftset_addattr_nest_end(n, nest_elem); + + /* interval attribute */ + if (data_interval && data_interval_len > 0) { + struct rtattr *nest_interval_end = _nftset_addattr_nest(n, PAYLOAD_MAX, NLA_F_NESTED | NFTA_LIST_ELEM); + _nftset_addattr_uint32(n, PAYLOAD_MAX, NFTA_SET_ELEM_FLAGS, htonl(NFT_SET_ELEM_INTERVAL_END)); + struct rtattr *nest_elem_interval_key = _nftset_addattr_nest(n, PAYLOAD_MAX, NLA_F_NESTED | NFTA_SET_ELEM_KEY); + + _nftset_addattr(n, PAYLOAD_MAX, NFTA_DATA_VALUE, data_interval, data_interval_len); + _nftset_addattr_nest_end(n, nest_elem_interval_key); + _nftset_addattr_nest_end(n, nest_interval_end); + } + _nftset_addattr_nest_end(n, nest_list); if (nextbuf) { @@ -350,25 +467,79 @@ static int _nftset_add_element(int nffamily, const char *table_name, const char return 0; } +static int _nftset_process_setflags(uint32_t flags, const unsigned char addr[], int addr_len, unsigned long *timeout, + uint8_t **interval_addr, int *interval_addr_len) +{ + uint8_t *addr_end = *interval_addr; + + if ((flags & NFT_SET_TIMEOUT) == 0 && timeout != NULL) { + *timeout = 0; + } + + if ((flags & NFT_SET_INTERVAL) && addr_end != NULL) { + if (addr_len == 4) { + addr_end[0] = addr[0]; + addr_end[1] = addr[1]; + addr_end[2] = addr[2]; + addr_end[3] = addr[3] + 1; + if (addr_end[3] == 0) { + return -1; + } + + *interval_addr_len = 4; + } else if (addr_len == 16) { + memcpy(addr_end, addr, 16); + addr_end[15] = addr[15] + 1; + if (addr_end[15] == 0) { + return -1; + } + *interval_addr_len = 16; + } + } else { + *interval_addr = NULL; + *interval_addr_len = 0; + } + + return 0; +} + static int _nftset_del(int nffamily, const char *tablename, const char *setname, const unsigned char addr[], - int addr_len) + int addr_len, const unsigned char addr_end[], int addr_end_len) { uint8_t buf[PAYLOAD_MAX]; void *next = buf; int buffer_len = 0; _nftset_start_batch(next, &next); - _nftset_del_element(nffamily, tablename, setname, addr, addr_len, next, &next); + _nftset_del_element(nffamily, tablename, setname, addr, addr_len, addr_end, addr_end_len, next, &next); _nftset_end_batch(next, &next); buffer_len = (uint8_t *)next - buf; - return _nftset_socket_send(buf, buffer_len);; + return _nftset_socket_send(buf, buffer_len); } int nftset_del(const char *familyname, const char *tablename, const char *setname, const unsigned char addr[], int addr_len) { int nffamily = _nftset_get_nffamily_from_str(familyname); - int ret = _nftset_del(nffamily, tablename, setname, addr, addr_len); + + uint8_t addr_end_buff[16] = {0}; + uint8_t *addr_end = addr_end_buff; + uint32_t flags = 0; + int addr_end_len = 0; + int ret = -1; + + ret = _nftset_get_flags(nffamily, tablename, setname, &flags); + if (ret == 0) { + ret = _nftset_process_setflags(flags, addr, addr_len, 0, &addr_end, &addr_end_len); + if (ret != 0) { + return -1; + } + } else { + addr_end = NULL; + addr_end_len = 0; + } + + ret = _nftset_del(nffamily, tablename, setname, addr, addr_len, addr_end, addr_end_len); if (ret != 0 && errno != ENOENT) { tlog(TLOG_ERROR, "nftset delete failed, family:%s, table:%s, set:%s, error:%s", familyname, tablename, setname, strerror(errno)); @@ -381,6 +552,10 @@ int nftset_add(const char *familyname, const char *tablename, const char *setnam int addr_len, unsigned long timeout) { uint8_t buf[PAYLOAD_MAX]; + uint8_t addr_end_buff[16] = {0}; + uint8_t *addr_end = addr_end_buff; + uint32_t flags = 0; + int addr_end_len = 0; void *next = buf; int buffer_len = 0; int ret = -1; @@ -390,9 +565,27 @@ int nftset_add(const char *familyname, const char *tablename, const char *setnam timeout = 0; } - _nftset_del(nffamily, tablename, setname, addr, addr_len); + ret = _nftset_get_flags(nffamily, tablename, setname, &flags); + if (ret == 0) { + ret = _nftset_process_setflags(flags, addr, addr_len, &timeout, &addr_end, &addr_end_len); + if (ret != 0) { + if (dns_conf_nftset_debug_enable) { + tlog(TLOG_ERROR, "nftset add failed, family:%s, table:%s, set:%s, error:%s", familyname, tablename, + setname, "ip is invalid"); + } + return -1; + } + } else { + addr_end = NULL; + addr_end_len = 0; + } + + if (timeout > 0) { + _nftset_del(nffamily, tablename, setname, addr, addr_len, addr_end, addr_end_len); + } + _nftset_start_batch(next, &next); - _nftset_add_element(nffamily, tablename, setname, addr, addr_len, timeout, next, &next); + _nftset_add_element(nffamily, tablename, setname, addr, addr_len, addr_end, addr_end_len, timeout, next, &next); _nftset_end_batch(next, &next); buffer_len = (uint8_t *)next - buf;