1 files changed, 14 insertions, 25 deletions

diff --git a/ssh/client.c b/ssh/client.c
index 9e60033..0adfffb 100644
--- a/ssh/client.c
+++ b/ssh/client.c
@@ -8,28 +8,8 @@
 #include <libssh/callbacks.h>
 #include <sys/select.h>
 #include <poll.h>
-#include "../global.h"
-
-
-static void parse_args(const char *arg, const char **server_host, int *port) {
-	const char *ptr = strchr(arg, ':');
-	if (ptr == NULL) {
-		*server_host = arg;
-	} else {
-		size_t length = ptr - arg;
-		char *host = malloc(length + 1, char);
-		memcpy(host, arg, length);
-		host[length] = '\0';
-		*server_host = host;
-
-		char *endp;
-		*port = strtol(ptr + 1, &endp, 10);
-		if (endp == ptr || *endp != '\0') {
-			fprintf(stderr, "Cannot parse server:port from argument '%s'\n", arg);
-			exit(1);
-		}
-	}
-}
+#include "util.h"
+
 
 static bool prompt_yn(const char *text) {
 	printf("%s", text);
@@ -159,7 +139,10 @@ int main(int argc, char **argv) {
 		return 1;
 	}
 
-	parse_args(argv[1], &server_host, &port);
+	if (!parse_host_port(argv[1], &server_host, &port)) {
+		fprintf(stderr, "Cannot parse host:port from argument '%s'\n", argv[1]);
+		return 1;
+	}
 
 	ssh_session session = ssh_new();
 	if (!session) {
@@ -169,7 +152,8 @@ int main(int argc, char **argv) {
 
 	const char *ciphers_str = "aes256-gcm@openssh.com,aes256-ctr,aes256-cbc";
 	bool procconfig = false;
-	bool ok = ssh_options_set(session, SSH_OPTIONS_PROCESS_CONFIG, &procconfig) == SSH_OK;
+	bool ok = true;
+	ok &= ssh_options_set(session, SSH_OPTIONS_PROCESS_CONFIG, &procconfig) == SSH_OK;
 	ok &= ssh_options_set(session, SSH_OPTIONS_USER, "tomsg") == SSH_OK;
 	ok &= ssh_options_set(session, SSH_OPTIONS_HOST, server_host) == SSH_OK;
 	ok &= ssh_options_set(session, SSH_OPTIONS_PORT, &port) == SSH_OK;
@@ -257,7 +241,12 @@ int main(int argc, char **argv) {
 
 	printf("Obtained tomsg subsystem on channel\n");
 
-	struct session_data *sesdata = malloc(1, struct session_data);
+	struct session_data *sesdata = malloc(sizeof(struct session_data));
+	if (!sesdata) {
+		fprintf(stderr, "Out of memory (allocating session_data)!\n");
+		return 1;
+	}
+
 	sesdata->session = session;
 	sesdata->channel = channel;
 	sesdata->should_close = false;