dns_cache.c 20 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809
  1. /*************************************************************************
  2. *
  3. * Copyright (C) 2018-2023 Ruilin Peng (Nick) <[email protected]>.
  4. *
  5. * smartdns is free software: you can redistribute it and/or modify
  6. * it under the terms of the GNU General Public License as published by
  7. * the Free Software Foundation, either version 3 of the License, or
  8. * (at your option) any later version.
  9. *
  10. * smartdns is distributed in the hope that it will be useful,
  11. * but WITHOUT ANY WARRANTY; without even the implied warranty of
  12. * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
  13. * GNU General Public License for more details.
  14. *
  15. * You should have received a copy of the GNU General Public License
  16. * along with this program. If not, see <http://www.gnu.org/licenses/>.
  17. */
  18. #include "dns_cache.h"
  19. #include "stringutil.h"
  20. #include "timer.h"
  21. #include "tlog.h"
  22. #include "util.h"
  23. #include <errno.h>
  24. #include <fcntl.h>
  25. #include <pthread.h>
  26. #include <stdio.h>
  27. #include <string.h>
  28. #include <sys/types.h>
  29. #define DNS_CACHE_MAX_HITNUM 6000
  30. #define DNS_CACHE_HITNUM_STEP 3
  31. #define DNS_CACHE_HITNUM_STEP_MAX 6
  32. #define DNS_CACHE_READ_TIMEOUT 60
  33. struct dns_cache_head {
  34. struct hash_table cache_hash;
  35. struct list_head cache_list;
  36. atomic_t num;
  37. int size;
  38. pthread_mutex_t lock;
  39. dns_cache_callback timeout_callback;
  40. };
  41. typedef int (*dns_cache_read_callback)(struct dns_cache_record *cache_record, struct dns_cache_data *cache_data);
  42. static int is_cache_init;
  43. static struct dns_cache_head dns_cache_head;
  44. int dns_cache_init(int size, dns_cache_callback timeout_callback)
  45. {
  46. int bits = 0;
  47. if (is_cache_init == 1) {
  48. return -1;
  49. }
  50. INIT_LIST_HEAD(&dns_cache_head.cache_list);
  51. bits = ilog2(size) - 1;
  52. if (bits >= 20) {
  53. bits = 20;
  54. } else if (bits < 12) {
  55. bits = 12;
  56. }
  57. hash_table_init(dns_cache_head.cache_hash, bits, malloc);
  58. atomic_set(&dns_cache_head.num, 0);
  59. dns_cache_head.size = size;
  60. dns_cache_head.timeout_callback = timeout_callback;
  61. pthread_mutex_init(&dns_cache_head.lock, NULL);
  62. is_cache_init = 1;
  63. return 0;
  64. }
  65. static struct dns_cache *_dns_cache_first(void)
  66. {
  67. return list_first_entry_or_null(&dns_cache_head.cache_list, struct dns_cache, list);
  68. }
  69. static void _dns_cache_delete(struct dns_cache *dns_cache)
  70. {
  71. hash_del(&dns_cache->node);
  72. list_del_init(&dns_cache->list);
  73. atomic_dec(&dns_cache_head.num);
  74. dns_cache_data_put(dns_cache->cache_data);
  75. dns_cache->cache_data = NULL;
  76. free(dns_cache);
  77. }
  78. void dns_cache_get(struct dns_cache *dns_cache)
  79. {
  80. if (atomic_inc_return(&dns_cache->ref) == 1) {
  81. tlog(TLOG_ERROR, "BUG: dns_cache is invalid.");
  82. return;
  83. }
  84. }
  85. void dns_cache_release(struct dns_cache *dns_cache)
  86. {
  87. if (dns_cache == NULL) {
  88. return;
  89. }
  90. if (!atomic_dec_and_test(&dns_cache->ref)) {
  91. return;
  92. }
  93. _dns_cache_delete(dns_cache);
  94. }
  95. static void _dns_cache_remove(struct dns_cache *dns_cache)
  96. {
  97. hash_del(&dns_cache->node);
  98. list_del_init(&dns_cache->list);
  99. dns_timer_del(&dns_cache->timer);
  100. dns_cache_release(dns_cache);
  101. }
  102. uint32_t dns_cache_get_query_flag(struct dns_cache *dns_cache)
  103. {
  104. return dns_cache->info.query_flag;
  105. }
  106. const char *dns_cache_get_dns_group_name(struct dns_cache *dns_cache)
  107. {
  108. return dns_cache->info.dns_group_name;
  109. }
  110. struct dns_cache_data *dns_cache_new_data_packet(void *packet, size_t packet_len)
  111. {
  112. struct dns_cache_packet *cache_packet = NULL;
  113. size_t data_size = 0;
  114. if (packet == NULL || packet_len <= 0) {
  115. return NULL;
  116. }
  117. data_size = sizeof(*cache_packet) + packet_len;
  118. cache_packet = malloc(data_size);
  119. if (cache_packet == NULL) {
  120. return NULL;
  121. }
  122. memcpy(cache_packet->data, packet, packet_len);
  123. memset(&cache_packet->head, 0, sizeof(cache_packet->head));
  124. cache_packet->head.size = packet_len;
  125. cache_packet->head.magic = MAGIC_CACHE_DATA;
  126. atomic_set(&cache_packet->head.ref, 1);
  127. return (struct dns_cache_data *)cache_packet;
  128. }
  129. static void dns_cache_timer_release(struct tw_timer_list *timer, void *data)
  130. {
  131. struct dns_cache *dns_cache = data;
  132. dns_cache_release(dns_cache);
  133. }
  134. static void dns_cache_expired(struct tw_timer_list *timer, void *data, unsigned long timestamp)
  135. {
  136. struct dns_cache *dns_cache = data;
  137. if (dns_cache->del_pending == 1) {
  138. dns_cache_release(dns_cache);
  139. return;
  140. }
  141. if (dns_cache_head.timeout_callback) {
  142. if (dns_cache_head.timeout_callback(dns_cache) != 0) {
  143. dns_cache_release(dns_cache);
  144. return;
  145. }
  146. }
  147. dns_cache->del_pending = 1;
  148. dns_timer_mod(&dns_cache->timer, 5);
  149. }
  150. static int _dns_cache_replace(struct dns_cache_key *cache_key, int rcode, int ttl, int speed, int timeout,
  151. int update_time, struct dns_cache_data *cache_data)
  152. {
  153. struct dns_cache *dns_cache = NULL;
  154. struct dns_cache_data *old_cache_data = NULL;
  155. if (dns_cache_head.size <= 0) {
  156. return 0;
  157. }
  158. /* lookup existing cache */
  159. dns_cache = dns_cache_lookup(cache_key);
  160. if (dns_cache == NULL) {
  161. return -1;
  162. }
  163. if (ttl < DNS_CACHE_TTL_MIN) {
  164. ttl = DNS_CACHE_TTL_MIN;
  165. }
  166. /* update cache data */
  167. pthread_mutex_lock(&dns_cache_head.lock);
  168. dns_cache->del_pending = 0;
  169. dns_cache->info.qtype = cache_key->qtype;
  170. dns_cache->info.query_flag = cache_key->query_flag;
  171. dns_cache->info.ttl = ttl;
  172. dns_cache->info.speed = speed;
  173. dns_cache->info.timeout = timeout;
  174. dns_cache->info.is_visited = 1;
  175. if (cache_data) {
  176. old_cache_data = dns_cache->cache_data;
  177. dns_cache->cache_data = cache_data;
  178. }
  179. if (update_time) {
  180. time(&dns_cache->info.insert_time);
  181. }
  182. time(&dns_cache->info.replace_time);
  183. list_del(&dns_cache->list);
  184. list_add_tail(&dns_cache->list, &dns_cache_head.cache_list);
  185. dns_timer_mod(&dns_cache->timer, timeout);
  186. pthread_mutex_unlock(&dns_cache_head.lock);
  187. if (old_cache_data) {
  188. dns_cache_data_put(old_cache_data);
  189. }
  190. dns_cache_release(dns_cache);
  191. return 0;
  192. }
  193. int dns_cache_replace(struct dns_cache_key *cache_key, int rcode, int ttl, int speed, int timeout, int update_time,
  194. struct dns_cache_data *cache_data)
  195. {
  196. return _dns_cache_replace(cache_key, rcode, ttl, speed, timeout, update_time, cache_data);
  197. }
  198. static void _dns_cache_remove_by_domain(struct dns_cache_key *cache_key)
  199. {
  200. uint32_t key = 0;
  201. struct dns_cache *dns_cache = NULL;
  202. key = hash_string(cache_key->domain);
  203. key = jhash(&cache_key->qtype, sizeof(cache_key->qtype), key);
  204. key = hash_string_initval(cache_key->dns_group_name, key);
  205. key = jhash(&cache_key->query_flag, sizeof(cache_key->query_flag), key);
  206. pthread_mutex_lock(&dns_cache_head.lock);
  207. hash_table_for_each_possible(dns_cache_head.cache_hash, dns_cache, node, key)
  208. {
  209. if (dns_cache->info.qtype != cache_key->qtype) {
  210. continue;
  211. }
  212. if (dns_cache->info.query_flag != cache_key->query_flag) {
  213. continue;
  214. }
  215. if (strncmp(cache_key->domain, dns_cache->info.domain, DNS_MAX_CNAME_LEN) != 0) {
  216. continue;
  217. }
  218. if (strncmp(dns_cache->info.dns_group_name, cache_key->dns_group_name, DNS_GROUP_NAME_LEN) != 0) {
  219. continue;
  220. }
  221. _dns_cache_remove(dns_cache);
  222. break;
  223. }
  224. pthread_mutex_unlock(&dns_cache_head.lock);
  225. }
  226. static int _dns_cache_insert(struct dns_cache_info *info, struct dns_cache_data *cache_data, struct list_head *head)
  227. {
  228. uint32_t key = 0;
  229. struct dns_cache *dns_cache = NULL;
  230. /* if cache already exists, free */
  231. struct dns_cache_key cache_key;
  232. cache_key.qtype = info->qtype;
  233. cache_key.query_flag = info->query_flag;
  234. cache_key.domain = info->domain;
  235. cache_key.dns_group_name = info->dns_group_name;
  236. _dns_cache_remove_by_domain(&cache_key);
  237. dns_cache = malloc(sizeof(*dns_cache));
  238. if (dns_cache == NULL) {
  239. goto errout;
  240. }
  241. memset(dns_cache, 0, sizeof(*dns_cache));
  242. key = hash_string(info->domain);
  243. key = jhash(&info->qtype, sizeof(info->qtype), key);
  244. key = hash_string_initval(info->dns_group_name, key);
  245. key = jhash(&info->query_flag, sizeof(info->query_flag), key);
  246. atomic_set(&dns_cache->ref, 1);
  247. memcpy(&dns_cache->info, info, sizeof(*info));
  248. dns_cache->del_pending = 0;
  249. dns_cache->cache_data = cache_data;
  250. dns_cache->timer.function = dns_cache_expired;
  251. dns_cache->timer.del_function = dns_cache_timer_release;
  252. dns_cache->timer.expires = info->timeout;
  253. dns_cache->timer.data = dns_cache;
  254. pthread_mutex_lock(&dns_cache_head.lock);
  255. hash_table_add(dns_cache_head.cache_hash, &dns_cache->node, key);
  256. list_add_tail(&dns_cache->list, head);
  257. INIT_LIST_HEAD(&dns_cache->check_list);
  258. /* Release extra cache, remove oldest cache record */
  259. if (atomic_inc_return(&dns_cache_head.num) > dns_cache_head.size) {
  260. struct dns_cache *del_cache = NULL;
  261. del_cache = _dns_cache_first();
  262. if (del_cache) {
  263. _dns_cache_remove(del_cache);
  264. }
  265. }
  266. dns_cache_get(dns_cache);
  267. dns_timer_add(&dns_cache->timer);
  268. pthread_mutex_unlock(&dns_cache_head.lock);
  269. return 0;
  270. errout:
  271. if (dns_cache) {
  272. dns_cache_release(dns_cache);
  273. }
  274. return -1;
  275. }
  276. int dns_cache_insert(struct dns_cache_key *cache_key, int rcode, int ttl, int speed, int timeout,
  277. struct dns_cache_data *cache_data)
  278. {
  279. struct dns_cache_info info;
  280. if (cache_data == NULL || cache_key == NULL || cache_key->dns_group_name == NULL || cache_key->domain == NULL) {
  281. return -1;
  282. }
  283. if (dns_cache_head.size <= 0) {
  284. dns_cache_data_put(cache_data);
  285. return 0;
  286. }
  287. if (ttl < DNS_CACHE_TTL_MIN) {
  288. ttl = DNS_CACHE_TTL_MIN;
  289. }
  290. memset(&info, 0, sizeof(info));
  291. info.hitnum = 3;
  292. safe_strncpy(info.domain, cache_key->domain, DNS_MAX_CNAME_LEN);
  293. info.qtype = cache_key->qtype;
  294. safe_strncpy(info.dns_group_name, cache_key->dns_group_name, DNS_GROUP_NAME_LEN);
  295. info.query_flag = cache_key->query_flag;
  296. info.ttl = ttl;
  297. info.hitnum_update_add = DNS_CACHE_HITNUM_STEP;
  298. info.speed = speed;
  299. info.timeout = timeout;
  300. info.is_visited = 1;
  301. info.rcode = rcode;
  302. time(&info.insert_time);
  303. time(&info.replace_time);
  304. return _dns_cache_insert(&info, cache_data, &dns_cache_head.cache_list);
  305. }
  306. struct dns_cache *dns_cache_lookup(struct dns_cache_key *cache_key)
  307. {
  308. uint32_t key = 0;
  309. struct dns_cache *dns_cache = NULL;
  310. struct dns_cache *dns_cache_ret = NULL;
  311. time_t now = 0;
  312. if (dns_cache_head.size <= 0) {
  313. return NULL;
  314. }
  315. key = hash_string(cache_key->domain);
  316. key = jhash(&cache_key->qtype, sizeof(cache_key->qtype), key);
  317. key = hash_string_initval(cache_key->dns_group_name, key);
  318. key = jhash(&cache_key->query_flag, sizeof(cache_key->query_flag), key);
  319. time(&now);
  320. /* find cache */
  321. pthread_mutex_lock(&dns_cache_head.lock);
  322. hash_table_for_each_possible(dns_cache_head.cache_hash, dns_cache, node, key)
  323. {
  324. if (dns_cache->info.qtype != cache_key->qtype) {
  325. continue;
  326. }
  327. if (strncmp(cache_key->domain, dns_cache->info.domain, DNS_MAX_CNAME_LEN) != 0) {
  328. continue;
  329. }
  330. if (strncmp(cache_key->dns_group_name, dns_cache->info.dns_group_name, DNS_GROUP_NAME_LEN) != 0) {
  331. continue;
  332. }
  333. if (cache_key->query_flag != dns_cache->info.query_flag) {
  334. continue;
  335. }
  336. dns_cache_ret = dns_cache;
  337. break;
  338. }
  339. if (dns_cache_ret) {
  340. dns_cache_get(dns_cache_ret);
  341. }
  342. pthread_mutex_unlock(&dns_cache_head.lock);
  343. return dns_cache_ret;
  344. }
  345. int dns_cache_get_ttl(struct dns_cache *dns_cache)
  346. {
  347. time_t now = 0;
  348. int ttl = 0;
  349. time(&now);
  350. ttl = dns_cache->info.insert_time + dns_cache->info.ttl - now;
  351. if (ttl < 0) {
  352. return 0;
  353. }
  354. return ttl;
  355. }
  356. struct dns_cache_data *dns_cache_get_data(struct dns_cache *dns_cache)
  357. {
  358. struct dns_cache_data *cache_data;
  359. pthread_mutex_lock(&dns_cache_head.lock);
  360. dns_cache_data_get(dns_cache->cache_data);
  361. cache_data = dns_cache->cache_data;
  362. pthread_mutex_unlock(&dns_cache_head.lock);
  363. return cache_data;
  364. }
  365. void dns_cache_data_get(struct dns_cache_data *cache_data)
  366. {
  367. if (atomic_inc_return(&cache_data->head.ref) == 1) {
  368. tlog(TLOG_ERROR, "BUG: dns_cache data is invalid.");
  369. return;
  370. }
  371. return;
  372. }
  373. void dns_cache_data_put(struct dns_cache_data *cache_data)
  374. {
  375. if (cache_data == NULL) {
  376. return;
  377. }
  378. if (!atomic_dec_and_test(&cache_data->head.ref)) {
  379. return;
  380. }
  381. free(cache_data);
  382. }
  383. int dns_cache_is_visited(struct dns_cache *dns_cache)
  384. {
  385. return dns_cache->info.is_visited;
  386. }
  387. void dns_cache_delete(struct dns_cache *dns_cache)
  388. {
  389. pthread_mutex_lock(&dns_cache_head.lock);
  390. _dns_cache_remove(dns_cache);
  391. pthread_mutex_unlock(&dns_cache_head.lock);
  392. }
  393. int dns_cache_hitnum_dec_get(struct dns_cache *dns_cache)
  394. {
  395. pthread_mutex_lock(&dns_cache_head.lock);
  396. dns_cache->info.hitnum--;
  397. if (dns_cache->info.hitnum_update_add > DNS_CACHE_HITNUM_STEP) {
  398. dns_cache->info.hitnum_update_add--;
  399. }
  400. pthread_mutex_unlock(&dns_cache_head.lock);
  401. return dns_cache->info.hitnum;
  402. }
  403. void dns_cache_update(struct dns_cache *dns_cache)
  404. {
  405. pthread_mutex_lock(&dns_cache_head.lock);
  406. if (!list_empty(&dns_cache->list)) {
  407. list_del_init(&dns_cache->list);
  408. list_add_tail(&dns_cache->list, &dns_cache_head.cache_list);
  409. dns_cache->info.hitnum += dns_cache->info.hitnum_update_add;
  410. if (dns_cache->info.hitnum > DNS_CACHE_MAX_HITNUM) {
  411. dns_cache->info.hitnum = DNS_CACHE_MAX_HITNUM;
  412. }
  413. if (dns_cache->info.hitnum_update_add < DNS_CACHE_HITNUM_STEP_MAX) {
  414. dns_cache->info.hitnum_update_add++;
  415. }
  416. dns_cache->info.is_visited = 1;
  417. }
  418. pthread_mutex_unlock(&dns_cache_head.lock);
  419. }
  420. static int _dns_cache_read_to_cache(struct dns_cache_record *cache_record, struct dns_cache_data *cache_data)
  421. {
  422. struct list_head *head = NULL;
  423. head = &dns_cache_head.cache_list;
  424. struct dns_cache_info *info = &cache_record->info;
  425. time_t now = time(NULL);
  426. unsigned int seed_tmp = now;
  427. int passed_time = now - info->replace_time;
  428. int timeout = info->timeout - passed_time;
  429. if (timeout < DNS_CACHE_READ_TIMEOUT * 2) {
  430. timeout = DNS_CACHE_READ_TIMEOUT + (rand_r(&seed_tmp) % DNS_CACHE_READ_TIMEOUT);
  431. }
  432. if (timeout > dns_conf_serve_expired_ttl && dns_conf_serve_expired_ttl >= 0) {
  433. timeout = dns_conf_serve_expired_ttl;
  434. }
  435. info->timeout = timeout;
  436. if (_dns_cache_insert(&cache_record->info, cache_data, head) != 0) {
  437. tlog(TLOG_ERROR, "insert cache data failed.");
  438. cache_data = NULL;
  439. goto errout;
  440. }
  441. dns_cache_data_get(cache_data);
  442. daemon_keepalive();
  443. return 0;
  444. errout:
  445. return -1;
  446. }
  447. static int _dns_cache_read_record(int fd, uint32_t cache_number, dns_cache_read_callback callback)
  448. {
  449. unsigned int i = 0;
  450. ssize_t ret = 0;
  451. struct dns_cache_record cache_record;
  452. struct dns_cache_data_head data_head;
  453. struct dns_cache_data *cache_data = NULL;
  454. for (i = 0; i < cache_number; i++) {
  455. ret = read(fd, &cache_record, sizeof(cache_record));
  456. if (ret != sizeof(cache_record)) {
  457. tlog(TLOG_ERROR, "read cache failed, %s", strerror(errno));
  458. goto errout;
  459. }
  460. if (cache_record.magic != MAGIC_RECORD) {
  461. tlog(TLOG_ERROR, "magic is invalid.");
  462. goto errout;
  463. }
  464. ret = read(fd, &data_head, sizeof(data_head));
  465. if (ret != sizeof(data_head)) {
  466. tlog(TLOG_ERROR, "read data head failed, %s", strerror(errno));
  467. goto errout;
  468. }
  469. if (data_head.magic != MAGIC_CACHE_DATA) {
  470. tlog(TLOG_ERROR, "data magic is invalid.");
  471. goto errout;
  472. }
  473. if (data_head.size > 1024 * 8) {
  474. tlog(TLOG_ERROR, "data may invalid, skip load cache.");
  475. goto errout;
  476. }
  477. cache_data = malloc(data_head.size + sizeof(data_head));
  478. if (cache_data == NULL) {
  479. tlog(TLOG_ERROR, "malloc cache data failed %s", strerror(errno));
  480. goto errout;
  481. }
  482. memcpy(&cache_data->head, &data_head, sizeof(data_head));
  483. atomic_set(&cache_data->head.ref, 1);
  484. ret = read(fd, cache_data->data, data_head.size);
  485. if (ret != data_head.size) {
  486. tlog(TLOG_ERROR, "read cache data failed, %s", strerror(errno));
  487. goto errout;
  488. }
  489. /* set cache unvisited, so that when refreshing ipset/nftset, reload ipset list by restarting smartdns */
  490. cache_record.info.is_visited = 0;
  491. cache_record.info.domain[DNS_MAX_CNAME_LEN - 1] = '\0';
  492. cache_record.info.dns_group_name[DNS_GROUP_NAME_LEN - 1] = '\0';
  493. ret = callback(&cache_record, cache_data);
  494. dns_cache_data_put(cache_data);
  495. cache_data = NULL;
  496. if (ret != 0) {
  497. goto errout;
  498. }
  499. }
  500. return 0;
  501. errout:
  502. if (cache_data) {
  503. dns_cache_data_put(cache_data);
  504. }
  505. return -1;
  506. }
  507. static int _dns_cache_file_read(const char *file, dns_cache_read_callback callback)
  508. {
  509. int fd = -1;
  510. ssize_t ret = 0;
  511. off_t filesize = 0;
  512. fd = open(file, O_RDONLY);
  513. if (fd < 0) {
  514. return 0;
  515. }
  516. filesize = lseek(fd, 0, SEEK_END);
  517. lseek(fd, 0, SEEK_SET);
  518. posix_fadvise(fd, 0, filesize, POSIX_FADV_WILLNEED | POSIX_FADV_SEQUENTIAL);
  519. struct dns_cache_file cache_file;
  520. ret = read(fd, &cache_file, sizeof(cache_file));
  521. if (ret != sizeof(cache_file)) {
  522. tlog(TLOG_ERROR, "read cache head failed.");
  523. goto errout;
  524. }
  525. if (cache_file.magic != MAGIC_NUMBER) {
  526. tlog(TLOG_ERROR, "cache file is invalid.");
  527. goto errout;
  528. }
  529. if (strncmp(cache_file.version, dns_cache_file_version(), DNS_CACHE_VERSION_LEN) != 0) {
  530. tlog(TLOG_WARN, "cache version is different, skip load cache.");
  531. goto errout;
  532. }
  533. tlog(TLOG_INFO, "load cache file %s, total %d records", file, cache_file.cache_number);
  534. if (_dns_cache_read_record(fd, cache_file.cache_number, callback) != 0) {
  535. goto errout;
  536. }
  537. close(fd);
  538. return 0;
  539. errout:
  540. if (fd > 0) {
  541. close(fd);
  542. }
  543. return -1;
  544. }
  545. int dns_cache_load(const char *file)
  546. {
  547. return _dns_cache_file_read(file, _dns_cache_read_to_cache);
  548. }
  549. static int _dns_cache_write_record(int fd, uint32_t *cache_number, struct list_head *head)
  550. {
  551. struct dns_cache *dns_cache = NULL;
  552. struct dns_cache *tmp = NULL;
  553. struct dns_cache_record cache_record;
  554. memset(&cache_record, 0, sizeof(cache_record));
  555. pthread_mutex_lock(&dns_cache_head.lock);
  556. list_for_each_entry_safe(dns_cache, tmp, head, list)
  557. {
  558. cache_record.magic = MAGIC_RECORD;
  559. memcpy(&cache_record.info, &dns_cache->info, sizeof(struct dns_cache_info));
  560. ssize_t ret = write(fd, &cache_record, sizeof(cache_record));
  561. if (ret != sizeof(cache_record)) {
  562. tlog(TLOG_ERROR, "write cache failed, %s", strerror(errno));
  563. goto errout;
  564. }
  565. struct dns_cache_data *cache_data = dns_cache->cache_data;
  566. ret = write(fd, cache_data, sizeof(*cache_data) + cache_data->head.size);
  567. if (ret != (int)sizeof(*cache_data) + cache_data->head.size) {
  568. tlog(TLOG_ERROR, "write cache data failed, %s", strerror(errno));
  569. goto errout;
  570. }
  571. (*cache_number)++;
  572. }
  573. pthread_mutex_unlock(&dns_cache_head.lock);
  574. return 0;
  575. errout:
  576. pthread_mutex_unlock(&dns_cache_head.lock);
  577. return -1;
  578. }
  579. static int _dns_cache_write_records(int fd, uint32_t *cache_number)
  580. {
  581. if (_dns_cache_write_record(fd, cache_number, &dns_cache_head.cache_list) != 0) {
  582. return -1;
  583. }
  584. return 0;
  585. }
  586. int dns_cache_save(const char *file, int check_lock)
  587. {
  588. int fd = -1;
  589. uint32_t cache_number = 0;
  590. tlog(TLOG_DEBUG, "write cache file %s", file);
  591. /* check lock */
  592. if (check_lock == 1) {
  593. if (pthread_mutex_trylock(&dns_cache_head.lock) != 0) {
  594. return -1;
  595. }
  596. pthread_mutex_unlock(&dns_cache_head.lock);
  597. }
  598. fd = open(file, O_TRUNC | O_CREAT | O_WRONLY, 0640);
  599. if (fd < 0) {
  600. tlog(TLOG_ERROR, "create file %s failed, %s", file, strerror(errno));
  601. goto errout;
  602. }
  603. struct dns_cache_file cache_file;
  604. memset(&cache_file, 0, sizeof(cache_file));
  605. cache_file.magic = MAGIC_NUMBER;
  606. safe_strncpy(cache_file.version, dns_cache_file_version(), DNS_CACHE_VERSION_LEN);
  607. cache_file.cache_number = 0;
  608. if (lseek(fd, sizeof(cache_file), SEEK_SET) < 0) {
  609. tlog(TLOG_ERROR, "seek file %s failed, %s", file, strerror(errno));
  610. goto errout;
  611. }
  612. if (_dns_cache_write_records(fd, &cache_number) != 0) {
  613. tlog(TLOG_ERROR, "write record to file %s failed.", file);
  614. goto errout;
  615. }
  616. if (lseek(fd, 0, SEEK_SET) < 0) {
  617. tlog(TLOG_ERROR, "seek file %s failed, %s", file, strerror(errno));
  618. goto errout;
  619. }
  620. cache_file.cache_number = cache_number;
  621. if (write(fd, &cache_file, sizeof(cache_file)) != sizeof(cache_file)) {
  622. tlog(TLOG_ERROR, "write file head %s failed, %s, %d", file, strerror(errno), fd);
  623. goto errout;
  624. }
  625. tlog(TLOG_DEBUG, "wrote total %d records.", cache_number);
  626. close(fd);
  627. return 0;
  628. errout:
  629. if (fd > 0) {
  630. close(fd);
  631. }
  632. return -1;
  633. }
  634. static int _dns_cache_print(struct dns_cache_record *cache_record, struct dns_cache_data *cache_data)
  635. {
  636. printf("domain: %s, qtype: %d, ttl: %d, speed: %.1fms\n", cache_record->info.domain, cache_record->info.qtype,
  637. cache_record->info.ttl, (float)cache_record->info.speed / 10);
  638. return 0;
  639. }
  640. int dns_cache_print(const char *file)
  641. {
  642. if (access(file, F_OK) != 0) {
  643. tlog(TLOG_ERROR, "cache file %s not exist.", file);
  644. return -1;
  645. }
  646. return _dns_cache_file_read(file, _dns_cache_print);
  647. }
  648. void dns_cache_destroy(void)
  649. {
  650. struct dns_cache *dns_cache = NULL;
  651. struct dns_cache *tmp = NULL;
  652. if (is_cache_init == 0) {
  653. return;
  654. }
  655. pthread_mutex_lock(&dns_cache_head.lock);
  656. list_for_each_entry_safe(dns_cache, tmp, &dns_cache_head.cache_list, list)
  657. {
  658. _dns_cache_remove(dns_cache);
  659. }
  660. pthread_mutex_unlock(&dns_cache_head.lock);
  661. pthread_mutex_destroy(&dns_cache_head.lock);
  662. hash_table_free(dns_cache_head.cache_hash, free);
  663. is_cache_init = 0;
  664. }
  665. const char *dns_cache_file_version(void)
  666. {
  667. const char *version = "cache ver 1.3";
  668. return version;
  669. }