queryperf.c

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

 *   Side effects:
 *   Rewinds the input and clears reached_end_input if we have reached the
 *   end of the input, but we are meant to run through it multiple times
 *   and have not hit the time limit yet (if any is set).
 */
int
keep_sending(int *reached_end_input) {
	static int stop = FALSE;

	if (stop == TRUE)
		return (FALSE);

	if ((*reached_end_input == FALSE) && (timelimit_reached() == FALSE))
		return (TRUE);
	else if ((*reached_end_input == TRUE) && (run_only_once == FALSE)
	         && (timelimit_reached() == FALSE)) {
		rewind(datafile_ptr);
		*reached_end_input = FALSE;
		runs_through_file++;
		return (TRUE);
	} else {
		if (*reached_end_input == TRUE)
			runs_through_file++;
		set_timenow(&time_of_stop_sending);
		stop = TRUE;
		return (FALSE);
	}
}

/*
 * queries_outstanding:
 *   How many queries are outstanding?
 *
 *   Returns the number of outstanding queries
 */
unsigned int
queries_outstanding(void) {
	return (num_queries_outstanding);
}

/*
 * next_input_line:
 *   Get the next non-comment line from the input file
 *
 *   Put text in line, up to max of n chars. Skip comment lines.
 *   Skip empty lines.
 *
 *   Return line length on success
 *   Return 0 if cannot read a non-comment line (EOF or error)
 */
int
next_input_line(char *line, int n) {
	char *result;

	do {
		result = fgets(line, n, datafile_ptr);
	} while ((result != NULL) &&
	         ((line[0] == COMMENT_CHAR) || (line[0] == '\n')));

	if (result == NULL)
		return (0);
	else
		return (strlen(line));
}

/*
 * identify_directive:
 *   Gives us a numerical value equivelant for a directive string
 *
 *   Returns the value for the directive
 *   Returns -1 if not a valid directive
 */
int
identify_directive(char *dir) {
	static char *directives[] = DIRECTIVES;
	static int dir_values[] = DIR_VALUES;
	unsigned int index, num_directives;

	num_directives = sizeof(directives) / sizeof(directives[0]);

	if (num_directives > (sizeof(dir_values) / sizeof(int)))
		num_directives = sizeof(dir_values) / sizeof(int);

	for (index = 0; index < num_directives; index++) {
		if (strcmp(dir, directives[index]) == 0)
			return (dir_values[index]);
	}

	return (-1);
}

/*
 * update_config:
 *   Update configuration options from a line from the input file
 */
void
update_config(char *config_change_desc) {
	char *directive, *config_value, *trailing_garbage;
	char conf_copy[MAX_INPUT_LEN + 1];
	unsigned int uint_val;
	int directive_number;
	int check;
	int old_af;

	if (ignore_config_changes == TRUE) {
		fprintf(stderr, "Ignoring configuration change: %s",
		        config_change_desc);
		return;
	}

	strcpy(conf_copy, config_change_desc);

	++config_change_desc;

	if (*config_change_desc == '\0') {
		fprintf(stderr, "Invalid config: No directive present: %s\n",
		        conf_copy);
		return;
	}

	if (index(WHITESPACE, *config_change_desc) != NULL) {
		fprintf(stderr, "Invalid config: Space before directive or "
		        "no directive present: %s\n", conf_copy);
		return;
	}

	directive = strtok(config_change_desc, WHITESPACE);
	config_value = strtok(NULL, WHITESPACE);
	trailing_garbage = strtok(NULL, WHITESPACE);

	if ((directive_number = identify_directive(directive)) == -1) {
		fprintf(stderr, "Invalid config: Bad directive: %s\n",
		        conf_copy);
		return;
	}

	if (config_value == NULL) {
		fprintf(stderr, "Invalid config: No value present: %s\n",
		        conf_copy);
		return;
	}

	if (trailing_garbage != NULL) {
		fprintf(stderr, "Config warning: "
		        "trailing garbage: %s\n", conf_copy);
	}

	switch(directive_number) {

	case V_SERVER:
		if (serverset && (setup_phase == TRUE)) {
			fprintf(stderr, "Config change overridden by command "
			        "line: %s\n", directive);
			return;
		}

		if (set_server(config_value) == -1) {
			fprintf(stderr, "Set server error: unable to change "
			        "the server name to '%s'\n", config_value);
			return;
		}

		old_af = server_ai->ai_family;
		if (set_server_sa() == -1) {
			fprintf(stderr, "Set server error: unable to resolve "
				"a new server '%s'\n",
				config_value);
			return;
		}
		if (old_af != server_ai->ai_family) {
			if ((query_socket = change_socket()) == -1) {
				/* XXX: this is fatal */
				fprintf(stderr, "Set server error: "
					"unable to open a new socket "
					"for '%s'\n", config_value);
				exit(1);
			}
		}

		break;

	case V_PORT:
		if (portset && (setup_phase == TRUE)) {
			fprintf(stderr, "Config change overridden by command "
			        "line: %s\n", directive);
			return;
		}

		check = is_uint(config_value, &uint_val);

		if ((check == TRUE) && (uint_val > 0)) {
			if (set_server_port(config_value) == -1) {
				fprintf(stderr, "Invalid config: Bad value for"
				        " %s: %s\n", directive, config_value);
			} else {
				if (set_server_sa() == -1) {
					fprintf(stderr,
						"Failed to set a new port\n");
					return;
				}
			}
		} else
			fprintf(stderr, "Invalid config: Bad value for "
			        "%s: %s\n", directive, config_value);
		break;

	case V_MAXQUERIES:
		if (queriesset && (setup_phase == TRUE)) {
			fprintf(stderr, "Config change overridden by command "
			        "line: %s\n", directive);
			return;
		}

		check = is_uint(config_value, &uint_val);

		if ((check == TRUE) && (uint_val > 0)) {
			set_max_queries(uint_val);
		} else
			fprintf(stderr, "Invalid config: Bad value for "
			        "%s: %s\n", directive, config_value);
		break;

	case V_MAXWAIT:
		if (timeoutset && (setup_phase == TRUE)) {
			fprintf(stderr, "Config change overridden by command "
			        "line: %s\n", directive);
			return;
		}

		check = is_uint(config_value, &uint_val);

		if ((check == TRUE) && (uint_val > 0)) {
			query_timeout = uint_val;
		} else
			fprintf(stderr, "Invalid config: Bad value for "
			        "%s: %s\n", directive, config_value);
		break;

	default:
		fprintf(stderr, "Invalid config: Bad directive: %s\n",
		        directive);
		break;
	}
}

/*
 * parse_query:
 *   Parse a query line from the input file
 *
 *   Set qname to be the domain to query (up to a max of qnlen chars)
 *   Set qtype to be the type of the query
 *
 *   Return -1 on failure
 *   Return a non-negative integer otherwise
 */
int
parse_query(char *input, char *qname, int qnlen, int *qtype) {
	static char *qtype_strings[] = QTYPE_STRINGS;
	static int qtype_codes[] = QTYPE_CODES;
	int num_types, index;
	int found = FALSE;
	char incopy[MAX_INPUT_LEN + 1];
	char *domain_str, *type_str;

	num_types = sizeof(qtype_strings) / sizeof(qtype_strings[0]);
	if (num_types > (sizeof(qtype_codes) / sizeof(int)))
		num_types = sizeof(qtype_codes) / sizeof(int);

	strcpy(incopy, input);

	domain_str = strtok(incopy, WHITESPACE);
	type_str = strtok(NULL, WHITESPACE);

	if ((domain_str == NULL) || (type_str == NULL)) {
		fprintf(stderr, "Invalid query input format: %s\n", input);
		return (-1);
	}

	if (strlen(domain_str) > qnlen) {
		fprintf(stderr, "Query domain too long: %s\n", domain_str);
		return (-1);
	}

	for (index = 0; (index < num_types) && (found == FALSE); index++) {
		if (strcasecmp(type_str, qtype_strings[index]) == 0) {
			*qtype = qtype_codes[index];
			found = TRUE;
		}
	}

	if (found == FALSE) {
		fprintf(stderr, "Query type not understood: %s\n", type_str);
		return (-1);
	}

	strcpy(qname, domain_str);

	return (0);
}

/*
 * dispatch_query:
 *   Send the query packet for the entry 
 *
 *   Return -1 on failure
 *   Return a non-negative integer otherwise
 */
int
dispatch_query(unsigned short int id, char *dom, int qt) {
	static u_char packet_buffer[PACKETSZ + 1];
	static socklen_t sockaddrlen = sizeof(struct sockaddr);
	int buffer_len = PACKETSZ;
	int bytes_sent;
	unsigned short int net_id = htons(id);
	char *id_ptr = (char *)&net_id;

	buffer_len = res_mkquery(QUERY, dom, C_IN, qt, NULL, 0,
				 NULL, packet_buffer, PACKETSZ);
	if (buffer_len == -1) {
		fprintf(stderr, "Failed to create query packet: %s %d\n",
		        dom, qt);
		return (-1);
	}
	if (edns) {
		unsigned char *p;
		if (buffer_len + EDNSLEN >= PACKETSZ) {
			fprintf(stderr, "Failed to add OPT to query packet\n");
			return (-1);
		}
		packet_buffer[11] = 1;
		p = &packet_buffer[buffer_len];
		*p++ = 0;	/* root name */
		*p++ = 0;
		*p++ = 41;	/* OPT */
		*p++ = 16;	
		*p++ = 0;	/* UDP payload size (4K) */
		*p++ = 0;	/* extended rcode */
		*p++ = 0;	/* version */
		if (dnssec)
			*p++ = 0x80;	/* upper flag bits - DO set */
		else
			*p++ = 0;	/* upper flag bits */
		*p++ = 0;	/* lower flag bit */
		*p++ = 0;
		*p++ = 0;	/* rdlen == 0 */
		buffer_len += EDNSLEN;
	}

	packet_buffer[0] = id_ptr[0];
	packet_buffer[1] = id_ptr[1];

	bytes_sent = sendto(query_socket, packet_buffer, buffer_len, 0,
			    server_ai->ai_addr, server_ai->ai_addrlen);
	if (bytes_sent == -1) {
		fprintf(stderr, "Failed to send query packet: %s %d\n",
		        dom, qt);
		return (-1);
	}

	if (bytes_sent != buffer_len)
		fprintf(stderr, "Warning: incomplete packet sent: %s %d\n",
		        dom, qt);

	return (0);
}

/*
 * send_query:
 *   Send a query based on a line of input
 */
void
send_query(char *query_desc) {
	static unsigned short int use_query_id = 0;
	static int qname_len = MAX_DOMAIN_LEN;
	static char domain[MAX_DOMAIN_LEN + 1];
	char serveraddr[NI_MAXHOST];
	int query_type;
	unsigned int count;

	use_query_id++;

	if (parse_query(query_desc, domain, qname_len, &query_type) == -1) {
		fprintf(stderr, "Error parsing query: %s\n", query_desc);
		return;
	}

	if (dispatch_query(use_query_id, domain, query_type) == -1) {
		char *addrstr;

		if (getnameinfo(server_ai->ai_addr, server_ai->ai_addrlen,
				serveraddr, sizeof(serveraddr), NULL, 0,
				NI_NUMERICHOST) == 0) {
			addrstr = serveraddr;
		} else
			addrstr = "???"; /* XXX: this should not happen */
		fprintf(stderr, "Error sending query to %s: %s\n",
			addrstr, query_desc);
		return;
	}

	if (setup_phase == TRUE) {
		set_timenow(&time_of_first_query);
		time_of_first_query_sec = (double)time_of_first_query.tv_sec +
			((double)time_of_first_query.tv_usec / 1000000.0);
		setup_phase = FALSE;
		if (getnameinfo(server_ai->ai_addr, server_ai->ai_addrlen,
				serveraddr, sizeof(serveraddr), NULL, 0,
				NI_NUMERICHOST) != 0) {
			fprintf(stderr, "Error printing server address\n");
			return;
		}
		printf("[Status] Sending queries (beginning with %s)\n",
		       serveraddr);
	}

	/* Find the first slot in status[] that is not in use */
	for (count = 0; (status[count].in_use == TRUE)
	     && (count < max_queries_outstanding); count++);

	if (status[count].in_use == TRUE) {
		fprintf(stderr, "Unexpected error: We have run out of "
			"status[] space!\n");
		return;
	}

	/* Register the query in status[] */
	status[count].in_use = TRUE;
	status[count].id = use_query_id;
	if (verbose)
		status[count].desc = strdup(query_desc);
	set_timenow(&status[count].sent_timestamp);

	if (num_queries_sent_interval == 0)
		set_timenow(&time_of_first_query_interval);

	num_queries_sent++;
	num_queries_sent_interval++;
	num_queries_outstanding++;
}

void
register_rtt(struct timeval *timestamp) {
	int i;
	int oldquery = FALSE;
	struct timeval now;
	double rtt;

	set_timenow(&now);
	rtt = difftv(now, *timestamp);

	if (difftv(*timestamp, time_of_first_query_interval) < 0)
		oldquery = TRUE;

	if (rtt_max < 0 || rtt_max < rtt)
		rtt_max = rtt;

	if (rtt_min < 0 || rtt_min > rtt)
		rtt_min = rtt;

	rtt_total += rtt;

	if (!oldquery) {
		if (rtt_max_interval < 0 || rtt_max_interval < rtt)
			rtt_max_interval = rtt;

		if (rtt_min_interval < 0 || rtt_min_interval > rtt)
			rtt_min_interval = rtt;

		rtt_total_interval += rtt;
	}

	if (rttarray == NULL)
		return;

	i = (int)(rtt * (1000000.0 / rttarray_unit));
	if (i < rttarray_size) {
		rttarray[i]++;
		if (!oldquery)
			rttarray_interval[i]++;
	} else {
		fprintf(stderr, "Warning: RTT is out of range: %.6lf\n",
			rtt);
		rtt_overflows++;
		if (!oldquery)
			rtt_overflows_interval++;
	}
}

/*
 * register_response:
 *   Register receipt of a query
 *
 *   Removes (sets in_use = FALSE) the record for the given query id in
 *   status[] if any exists.
 */
void
register_response(unsigned short int id, unsigned int rcode) {
	unsigned int ct = 0;
	int found = FALSE;
	struct timeval now;
	double rtt;

	for (; (ct < query_status_allocated) && (found == FALSE); ct++) {
		if ((status[ct].in_use == TRUE) && (status[ct].id == id)) {
			status[ct].in_use = FALSE;
			num_queries_outstanding--;
			found = TRUE;

			register_rtt(&status[ct].sent_timestamp);

			if (status[ct].desc) {
				printf("> %s %s\n", rcode_strings[rcode],
				       status[ct].desc);
				free(status[ct].desc);
			}
			if (countrcodes)
				rcodecounts[rcode]++;
		}
	}

	if (found == FALSE) {
		if (target_qps > 0) {
			num_queries_possiblydelayed++;
			num_queries_possiblydelayed_interval++;
		} else {
			fprintf(stderr, "Warning: Received a response with an "
				"unexpected (maybe timed out) id: %u\n", id);
		}