dlz_postgres_driver.c

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

				len += strlen(PQgetvalue(rs, i, j)) + 1;
			}
			/* allocate memory, allow for NULL to term string */
			tmpString = isc_mem_allocate(ns_g_mctx, len + 1);
			if (tmpString == NULL) {	/* we need more ram. */
				isc_log_write(dns_lctx,
					      DNS_LOGCATEGORY_DATABASE,
					      DNS_LOGMODULE_DLZ, ISC_LOG_ERROR,
					      "Postgres driver unable to "
					      "allocate memory for "
					      "temporary string");
				PQclear(rs);
				return (ISC_R_FAILURE);
			}
			/* copy this field to tmpString */
			strcpy(tmpString, PQgetvalue(rs, i, 3));
			/* concatonate the rest, with spaces between */
			for (j=4; j < fields; j++) {
				strcat(tmpString, " ");
				strcat(tmpString, PQgetvalue(rs, i, j));
			}
			/* tell Bind about it. */
			result = dns_sdlz_putnamedrr(allnodes,
						     PQgetvalue(rs, i, 2),
						     PQgetvalue(rs, i, 1),
						     ttl, tmpString);
			isc_mem_free(ns_g_mctx, tmpString);
		}
		/* if we weren't successful, log err msg */
		if (result != ISC_R_SUCCESS) {
			PQclear(rs);
			isc_log_write(dns_lctx, DNS_LOGCATEGORY_DATABASE,
				      DNS_LOGMODULE_DLZ, ISC_LOG_ERROR,
				      "dns_sdlz_putnamedrr returned error. "
				      "Error code was: %s",
				      isc_result_totext(result));
			return (ISC_R_FAILURE);
		}
	}

	/* free result set memory */
	PQclear(rs);

	/* if we did return results, we are successful */
	if (rows > 0)
		return (ISC_R_SUCCESS);

	/* empty result set, no data found */
	return (ISC_R_NOTFOUND);
}

/*%
 * if the lookup function does not return SOA or NS records for the zone,
 * use this function to get that information for Bind.
 */

static isc_result_t
postgres_authority(const char *zone, void *driverarg, void *dbdata,
		   dns_sdlzlookup_t *lookup)
{
	isc_result_t result;
	PGresult *rs = NULL;

	UNUSED(driverarg);

	/* run the query and get the result set from the database. */
	result = postgres_get_resultset(zone, NULL, NULL,
					AUTHORITY, dbdata, &rs);
	/* if we get "not implemented", send it along */
	if (result == ISC_R_NOTIMPLEMENTED)
		return result;
	/* if we didn't get a result set, log an err msg. */
	if (result != ISC_R_SUCCESS) {
		if (rs != NULL)
			PQclear(rs);
		isc_log_write(dns_lctx, DNS_LOGCATEGORY_DATABASE,
			      DNS_LOGMODULE_DLZ, ISC_LOG_ERROR,
			      "Postgres driver unable to return "
			      "result set for authority query");
		return (ISC_R_FAILURE);
	}
	/*
	 * lookup and authority result sets are processed in the same
	 * manner postgres_process_rs does the job for both
	 * functions.
	 */
	return postgres_process_rs(lookup, rs);
}

/*% if zone is supported, lookup up a (or multiple) record(s) in it */
static isc_result_t
postgres_lookup(const char *zone, const char *name, void *driverarg,
		void *dbdata, dns_sdlzlookup_t *lookup)
{
	isc_result_t result;
	PGresult *rs = NULL;

	UNUSED(driverarg);

	/* run the query and get the result set from the database. */
	result = postgres_get_resultset(zone, name, NULL, LOOKUP, dbdata, &rs);
	/* if we didn't get a result set, log an err msg. */
	if (result != ISC_R_SUCCESS) {
		if (rs != NULL)
			PQclear(rs);
		isc_log_write(dns_lctx, DNS_LOGCATEGORY_DATABASE,
			      DNS_LOGMODULE_DLZ, ISC_LOG_ERROR,
			      "Postgres driver unable to "
			      "return result set for lookup query");
		return (ISC_R_FAILURE);
	}
	/*
	 * lookup and authority result sets are processed in the same
	 * manner postgres_process_rs does the job for both functions.
	 */
	return postgres_process_rs(lookup, rs);
}

/*%
 * create an instance of the driver.  Remember, only 1 copy of the driver's
 * code is ever loaded, the driver has to remember which context it's
 * operating in.  This is done via use of the dbdata argument which is
 * passed into all query functions.
 */
static isc_result_t
postgres_create(const char *dlzname, unsigned int argc, char *argv[],
		void *driverarg, void **dbdata)
{
	isc_result_t result;
	dbinstance_t *dbi = NULL;
	unsigned int j;

#ifdef ISC_PLATFORM_USETHREADS
	/* if multi-threaded, we need a few extra variables. */
	int dbcount;
	db_list_t *dblist = NULL;
	int i;
	char *endp;

#endif /* ISC_PLATFORM_USETHREADS */

	UNUSED(driverarg);
	UNUSED(dlzname);

/* seed random # generator */
	srand( (unsigned)time( NULL ) );


#ifdef ISC_PLATFORM_USETHREADS
	/* if debugging, let user know we are multithreaded. */
	isc_log_write(dns_lctx, DNS_LOGCATEGORY_DATABASE,
		      DNS_LOGMODULE_DLZ, ISC_LOG_DEBUG(1),
		      "Postgres driver running multithreaded");
#else /* ISC_PLATFORM_USETHREADS */
	/* if debugging, let user know we are single threaded. */
	isc_log_write(dns_lctx, DNS_LOGCATEGORY_DATABASE,
		      DNS_LOGMODULE_DLZ, ISC_LOG_DEBUG(1),
		      "Postgres driver running single threaded");
#endif /* ISC_PLATFORM_USETHREADS */

	/* verify we have at least 5 arg's passed to the driver */
	if (argc < 5) {
		isc_log_write(dns_lctx, DNS_LOGCATEGORY_DATABASE,
			      DNS_LOGMODULE_DLZ, ISC_LOG_ERROR,
			      "Postgres driver requires at least "
			      "4 command line args.");
		return (ISC_R_FAILURE);
	}

	/* no more than 8 arg's should be passed to the driver */
	if (argc > 8) {
		isc_log_write(dns_lctx, DNS_LOGCATEGORY_DATABASE,
			      DNS_LOGMODULE_DLZ, ISC_LOG_ERROR,
			      "Postgres driver cannot accept more than "
			      "7 command line args.");
		return (ISC_R_FAILURE);
	}

	/* multithreaded build can have multiple DB connections */
#ifdef ISC_PLATFORM_USETHREADS

	/* check how many db connections we should create */
	dbcount = strtol(argv[1], &endp, 10);
	if (*endp != '\0' || dbcount < 0) {
		isc_log_write(dns_lctx, DNS_LOGCATEGORY_DATABASE,
			      DNS_LOGMODULE_DLZ, ISC_LOG_ERROR,
			      "Postgres driver database connection count "
			      "must be positive.");
		return (ISC_R_FAILURE);
	}

	/* allocate memory for database connection list */
	dblist = isc_mem_get(ns_g_mctx, sizeof(db_list_t));
	if (dblist == NULL)
		return (ISC_R_NOMEMORY);

	/* initialize DB connection list */
	ISC_LIST_INIT(*dblist);

	/*
	 * create the appropriate number of database instances (DBI)
	 * append each new DBI to the end of the list
	 */
	for (i=0; i < dbcount; i++) {

#endif /* ISC_PLATFORM_USETHREADS */

		/* how many queries were passed in from config file? */
		switch(argc) {
		case 5:
			result = build_sqldbinstance(ns_g_mctx, NULL, NULL,
						     NULL, argv[3], argv[4],
						     NULL, &dbi);
			break;
		case 6:
			result = build_sqldbinstance(ns_g_mctx, NULL, NULL,
						     argv[5], argv[3], argv[4],
						     NULL, &dbi);
			break;
		case 7:
			result = build_sqldbinstance(ns_g_mctx, argv[6], NULL,
						     argv[5], argv[3], argv[4],
						     NULL, &dbi);
			break;
		case 8:
			result = build_sqldbinstance(ns_g_mctx, argv[6],
						     argv[7], argv[5], argv[3],
						     argv[4], NULL, &dbi);
			break;
		default:
			/* not really needed, should shut up compiler. */
			result = ISC_R_FAILURE;
		}


		if (result == ISC_R_SUCCESS) {
			isc_log_write(dns_lctx, DNS_LOGCATEGORY_DATABASE,
				      DNS_LOGMODULE_DLZ, ISC_LOG_DEBUG(2),
				      "Postgres driver created database "
				      "instance object.");
		} else { /* unsuccessful?, log err msg and cleanup. */
			isc_log_write(dns_lctx, DNS_LOGCATEGORY_DATABASE,
				      DNS_LOGMODULE_DLZ, ISC_LOG_ERROR,
				      "Postgres driver could not create "
				      "database instance object.");
			goto cleanup;
		}

#ifdef ISC_PLATFORM_USETHREADS

		/* when multithreaded, build a list of DBI's */
		ISC_LINK_INIT(dbi, link);
		ISC_LIST_APPEND(*dblist, dbi, link);

#endif

		/* create and set db connection */
		dbi->dbconn = PQconnectdb(argv[2]);
		/*
		 * if db connection cannot be created, log err msg and
		 * cleanup.
		 */
		if (dbi->dbconn == NULL) {
			isc_log_write(dns_lctx, DNS_LOGCATEGORY_DATABASE,
				      DNS_LOGMODULE_DLZ, ISC_LOG_ERROR,
				      "Postgres driver could not allocate "
				      "memory for database connection");
			goto cleanup;
		}

		/* if we cannot connect the first time, try 3 more times. */
		for (j = 0;
		     PQstatus((PGconn *) dbi->dbconn) != CONNECTION_OK &&
			     j < 3;
		     j++)
			PQreset((PGconn *) dbi->dbconn);


#ifdef ISC_PLATFORM_USETHREADS

		/*
		 * if multi threaded, let user know which connection
		 * failed.  user could be attempting to create 10 db
		 * connections and for some reason the db backend only
		 * allows 9
		 */
		if (PQstatus((PGconn *) dbi->dbconn) != CONNECTION_OK) {
			isc_log_write(dns_lctx, DNS_LOGCATEGORY_DATABASE,
				      DNS_LOGMODULE_DLZ, ISC_LOG_ERROR,
				      "Postgres driver failed to create "
				      "database connection number %u "
				      "after 4 attempts",
				      i + 1);
			goto cleanup;
		}

		/* set DBI = null for next loop through. */
		dbi = NULL;
	}	/* end for loop */

		/* set dbdata to the list we created. */
	*dbdata = dblist;

#else /* ISC_PLATFORM_USETHREADS */
	/* if single threaded, just let user know we couldn't connect. */
	if (PQstatus((PGconn *) dbi->dbconn) != CONNECTION_OK) {
		isc_log_write(dns_lctx, DNS_LOGCATEGORY_DATABASE,
			      DNS_LOGMODULE_DLZ, ISC_LOG_ERROR,
			      "Postgres driver failed to create database "
			      "connection after 4 attempts");
		goto cleanup;
	}

	/*
	 * single threaded build can only use 1 db connection, return
	 * it via dbdata
	 */
	*dbdata = dbi;

#endif /* ISC_PLATFORM_USETHREADS */

	/* hey, we got through all of that ok, return success. */
	return(ISC_R_SUCCESS);

 cleanup:

#ifdef ISC_PLATFORM_USETHREADS
	/*
	 * if multithreaded, we could fail because only 1 connection
	 * couldn't be made.  We should cleanup the other successful
	 * connections properly.
	 */
	postgres_destroy_dblist(dblist);

#else /* ISC_PLATFORM_USETHREADS */
	if (dbi != NULL)
		destroy_sqldbinstance(dbi);

#endif /* ISC_PLATFORM_USETHREADS */
	return(ISC_R_FAILURE);
}

/*%
 * destroy an instance of the driver.  Remember, only 1 copy of the driver's
 * code is ever loaded, the driver has to remember which context it's
 * operating in.  This is done via use of the dbdata argument.
 * so we really only need to clean it up since we are not using driverarg.
 */
static void
postgres_destroy(void *driverarg, void *dbdata)
{

#ifdef ISC_PLATFORM_USETHREADS

	UNUSED(driverarg);
	/* cleanup the list of DBI's */
	postgres_destroy_dblist((db_list_t *) dbdata);

#else /* ISC_PLATFORM_USETHREADS */

	dbinstance_t *dbi;

	UNUSED(driverarg);

	dbi = (dbinstance_t *) dbdata;

	/* release DB connection */
	if (dbi->dbconn != NULL)
		PQfinish((PGconn *) dbi->dbconn);

	/* destroy single DB instance */
	destroy_sqldbinstance(dbi);

#endif /* ISC_PLATFORM_USETHREADS */
}

/* pointers to all our runtime methods. */
/* this is used during driver registration */
/* i.e. in dlz_postgres_init below. */
static dns_sdlzmethods_t dlz_postgres_methods = {
	postgres_create,
	postgres_destroy,
	postgres_findzone,
	postgres_lookup,
	postgres_authority,
	postgres_allnodes,
	postgres_allowzonexfr
};

/*%
 * Wrapper around dns_sdlzregister().
 */
isc_result_t
dlz_postgres_init(void) {
	isc_result_t result;

	/*
	 * Write debugging message to log
	 */
	isc_log_write(dns_lctx, DNS_LOGCATEGORY_DATABASE,
		      DNS_LOGMODULE_DLZ, ISC_LOG_DEBUG(2),
		      "Registering DLZ postgres driver.");

	/*
	 * Driver is always threadsafe.  When multithreaded all
	 * functions use multithreaded code.  When not multithreaded,
	 * all functions can only be entered once, but only 1 thread
	 * of operation is available in Bind.  So everything is still
	 * threadsafe.
	 */
	result = dns_sdlzregister("postgres", &dlz_postgres_methods, NULL,
				  DNS_SDLZFLAG_RELATIVEOWNER |
				  DNS_SDLZFLAG_RELATIVERDATA |
				  DNS_SDLZFLAG_THREADSAFE,
				  ns_g_mctx, &dlz_postgres);
	/* if we can't register the driver, there are big problems. */
	if (result != ISC_R_SUCCESS) {
		UNEXPECTED_ERROR(__FILE__, __LINE__,
				 "dns_sdlzregister() failed: %s",
				 isc_result_totext(result));
		result = ISC_R_UNEXPECTED;
	}


	return result;
}

/*%
 * Wrapper around dns_sdlzunregister().
 */
void
dlz_postgres_clear(void) {

	/*
	 * Write debugging message to log
	 */
	isc_log_write(dns_lctx, DNS_LOGCATEGORY_DATABASE,
		      DNS_LOGMODULE_DLZ, ISC_LOG_DEBUG(2),
		      "Unregistering DLZ postgres driver.");

	/* unregister the driver. */
	if (dlz_postgres != NULL)
		dns_sdlzunregister(&dlz_postgres);
}

#endif