2 files changed, 43 insertions, 15 deletions

diff --git a/ecp/util/keygen.c b/ecp/util/keygen.c
index a0a1507..c7cee18 100644
--- a/ecp/util/keygen.c
+++ b/ecp/util/keygen.c
@@ -18,7 +18,6 @@ int main(int argc, char *argv[]) {
     int rv;
 
     if (argc != 3) usage(argv[0]);
-    printf("keygen: public=%p, private=%p\n", key.public, key.private);
 
     rv = ecp_dhkey_gen(&key);
     if (rv) goto err;
diff --git a/ecp/util/util.c b/ecp/util/util.c
index abe562d..0f7f066 100644
--- a/ecp/util/util.c
+++ b/ecp/util/util.c
@@ -10,27 +10,56 @@
 
 #include "util.h"
 
-int ecp_util_read_key(int fd, ecp_ecdh_public_t *public, ecp_ecdh_private_t *private) {
+static int _read_key(int fd, uint8_t *key) {
     char buffer[ECP_SIZE_ECDH_KEY_BUF];
+    char ch;
+    int cmt;
     ssize_t rv;
     int _rv;
 
-    memset(buffer, 0, sizeof(buffer));
-    rv = read(fd, buffer, sizeof(buffer));
-    if (rv < (ECP_SIZE_ECDH_KEY_BUF - 1)) return ECP_ERR;
-    if (buffer[ECP_SIZE_ECDH_KEY_BUF - 1] == '\n') buffer[ECP_SIZE_ECDH_KEY_BUF - 1] = '\0';
+    do {
+        cmt = 0;
+        rv = read(fd, &ch, 1);
+        if (rv != 1) return ECP_ERR;
+        switch (ch) {
+            case '#':
+            case '\n': {
+                cmt = 1;
+                while (ch != '\n') {
+                    rv = read(fd, &ch, 1);
+                    if (rv != 1) return ECP_ERR;
+                }
+                break;
+            }
+
+            default: {
+                buffer[0] = ch;
+                break;
+            }
+        }
+    } while(cmt);
+
+    rv = read(fd, buffer + 1, sizeof(buffer) - 2);
+    if (rv < (sizeof(buffer) - 2)) return ECP_ERR;
+    buffer[ECP_SIZE_ECDH_KEY_BUF - 1] = '\0';
+
+    do {
+        rv = read(fd, &ch, 1);
+    } while ((rv == 1) && (ch != '\n'));
+
+    _rv = ecp_str2key(key, buffer);
+    return _rv;
+}
 
-    _rv = ecp_str2key((uint8_t *)public, buffer);
-    if (_rv) return _rv;
+int ecp_util_read_key(int fd, ecp_ecdh_public_t *public, ecp_ecdh_private_t *private) {
+    int rv;
 
-    if (private) {
-        memset(buffer, 0, sizeof(buffer));
-        rv = read(fd, buffer, sizeof(buffer));
-        if (rv < (ECP_SIZE_ECDH_KEY_BUF - 1)) return ECP_ERR;
-        if (buffer[ECP_SIZE_ECDH_KEY_BUF - 1] == '\n') buffer[ECP_SIZE_ECDH_KEY_BUF - 1] = '\0';
+    rv = _read_key(fd, (uint8_t *)public);
+    if (rv) return rv;
 
-        _rv = ecp_str2key((uint8_t *)private, buffer);
-        if (_rv) return _rv;
+    if (private) {
+        rv = _read_key(fd, (uint8_t *)private);
+        if (rv) return rv;
     }
 
     return ECP_OK;