queryperf.c

来自「非常好的dns解析软件」· C语言 代码 · 共 2,111 行 · 第 1/4 页

			}
			break;
		case 'q':
			if (is_uint(optarg, &uint_arg_val) == TRUE) {
				set_max_queries(uint_arg_val);
				queriesset = TRUE;
			} else {
				fprintf(stderr, "Option requires a positive "
				        "integer value: -%c %s\n",
					c, optarg);
				return (-1);
			}
			break;

		case 't':
			if (is_uint(optarg, &uint_arg_val) == TRUE) {
				query_timeout = uint_arg_val;
				timeoutset = TRUE;
			} else {
				fprintf(stderr, "Option requires a positive "
				        "integer value: -%c %s\n",
				        c, optarg);
				return (-1);
			}
			break;

		case 'n':
			ignore_config_changes = TRUE;
			break;

		case 'd':
			if (set_datafile(optarg) == -1) {
				fprintf(stderr, "Error setting datafile "
					"name: %s\n", optarg);
				return (-1);
			}
			break;

		case 's':
			if (set_server(optarg) == -1) {
				fprintf(stderr, "Error setting server "
					"name: %s\n", optarg);
				return (-1);
			}
			serverset = TRUE;
			break;

		case 'p':
			if (is_uint(optarg, &uint_arg_val) == TRUE &&
			    uint_arg_val < MAX_PORT)
			{
				set_server_port(optarg);
				portset = TRUE;
			} else {
				fprintf(stderr, "Option requires a positive "
				        "integer between 0 and %d: -%c %s\n",
					MAX_PORT - 1, c, optarg);
				return (-1);
			}
			break;

		case '1':
			run_only_once = TRUE;
			break;

		case 'l':
			if (is_uint(optarg, &uint_arg_val) == TRUE) {
				use_timelimit = TRUE;
				run_timelimit = uint_arg_val;
			} else {
				fprintf(stderr, "Option requires a positive "
				        "integer: -%c %s\n",
					c, optarg);
				return (-1);
			}
			break;

		case 'b':
			if (is_uint(optarg, &uint_arg_val) == TRUE) {
				socket_bufsize = uint_arg_val;
			} else {
				fprintf(stderr, "Option requires a positive "
					"integer: -%c %s\n",
					c, optarg);
				return (-1);
			}
			break;
		case 'e':
			edns = TRUE;
			break;
		case 'D':
			dnssec = TRUE;
			edns = TRUE;
			break;
		case 'c':
			countrcodes = TRUE;
			break;
		case 'v':
			verbose = 1;
			break;
		case 'i':
			if (is_uint(optarg, &uint_arg_val) == TRUE)
				print_interval = uint_arg_val;
			else {
				fprintf(stderr, "Invalid interval: %s\n",
					optarg);
				return (-1);
			}
			break;
		case 'r':
			if (is_uint(optarg, &uint_arg_val) == TRUE)
				rttarray_size = uint_arg_val;
			else {
				fprintf(stderr, "Invalid RTT array size: %s\n",
					optarg);
				return (-1);
			}
			break;
		case 'u':
			if (is_uint(optarg, &uint_arg_val) == TRUE)
				rttarray_unit = uint_arg_val;
			else {
				fprintf(stderr, "Invalid RTT unit: %s\n",
					optarg);
				return (-1);
			}
			break;
		case 'H':
			rtt_histogram_file = optarg;
			break;
		case 'T':
			if (is_uint(optarg, &uint_arg_val) == TRUE)
				target_qps = uint_arg_val;
			else {
				fprintf(stderr, "Invalid target qps: %s\n",
					optarg);
				return (-1);
			}
			break;
		case 'h':
			return (-1);
		default:
			fprintf(stderr, "Invalid option: -%c\n", optopt);
			return (-1);
		}
	}

	if (run_only_once == FALSE && use_timelimit == FALSE)
		run_only_once = TRUE;

	return (0);
}

/*
 * open_datafile:
 *   Open the data file ready for reading
 *
 *   Return -1 on failure
 *   Return non-negative integer on success
 */
int
open_datafile(void) {
	if (use_stdin == TRUE) {
		datafile_ptr = stdin;
		return (0);
	} else {
		if ((datafile_ptr = fopen(datafile_name, "r")) == NULL) {
			fprintf(stderr, "Error: unable to open datafile: %s\n",
			        datafile_name);
			return (-1);
		} else
			return (0);
	}
}

/*
 * close_datafile:
 *   Close the data file if any is open
 *
 *   Return -1 on failure
 *   Return non-negative integer on success, including if not needed
 */
int
close_datafile(void) {
	if ((use_stdin == FALSE) && (datafile_ptr != NULL)) {
		if (fclose(datafile_ptr) != 0) {
			fprintf(stderr, "Error: unable to close datafile\n");
			return (-1);
		}
	}

	return (0);
}

/*
 * open_socket:
 *   Open a socket for the queries.  When we have an active socket already,
 *   close it and open a new one.
 *
 *   Return -1 on failure
 *   Return the socket identifier
 */
int
open_socket(void) {
	int sock;
	int ret;
	int bufsize;
	struct addrinfo hints, *res;

	memset(&hints, 0, sizeof(hints));
	hints.ai_family = server_ai->ai_family;
	hints.ai_socktype = server_ai->ai_socktype;
	hints.ai_protocol = server_ai->ai_protocol;
	hints.ai_flags = AI_PASSIVE;

	if ((ret = getaddrinfo(NULL, "0", &hints, &res)) != 0) {
		fprintf(stderr,
			"Error: getaddrinfo for bind socket failed: %s\n",
			gai_strerror(ret));
		return (-1);
	}

	if ((sock = socket(res->ai_family, SOCK_DGRAM,
			   res->ai_protocol)) == -1) {
		fprintf(stderr, "Error: socket call failed");
		goto fail;
	}

#if defined(AF_INET6) && defined(IPV6_V6ONLY)
	if (res->ai_family == AF_INET6) {
		int on = 1;

		if (setsockopt(sock, IPPROTO_IPV6, IPV6_V6ONLY,
			       &on, sizeof(on)) == -1) {
			fprintf(stderr,
				"Warning: setsockopt(IPV6_V6ONLY) failed\n");
		}
	}
#endif

	if (bind(sock, res->ai_addr, res->ai_addrlen) == -1)
		fprintf(stderr, "Error: bind call failed");

	freeaddrinfo(res);

	bufsize = 1024 * socket_bufsize;

	ret = setsockopt(sock, SOL_SOCKET, SO_RCVBUF,
			 (char *) &bufsize, sizeof(bufsize));
	if (ret < 0)
		fprintf(stderr, "Warning:  setsockbuf(SO_RCVBUF) failed\n");

	ret = setsockopt(sock, SOL_SOCKET, SO_SNDBUF,
			 (char *) &bufsize, sizeof(bufsize));
	if (ret < 0)
		fprintf(stderr, "Warning:  setsockbuf(SO_SNDBUF) failed\n");

	return (sock);

 fail:
	if (res)
		freeaddrinfo(res);
	return (-1);
}

/*
 * close_socket:
 *   Close the query socket(s)
 *
 *   Return -1 on failure
 *   Return a non-negative integer otherwise
 */
int
close_socket(void) {
	if (socket4 != -1) {
		if (close(socket4) != 0) {
			fprintf(stderr,
				"Error: unable to close IPv4 socket\n");
			return (-1);
		}
	}

	if (socket6 != -1) {
		if (close(socket6) != 0) {
			fprintf(stderr,
				"Error: unable to close IPv6 socket\n");
			return (-1);
		}
	}

	query_socket = -1;

	return (0);
}

/*
 * change_socket:
 *   Choose an appropriate socket according to the address family of the
 *   current server.  Open a new socket if necessary.
 *
 *   Return -1 on failure
 *   Return the socket identifier
 */
int
change_socket(void) {
	int s, *sockp;

	switch (server_ai->ai_family) {
	case AF_INET:
		sockp = &socket4;
		break;
#ifdef AF_INET6
	case AF_INET6:
		sockp = &socket6;
		break;
#endif
	default:
		fprintf(stderr, "unexpected address family: %d\n",
			server_ai->ai_family);
		exit(1);
	}

	if (*sockp == -1) {
		if ((s = open_socket()) == -1)
			return (-1);
		*sockp = s;
	}

	return (*sockp);
}

/*
 * reset_rttarray:
 *   (re)allocate RTT array and zero-clear the whole buffer.
 *   if array is being used, it is freed.
 *   Returns -1 on failure
 *   Returns a non-negative integer otherwise
 */
int
reset_rttarray(int size) {
	if (rttarray != NULL)
		free(rttarray);
	if (rttarray_interval != NULL)
		free(rttarray_interval);

	rttarray = NULL;
	rttarray_interval = NULL;
	rtt_max = -1;
	rtt_min = -1;

	if (size > 0) {
		rttarray = malloc(size * sizeof(rttarray[0]));
		if (rttarray == NULL) {
			fprintf(stderr,
				"Error: allocating memory for RTT array\n");
			return (-1);
		}
		memset(rttarray, 0, size * sizeof(rttarray[0]));

		rttarray_interval = malloc(size *
					   sizeof(rttarray_interval[0]));
		if (rttarray_interval == NULL) {
			fprintf(stderr,
				"Error: allocating memory for RTT array\n");
			return (-1);
		}

		memset(rttarray_interval, 0,
		       size * sizeof(rttarray_interval[0]));
	}

	return (0);
}

/*
 * set_query_interval:
 *   set the interval of consecutive queries if the target qps are specified.
 *   Returns -1 on failure
 *   Returns a non-negative integer otherwise
 */
int
set_query_interval(unsigned int qps) {
	if (qps == 0)
		return (0);

	query_interval = (1.0 / (double)qps);

	return (0);
}

/*
 * setup:
 *   Set configuration options from command line arguments
 *   Open datafile ready for reading
 *
 *   Return -1 on failure
 *   Return non-negative integer on success
 */
int
setup(int argc, char **argv) {
	set_input_stdin();

	if (set_max_queries(DEF_MAX_QUERIES_OUTSTANDING) == -1) {
		fprintf(stderr, "%s: Unable to set default max outstanding "
		        "queries\n", argv[0]);
		return (-1);
	}

	if (set_server(DEF_SERVER_TO_QUERY) == -1) {
		fprintf(stderr, "%s: Error setting default server name\n",
		        argv[0]);
		return (-1);
	}

	if (set_server_port(DEF_SERVER_PORT) == -1) {
		fprintf(stderr, "%s: Error setting default server port\n",
		        argv[0]);
		return (-1);
	}

	if (parse_args(argc, argv) == -1) {
		show_usage();
		return (-1);
	}

	if (open_datafile() == -1)
		return (-1);

	if (set_server_sa() == -1)
		return (-1);

	if ((query_socket = change_socket()) == -1)
		return (-1);

	if (reset_rttarray(rttarray_size) == -1)
		return (-1);

	if (set_query_interval(target_qps) == -1)
		return (-1);

	return (0);
}

/*
 * set_timenow:
 *   Set a timeval struct to indicate the current time
 */
void
set_timenow(struct timeval *tv) {
	if (gettimeofday(tv, NULL) == -1) {
		fprintf(stderr, "Error in gettimeofday(). Using inaccurate "
		        "time() instead\n");
		tv->tv_sec = time(NULL);
		tv->tv_usec = 0;
	}
}

/*
 * addtv:
 *   add tv1 and tv2, store the result in tv_result.
 */
void
addtv(struct timeval *tv1, struct timeval *tv2, struct timeval *tv_result) {
	tv_result->tv_sec = tv1->tv_sec + tv2->tv_sec;
	tv_result->tv_usec = tv1->tv_usec + tv2->tv_usec;
	if (tv_result->tv_usec > 1000000) {
		tv_result->tv_sec++;
		tv_result->tv_usec -= 1000000;
	}
}

/*
 * difftv:
 *   Find the difference in seconds between two timeval structs.
 *
 *   Return the difference between tv1 and tv2 in seconds in a double.
 */
double
difftv(struct timeval tv1, struct timeval tv2) {
	long diff_sec, diff_usec;
	double diff;

	diff_sec = tv1.tv_sec - tv2.tv_sec;
	diff_usec = tv1.tv_usec - tv2.tv_usec;

	diff = (double)diff_sec + ((double)diff_usec / 1000000.0);

	return (diff);
}

/*
 * timelimit_reached:
 *   Have we reached the time limit (if any)?
 *
 *   Returns FALSE if there is no time limit or if we have not reached it
 *   Returns TRUE otherwise
 */
int
timelimit_reached(void) {
	struct timeval time_now;

	set_timenow(&time_now);

	if (use_timelimit == FALSE)
		return (FALSE);

	if (setup_phase == TRUE) {
		if (difftv(time_now, time_of_program_start)
		    < (double)(run_timelimit + HARD_TIMEOUT_EXTRA))
			return (FALSE);
		else
			return (TRUE);
	} else {
		if (difftv(time_now, time_of_first_query)
		    < (double)run_timelimit)
			return (FALSE);
		else
			return (TRUE);
	}
}

/*
 * keep_sending:
 *   Should we keep sending queries or stop here?
 *
 *   Return TRUE if we should keep on sending queries
 *   Return FALSE if we should stop
 *