datastore_sqlite.c

一个C语言写的快速贝叶斯垃圾邮件过滤工具

    rc = sqlite3_open(dbh->name, &dbh->db);
    if (rc) {
	print_error(__FILE__, __LINE__, "Can't open database %s: %s\n",
		dbh->name, sqlite3_errmsg(dbh->db));
	goto barf;
    }

    /* set trace mode */
    if (DEBUG_DATABASE(1) || getenv("BF_DEBUG_DB"))
	sqlite3_trace(dbh->db, db_trace, NULL);

    /* set busy handler */
    if (sqlite3_busy_handler(dbh->db, busyhandler, NULL)) {
	print_error(__FILE__, __LINE__, "Can't set busy handler: %s\n",
		sqlite3_errmsg(dbh->db));
	goto barf;
    }

    /* check/set endianness marker and create table if needed */
    if (mode != DS_READ) {
	/* using IMMEDIATE or DEFERRED here locks up in t.lock3
	 * or t.bulkmode
	 * using EXCLUSIVE locks up in t.lock3 on MAC OSX
	 */
	if (sqlexec(dbh->db, BEGIN)) goto barf;
	/*
	 * trick: the sqlite_master table (see SQLite FAQ) is read-only
	 * and lists all table, indexes etc. so we use it to check if
	 * the bogofilter table is already there, the error codes are
	 * too vague either way, for "no such table" and "table already
	 * exists" we always get SQLITE_ERROR, which we'll also get for
	 * syntax errors, such as "EXCLUSIVE" not supported on older
	 * versions :-(
	 */
	rc = db_loop(dbh->db, "SELECT name FROM sqlite_master "
		"WHERE type='table' AND name='bogofilter';",
		NULL, NULL);
	switch (rc) {
	    case 0:
		if (sqlexec(dbh->db, "COMMIT;")) goto barf;
		break;
	    case DS_NOTFOUND:
		{
		    u_int32_t p[2] = { 0x01020304, 0x01020304 };

		    if (sqlexec(dbh->db, LAYOUT)) goto barf;

		    /* set endianness marker */
		    k.data = xstrdup(ENDIAN32);
		    k.leng = strlen(k.data);
		    v.data = p;
		    v.leng = sizeof(p);
		    rc = db_set_dbvalue(dbh, &k, &v);
		    xfree(k.data);
		    if (rc)
			goto barf;

		    if (sqlexec(dbh->db, "COMMIT;")) goto barf;
		    dbh->created = true;
		}
		break;
	    default:
		goto barf;
	}
    }

    /*
     * initialize common statements
     * dbh->insert is not here as it's needed earlier,
     * so it sets itself up lazily
     */
    dbh->select = sqlprep(dbh, "SELECT value FROM bogofilter WHERE key=? LIMIT 1;", false);
    if (dbh->select == NULL)
    {
	fprintf(stderr,
		"\nRemember to register some spam and ham messages before you\n"
		"use bogofilter to evaluate mail for its probable spam status!\n\n");
	exit(EX_ERROR);
    }

    dbh->delete = sqlprep(dbh, "DELETE FROM bogofilter WHERE(key = ?);", true);

    /* check if byteswapped */
    {
	u_int32_t t, b[2];
	int ee;

	k.data = xstrdup(ENDIAN32);
	k.leng = strlen(k.data);
	v.data = b;
	v.leng = sizeof(b);

	ee = db_get_dbvalue(dbh, &k, &v);
	xfree(k.data);
	switch (ee) {
	    case 0: /* found endian marker token, read it */
		if (v.leng < 4)
		    goto barf;
		t = ((u_int32_t *)v.data)[0];
		switch (t) {
		    case 0x01020304: /* same endian, "UNIX" */
			dbh->swapped = false;
			break;
		    case 0x04030201: /* swapped, "XINU" */
			dbh->swapped = true;
			break;
		    default: /* NUXI or IXUN or crap */
			print_error(__FILE__, __LINE__,
				"Unknown endianness on %s: %08x.\n",
				dbh->name, ((u_int32_t *)v.data)[0]);
			goto barf2;
		}
		break;
	    case DS_NOTFOUND: /* no marker token, assume not swapped */
		dbh->swapped = false;
		break;
	    default:
		goto barf;
	}
    }

    return dbh;
barf:
    print_error(__FILE__, __LINE__, "Error on database %s: %s\n",
	    dbh->name, sqlite3_errmsg(dbh->db));
barf2:
    db_close(dbh);
    return NULL;
}

void db_close(void *handle) {
    int rc;
    dbh_t *dbh = handle;
    if (dbh->delete) sqlite3_finalize(dbh->delete);
    if (dbh->insert) sqlite3_finalize(dbh->insert);
    if (dbh->select) sqlite3_finalize(dbh->select);
    rc = sqlite3_close(dbh->db);
    if (rc) {
	print_error(__FILE__, __LINE__, "Can't close database %s: %d",
		dbh->name, rc);
	exit(EX_ERROR);
    }
    free_dbh(dbh);
}

const char *db_version_str(void) {
    static char buf[80];

    if (!buf[0])
	snprintf(buf, sizeof(buf), "SQLite %s", sqlite3_libversion());
    return buf;
}

static int sql_txn_begin(void *vhandle) {
    dbh_t *dbh = vhandle;
    return sqlexec(dbh->db,  BEGIN );
}

static int sql_txn_abort(void *vhandle) {
    dbh_t *dbh = vhandle;
    return sqlexec(dbh->db, "ROLLBACK;");
}

static int sql_txn_commit(void *vhandle) {
    dbh_t *dbh = vhandle;
    return sqlexec(dbh->db, "COMMIT;");
}

/** common code for db_delete, db_(get|set)_dbvalue.
 * This works by setting variables in precompiled statements (see PREP,
 * sqlite3_prepare, sqlite3_bind_*, sqlite3_reset) and avoids encoding
 * binary data into SQL's hex representation as well as compiling the
 * same SQL statement over and over again. */
static int sql_fastpath(
	dbh_t *dbh,		/**< database handle */
	const char *func,	/**< function name to report in errors */
	sqlite3_stmt *stmt,	/**< SQLite3 statement to execute/reset */
	dbv_t *val,		/**< OUT value from first row, NULL ok */
	int retnotfound		/**  return value if no rows found */
	)
{
    int rc;
    bool found = false;

    while (1) {
	rc = sqlite3_step(stmt);
	switch (rc) {
	    case SQLITE_ROW:	/* this is the only branch that loops */
		if (val) {
		    int len = min(INT_MAX, val->leng);
		    val->leng = min(len, sqlite3_column_bytes(stmt, 0));
		    memcpy(val->data, sqlite3_column_blob(stmt, 0), val->leng);
		}
		found = 1;
		break;
		/* all other branches below return control to the caller */
	    case SQLITE_BUSY:
		sqlite3_reset(stmt);
		sql_txn_abort(dbh);
		return DS_ABORT_RETRY;

	    case SQLITE_DONE:
		sqlite3_reset(stmt);
		return found ? 0 : retnotfound;

	    default:
		print_error(__FILE__, __LINE__,
			"%s: error executing statement on %s: %s (%d)\n",
			func, dbh->name, sqlite3_errmsg(dbh->db), rc);
		sqlite3_reset(stmt);
		return rc;
	}
    }
}

int db_delete(void *vhandle, const dbv_t *key) {
    dbh_t *dbh = vhandle;

    sqlite3_bind_blob(dbh->delete, 1, key->data, key->leng, SQLITE_STATIC);
    return sql_fastpath(dbh, "db_delete", dbh->delete, NULL, 0);
}

int db_set_dbvalue(void *vhandle, const dbv_t *key, const dbv_t *val) {
    dbh_t *dbh = vhandle;

    if (!dbh->insert)
	dbh->insert = sqlprep(dbh, "INSERT OR REPLACE INTO bogofilter VALUES(?,?);", true);

    sqlite3_bind_blob(dbh->insert, 1, key->data, key->leng, SQLITE_STATIC);
    sqlite3_bind_blob(dbh->insert, 2, val->data, val->leng, SQLITE_STATIC);
    return sql_fastpath(dbh, "db_set_dbvalue", dbh->insert, NULL, 0);
}

int db_get_dbvalue(void *vhandle, const dbv_t* key, /*@out@*/ dbv_t *val) {
    dbh_t *dbh = vhandle;

    sqlite3_bind_blob(dbh->select, 1, key->data, key->leng, SQLITE_STATIC);
    return sql_fastpath(dbh, "db_get_dbvalue", dbh->select, val, DS_NOTFOUND);
}

ex_t db_foreach(void *vhandle, db_foreach_t hook, void *userdata) {
    dbh_t *dbh = vhandle;
    const char *cmd = "SELECT key, value FROM bogofilter;";
    return db_loop(dbh->db, cmd, hook, userdata);
}

const char *db_str_err(int e) {
    return e == 0 ? "no error" : "unknown condition (not yet implemented)";
}

bool db_created(void *vhandle) {
    dbh_t *dbh = vhandle;
    return dbh->created;
}

bool db_is_swapped(void *vhandle) {
    dbh_t *dbh = vhandle;
    return dbh->swapped;
}

static int pagesize_cb(void *ptr, int argc, char **argv, char **dummy) {
    u_int32_t *uptr = ptr;

    (void)dummy;

    if (argc != 1)
	return -1;
    errno = 0;
    *uptr = strtoul(argv[0], NULL, 0);
    return errno;
}

static u_int32_t sql_pagesize(bfpath *bfp)
{
    dbh_t *dbh;
    int rc;
    u_int32_t size;

    dbh = db_open(NULL, bfp, DS_READ);
    if (!dbh)
	return 0xffffffff;
    rc = sqlite3_exec(dbh->db, "PRAGMA page_size;", pagesize_cb, &size, NULL);
    if (rc != SQLITE_OK) {
	return 0xffffffff;
    }
    db_close(dbh);
    return size;
}