[people/ms/libloc.git] / src / writer.c

/*
	libloc - A library to determine the location of someone on the Internet

	Copyright (C) 2017 IPFire Development Team <info@ipfire.org>

	This library is free software; you can redistribute it and/or
	modify it under the terms of the GNU Lesser General Public
	License as published by the Free Software Foundation; either
	version 2.1 of the License, or (at your option) any later version.

	This library is distributed in the hope that it will be useful,
	but WITHOUT ANY WARRANTY; without even the implied warranty of
	MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
	Lesser General Public License for more details.
*/

#include <assert.h>
#include <errno.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/queue.h>
#include <time.h>

#ifdef HAVE_ENDIAN_H
#  include <endian.h>
#endif

#include <openssl/bio.h>
#include <openssl/err.h>
#include <openssl/evp.h>
#include <openssl/pem.h>

#include <loc/libloc.h>
#include <loc/as.h>
#include <loc/compat.h>
#include <loc/country.h>
#include <loc/database.h>
#include <loc/format.h>
#include <loc/network.h>
#include <loc/private.h>
#include <loc/writer.h>

struct loc_writer {
	struct loc_ctx* ctx;
	int refcount;

	struct loc_stringpool* pool;
	off_t vendor;
	off_t description;
	off_t license;

	// Private key to sign any databases
	EVP_PKEY* private_key;

	struct loc_as** as;
	size_t as_count;

	struct loc_country** countries;
	size_t countries_count;

	struct loc_network_tree* networks;
};

static int parse_private_key(struct loc_writer* writer, FILE* f) {
	// Free any previously loaded keys
	if (writer->private_key)
		EVP_PKEY_free(writer->private_key);

	// Read the key
	writer->private_key = PEM_read_PrivateKey(f, NULL, NULL, NULL);

	// Log any errors
	if (!writer->private_key) {
		char* error = ERR_error_string(ERR_get_error(), NULL);
		ERROR(writer->ctx, "Could not parse private key: %s\n", error);

		return -1;
	}

	return 0;
}

LOC_EXPORT int loc_writer_new(struct loc_ctx* ctx, struct loc_writer** writer, FILE* fkey) {
	struct loc_writer* w = calloc(1, sizeof(*w));
	if (!w)
		return -ENOMEM;

	w->ctx = loc_ref(ctx);
	w->refcount = 1;

	int r = loc_stringpool_new(ctx, &w->pool);
	if (r) {
		loc_writer_unref(w);
		return r;
	}

	// Initialize the network tree
	r = loc_network_tree_new(ctx, &w->networks);
	if (r) {
		loc_writer_unref(w);
		return r;
	}

	// Load the private key to sign databases
	if (fkey) {
		r = parse_private_key(w, fkey);
		if (r) {
			loc_writer_unref(w);
			return r;
		}
	}

	*writer = w;
	return 0;
}

LOC_EXPORT struct loc_writer* loc_writer_ref(struct loc_writer* writer) {
	writer->refcount++;

	return writer;
}

static void loc_writer_free(struct loc_writer* writer) {
	DEBUG(writer->ctx, "Releasing writer at %p\n", writer);

	// Free private key
	if (writer->private_key)
		EVP_PKEY_free(writer->private_key);

	// Unref all AS
	for (unsigned int i = 0; i < writer->as_count; i++) {
		loc_as_unref(writer->as[i]);
	}

	// Release network tree
	if (writer->networks)
		loc_network_tree_unref(writer->networks);

	// Unref the string pool
	loc_stringpool_unref(writer->pool);

	loc_unref(writer->ctx);
	free(writer);
}

LOC_EXPORT struct loc_writer* loc_writer_unref(struct loc_writer* writer) {
	if (--writer->refcount > 0)
		return writer;

	loc_writer_free(writer);

	return NULL;
}

LOC_EXPORT const char* loc_writer_get_vendor(struct loc_writer* writer) {
	return loc_stringpool_get(writer->pool, writer->vendor);
}

LOC_EXPORT int loc_writer_set_vendor(struct loc_writer* writer, const char* vendor) {
	// Add the string to the string pool
	off_t offset = loc_stringpool_add(writer->pool, vendor);
	if (offset < 0)
		return offset;

	writer->vendor = offset;
	return 0;
}

LOC_EXPORT const char* loc_writer_get_description(struct loc_writer* writer) {
	return loc_stringpool_get(writer->pool, writer->description);
}

LOC_EXPORT int loc_writer_set_description(struct loc_writer* writer, const char* description) {
	// Add the string to the string pool
	off_t offset = loc_stringpool_add(writer->pool, description);
	if (offset < 0)
		return offset;

	writer->description = offset;
	return 0;
}

LOC_EXPORT const char* loc_writer_get_license(struct loc_writer* writer) {
	return loc_stringpool_get(writer->pool, writer->license);
}

LOC_EXPORT int loc_writer_set_license(struct loc_writer* writer, const char* license) {
	// Add the string to the string pool
	off_t offset = loc_stringpool_add(writer->pool, license);
	if (offset < 0)
		return offset;

	writer->license = offset;
	return 0;
}

static int __loc_as_cmp(const void* as1, const void* as2) {
	return loc_as_cmp(*(struct loc_as**)as1, *(struct loc_as**)as2);
}

LOC_EXPORT int loc_writer_add_as(struct loc_writer* writer, struct loc_as** as, uint32_t number) {
	int r = loc_as_new(writer->ctx, as, number);
	if (r)
		return r;

	// We have a new AS to add
	writer->as_count++;

	// Make space
	writer->as = realloc(writer->as, sizeof(*writer->as) * writer->as_count);
	if (!writer->as)
		return -ENOMEM;

	// Add as last element
	writer->as[writer->as_count - 1] = loc_as_ref(*as);

	// Sort everything
	qsort(writer->as, writer->as_count, sizeof(*writer->as), __loc_as_cmp);

	return 0;
}

LOC_EXPORT int loc_writer_add_network(struct loc_writer* writer, struct loc_network** network, const char* string) {
	int r;

	// Create a new network object
	r = loc_network_new_from_string(writer->ctx, network, string);
	if (r)
		return r;

	// Add it to the local tree
	return loc_network_tree_add_network(writer->networks, *network);
}

static int __loc_country_cmp(const void* country1, const void* country2) {
	return loc_country_cmp(*(struct loc_country**)country1, *(struct loc_country**)country2);
}

LOC_EXPORT int loc_writer_add_country(struct loc_writer* writer, struct loc_country** country, const char* country_code) {
	int r = loc_country_new(writer->ctx, country, country_code);
	if (r)
		return r;

	// We have a new country to add
	writer->countries_count++;

	// Make space
	writer->countries = realloc(writer->countries, sizeof(*writer->countries) * writer->countries_count);
	if (!writer->countries)
		return -ENOMEM;

	// Add as last element
	writer->countries[writer->countries_count - 1] = loc_country_ref(*country);

	// Sort everything
	qsort(writer->countries, writer->countries_count, sizeof(*writer->countries), __loc_country_cmp);

	return 0;
}

static void make_magic(struct loc_writer* writer, struct loc_database_magic* magic,
		enum loc_database_version version) {
	// Copy magic bytes
	for (unsigned int i = 0; i < strlen(LOC_DATABASE_MAGIC); i++)
		magic->magic[i] = LOC_DATABASE_MAGIC[i];

	// Set version
	magic->version = version;
}

static void align_page_boundary(off_t* offset, FILE* f) {
	// Move to next page boundary
	while (*offset % LOC_DATABASE_PAGE_SIZE > 0)
		*offset += fwrite("", 1, 1, f);
}

static int loc_database_write_pool(struct loc_writer* writer,
		struct loc_database_header_v1* header, off_t* offset, FILE* f) {
	// Save the offset of the pool section
	DEBUG(writer->ctx, "Pool starts at %jd bytes\n", (intmax_t)*offset);
	header->pool_offset = htobe32(*offset);

	// Write the pool
	size_t pool_length = loc_stringpool_write(writer->pool, f);
	*offset += pool_length;

	DEBUG(writer->ctx, "Pool has a length of %zu bytes\n", pool_length);
	header->pool_length = htobe32(pool_length);

	return 0;
}

static int loc_database_write_as_section(struct loc_writer* writer,
		struct loc_database_header_v1* header, off_t* offset, FILE* f) {
	DEBUG(writer->ctx, "AS section starts at %jd bytes\n", (intmax_t)*offset);
	header->as_offset = htobe32(*offset);

	size_t as_length = 0;

	struct loc_database_as_v1 as;
	for (unsigned int i = 0; i < writer->as_count; i++) {
		// Convert AS into database format
		loc_as_to_database_v1(writer->as[i], writer->pool, &as);

		// Write to disk
		*offset += fwrite(&as, 1, sizeof(as), f);
		as_length += sizeof(as);
	}

	DEBUG(writer->ctx, "AS section has a length of %zu bytes\n", as_length);
	header->as_length = htobe32(as_length);

	align_page_boundary(offset, f);

	return 0;
}

struct node {
	TAILQ_ENTRY(node) nodes;

	struct loc_network_tree_node* node;

	// Indices of the child nodes
	uint32_t index_zero;
	uint32_t index_one;
};

static struct node* make_node(struct loc_network_tree_node* node) {
	struct node* n = malloc(sizeof(*n));
	if (!n)
		return NULL;

	n->node  = loc_network_tree_node_ref(node);
	n->index_zero = n->index_one = 0;

	return n;
}

static void free_node(struct node* node) {
	loc_network_tree_node_unref(node->node);

	free(node);
}

struct network {
	TAILQ_ENTRY(network) networks;

	struct loc_network* network;
};

static struct network* make_network(struct loc_network* network) {
	struct network* n = malloc(sizeof(*n));
	if (!n)
		return NULL;

	n->network = loc_network_ref(network);

	return n;
}

static void free_network(struct network* network) {
	loc_network_unref(network->network);

	free(network);
}

static int loc_database_write_networks(struct loc_writer* writer,
		struct loc_database_header_v1* header, off_t* offset, FILE* f) {
	// Write the network tree
	DEBUG(writer->ctx, "Network tree starts at %jd bytes\n", (intmax_t)*offset);
	header->network_tree_offset = htobe32(*offset);

	size_t network_tree_length = 0;
	size_t network_data_length = 0;

	struct node* node;
	struct node* child_node;

	uint32_t index = 0;
	uint32_t network_index = 0;

	struct loc_database_network_v1 db_network;
	struct loc_database_network_node_v1 db_node;

	// Initialize queue for nodes
	TAILQ_HEAD(node_t, node) nodes;
	TAILQ_INIT(&nodes);

	// Initialize queue for networks
	TAILQ_HEAD(network_t, network) networks;
	TAILQ_INIT(&networks);

	// Add root
	struct loc_network_tree_node* root = loc_network_tree_get_root(writer->networks);
	node = make_node(root);

	TAILQ_INSERT_TAIL(&nodes, node, nodes);

	while (!TAILQ_EMPTY(&nodes)) {
		// Pop first node in list
		node = TAILQ_FIRST(&nodes);
		TAILQ_REMOVE(&nodes, node, nodes);

		DEBUG(writer->ctx, "Processing node %p\n", node);

		// Get child nodes
		struct loc_network_tree_node* node_zero = loc_network_tree_node_get(node->node, 0);
		if (node_zero) {
			node->index_zero = ++index;

			child_node = make_node(node_zero);
			loc_network_tree_node_unref(node_zero);

			TAILQ_INSERT_TAIL(&nodes, child_node, nodes);
		}

		struct loc_network_tree_node* node_one = loc_network_tree_node_get(node->node, 1);
		if (node_one) {
			node->index_one = ++index;

			child_node = make_node(node_one);
			loc_network_tree_node_unref(node_one);

			TAILQ_INSERT_TAIL(&nodes, child_node, nodes);
		}

		// Prepare what we are writing to disk
		db_node.zero = htobe32(node->index_zero);
		db_node.one  = htobe32(node->index_one);

		if (loc_network_tree_node_is_leaf(node->node)) {
			struct loc_network* network = loc_network_tree_node_get_network(node->node);

			// Append network to be written out later
			struct network* nw = make_network(network);
			TAILQ_INSERT_TAIL(&networks, nw, networks);

			db_node.network = htobe32(network_index++);
			loc_network_unref(network);
		} else {
			db_node.network = htobe32(0xffffffff);
		}

		// Write the current node
		DEBUG(writer->ctx, "Writing node %p (0 = %d, 1 = %d)\n",
			node, node->index_zero, node->index_one);

		*offset += fwrite(&db_node, 1, sizeof(db_node), f);
		network_tree_length += sizeof(db_node);

		free_node(node);
	}

	loc_network_tree_node_unref(root);

	header->network_tree_length = htobe32(network_tree_length);

	align_page_boundary(offset, f);

	DEBUG(writer->ctx, "Networks data section starts at %jd bytes\n", (intmax_t)*offset);
	header->network_data_offset = htobe32(*offset);

	// We have now written the entire tree and have all networks
	// in a queue in order as they are indexed
	while (!TAILQ_EMPTY(&networks)) {
		struct network* nw = TAILQ_FIRST(&networks);
		TAILQ_REMOVE(&networks, nw, networks);

		// Prepare what we are writing to disk
		int r = loc_network_to_database_v1(nw->network, &db_network);
		if (r)
			return r;

		*offset += fwrite(&db_network, 1, sizeof(db_network), f);
		network_data_length += sizeof(db_network);

		free_network(nw);
	}

	header->network_data_length = htobe32(network_data_length);

	align_page_boundary(offset, f);

	return 0;
}

static int loc_database_write_countries(struct loc_writer* writer,
		struct loc_database_header_v1* header, off_t* offset, FILE* f) {
	DEBUG(writer->ctx, "Countries section starts at %jd bytes\n", (intmax_t)*offset);
	header->countries_offset = htobe32(*offset);

	size_t countries_length = 0;

	struct loc_database_country_v1 country;
	for (unsigned int i = 0; i < writer->countries_count; i++) {
		// Convert country into database format
		loc_country_to_database_v1(writer->countries[i], writer->pool, &country);

		// Write to disk
		*offset += fwrite(&country, 1, sizeof(country), f);
		countries_length += sizeof(country);
	}

	DEBUG(writer->ctx, "Countries section has a length of %zu bytes\n", countries_length);
	header->countries_length = htobe32(countries_length);

	align_page_boundary(offset, f);

	return 0;
}

static int loc_writer_create_signature(struct loc_writer* writer,
		struct loc_database_header_v1* header, FILE* f) {
	DEBUG(writer->ctx, "Signing database...\n");

	// Read file from the beginning
	rewind(f);

	// Create a new context for signing
	EVP_MD_CTX* mdctx = EVP_MD_CTX_new();

	// Initialise the context
	int r = EVP_DigestSignInit(mdctx, NULL, NULL, NULL, writer->private_key);
	if (r != 1) {
		ERROR(writer->ctx, "%s\n", ERR_error_string(ERR_get_error(), NULL));
		goto END;
	}

	// Read magic
	struct loc_database_magic magic;
	fread(&magic, 1, sizeof(magic), f);

	hexdump(writer->ctx, &magic, sizeof(magic));

	// Feed magic into the signature
	r = EVP_DigestSignUpdate(mdctx, &magic, sizeof(magic));
	if (r != 1) {
		ERROR(writer->ctx, "%s\n", ERR_error_string(ERR_get_error(), NULL));
		goto END;
	}

	hexdump(writer->ctx, header, sizeof(*header));

	// Feed the header into the signature
	r = EVP_DigestSignUpdate(mdctx, header, sizeof(*header));
	if (r != 1) {
		ERROR(writer->ctx, "%s\n", ERR_error_string(ERR_get_error(), NULL));
		goto END;
	}

	// Skip header
	fseek(f, sizeof(*header), SEEK_CUR);

	// Walk through the file in chunks of 64kB
	char buffer[64 * 1024];
	while (!feof(f)) {
		size_t bytes_read = fread(buffer, 1, sizeof(buffer), f);

		if (ferror(f)) {
			ERROR(writer->ctx, "Error reading from file: %s\n", strerror(errno));
			r = errno;
			goto END;
		}

		hexdump(writer->ctx, buffer, bytes_read);

		r = EVP_DigestSignUpdate(mdctx, buffer, bytes_read);
		if (r != 1) {
			ERROR(writer->ctx, "%s\n", ERR_error_string(ERR_get_error(), NULL));
			goto END;
		}
	}

	// Compute the signature
	size_t signature_length = sizeof(header->signature);

	r = EVP_DigestSignFinal(mdctx,
		(unsigned char*)header->signature, &signature_length);
	if (r != 1) {
		ERROR(writer->ctx, "%s\n", ERR_error_string(ERR_get_error(), NULL));
		goto END;
	}

	// Save length of the signature
	header->signature_length = htobe32(signature_length);

	DEBUG(writer->ctx, "Successfully generated signature of %lu bytes\n",
		signature_length);
	r = 0;

	// Dump signature
	hexdump(writer->ctx, header->signature, signature_length);

END:
	EVP_MD_CTX_free(mdctx);

	return r;
}

LOC_EXPORT int loc_writer_write(struct loc_writer* writer, FILE* f, enum loc_database_version version) {
	// Check version
	switch (version) {
		case LOC_DATABASE_VERSION_UNSET:
			version = LOC_DATABASE_VERSION_LATEST;
			break;

		case LOC_DATABASE_VERSION_1:
			break;

		default:
			ERROR(writer->ctx, "Invalid database version: %d\n", version);
			return -1;
	}

	DEBUG(writer->ctx, "Writing database in version %d\n", version);

	struct loc_database_magic magic;
	make_magic(writer, &magic, version);

	// Make the header
	struct loc_database_header_v1 header;
	header.vendor      = htobe32(writer->vendor);
	header.description = htobe32(writer->description);
	header.license     = htobe32(writer->license);

	time_t now = time(NULL);
	header.created_at = htobe64(now);

	// Clear the signature
	header.signature_length = 0;
	for (unsigned int i = 0; i < sizeof(header.signature); i++)
		header.signature[i] = '\0';

	// Clear the padding
	for (unsigned int i = 0; i < sizeof(header.padding); i++)
		header.padding[i] = '\0';

	int r;
	off_t offset = 0;

	// Start writing at the beginning of the file
	r = fseek(f, 0, SEEK_SET);
	if (r)
		return r;

	// Write the magic
	offset += fwrite(&magic, 1, sizeof(magic), f);

	// Skip the space we need to write the header later
	r = fseek(f, sizeof(header), SEEK_CUR);
	if (r) {
		DEBUG(writer->ctx, "Could not seek to position after header\n");
		return r;
	}
	offset += sizeof(header);

	align_page_boundary(&offset, f);

	// Write all ASes
	r = loc_database_write_as_section(writer, &header, &offset, f);
	if (r)
		return r;

	// Write all networks
	r = loc_database_write_networks(writer, &header, &offset, f);
	if (r)
		return r;

	// Write countries
	r = loc_database_write_countries(writer, &header, &offset, f);
	if (r)
		return r;

	// Write pool
	r = loc_database_write_pool(writer, &header, &offset, f);
	if (r)
		return r;

	// Create the signature
	if (writer->private_key) {
		r = loc_writer_create_signature(writer, &header, f);
		if (r)
			return r;
	}

	// Write the header
	r = fseek(f, sizeof(magic), SEEK_SET);
	if (r)
		return r;

	fwrite(&header, 1, sizeof(header), f);

	return r;
}
Commit	Line	Data
c182393f MT	1	/*
	2	libloc - A library to determine the location of someone on the Internet
	3
	4	Copyright (C) 2017 IPFire Development Team <info@ipfire.org>
	5
	6	This library is free software; you can redistribute it and/or
	7	modify it under the terms of the GNU Lesser General Public
	8	License as published by the Free Software Foundation; either
	9	version 2.1 of the License, or (at your option) any later version.
	10
	11	This library is distributed in the hope that it will be useful,
	12	but WITHOUT ANY WARRANTY; without even the implied warranty of
	13	MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
	14	Lesser General Public License for more details.
	15	*/
	16
726f9984	17	#include <assert.h>
c182393f MT	18	#include <errno.h>
	19	#include <stdio.h>
	20	#include <stdlib.h>
	21	#include <string.h>
438db08c	22	#include <sys/queue.h>
c182393f MT	23	#include <time.h>
c182393f MT	24
42f3ccd7 MT	25	#ifdef HAVE_ENDIAN_H
	26	# include <endian.h>
	27	#endif
	28
726f9984 MT	29	#include <openssl/bio.h>
	30	#include <openssl/err.h>
	31	#include <openssl/evp.h>
	32	#include <openssl/pem.h>
	33
9fc7f001 MT	34	#include <loc/libloc.h>
9fc7f001 MT	35	#include <loc/as.h>
42f3ccd7	36	#include <loc/compat.h>
ec684c1a	37	#include <loc/country.h>
726f9984	38	#include <loc/database.h>
c182393f	39	#include <loc/format.h>
3b5f4af2	40	#include <loc/network.h>
9fc7f001	41	#include <loc/private.h>
c182393f MT	42	#include <loc/writer.h>
c182393f MT	43
c182393f MT	44	struct loc_writer {
	45	struct loc_ctx* ctx;
	46	int refcount;
	47
	48	struct loc_stringpool* pool;
	49	off_t vendor;
	50	off_t description;
4bf49d00	51	off_t license;
c182393f	52
726f9984 MT	53	// Private key to sign any databases
	54	EVP_PKEY* private_key;
	55
c182393f MT	56	struct loc_as** as;
c182393f MT	57	size_t as_count;
3b5f4af2	58
ec684c1a MT	59	struct loc_country** countries;
	60	size_t countries_count;
	61
3b5f4af2	62	struct loc_network_tree* networks;
c182393f MT	63	};
c182393f MT	64
726f9984 MT	65	static int parse_private_key(struct loc_writer* writer, FILE* f) {
	66	// Free any previously loaded keys
	67	if (writer->private_key)
	68	EVP_PKEY_free(writer->private_key);
	69
	70	// Read the key
	71	writer->private_key = PEM_read_PrivateKey(f, NULL, NULL, NULL);
	72
	73	// Log any errors
	74	if (!writer->private_key) {
	75	char* error = ERR_error_string(ERR_get_error(), NULL);
	76	ERROR(writer->ctx, "Could not parse private key: %s\n", error);
	77
	78	return -1;
	79	}
	80
	81	return 0;
	82	}
	83
22c7b98b	84	LOC_EXPORT int loc_writer_new(struct loc_ctx* ctx, struct loc_writer** writer, FILE* fkey) {
c182393f MT	85	struct loc_writer* w = calloc(1, sizeof(*w));
	86	if (!w)
	87	return -ENOMEM;
	88
	89	w->ctx = loc_ref(ctx);
	90	w->refcount = 1;
	91
0e974d4b	92	int r = loc_stringpool_new(ctx, &w->pool);
c182393f MT	93	if (r) {
	94	loc_writer_unref(w);
	95	return r;
	96	}
	97
3b5f4af2 MT	98	// Initialize the network tree
	99	r = loc_network_tree_new(ctx, &w->networks);
	100	if (r) {
	101	loc_writer_unref(w);
	102	return r;
	103	}
	104
726f9984 MT	105	// Load the private key to sign databases
	106	if (fkey) {
	107	r = parse_private_key(w, fkey);
	108	if (r) {
	109	loc_writer_unref(w);
	110	return r;
	111	}
	112	}
	113
c182393f MT	114	*writer = w;
	115	return 0;
	116	}
	117
	118	LOC_EXPORT struct loc_writer* loc_writer_ref(struct loc_writer* writer) {
	119	writer->refcount++;
	120
	121	return writer;
	122	}
	123
	124	static void loc_writer_free(struct loc_writer* writer) {
	125	DEBUG(writer->ctx, "Releasing writer at %p\n", writer);
	126
726f9984 MT	127	// Free private key
	128	if (writer->private_key)
	129	EVP_PKEY_free(writer->private_key);
	130
c182393f MT	131	// Unref all AS
	132	for (unsigned int i = 0; i < writer->as_count; i++) {
	133	loc_as_unref(writer->as[i]);
	134	}
	135
3b5f4af2 MT	136	// Release network tree
	137	if (writer->networks)
	138	loc_network_tree_unref(writer->networks);
	139
c182393f MT	140	// Unref the string pool
	141	loc_stringpool_unref(writer->pool);
	142
	143	loc_unref(writer->ctx);
	144	free(writer);
	145	}
	146
	147	LOC_EXPORT struct loc_writer* loc_writer_unref(struct loc_writer* writer) {
	148	if (--writer->refcount > 0)
	149	return writer;
	150
	151	loc_writer_free(writer);
	152
	153	return NULL;
	154	}
	155
	156	LOC_EXPORT const char* loc_writer_get_vendor(struct loc_writer* writer) {
	157	return loc_stringpool_get(writer->pool, writer->vendor);
	158	}
	159
	160	LOC_EXPORT int loc_writer_set_vendor(struct loc_writer* writer, const char* vendor) {
	161	// Add the string to the string pool
	162	off_t offset = loc_stringpool_add(writer->pool, vendor);
	163	if (offset < 0)
	164	return offset;
	165
	166	writer->vendor = offset;
	167	return 0;
	168	}
	169
	170	LOC_EXPORT const char* loc_writer_get_description(struct loc_writer* writer) {
	171	return loc_stringpool_get(writer->pool, writer->description);
	172	}
	173
	174	LOC_EXPORT int loc_writer_set_description(struct loc_writer* writer, const char* description) {
	175	// Add the string to the string pool
	176	off_t offset = loc_stringpool_add(writer->pool, description);
	177	if (offset < 0)
	178	return offset;
	179
	180	writer->description = offset;
	181	return 0;
	182	}
	183
4bf49d00 MT	184	LOC_EXPORT const char* loc_writer_get_license(struct loc_writer* writer) {
	185	return loc_stringpool_get(writer->pool, writer->license);
	186	}
	187
	188	LOC_EXPORT int loc_writer_set_license(struct loc_writer* writer, const char* license) {
	189	// Add the string to the string pool
	190	off_t offset = loc_stringpool_add(writer->pool, license);
	191	if (offset < 0)
	192	return offset;
	193
	194	writer->license = offset;
	195	return 0;
	196	}
	197
c182393f MT	198	static int __loc_as_cmp(const void* as1, const void* as2) {
	199	return loc_as_cmp((struct loc_as)as1, (struct loc_as**)as2);
	200	}
	201
	202	LOC_EXPORT int loc_writer_add_as(struct loc_writer* writer, struct loc_as** as, uint32_t number) {
29bd8211	203	int r = loc_as_new(writer->ctx, as, number);
c182393f MT	204	if (r)
	205	return r;
	206
	207	// We have a new AS to add
	208	writer->as_count++;
	209
	210	// Make space
	211	writer->as = realloc(writer->as, sizeof(writer->as) writer->as_count);
	212	if (!writer->as)
	213	return -ENOMEM;
	214
	215	// Add as last element
	216	writer->as[writer->as_count - 1] = loc_as_ref(*as);
	217
	218	// Sort everything
	219	qsort(writer->as, writer->as_count, sizeof(*writer->as), __loc_as_cmp);
	220
	221	return 0;
	222	}
	223
3b5f4af2 MT	224	LOC_EXPORT int loc_writer_add_network(struct loc_writer* writer, struct loc_network** network, const char* string) {
	225	int r;
	226
	227	// Create a new network object
	228	r = loc_network_new_from_string(writer->ctx, network, string);
	229	if (r)
	230	return r;
	231
	232	// Add it to the local tree
	233	return loc_network_tree_add_network(writer->networks, *network);
	234	}
	235
ec684c1a MT	236	static int __loc_country_cmp(const void* country1, const void* country2) {
	237	return loc_country_cmp((struct loc_country)country1, (struct loc_country**)country2);
	238	}
	239
	240	LOC_EXPORT int loc_writer_add_country(struct loc_writer* writer, struct loc_country** country, const char* country_code) {
	241	int r = loc_country_new(writer->ctx, country, country_code);
	242	if (r)
	243	return r;
	244
	245	// We have a new country to add
	246	writer->countries_count++;
	247
	248	// Make space
	249	writer->countries = realloc(writer->countries, sizeof(writer->countries) writer->countries_count);
	250	if (!writer->countries)
	251	return -ENOMEM;
	252
	253	// Add as last element
	254	writer->countries[writer->countries_count - 1] = loc_country_ref(*country);
	255
	256	// Sort everything
	257	qsort(writer->countries, writer->countries_count, sizeof(*writer->countries), __loc_country_cmp);
	258
	259	return 0;
	260	}
	261
22c7b98b MT	262	static void make_magic(struct loc_writer* writer, struct loc_database_magic* magic,
22c7b98b MT	263	enum loc_database_version version) {
c182393f MT	264	// Copy magic bytes
	265	for (unsigned int i = 0; i < strlen(LOC_DATABASE_MAGIC); i++)
	266	magic->magic[i] = LOC_DATABASE_MAGIC[i];
	267
	268	// Set version
22c7b98b	269	magic->version = version;
c182393f MT	270	}
	271
	272	static void align_page_boundary(off_t* offset, FILE* f) {
	273	// Move to next page boundary
	274	while (*offset % LOC_DATABASE_PAGE_SIZE > 0)
	275	*offset += fwrite("", 1, 1, f);
	276	}
	277
	278	static int loc_database_write_pool(struct loc_writer* writer,
b904896a	279	struct loc_database_header_v1* header, off_t* offset, FILE* f) {
c182393f	280	// Save the offset of the pool section
5c57de03	281	DEBUG(writer->ctx, "Pool starts at %jd bytes\n", (intmax_t)*offset);
0676cd80	282	header->pool_offset = htobe32(*offset);
c182393f MT	283
	284	// Write the pool
	285	size_t pool_length = loc_stringpool_write(writer->pool, f);
	286	*offset += pool_length;
	287
	288	DEBUG(writer->ctx, "Pool has a length of %zu bytes\n", pool_length);
0676cd80	289	header->pool_length = htobe32(pool_length);
c182393f MT	290
	291	return 0;
	292	}
	293
	294	static int loc_database_write_as_section(struct loc_writer* writer,
b904896a	295	struct loc_database_header_v1* header, off_t* offset, FILE* f) {
5c57de03	296	DEBUG(writer->ctx, "AS section starts at %jd bytes\n", (intmax_t)*offset);
0676cd80	297	header->as_offset = htobe32(*offset);
c182393f MT	298
	299	size_t as_length = 0;
	300
b904896a	301	struct loc_database_as_v1 as;
c182393f MT	302	for (unsigned int i = 0; i < writer->as_count; i++) {
c182393f MT	303	// Convert AS into database format
b904896a	304	loc_as_to_database_v1(writer->as[i], writer->pool, &as);
c182393f MT	305
c182393f MT	306	// Write to disk
65862405	307	*offset += fwrite(&as, 1, sizeof(as), f);
c182393f MT	308	as_length += sizeof(as);
	309	}
	310
	311	DEBUG(writer->ctx, "AS section has a length of %zu bytes\n", as_length);
0676cd80	312	header->as_length = htobe32(as_length);
c182393f	313
fda89f24 MT	314	align_page_boundary(offset, f);
fda89f24 MT	315
c182393f MT	316	return 0;
	317	}
	318
438db08c MT	319	struct node {
438db08c MT	320	TAILQ_ENTRY(node) nodes;
f3e02bc5	321
438db08c	322	struct loc_network_tree_node* node;
f3e02bc5	323
438db08c MT	324	// Indices of the child nodes
	325	uint32_t index_zero;
	326	uint32_t index_one;
	327	};
f3e02bc5	328
438db08c MT	329	static struct node* make_node(struct loc_network_tree_node* node) {
	330	struct node* n = malloc(sizeof(*n));
	331	if (!n)
	332	return NULL;
f3e02bc5	333
438db08c MT	334	n->node = loc_network_tree_node_ref(node);
	335	n->index_zero = n->index_one = 0;
	336
	337	return n;
	338	}
	339
	340	static void free_node(struct node* node) {
	341	loc_network_tree_node_unref(node->node);
	342
	343	free(node);
	344	}
	345
	346	struct network {
	347	TAILQ_ENTRY(network) networks;
	348
	349	struct loc_network* network;
	350	};
	351
	352	static struct network* make_network(struct loc_network* network) {
	353	struct network* n = malloc(sizeof(*n));
	354	if (!n)
	355	return NULL;
	356
	357	n->network = loc_network_ref(network);
	358
	359	return n;
f3e02bc5 MT	360	}
f3e02bc5 MT	361
438db08c MT	362	static void free_network(struct network* network) {
	363	loc_network_unref(network->network);
	364
	365	free(network);
	366	}
	367
	368	static int loc_database_write_networks(struct loc_writer* writer,
b904896a	369	struct loc_database_header_v1* header, off_t* offset, FILE* f) {
438db08c	370	// Write the network tree
5c57de03	371	DEBUG(writer->ctx, "Network tree starts at %jd bytes\n", (intmax_t)*offset);
438db08c	372	header->network_tree_offset = htobe32(*offset);
f3e02bc5	373
438db08c MT	374	size_t network_tree_length = 0;
438db08c MT	375	size_t network_data_length = 0;
f3e02bc5	376
438db08c MT	377	struct node* node;
	378	struct node* child_node;
	379
	380	uint32_t index = 0;
	381	uint32_t network_index = 0;
	382
b904896a MT	383	struct loc_database_network_v1 db_network;
b904896a MT	384	struct loc_database_network_node_v1 db_node;
438db08c MT	385
	386	// Initialize queue for nodes
	387	TAILQ_HEAD(node_t, node) nodes;
	388	TAILQ_INIT(&nodes);
	389
	390	// Initialize queue for networks
	391	TAILQ_HEAD(network_t, network) networks;
	392	TAILQ_INIT(&networks);
	393
	394	// Add root
	395	struct loc_network_tree_node* root = loc_network_tree_get_root(writer->networks);
	396	node = make_node(root);
	397
	398	TAILQ_INSERT_TAIL(&nodes, node, nodes);
	399
	400	while (!TAILQ_EMPTY(&nodes)) {
	401	// Pop first node in list
	402	node = TAILQ_FIRST(&nodes);
	403	TAILQ_REMOVE(&nodes, node, nodes);
	404
	405	DEBUG(writer->ctx, "Processing node %p\n", node);
	406
	407	// Get child nodes
	408	struct loc_network_tree_node* node_zero = loc_network_tree_node_get(node->node, 0);
	409	if (node_zero) {
	410	node->index_zero = ++index;
	411
	412	child_node = make_node(node_zero);
	413	loc_network_tree_node_unref(node_zero);
	414
	415	TAILQ_INSERT_TAIL(&nodes, child_node, nodes);
	416	}
	417
	418	struct loc_network_tree_node* node_one = loc_network_tree_node_get(node->node, 1);
	419	if (node_one) {
	420	node->index_one = ++index;
	421
	422	child_node = make_node(node_one);
	423	loc_network_tree_node_unref(node_one);
	424
	425	TAILQ_INSERT_TAIL(&nodes, child_node, nodes);
	426	}
	427
	428	// Prepare what we are writing to disk
39a55353 MT	429	db_node.zero = htobe32(node->index_zero);
	430	db_node.one = htobe32(node->index_one);
	431
438db08c MT	432	if (loc_network_tree_node_is_leaf(node->node)) {
	433	struct loc_network* network = loc_network_tree_node_get_network(node->node);
	434
	435	// Append network to be written out later
	436	struct network* nw = make_network(network);
	437	TAILQ_INSERT_TAIL(&networks, nw, networks);
	438
39a55353	439	db_node.network = htobe32(network_index++);
438db08c MT	440	loc_network_unref(network);
438db08c MT	441	} else {
39a55353	442	db_node.network = htobe32(0xffffffff);
438db08c MT	443	}
	444
	445	// Write the current node
	446	DEBUG(writer->ctx, "Writing node %p (0 = %d, 1 = %d)\n",
	447	node, node->index_zero, node->index_one);
	448
	449	*offset += fwrite(&db_node, 1, sizeof(db_node), f);
	450	network_tree_length += sizeof(db_node);
	451
	452	free_node(node);
	453	}
	454
	455	loc_network_tree_node_unref(root);
	456
	457	header->network_tree_length = htobe32(network_tree_length);
	458
	459	align_page_boundary(offset, f);
	460
5c57de03	461	DEBUG(writer->ctx, "Networks data section starts at %jd bytes\n", (intmax_t)*offset);
438db08c MT	462	header->network_data_offset = htobe32(*offset);
	463
	464	// We have now written the entire tree and have all networks
	465	// in a queue in order as they are indexed
	466	while (!TAILQ_EMPTY(&networks)) {
	467	struct network* nw = TAILQ_FIRST(&networks);
	468	TAILQ_REMOVE(&networks, nw, networks);
	469
	470	// Prepare what we are writing to disk
b904896a	471	int r = loc_network_to_database_v1(nw->network, &db_network);
438db08c MT	472	if (r)
	473	return r;
	474
	475	*offset += fwrite(&db_network, 1, sizeof(db_network), f);
	476	network_data_length += sizeof(db_network);
	477
	478	free_network(nw);
	479	}
f3e02bc5	480
438db08c	481	header->network_data_length = htobe32(network_data_length);
f3e02bc5	482
fda89f24 MT	483	align_page_boundary(offset, f);
fda89f24 MT	484
f3e02bc5 MT	485	return 0;
	486	}
	487
ec684c1a	488	static int loc_database_write_countries(struct loc_writer* writer,
b904896a	489	struct loc_database_header_v1* header, off_t* offset, FILE* f) {
2e2325a9	490	DEBUG(writer->ctx, "Countries section starts at %jd bytes\n", (intmax_t)*offset);
ec684c1a MT	491	header->countries_offset = htobe32(*offset);
	492
	493	size_t countries_length = 0;
	494
b904896a	495	struct loc_database_country_v1 country;
ec684c1a MT	496	for (unsigned int i = 0; i < writer->countries_count; i++) {
ec684c1a MT	497	// Convert country into database format
b904896a	498	loc_country_to_database_v1(writer->countries[i], writer->pool, &country);
ec684c1a MT	499
	500	// Write to disk
	501	*offset += fwrite(&country, 1, sizeof(country), f);
	502	countries_length += sizeof(country);
	503	}
	504
	505	DEBUG(writer->ctx, "Countries section has a length of %zu bytes\n", countries_length);
	506	header->countries_length = htobe32(countries_length);
	507
	508	align_page_boundary(offset, f);
	509
	510	return 0;
	511	}
	512
726f9984	513	static int loc_writer_create_signature(struct loc_writer* writer,
b904896a	514	struct loc_database_header_v1* header, FILE* f) {
726f9984 MT	515	DEBUG(writer->ctx, "Signing database...\n");
	516
	517	// Read file from the beginning
	518	rewind(f);
	519
	520	// Create a new context for signing
	521	EVP_MD_CTX* mdctx = EVP_MD_CTX_new();
	522
	523	// Initialise the context
	524	int r = EVP_DigestSignInit(mdctx, NULL, NULL, NULL, writer->private_key);
	525	if (r != 1) {
	526	ERROR(writer->ctx, "%s\n", ERR_error_string(ERR_get_error(), NULL));
	527	goto END;
	528	}
	529
	530	// Read magic
	531	struct loc_database_magic magic;
	532	fread(&magic, 1, sizeof(magic), f);
	533
a0cff45d MT	534	hexdump(writer->ctx, &magic, sizeof(magic));
a0cff45d MT	535
726f9984 MT	536	// Feed magic into the signature
	537	r = EVP_DigestSignUpdate(mdctx, &magic, sizeof(magic));
	538	if (r != 1) {
	539	ERROR(writer->ctx, "%s\n", ERR_error_string(ERR_get_error(), NULL));
	540	goto END;
	541	}
	542
a0cff45d MT	543	hexdump(writer->ctx, header, sizeof(*header));
a0cff45d MT	544
726f9984	545	// Feed the header into the signature
72633a85	546	r = EVP_DigestSignUpdate(mdctx, header, sizeof(*header));
726f9984 MT	547	if (r != 1) {
	548	ERROR(writer->ctx, "%s\n", ERR_error_string(ERR_get_error(), NULL));
	549	goto END;
	550	}
	551
	552	// Skip header
72633a85	553	fseek(f, sizeof(*header), SEEK_CUR);
726f9984 MT	554
	555	// Walk through the file in chunks of 64kB
	556	char buffer[64 * 1024];
	557	while (!feof(f)) {
	558	size_t bytes_read = fread(buffer, 1, sizeof(buffer), f);
	559
	560	if (ferror(f)) {
	561	ERROR(writer->ctx, "Error reading from file: %s\n", strerror(errno));
	562	r = errno;
	563	goto END;
	564	}
	565
a0cff45d MT	566	hexdump(writer->ctx, buffer, bytes_read);
a0cff45d MT	567
726f9984 MT	568	r = EVP_DigestSignUpdate(mdctx, buffer, bytes_read);
	569	if (r != 1) {
	570	ERROR(writer->ctx, "%s\n", ERR_error_string(ERR_get_error(), NULL));
	571	goto END;
	572	}
	573	}
	574
	575	// Compute the signature
	576	size_t signature_length = sizeof(header->signature);
	577
	578	r = EVP_DigestSignFinal(mdctx,
	579	(unsigned char*)header->signature, &signature_length);
	580	if (r != 1) {
	581	ERROR(writer->ctx, "%s\n", ERR_error_string(ERR_get_error(), NULL));
	582	goto END;
	583	}
	584
	585	// Save length of the signature
	586	header->signature_length = htobe32(signature_length);
	587
257626f5 MT	588	DEBUG(writer->ctx, "Successfully generated signature of %lu bytes\n",
257626f5 MT	589	signature_length);
726f9984 MT	590	r = 0;
726f9984 MT	591
257626f5 MT	592	// Dump signature
	593	hexdump(writer->ctx, header->signature, signature_length);
	594
726f9984 MT	595	END:
	596	EVP_MD_CTX_free(mdctx);
	597
	598	return r;
	599	}
	600
22c7b98b MT	601	LOC_EXPORT int loc_writer_write(struct loc_writer* writer, FILE* f, enum loc_database_version version) {
	602	// Check version
	603	switch (version) {
	604	case LOC_DATABASE_VERSION_UNSET:
	605	version = LOC_DATABASE_VERSION_LATEST;
	606	break;
	607
	608	case LOC_DATABASE_VERSION_1:
	609	break;
	610
	611	default:
	612	ERROR(writer->ctx, "Invalid database version: %d\n", version);
	613	return -1;
	614	}
	615
	616	DEBUG(writer->ctx, "Writing database in version %d\n", version);
	617
c182393f	618	struct loc_database_magic magic;
22c7b98b	619	make_magic(writer, &magic, version);
c182393f MT	620
c182393f MT	621	// Make the header
b904896a	622	struct loc_database_header_v1 header;
0676cd80 MT	623	header.vendor = htobe32(writer->vendor);
0676cd80 MT	624	header.description = htobe32(writer->description);
4bf49d00	625	header.license = htobe32(writer->license);
c182393f MT	626
	627	time_t now = time(NULL);
	628	header.created_at = htobe64(now);
	629
b1720435	630	// Clear the signature
726f9984	631	header.signature_length = 0;
b1720435 MT	632	for (unsigned int i = 0; i < sizeof(header.signature); i++)
	633	header.signature[i] = '\0';
	634
5e164830 MT	635	// Clear the padding
	636	for (unsigned int i = 0; i < sizeof(header.padding); i++)
	637	header.padding[i] = '\0';
	638
c182393f MT	639	int r;
	640	off_t offset = 0;
	641
	642	// Start writing at the beginning of the file
	643	r = fseek(f, 0, SEEK_SET);
	644	if (r)
	645	return r;
	646
	647	// Write the magic
	648	offset += fwrite(&magic, 1, sizeof(magic), f);
	649
	650	// Skip the space we need to write the header later
	651	r = fseek(f, sizeof(header), SEEK_CUR);
	652	if (r) {
	653	DEBUG(writer->ctx, "Could not seek to position after header\n");
	654	return r;
	655	}
	656	offset += sizeof(header);
	657
	658	align_page_boundary(&offset, f);
	659
	660	// Write all ASes
	661	r = loc_database_write_as_section(writer, &header, &offset, f);
	662	if (r)
	663	return r;
	664
f3e02bc5	665	// Write all networks
438db08c	666	r = loc_database_write_networks(writer, &header, &offset, f);
f3e02bc5 MT	667	if (r)
	668	return r;
	669
47bea941 MT	670	// Write countries
47bea941 MT	671	r = loc_database_write_countries(writer, &header, &offset, f);
c182393f MT	672	if (r)
	673	return r;
	674
47bea941 MT	675	// Write pool
47bea941 MT	676	r = loc_database_write_pool(writer, &header, &offset, f);
ec684c1a MT	677	if (r)
	678	return r;
	679
726f9984 MT	680	// Create the signature
	681	if (writer->private_key) {
	682	r = loc_writer_create_signature(writer, &header, f);
	683	if (r)
	684	return r;
	685	}
	686
c182393f MT	687	// Write the header
	688	r = fseek(f, sizeof(magic), SEEK_SET);
	689	if (r)
	690	return r;
	691
726f9984	692	fwrite(&header, 1, sizeof(header), f);
c182393f	693
726f9984	694	return r;
c182393f	695	}