1 files changed, 21 insertions, 11 deletions

diff --git a/ecp/util/util.c b/ecp/util/util.c
index 0f7f066..e11cab9 100644
--- a/ecp/util/util.c
+++ b/ecp/util/util.c
@@ -10,16 +10,18 @@
 
 #include "util.h"
 
-static int _read_key(int fd, uint8_t *key) {
+static int _read_key(int fd, uint8_t *key, int *line_cnt) {
     char buffer[ECP_SIZE_ECDH_KEY_BUF];
     char ch;
     int cmt;
     ssize_t rv;
     int _rv;
 
+    if (line_cnt) *line_cnt = 0;
     do {
         cmt = 0;
         rv = read(fd, &ch, 1);
+        if (rv == 0) return ECP_ERR_EOF;
         if (rv != 1) return ECP_ERR;
         switch (ch) {
             case '#':
@@ -27,8 +29,10 @@ static int _read_key(int fd, uint8_t *key) {
                 cmt = 1;
                 while (ch != '\n') {
                     rv = read(fd, &ch, 1);
+                    if (rv == 0) return ECP_ERR_EOF;
                     if (rv != 1) return ECP_ERR;
                 }
+                if (line_cnt) *line_cnt += 1;
                 break;
             }
 
@@ -38,6 +42,7 @@ static int _read_key(int fd, uint8_t *key) {
             }
         }
     } while(cmt);
+    if (line_cnt) *line_cnt += 1;
 
     rv = read(fd, buffer + 1, sizeof(buffer) - 2);
     if (rv < (sizeof(buffer) - 2)) return ECP_ERR;
@@ -51,29 +56,31 @@ static int _read_key(int fd, uint8_t *key) {
     return _rv;
 }
 
-int ecp_util_read_key(int fd, ecp_ecdh_public_t *public, ecp_ecdh_private_t *private) {
+int ecp_util_read_key(int fd, ecp_ecdh_public_t *public, ecp_ecdh_private_t *private, int *line_cnt) {
     int rv;
 
-    rv = _read_key(fd, (uint8_t *)public);
+    rv = _read_key(fd, (uint8_t *)public, line_cnt);
     if (rv) return rv;
 
     if (private) {
-        rv = _read_key(fd, (uint8_t *)private);
+        rv = _read_key(fd, (uint8_t *)private, line_cnt);
         if (rv) return rv;
     }
 
     return ECP_OK;
 }
 
-int ecp_util_write_key(int fd, ecp_ecdh_public_t *public, ecp_ecdh_private_t *private) {
+int ecp_util_write_key(int fd, ecp_ecdh_public_t *public, ecp_ecdh_private_t *private, int *line_cnt) {
     char buffer[ECP_SIZE_ECDH_KEY_BUF];
     ssize_t rv;
 
+    if (line_cnt) *line_cnt = 0;
     ecp_key2str(buffer, (uint8_t *)public);
     buffer[ECP_SIZE_ECDH_KEY_BUF - 1] = '\n';
 
     rv = write(fd, buffer, sizeof(buffer));
     if (rv != sizeof(buffer)) return ECP_ERR;
+    if (line_cnt) *line_cnt += 1;
 
     if (private) {
         ecp_key2str(buffer, (uint8_t *)private);
@@ -81,6 +88,7 @@ int ecp_util_write_key(int fd, ecp_ecdh_public_t *public, ecp_ecdh_private_t *pr
 
         rv = write(fd, buffer, sizeof(buffer));
         if (rv != sizeof(buffer)) return ECP_ERR;
+        if (line_cnt) *line_cnt += 1;
     }
 
     return ECP_OK;
@@ -89,21 +97,23 @@ int ecp_util_write_key(int fd, ecp_ecdh_public_t *public, ecp_ecdh_private_t *pr
 int ecp_util_load_key(char *filename, ecp_ecdh_public_t *public, ecp_ecdh_private_t *private) {
     int rv, fd;
 
-    if ((fd = open(filename, O_RDONLY)) < 0) return ECP_ERR;
+    if ((fd = open(filename, O_RDONLY)) < 0) return ECP_ERR_OPEN;
 
-    rv = ecp_util_read_key(fd, public, private);
+    rv = ecp_util_read_key(fd, public, private, NULL);
     close(fd);
 
-    return rv;
+    if (rv < 0) return rv;
+    return ECP_OK;
 }
 
 int ecp_util_save_key(char *filename, ecp_ecdh_public_t *public, ecp_ecdh_private_t *private) {
     int rv, fd;
 
-    if ((fd = open(filename, O_WRONLY | O_CREAT, S_IRUSR | S_IWUSR)) < 0) return ECP_ERR;
+    if ((fd = open(filename, O_WRONLY | O_CREAT, S_IRUSR | S_IWUSR)) < 0) return ECP_ERR_OPEN;
 
-    rv = ecp_util_write_key(fd, public, private);
+    rv = ecp_util_write_key(fd, public, private, NULL);
     close(fd);
 
-    return rv;
+    if (rv < 0) return rv;
+    return ECP_OK;
 }