knownhosts writing and unit test

3 files changed, 138 insertions, 10 deletions

diff --git a/libssh/keyfiles.c b/libssh/keyfiles.c
index 63b6321a..74ec7ad9 100644
--- a/libssh/keyfiles.c
+++ b/libssh/keyfiles.c
@@ -1695,7 +1695,7 @@ int ssh_is_server_known(ssh_session session) {
  *
  * @param[in]  session  The ssh session to use.
  *
- * @return              0 on success, -1 on error.
+ * @return              SSH_OK on success, SSH_ERROR on error.
  */
 int ssh_write_knownhost(ssh_session session) {
   ssh_string pubkey = session->current_crypto->server_pubkey;
@@ -1703,8 +1703,25 @@ int ssh_write_knownhost(ssh_session session) {
   char buffer[4096] = {0};
   FILE *file;
   char *dir;
+  char *host;
+  char *hostport;
   size_t len = 0;
 
+  if (session->host == NULL) {
+    ssh_set_error(session, SSH_FATAL,
+        "Can't write host in known hosts if the hostname isn't known");
+    return SSH_ERROR;
+  }
+
+  host = ssh_lowercase(session->host);
+  /* If using a nonstandard port, save the host in the [host]:port format */
+  if(session->port != 22){
+    hostport = ssh_hostport(host,session->port);
+    SAFE_FREE(host);
+    host=hostport;
+    hostport=NULL;
+  }
+
   if (session->knownhosts == NULL) {
     if (ssh_options_set(session, SSH_OPTIONS_KNOWNHOSTS, NULL) < 0) {
       ssh_set_error(session, SSH_FATAL, "Can't find a known_hosts file");
@@ -1712,12 +1729,6 @@ int ssh_write_knownhost(ssh_session session) {
     }
   }
 
-  if (session->host == NULL) {
-    ssh_set_error(session, SSH_FATAL,
-        "Cannot write host in known hosts if the hostname is unknown");
-    return -1;
-  }
-
   /* Check if ~/.ssh exists and create it if not */
   dir = ssh_dirname(session->knownhosts);
   if (dir == NULL) {
@@ -1739,6 +1750,7 @@ int ssh_write_knownhost(ssh_session session) {
     ssh_set_error(session, SSH_FATAL,
         "Couldn't open known_hosts file %s for appending: %s",
         session->knownhosts, strerror(errno));
+    SAFE_FREE(host);
     return -1;
   }
 
@@ -1758,6 +1770,7 @@ int ssh_write_knownhost(ssh_session session) {
     key = publickey_from_string(session, pubkey);
     if (key == NULL) {
       fclose(file);
+      SAFE_FREE(host);
       return -1;
     }
 
@@ -1766,6 +1779,7 @@ int ssh_write_knownhost(ssh_session session) {
     if (sexp == NULL) {
       publickey_free(key);
       fclose(file);
+      SAFE_FREE(host);
       return -1;
     }
     e = gcry_sexp_nth_mpi(sexp, 1, GCRYMPI_FMT_USG);
@@ -1773,6 +1787,7 @@ int ssh_write_knownhost(ssh_session session) {
     if (e == NULL) {
       publickey_free(key);
       fclose(file);
+      SAFE_FREE(host);
       return -1;
     }
 
@@ -1781,6 +1796,7 @@ int ssh_write_knownhost(ssh_session session) {
       publickey_free(key);
       bignum_free(e);
       fclose(file);
+      SAFE_FREE(host);
       return -1;
     }
     n = gcry_sexp_nth_mpi(sexp, 1, GCRYMPI_FMT_USG);
@@ -1789,6 +1805,7 @@ int ssh_write_knownhost(ssh_session session) {
       publickey_free(key);
       bignum_free(e);
       fclose(file);
+      SAFE_FREE(host);
       return -1;
     }
 
@@ -1813,12 +1830,13 @@ int ssh_write_knownhost(ssh_session session) {
 #endif
       publickey_free(key);
       fclose(file);
+      SAFE_FREE(host);
       return -1;
     }
 
     snprintf(buffer, sizeof(buffer),
         "%s %d %s %s\n",
-        session->host,
+        host,
         rsa_size << 3,
         e_string,
         n_string);
@@ -1838,18 +1856,19 @@ int ssh_write_knownhost(ssh_session session) {
     pubkey_64 = bin_to_base64(pubkey->string, string_len(pubkey));
     if (pubkey_64 == NULL) {
       fclose(file);
+      SAFE_FREE(host);
       return -1;
     }
 
     snprintf(buffer, sizeof(buffer),
         "%s %s %s\n",
-        session->host,
+        host,
         session->current_crypto->server_pubkey_type,
         pubkey_64);
 
     SAFE_FREE(pubkey_64);
   }
-
+  SAFE_FREE(host);
   len = strlen(buffer);
   if (fwrite(buffer, len, 1, file) != 1 || ferror(file)) {
     fclose(file);
diff --git a/tests/unittests/CMakeLists.txt b/tests/unittests/CMakeLists.txt
index 2d0a8607..855c027e 100644
--- a/tests/unittests/CMakeLists.txt
+++ b/tests/unittests/CMakeLists.txt
@@ -4,3 +4,4 @@ add_check_test(torture_list torture_list.c ${TORTURE_LIBRARY})
 add_check_test(torture_misc torture_misc.c ${TORTURE_LIBRARY})
 add_check_test(torture_keyfiles torture_keyfiles.c ${TORTURE_LIBRARY})
 add_check_test(torture_options torture_options.c ${TORTURE_LIBRARY})
+add_check_test(torture_knownhosts torture_knownhosts.c ${TORTURE_LIBRARY})
diff --git a/tests/unittests/torture_knownhosts.c b/tests/unittests/torture_knownhosts.c
new file mode 100644
index 00000000..29b8f897
--- /dev/null
+++ b/tests/unittests/torture_knownhosts.c
@@ -0,0 +1,108 @@
+/*
+ * This file is part of the SSH Library
+ *
+ * Copyright (c) 2010 by Aris Adamantiadis
+ *
+ * The SSH 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.
+ *
+ * The SSH 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.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with the SSH Library; see the file COPYING.  If not, write to
+ * the Free Software Foundation, Inc., 59 Temple Place - Suite 330, Boston,
+ * MA 02111-1307, USA.
+ */
+
+#define LIBSSH_STATIC
+
+#include "torture.h"
+#include "libssh/libssh.h"
+#include "libssh/priv.h"
+#include "libssh/session.h"
+ssh_session session;
+
+#define KNOWNHOSTFILES "libssh_torture_knownhosts"
+
+static void setup(void) {
+    session = ssh_new();
+}
+
+static void teardown(void) {
+    ssh_free(session);
+    unlink(KNOWNHOSTFILES);
+}
+
+START_TEST (torture_knownhosts_port)
+{
+  int rc;
+  char buffer[200];
+  FILE *file;
+  /* Connect to localhost:22, force the port to 1234 and then write
+   * the known hosts file. Then check that the entry written is
+   * [localhost]:1234
+   */
+  ssh_options_set(session,SSH_OPTIONS_HOST,"localhost");
+  ssh_options_set(session,SSH_OPTIONS_KNOWNHOSTS,KNOWNHOSTFILES);
+  rc=ssh_connect(session);
+  ck_assert_int_eq(rc,SSH_OK);
+  session->port=1234;
+  rc=ssh_write_knownhost(session);
+  ck_assert_int_eq(rc,SSH_OK);
+  ssh_disconnect(session);
+  ssh_free(session);
+  file=fopen(KNOWNHOSTFILES,"r");
+  ck_assert(file != NULL);
+  fgets(buffer,sizeof(buffer),file);
+  buffer[sizeof(buffer)-1]='\0';
+  ck_assert(strstr(buffer,"[localhost]:1234 ") != NULL);
+  fclose(file);
+
+  /* now, connect back to the ssh server and verify the known host line */
+  session=ssh_new();
+  ssh_options_set(session,SSH_OPTIONS_HOST,"localhost");
+  ssh_options_set(session,SSH_OPTIONS_KNOWNHOSTS,KNOWNHOSTFILES);
+  rc=ssh_connect(session);
+  ck_assert_int_eq(rc,SSH_OK);
+  session->port=1234;
+  rc=ssh_is_server_known(session);
+  ck_assert_int_eq(rc,SSH_SERVER_KNOWN_OK);
+  ssh_disconnect(session);
+}
+END_TEST
+
+static Suite *torture_make_suite(void) {
+  Suite *s = suite_create("libssh_knownhosts");
+
+  torture_create_case_fixture(s, "torture_knownhosts_port",
+          torture_knownhosts_port, setup, teardown);
+  return s;
+}
+
+int main(int argc, char **argv) {
+  Suite *s = NULL;
+  SRunner *sr = NULL;
+  struct argument_s arguments;
+  int nf;
+
+  ZERO_STRUCT(arguments);
+
+  torture_cmdline_parse(argc, argv, &arguments);
+
+  s = torture_make_suite();
+
+  sr = srunner_create(s);
+  if (arguments.nofork) {
+    srunner_set_fork_status(sr, CK_NOFORK);
+  }
+  srunner_run_all(sr, CK_VERBOSE);
+  nf = srunner_ntests_failed(sr);
+  srunner_free(sr);
+
+  return (nf == 0) ? EXIT_SUCCESS : EXIT_FAILURE;
+}