dlzbdb.c

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

		db.cursor = NULL;
	}

	/* close databases */
	if (db.data != NULL) {
		db.data->close(db.data, 0);
		db.data = NULL;
	}
	if (db.host != NULL) {
		db.host->close(db.host, 0);
		db.host = NULL;
	}
	if (db.zone != NULL) {
		db.zone->close(db.zone, 0);
		db.zone = NULL;
	}
	if (db.client != NULL) {
		db.client->close(db.client, 0);
		db.client = NULL;
	}

	/* close environment */
	if (db.dbenv != NULL) {
		db.dbenv->close(db.dbenv, 0);
		db.dbenv = NULL;
	}
}

/*% Initializes, sets flags and then opens Berkeley databases. */

isc_result_t
bdb_opendb(DBTYPE db_type, DB **db_out, const char *db_name, int flags) {

	int result;
	int createFlag = 0;

	/* Initialize the database. */
	if ((result = db_create(db_out, db.dbenv, 0)) != 0) {
		fprintf(stderr, "BDB could not initialize %s database. BDB error: %s",
			db_name, db_strerror(result));
		return ISC_R_FAILURE;
	}

	/* set database flags. */
	if ((result = (*db_out)->set_flags(*db_out, flags)) != 0) {
		fprintf(stderr, "BDB could not set flags for %s database. BDB error: %s",
			db_name, db_strerror(result));
		return ISC_R_FAILURE;
	}

	if (create_allowed == isc_boolean_true) {
		createFlag = DB_CREATE;
	}
	/* open the database. */
	if ((result = (*db_out)->open(*db_out, NULL, db_file, db_name, db_type,
				      createFlag, 0)) != 0) {
		fprintf(stderr, "BDB could not open %s database in %s. BDB error: %s",
			db_name, db_file, db_strerror(result));
		return ISC_R_FAILURE;
	}

	return ISC_R_SUCCESS;
}

/*%
 * parses input and adds it to the BDB database
 * Lexer should be instantiated, and either a file or buffer opened for it.
 * The insert_data function is used by both the add, and bulk insert
 * operations
 */

void
put_data(isc_boolean_t dns_data, char *input_key, char *input_data) {

	int bdbres;
	DBT key, data;

	/* make sure key & data are completely empty */
	memset(&key, 0, sizeof(key));
	memset(&data, 0, sizeof(data));

	/* if client data, setup key for insertion */
	if (!dns_data && input_key != NULL) {
		key.data = input_key;
		key.size = strlen(input_key);
		key.flags = 0;
	}
	/* always setup data for insertion */
	data.data = input_data;
	data.size = strlen(input_data);
	data.flags = 0;

	/* execute insert against appropriate database. */
	if (dns_data) {
		bdbres = db.data->put(db.data, NULL, &key, &data, DB_APPEND);
	} else {
		bdbres = db.client->put(db.client, NULL, &key, &data, 0);
	}

	/* if something went wrong, log error and quit */
	if (bdbres != 0) {
		fprintf(stderr, "BDB could not insert data.  Error: %s",
			db_strerror(bdbres));
		quit(5);
	}
}

void
insert_data(void) {
	unsigned int opt =
		ISC_LEXOPT_EOL |		/* Want end-of-line token. */
		ISC_LEXOPT_EOF | 		/* Want end-of-file token. */
		ISC_LEXOPT_QSTRING |		/* Recognize qstrings. */
		ISC_LEXOPT_QSTRINGMULTILINE;	/* Allow multiline "" strings */

	isc_result_t result;
	isc_token_t token;	/* token from lexer */
	isc_boolean_t loop = isc_boolean_true;
	isc_boolean_t have_czone = isc_boolean_false;
	char data_arr[max_data_len];
	isc_buffer_t buf;
	char data_arr2[max_data_len];
	isc_buffer_t buf2;
	char data_type = 'u'; /* u =unknown, b =bad token, d/D =DNS, c/C =client IP */

	/* Initialize buffers */
	isc_buffer_init(&buf, &data_arr, max_data_len);
	isc_buffer_init(&buf2, &data_arr2, max_data_len);

	while (loop) {
		result = isc_lex_gettoken(lexer, opt, &token); 
		if (result != ISC_R_SUCCESS) 
			goto data_cleanup;

		switch(token.type) {
		case isc_tokentype_string:
			if (data_type == 'u') {
				/* store data_type */
				strncpy(&data_type, token.value.as_pointer, 1);
				/* verify data_type was specified correctly on input */
				if (strlen(token.value.as_pointer) > 1 || (
					    data_type != 'd' && data_type != 'D' &&
					    data_type != 'c' && data_type != 'C') ) {
					/* if not, set to 'b' so this line is ignored. */
					data_type = 'b';
				}
			} else if (data_type == 'c' || data_type == 'C') {
				if (have_czone == isc_boolean_true) {
					isc_buffer_putstr(&buf2, token.value.as_pointer);
					/* add string terminator to buffer */
					isc_buffer_putmem(&buf2, "\0", 1);
				} else {
					isc_buffer_putstr(&buf, token.value.as_pointer);
					/* add string terminator to buffer */
					isc_buffer_putmem(&buf, "\0", 1);
					have_czone = isc_boolean_true;
				}
			} else {
				isc_buffer_putstr(&buf, token.value.as_pointer);
				isc_buffer_putstr(&buf, " ");
			}
			break;
		case isc_tokentype_qstring:
			isc_buffer_putstr(&buf, "\"");
			isc_buffer_putstr(&buf, token.value.as_pointer);
			isc_buffer_putstr(&buf, "\" ");
			break;
		case isc_tokentype_eol:
		case isc_tokentype_eof:
			
			if ((data_type != 'u' && isc_buffer_usedlength(&buf) > 0) || data_type == 'b') {
				/* perform insert operation */
				if (data_type == 'd' || data_type == 'D') {
					/* add string terminator to buffer */
					isc_buffer_putmem(&buf, "\0", 1);
					put_data(isc_boolean_true, NULL, (char *) &data_arr);
				} else if (data_type == 'c' || data_type == 'C') {
					put_data(isc_boolean_false, (char *) &data_arr, 
						 (char *) &data_arr2);
				} else if (data_type == 'b') {
					fprintf(stderr, "Bad / unknown token encountered on line %lu."\
						"  Skipping line.",	isc_lex_getsourceline(lexer) - 1);
				} else {
					fprintf(stderr, "Bad / unknown db data type encountered on " \
						"line %lu.  Skipping line\n", isc_lex_getsourceline(lexer) - 1);
				}
			}

			if (token.type == isc_tokentype_eof) {
				loop = isc_boolean_false;
			}	

			/* reset buffer for next insert */
			isc_buffer_clear(&buf);
			isc_buffer_clear(&buf2);
			have_czone = isc_boolean_false;
			data_type ='u';
			break;
		default:
			data_type = 'b';
			break;
		}
	}

	return;

 data_cleanup:
	/* let user know we had problems */
	fprintf(stderr, "Unknown error processing tokens during \"add\" or " \
		"\"bulk\" operation.\nStoped processing on line %lu.", 
		isc_lex_getsourceline(lexer));
}


isc_result_t
openBDB(void) {

	int bdbres;
	isc_result_t result;

	/* create BDB environment  */
	/* Basically BDB allocates and assigns memory to db->dbenv */
	bdbres = db_env_create(&db.dbenv, 0);
	if (bdbres != 0) {
		fprintf(stderr, "BDB environment could not be created. BDB error: %s", 
			db_strerror(bdbres));
		result = ISC_R_FAILURE;
		goto openBDB_cleanup;
	}

	/* open BDB environment */
	if (create_allowed == isc_boolean_true) {
		/* allowed to create new files */
		bdbres = db.dbenv->open(db.dbenv, db_envdir, 
					DB_INIT_CDB | DB_INIT_MPOOL | DB_CREATE, 0);
	} else {	/* not allowed to create new files. */
		bdbres = db.dbenv->open(db.dbenv, db_envdir, 
					DB_INIT_CDB | DB_INIT_MPOOL, 0);
	}
	if (bdbres != 0) {
		fprintf(stderr, "BDB environment at '%s' could not be opened. BDB " \
			"error: %s", db_envdir, db_strerror(bdbres));
		result = ISC_R_FAILURE;
		goto openBDB_cleanup;
	}

	/* open dlz_data database. */

	result = bdb_opendb(DB_RECNO, &db.data, dlz_data, 0);
	if (result != ISC_R_SUCCESS)
		goto openBDB_cleanup;
	
	/* open dlz_host database */
	result = bdb_opendb(DB_BTREE, &db.host, dlz_host, DB_DUP | DB_DUPSORT);
	if (result != ISC_R_SUCCESS)
		goto openBDB_cleanup;

	/* open dlz_zone database. */
	result = bdb_opendb(DB_BTREE, &db.zone, dlz_zone, DB_DUP | DB_DUPSORT);
	if (result != ISC_R_SUCCESS)
		goto openBDB_cleanup;

	/* open dlz_client database. */
	result = bdb_opendb(DB_BTREE, &db.client, dlz_client, DB_DUP | DB_DUPSORT);
	if (result != ISC_R_SUCCESS)
		goto openBDB_cleanup;

	/* associate the host secondary database with the primary database */
	bdbres = db.data->associate(db.data, NULL, db.host, gethost, 0);
	if (bdbres != 0) {
		fprintf(stderr, "BDB could not associate %s database with %s. BDB "\
			"error: %s", dlz_host, dlz_data, db_strerror(bdbres));
		result = ISC_R_FAILURE;
		goto openBDB_cleanup;
	}

	/* associate the zone secondary database with the primary database */
	bdbres = db.data->associate(db.data, NULL, db.zone, getzone, 0);
	if (bdbres != 0) {
		fprintf(stderr, "BDB could not associate %s database with %s. BDB "\
			"error: %s", dlz_zone, dlz_data, db_strerror(bdbres));
		result = ISC_R_FAILURE;
		goto openBDB_cleanup;
	}

	return result;

 openBDB_cleanup:

	bdb_cleanup();
	return result;
}

/*% Create & open lexer to parse input data */

isc_result_t
open_lexer(void) {
	isc_result_t result;

	/* check if we already opened the lexer, if we did, return success */
	if (lexer != NULL)
		return ISC_R_SUCCESS;

	/* allocate memory for lexer, and verify it was allocated */
	result = isc_mem_create(0, 0, &lex_mctx);
	if (result != ISC_R_SUCCESS) {
		fprintf(stderr, "unexpected error creating lexer\n");
		return result;
	}

	/* create lexer */
	result = isc_lex_create(lex_mctx, 1500, &lexer);
	if (result != ISC_R_SUCCESS)
		fprintf(stderr, "unexpected error creating lexer\n");

	/* set allowed commenting style */
	isc_lex_setcomments(lexer, ISC_LEXCOMMENT_C |	/* Allow C comments */
			    ISC_LEXCOMMENT_CPLUSPLUS |	/* Allow C++ comments */
			    ISC_LEXCOMMENT_SHELL);		/* Allow shellcomments */

	isc_buffer_init(&lex_buffer, &lex_data_buf, max_data_len);

	return result;
}

/*% Close the lexer, and cleanup memory */

void
close_lexer(void) {
		
	/* If lexer is still open, close it & destroy it. */
	if (lexer != NULL) {
		isc_lex_close(lexer);
		isc_lex_destroy(&lexer);
	}

	/* if lexer memory is still allocated, destroy it. */
	if (lex_mctx != NULL)
		isc_mem_destroy(&lex_mctx);
}

/*% Perform add operation */

void
operation_add(void) {
	/* check for any parameters that are not allowed during add */
	checkInvalidParam(key, "k", "for add operation");
	checkInvalidParam(zone, "z", "for add operation");
	checkInvalidParam(host, "h", "for add operation");
	checkInvalidParam(c_zone, "c", "for add operation");
	checkInvalidParam(c_ip, "i", "for add operation");
	checkInvalidOption(list_everything, isc_boolean_true, "e",
			   "for add operation");

	/* if open lexer fails it alread prints error messages. */
	if (open_lexer() != ISC_R_SUCCESS) {
		quit(4);
	}
	
	/* copy input data to buffer */
	isc_buffer_putstr(&lex_buffer, a_data);
		
	/* tell lexer to use buffer as input  */
	if (isc_lex_openbuffer(lexer, &lex_buffer) != ISC_R_SUCCESS) {
		fprintf(stderr, "unexpected error opening lexer buffer");
		quit(4);
	}

	/*common logic for "add" & "bulk" operations are handled by insert_data */
	insert_data();
}

/*% Perform bulk insert operation */

void
operation_bulk(void) {
	/* check for any parameters that are not allowed during bulk */
	checkInvalidParam(key, "k", "for bulk load operation");
	checkInvalidParam(zone, "z", "for bulk load operation");
	checkInvalidParam(host, "h", "for bulk load operation");
	checkInvalidParam(c_zone, "c", "for bulk load operation");
	checkInvalidParam(c_ip, "i", "for bulk load operation");
	checkInvalidOption(list_everything, isc_boolean_true, "e",
			   "for bulk load operation");

	/* if open lexer fails it already prints error messages. */
	if (open_lexer() != ISC_R_SUCCESS) {
		quit(4);
	}

	if (bulk_file == NULL) {
		if (isc_lex_openstream(lexer, stdin) != ISC_R_SUCCESS) {
			fprintf(stderr, "unexpected error opening stdin by lexer.");
			quit(4);				
		}
	} else if (isc_lex_openfile(lexer, bulk_file) != ISC_R_SUCCESS) {
		fprintf(stderr, "unexpected error opening %s by lexer.", bulk_file);
		quit(4);
	}

	/* common logic for "add" & "bulk" operations are handled by insert_data */
	insert_data();	
}

isc_result_t
bulk_write(char type, DB *database, DBC *dbcursor, DBT *bdbkey, DBT *bdbdata) {

	int bdbres;
	db_recno_t recNum;
	char *retkey = NULL, *retdata;
	size_t retklen = 0, retdlen;
	void *p;

	/* use a 5MB buffer for the bulk dump */
	int buffer_size = 5 * 1024 * 1024;
         
	/* try to allocate a 5 MB buffer, if we fail write err msg, die. */
	bdbdata->data = malloc(buffer_size);
	if (bdbdata->data == NULL) {
		fprintf(stderr,
			"Unable to allocate 5 MB buffer for bulk database dump\n");
		return ISC_R_FAILURE;
	}
	bdbdata->ulen = buffer_size;
	bdbdata->flags = DB_DBT_USERMEM;