From 71424899eed0bb94c2adb029b4b49f5d1a8f4ce4 Mon Sep 17 00:00:00 2001
From: Ali Gholami Rudi <ali@rudi.ir>
Date: Tue, 2 Feb 2016 08:58:39 +0330
Subject: [PATCH] conn: move tls handshake to conn_tls()

---
 conn.h         |  3 ++-
 conn_mbedtls.c | 25 ++++++++++++++-----------
 conn_openssl.c | 22 ++++++++++++----------
 pop3.c         |  6 +++++-
 4 files changed, 33 insertions(+), 23 deletions(-)

diff --git a/conn.h b/conn.h
index 40d4dbb..3106979 100644
--- a/conn.h
+++ b/conn.h
@@ -1,4 +1,5 @@
-struct conn *conn_connect(char *addr, char *port, char *certfile);
+struct conn *conn_connect(char *addr, char *port);
+int conn_tls(struct conn *conn, char *certfile);
 int conn_write(struct conn *conn, char *buf, int len);
 int conn_read(struct conn *conn, char *buf, int len);
 int conn_close(struct conn *conn);
diff --git a/conn_mbedtls.c b/conn_mbedtls.c
index af148f7..64d892c 100644
--- a/conn_mbedtls.c
+++ b/conn_mbedtls.c
@@ -16,6 +16,7 @@
 
 struct conn {
 	int fd;
+	int tls;
 	ssl_context ssl;
 	ssl_session ssn;
 	ctr_drbg_context ctr_drbg;
@@ -34,15 +35,19 @@ static int ps_recv(void *ctx, unsigned char *buf, size_t len)
 
 int conn_read(struct conn *conn, char *buf, int len)
 {
-	return ssl_read(&conn->ssl, (unsigned char *) buf, sizeof(buf));
+	if (conn->tls)
+		return ssl_read(&conn->ssl, (unsigned char *) buf, len);
+	return read(conn->fd, buf, len);
 }
 
 int conn_write(struct conn *conn, char *buf, int len)
 {
-	return ssl_write(&conn->ssl, (unsigned char *) buf, len);
+	if (conn->tls)
+		return ssl_write(&conn->ssl, (unsigned char *) buf, len);
+	return write(conn->fd, buf, len);
 }
 
-static int conns_init(struct conn *conn, char *certfile)
+int conn_tls(struct conn *conn, char *certfile)
 {
 	entropy_context entropy;
 	entropy_init(&entropy);
@@ -61,6 +66,7 @@ static int conns_init(struct conn *conn, char *certfile)
 	ssl_set_bio(&conn->ssl, ps_recv, &conn->fd, ps_send, &conn->fd);
 	ssl_set_ciphersuites(&conn->ssl, ssl_list_ciphersuites());
 	ssl_set_session(&conn->ssl, &conn->ssn);
+	conn->tls = 1;
 	return ssl_handshake(&conn->ssl);
 }
 
@@ -90,19 +96,16 @@ struct conn *conn_connect(char *addr, char *port, char *certfile)
 	conn = malloc(sizeof(*conn));
 	memset(conn, 0, sizeof(*conn));
 	conn->fd = fd;
-	if (conns_init(conn, certfile)) {
-		free(conn);
-		return NULL;
-	}
 	return conn;
 }
 
 int conn_close(struct conn *conn)
 {
-	ssl_close_notify(&conn->ssl);
-	x509_crt_free(&conn->cert);
-	ssl_free(&conn->ssl);
-
+	if (conn->tls) {
+		ssl_close_notify(&conn->ssl);
+		x509_crt_free(&conn->cert);
+		ssl_free(&conn->ssl);
+	}
 	close(conn->fd);
 	free(conn);
 	return 0;
diff --git a/conn_openssl.c b/conn_openssl.c
index cc5adce..2b1db0b 100644
--- a/conn_openssl.c
+++ b/conn_openssl.c
@@ -25,15 +25,19 @@ struct conn {
 
 int conn_read(struct conn *conn, char *buf, int len)
 {
-	return SSL_read(conn->ssl, buf, sizeof(buf) - 1);
+	if (conn->ssl)
+		return SSL_read(conn->ssl, buf, len);
+	return read(conn->fd, buf, len);
 }
 
 int conn_write(struct conn *conn, char *buf, int len)
 {
-	return SSL_write(conn->ssl, buf, len);
+	if (conn->ssl)
+		return SSL_write(conn->ssl, buf, len);
+	return write(conn->fd, buf, len);
 }
 
-static int conns_init(struct conn *conn, char *certfile)
+int conn_tls(struct conn *conn, char *certfile)
 {
 	SSLeay_add_ssl_algorithms();
 	SSL_load_error_strings();
@@ -81,18 +85,16 @@ struct conn *conn_connect(char *addr, char *port, char *certfile)
 	conn = malloc(sizeof(*conn));
 	memset(conn, 0, sizeof(*conn));
 	conn->fd = fd;
-	if (conns_init(conn, certfile)) {
-		free(conn);
-		return NULL;
-	}
 	return conn;
 }
 
 int conn_close(struct conn *conn)
 {
-	SSL_shutdown(conn->ssl);
-	SSL_free(conn->ssl);
-	SSL_CTX_free(conn->ctx);
+	if (conn->ssl) {
+		SSL_shutdown(conn->ssl);
+		SSL_free(conn->ssl);
+		SSL_CTX_free(conn->ctx);
+	}
 	close(conn->fd);
 	free(conn);
 	return 0;
diff --git a/pop3.c b/pop3.c
index 7ad01dc..c9582f2 100644
--- a/pop3.c
+++ b/pop3.c
@@ -262,9 +262,13 @@ static int fetch(struct account *account)
 	int failed = 0;
 	int i;
 	nmails = 0;
-	conn = conn_connect(account->server, account->port, account->cert);
+	conn = conn_connect(account->server, account->port);
 	if (!conn)
 		return 1;
+	if (conn_tls(conn, account->cert)) {
+		conn_close(conn);
+		return 1;
+	}
 	buf_pos = 0;
 	buf_len = 0;
 	if (account->uidl)
-- 
2.11.4.GIT