[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;
	enum loc_database_version version;

	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,
		enum loc_database_version version, FILE* fkey) {
	struct loc_writer* w = calloc(1, sizeof(*w));
	if (!w)
		return -ENOMEM;

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

	// Check version
	switch (version) {
		case LOC_DATABASE_VERSION_1:
			w->version = version;
			break;

		case LOC_DATABASE_VERSION_UNSET:
			w->version = LOC_DATABASE_VERSION_LATEST;
			break;

		default:
			ERROR(ctx, "Invalid database version: %d\n", version);
			loc_writer_unref(w);
			return -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) {
	// 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 = htobe16(writer->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) {
	struct loc_database_magic magic;
	make_magic(writer, &magic);

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