int keys_num(const char *pattern)
{
int i;
int num = 0;
redisReply *reply;
reply = redisCommand(ctx, "KEYS %s", pattern);
if (reply == NULL)
{
log_error("keys fails: %s", ctx->errstr);
// handle error
return 0;
}
if (reply->type == REDIS_REPLY_ARRAY) {
for (i = 0; i < reply->elements; i++) {
log_debugv("%d) %s", i, reply->element[i]->str);
}
num = reply->elements;
log_debugv("keys num: %s %d", pattern, num);
}
else {
log_warn("keys: unexpected reply type: %d", reply->type);
}
freeReplyObject(reply);
return num;
}
int get_int(const char *key)
{
int val = 0;
redisReply *reply;
reply = redisCommand(ctx, "GET %s", key);
if (reply == NULL)
{
log_error("get_int fails: %s", ctx->errstr);
// handle error
return 0;
}
if (reply->type == REDIS_REPLY_NIL) {
val = 0;
log_debugv("get_int: key not found: %s", key);
}
else if (reply->type == REDIS_REPLY_INTEGER) {
val = reply->integer;
log_debugv("get_int: %s %d", key, val);
}
else if (reply->type == REDIS_REPLY_STRING) {
val = strtol(reply->str, NULL, 10);
log_debugv("get_int: %s %d", key, val);
}
else {
log_warn("get_int: unexpected reply type: %d", reply->type);
}
freeReplyObject(reply);
return val;
}
void get_str(const char *key, char *val, size_t len)
{
redisReply *reply;
reply = redisCommand(ctx, "GET %s", key);
val[0] = 0;
if (reply == NULL)
{
log_error("get_str fails: %s", ctx->errstr);
// handle error
return;
}
if (reply->type == REDIS_REPLY_NIL) {
log_debugv("get_str: key not found: %s", key);
}
else if (reply->type == REDIS_REPLY_STRING) {
strncat(val, reply->str, len);
log_debugv("get_str: %s %s", key, val);
}
else {
log_warn("get_str: unexpected reply type: %d", reply->type);
}
freeReplyObject(reply);
}
// get all the keys matching a certain pattern
int scan_match(const char *pattern, char keys[][64], int max_size)
{
int i, j;
int count = 0;
int cursor = 0;
do {
redisReply *reply;
reply = redisCommand(ctx, "SCAN %d MATCH %s", cursor, pattern);
if (reply == NULL)
{
log_error("scan_match fails: %s", ctx->errstr);
// handle error
return 0;
}
//if (reply->type == REDIS_REPLY_NIL) {
// log_debugv("scan_match: key not found: %s", pattern);
//}
if (reply->type == REDIS_REPLY_ARRAY) {
log_debugv("scan_match: %s", pattern);
for (i = 0; i < reply->elements; i++) {
if (reply->element[i]->type == REDIS_REPLY_STRING) {
log_debugv("%u) %s", i+1, reply->element[i]->str);
}
else if (reply->element[i]->type == REDIS_REPLY_ARRAY) {
for (j = 0; j < reply->element[i]->elements; j++) {
if (reply->element[i]->element[j]->type == REDIS_REPLY_STRING) {
log_debugv("%u) %u) %s", i+1, j+1, reply->element[i]->element[j]->str);
if (count >= max_size) {
log_warn("scan_match: exceeds max size of %d", max_size);
return max_size;
}
keys[count][0] = 0;
strncat(keys[count], reply->element[i]->element[j]->str, 64);
count++;
}
else
log_warn("scan_match: unexpected reply type: %d",
reply->element[i]->element[j]->type);
}
}
else
log_warn("scan_match: unexpected reply type: %d",
reply->element[i]->type);
}
}
else {
log_warn("scan_match: unexpected reply type: %d", reply->type);
}
cursor = strtol(reply->element[0]->str, NULL, 10);
freeReplyObject(reply);
}
while (cursor != 0);
log_debugv("done");
return count;
}
void log_debug(const char *format, ...)
{
va_list ap;
va_start(ap, format);
log_debugv(format, ap);
va_end(ap);
}
void set_int_ex_nx(const char *key, int val, int timeout)
{
redisReply *reply;
reply = redisCommand(ctx, "SET %s %d EX %d NX", key, val, timeout);
if (reply == NULL)
{
log_error("set_int_ex_nx fails: %s", ctx->errstr);
// handle error
return;
}
if (reply->type == REDIS_REPLY_STATUS) {
log_debugv("set_int_ex_nx %s %d %s", key, val, reply->str);
}
else if (reply->type == REDIS_REPLY_NIL) {
log_debugv("set_int_ex_nx: %s exists", key);
}
else {
log_warn("set_int_ex_nx: unexpected reply type: %d", reply->type);
}
freeReplyObject(reply);
}
void on_expire(redisAsyncContext *ac, void *r, void *privdata) {
int i;
redisReply *reply = r;
if (reply == NULL) return;
if (reply->type == REDIS_REPLY_ARRAY && strcmp(reply->element[0]->str, "message") == 0) {
if (reply->elements == 3) {
log_debugv("expired: %s", reply->element[2]->str);
on_request_expire(reply->element[2]->str);
}
else {
log_warn("on_expire: number of element is not 3. (%d)", reply->elements);
}
}
}
void expire(const char *key, int timeout)
{
redisReply *reply;
reply = redisCommand(ctx, "EXPIRE %s %d", key, timeout);
if (reply == NULL)
{
log_error("expire fails: %s", ctx->errstr);
// handle error
return;
}
if (reply->type == REDIS_REPLY_INTEGER) {
log_debugv("expire: %s %d", key, reply->integer);
}
else {
log_warn("expire: unexpected reply type: %d", reply->type);
}
freeReplyObject(reply);
}
void set_str_ex(const char *key, const char *val, int timeout)
{
redisReply *reply;
reply = redisCommand(ctx, "SETEX %s %d %s", key, timeout, val);
if (reply == NULL)
{
log_error("set_str_ex fails: %s", ctx->errstr);
// handle error
return;
}
if (reply->type == REDIS_REPLY_STATUS) {
log_debugv("set_str_ex %s %s %s", key, val, reply->str);
}
else {
log_warn("set_str_ex: unexpected reply type: %d", reply->type);
}
freeReplyObject(reply);
}
void del_key(const char *key)
{
int i;
redisReply *reply;
reply = redisCommand(ctx, "DEL %s", key);
if (reply == NULL)
{
log_error("del_key fails: %s", ctx->errstr);
// handle error
return;
}
if (reply->type == REDIS_REPLY_INTEGER) {
log_debugv("del_key: %s %d", key, reply->integer);
}
else {
log_warn("del_key: unexpected reply type: %d", reply->type);
}
freeReplyObject(reply);
}
int incr(const char *key)
{
int val = 0;
redisReply *reply;
reply = redisCommand(ctx, "INCR %s", key);
if (reply == NULL)
{
log_error("incr fails: %s", ctx->errstr);
// handle error
return 0;
}
if (reply->type == REDIS_REPLY_INTEGER) {
val = reply->integer;
log_debugv("incr %s %d", key, val);
}
else {
log_warn("incr: unexpected reply type: %d", reply->type);
}
freeReplyObject(reply);
return val;
}
int dns_proxy_loop()
{
int sockfd;
struct sockaddr_in local_addr, remote_addr;
socklen_t addrlen = sizeof(remote_addr);
int retval;
int recvlen;
char buf[MAX_PACKET_SIZE];
fd_set rfds;
struct timeval tv;
time_t now, last_ack_time;
load_dns_blacklist();
//print_blacklist();
// create the listening socket
if ((sockfd = socket(AF_INET, SOCK_DGRAM, 0)) < 0) {
log_error("dns_main_loop: Cannot create DNS socket.");
return -1;
}
memset((char*)&local_addr, 0, sizeof(local_addr));
local_addr.sin_family = AF_INET;
local_addr.sin_port = htons(LOCAL_DNS_PORT);
local_addr.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
if (bind(sockfd, (struct sockaddr*)&local_addr, sizeof(local_addr)) < 0) {
log_error("dns_main_loop: Bind failed.");
return -1;
}
// init connection to public resolver using TCP
if (init_dns_tcp_conn() < 0) {
log_error("dns_main_loop: Init DNS TCP conn failed.");
return -1;
}
last_ack_time = time(NULL);
while (1) {
FD_ZERO(&rfds);
FD_SET(sockfd, &rfds);
tv.tv_sec = 1;
tv.tv_usec = 0;
retval = select(sockfd+1, &rfds, NULL, NULL, &tv);
if (retval == -1)
log_error("select() failed. errno: %d", errno);
else if (retval) {
recvlen = recvfrom(sockfd, buf, MAX_PACKET_SIZE, 0, (struct sockaddr*)&remote_addr, &addrlen);
log_debugv("received %d bytes", recvlen);
if (recvlen > 0) {
send_dns_req_over_tcp(buf, recvlen);
}
}
// need to send a normal DNS request in order to keep connection alive, does it worth?
now = time(NULL);
if (now - last_ack_time > KEEP_ALIVE_INTERVAL) {
send_keep_alive_packet();
last_ack_time = now;
}
}
return 0;
}
/* Process UDP packets
* Return 0 to accept packet, otherwise to drop packet
*/
int process_udp_packet(struct mypacket *packet, char inout)
{
int i, ret;
struct myiphdr *iphdr = packet->iphdr;
struct myudphdr *udphdr = packet->udphdr;
char sip[16], dip[16];
ip2str(iphdr->saddr, sip);
ip2str(iphdr->daddr, dip);
unsigned short sport, dport;
sport = ntohs(udphdr->uh_sport);
dport = ntohs(udphdr->uh_dport);
//log_debug("[UDP] This packet goes from %s:%d to %s:%d", sip, sport, dip, dport);
struct fourtuple fourtp;
fourtp.saddr = iphdr->saddr;
fourtp.daddr = iphdr->daddr;
fourtp.sport = udphdr->uh_sport;
fourtp.dport = udphdr->uh_dport;
if (dport == 53 || sport == 53) {
// Parse DNS header
struct mydnshdr *dnshdr = (struct mydnshdr*)packet->payload;
unsigned short txn_id = ntohs(dnshdr->txn_id);
int qdcount = ntohs(dnshdr->questions);
char qname[MAX_QNAME_LEN];
if (qdcount > 0) {
int flag = 0;
//log_debug("Questions: %d", qdcount);
unsigned char *ptr = (unsigned char *)(dnshdr + 1);
for (i = 0; i < qdcount; i++) {
struct mydnsquery query;
int j = 0, l;
while (1) {
for (l = *ptr++; l != 0; l--) {
query.qname[j++] = *ptr++;
if (j >= MAX_QNAME_LEN) {
while (*ptr != 0) ptr++;
break;
}
}
if (*ptr == 0) {
query.qname[j] = 0;
ptr++;
break;
}
query.qname[j++] = '.';
}
query.qtype = (ptr[0] << 8) + ptr[1];
query.qclass = (ptr[2] << 8) + ptr[3];
log_debug("DNS Query: %s %d %d", query.qname, query.qtype, query.qclass);
// save the first query name for later usage
if (i == 0) {
qname[0] = 0;
strncat(qname, query.qname, MAX_QNAME_LEN - 1);
}
// check if qname in blacklist
if (is_poisoned_domain(qname))
flag = 1;
}
if (dport == 53) {
// Process outgoing DNS requests
log_debug("[UDP] Sent a DNS request from %s:%d to %s:%d.", sip, sport, dip, dport);
if (opt_protect_dns_protocol == 1 &&
(opt_dns_only_blacklisted == 0 || opt_dns_only_blacklisted == 1 && flag)) {
log_debug("Redirecting to TCP.");
log_debugv("[EVAL] DNS TCP request %d", txn_id);
// Tell the caching thread to cache the request
// use DNS transaction ID and first query name as unique ID
// transaction ID alone may cause collision
cache_dns_udp_request(txn_id, qname, &fourtp);
// send the request over TCP
ret = send_dns_req(packet->payload, packet->payload_len);
if (ret == 0) {
// drop the packet
return -1;
} else {
log_error("DNS redirect failed.");
}
//send_udp2(packet);
}
else {
log_debugv("[EVAL] DNS UDP request %d", txn_id);
}
}
else if (sport == 53) {
if (opt_protect_dns_protocol == 1 &&
(opt_dns_only_blacklisted == 0 || opt_dns_only_blacklisted == 1 && flag)) {
// the response must be sent by ourself
return 0;
}
// Process incoming DNS responses
log_debug("[UDP] Got a DNS response from %s:%d to %s:%d.", sip, sport, dip, dport);
log_debugv("[EVAL] DNS UDP response %d", txn_id);
// Tell the caching thread to process the dns udp response
process_dns_udp_response(txn_id, qname, &fourtp, iphdr->ttl);
}
}
else {
log_error("[UDP] DNS request has 0 question.");
}
}
return 0;
}