From 8fcb5bc7f0b6f426a091256ad7758181d71734e3 Mon Sep 17 00:00:00 2001
From: Michael Tremer <michael.tremer@ipfire.org>
Date: Tue, 23 Aug 2022 09:35:54 +0000
Subject: [PATCH] database: country: Return better error codes

Signed-off-by: Michael Tremer <michael.tremer@ipfire.org>
---
 src/database.c        |  8 +++++++-
 src/python/database.c | 19 +++++++++++++++++--
 2 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/src/database.c b/src/database.c
index bc53f41..374da57 100644
--- a/src/database.c
+++ b/src/database.c
@@ -954,6 +954,12 @@ LOC_EXPORT int loc_database_get_country(struct loc_database* db,
 	off_t lo = 0;
 	off_t hi = db->country_objects.count - 1;
 
+	// Check if the country code is valid
+	if (!loc_country_code_is_valid(code)) {
+		errno = EINVAL;
+		return 1;
+	}
+
 #ifdef ENABLE_DEBUG
 	// Save start time
 	clock_t start = clock();
@@ -996,7 +1002,7 @@ LOC_EXPORT int loc_database_get_country(struct loc_database* db,
 	// Nothing found
 	*country = NULL;
 
-	return 1;
+	return 0;
 }
 
 // Enumerator
diff --git a/src/python/database.c b/src/python/database.c
index 250750d..d6ee4d0 100644
--- a/src/python/database.c
+++ b/src/python/database.c
@@ -166,17 +166,32 @@ static PyObject* Database_get_as(DatabaseObject* self, PyObject* args) {
 }
 
 static PyObject* Database_get_country(DatabaseObject* self, PyObject* args) {
+	struct loc_country* country = NULL;
 	const char* country_code = NULL;
 
 	if (!PyArg_ParseTuple(args, "s", &country_code))
 		return NULL;
 
-	struct loc_country* country;
+	// Fetch the country
 	int r = loc_database_get_country(self->db, &country, country_code);
 	if (r) {
-		Py_RETURN_NONE;
+		switch (errno) {
+			case EINVAL:
+				PyErr_SetString(PyExc_ValueError, "Invalid country code");
+				break;
+
+			default:
+				PyErr_SetFromErrno(PyExc_OSError);
+				break;
+		}
+
+		return NULL;
 	}
 
+	// No result
+	if (!country)
+		Py_RETURN_NONE;
+
 	PyObject* obj = new_country(&CountryType, country);
 	loc_country_unref(country);
 
-- 
2.39.5