From 85e745121fd94a92cb1c282c2362839ccd2bd77e Mon Sep 17 00:00:00 2001
From: Zhi Guan <guanzhi1980@gmail.com>
Date: Tue, 1 Nov 2022 17:49:48 +0800
Subject: [PATCH] Add socket wrapper

---
 include/gmssl/socket.h | 70 ++++++++++++++++++++++++++++++++++++++++++
 include/gmssl/tls.h    | 33 ++++----------------
 src/tlcp.c             |  9 ------
 src/tls.c              | 66 ++++++---------------------------------
 src/tls12.c            | 10 ------
 src/tls13.c            |  9 ------
 src/tls_ext.c          | 12 --------
 tools/tlcp_client.c    | 24 ++-------------
 tools/tlcp_server.c    | 35 +++------------------
 tools/tls12_client.c   | 25 ++-------------
 tools/tls12_server.c   | 36 +++-------------------
 tools/tls13_client.c   | 22 ++-----------
 tools/tls13_server.c   | 35 +++------------------
 13 files changed, 105 insertions(+), 281 deletions(-)
 create mode 100644 include/gmssl/socket.h

diff --git a/include/gmssl/socket.h b/include/gmssl/socket.h
new file mode 100644
index 00000000..b0a34569
--- /dev/null
+++ b/include/gmssl/socket.h
@@ -0,0 +1,70 @@
+/*
+ *  Copyright 2014-2022 The GmSSL Project. All Rights Reserved.
+ *
+ *  Licensed under the Apache License, Version 2.0 (the License); you may
+ *  not use this file except in compliance with the License.
+ *
+ *  http://www.apache.org/licenses/LICENSE-2.0
+ */
+
+
+#ifndef GMSSL_SOCKET_H
+#define GMSSL_SOCKET_H
+
+#include <string.h>
+#include <stdint.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+
+#ifdef WIN32
+#pragma comment (lib, "Ws2_32.lib")
+#pragma comment (lib, "Mswsock.lib")
+#pragma comment (lib, "AdvApi32.lib")
+
+#include <winsock2.h>
+
+typedef SOCKET tls_socket_t;
+typedef int tls_ret_t;
+typedef int tls_socklen_t;
+
+
+#define tls_socket_send(sock,buf,len,flags)	send(sock,buf,(int)(len),flags)
+#define tls_socket_recv(sock,buf,len,flags)	recv(sock,buf,(int)(len),flags)
+#define tls_socket_close(sock)			closesocket(sock)
+
+
+#else
+
+#include <fcntl.h>
+#include <netdb.h>
+#include <arpa/inet.h>
+#include <sys/types.h>
+#include <sys/socket.h>
+#include <netinet/in.h>
+#include <unistd.h>
+
+typedef int tls_socket_t;
+typedef ssize_t tls_ret_t;
+typedef socklen_t tls_socklen_t;
+
+
+#define tls_socket_send(sock,buf,len,flags)	send(sock,buf,len,flags)
+#define tls_socket_recv(sock,buf,len,flags)	recv(sock,buf,len,flags)
+#define tls_socket_close(sock)			close(sock)
+
+
+
+#endif
+
+
+
+
+
+
+#ifdef __cplusplus
+}
+#endif
+#endif
diff --git a/include/gmssl/tls.h b/include/gmssl/tls.h
index f9cb8a19..30d93bf7 100644
--- a/include/gmssl/tls.h
+++ b/include/gmssl/tls.h
@@ -12,13 +12,6 @@
 #ifndef GMSSL_TLS_H
 #define GMSSL_TLS_H
 
-#ifdef WIN32
-#pragma comment (lib, "Ws2_32.lib")
-#pragma comment (lib, "Mswsock.lib")
-#pragma comment (lib, "AdvApi32.lib")
-#include <winsock2.h>
-#endif
-
 
 #include <stdint.h>
 #include <gmssl/sm2.h>
@@ -26,6 +19,7 @@
 #include <gmssl/sm4.h>
 #include <gmssl/digest.h>
 #include <gmssl/block_cipher.h>
+#include <gmssl/socket.h>
 
 
 #ifdef __cplusplus
@@ -458,15 +452,9 @@ int tls_record_set_data(uint8_t *record, const uint8_t *data, size_t datalen);
 int tls_record_print(FILE *fp, const uint8_t *record,  size_t recordlen, int format, int indent);
 int tlcp_record_print(FILE *fp, const uint8_t *record,  size_t recordlen, int format, int indent);
 
-#ifdef WIN32
-int tls_record_send(const uint8_t* record, size_t recordlen, SOCKET sock);
-int tls_record_recv(uint8_t* record, size_t* recordlen, SOCKET sock);
-int tls12_record_recv(uint8_t* record, size_t* recordlen, SOCKET sock);
-#else
-int tls_record_send(const uint8_t *record, size_t recordlen, int sock);
-int tls_record_recv(uint8_t *record, size_t *recordlen, int sock);
-int tls12_record_recv(uint8_t *record, size_t *recordlen, int sock);
-#endif
+int tls_record_send(const uint8_t *record, size_t recordlen, tls_socket_t sock);
+int tls_record_recv(uint8_t *record, size_t *recordlen, tls_socket_t sock);
+int tls12_record_recv(uint8_t *record, size_t *recordlen, tls_socket_t sock);
 
 
 // Handshake
@@ -752,12 +740,7 @@ typedef struct {
 	int is_client;
 	int cipher_suites[TLS_MAX_CIPHER_SUITES_COUNT];
 	size_t cipher_suites_cnt;
-
-#ifdef WIN32
-	SOCKET sock;
-#else
-	int sock;
-#endif
+	tls_socket_t sock;
 
 	uint8_t enced_record[TLS_MAX_RECORD_SIZE];
 	size_t enced_record_len;
@@ -807,11 +790,7 @@ typedef struct {
 
 
 int tls_init(TLS_CONNECT *conn, const TLS_CTX *ctx);
-#ifdef WIN32
-int tls_set_socket(TLS_CONNECT* conn, SOCKET sock);
-#else
-int tls_set_socket(TLS_CONNECT *conn, int sock);
-#endif
+int tls_set_socket(TLS_CONNECT *conn, tls_socket_t sock);
 int tls_do_handshake(TLS_CONNECT *conn);
 int tls_send(TLS_CONNECT *conn, const uint8_t *in, size_t inlen, size_t *sentlen);
 int tls_recv(TLS_CONNECT *conn, uint8_t *out, size_t outlen, size_t *recvlen);
diff --git a/src/tlcp.c b/src/tlcp.c
index dcdfadc2..34be96d3 100644
--- a/src/tlcp.c
+++ b/src/tlcp.c
@@ -14,15 +14,6 @@
 #include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
-#ifdef WIN32
-#include <winsock2.h>
-#else
-#include <unistd.h>
-#include <sys/types.h>
-#include <arpa/inet.h>
-#include <sys/socket.h>
-#include <netinet/in.h>
-#endif
 #include <gmssl/rand.h>
 #include <gmssl/x509.h>
 #include <gmssl/error.h>
diff --git a/src/tls.c b/src/tls.c
index 206f0976..8d586bbc 100644
--- a/src/tls.c
+++ b/src/tls.c
@@ -14,19 +14,6 @@
 #include <assert.h>
 #include <stdlib.h>
 #include <string.h>
-
-#ifdef WIN32
-#include <winsock2.h>
-#else
-
-#include <unistd.h>
-#include <fcntl.h>
-#include <sys/types.h>
-#include <arpa/inet.h>
-#include <sys/socket.h>
-#include <netinet/in.h>
-#endif
-
 #include <gmssl/rand.h>
 #include <gmssl/x509.h>
 #include <gmssl/error.h>
@@ -1460,17 +1447,10 @@ int tls_cipher_suite_in_list(int cipher, const int *list, size_t list_count)
 	return 0;
 }
 
-#ifdef WIN32
-int tls_record_send(const uint8_t *record, size_t recordlen, SOCKET sock)
-#else
-int tls_record_send(const uint8_t *record, size_t recordlen, int sock)
-#endif
+int tls_record_send(const uint8_t *record, size_t recordlen, tls_socket_t sock)
 {
-#ifdef WIN32
-	int r;
-#else
-	ssize_t r;
-#endif
+	tls_ret_t r;
+
 	if (!record) {
 		error_print();
 		return -1;
@@ -1483,11 +1463,7 @@ int tls_record_send(const uint8_t *record, size_t recordlen, int sock)
 		error_print();
 		return -1;
 	}
-#ifdef WIN32
-	if ((r = send(sock, record, (int)recordlen, 0)) < 0) {
-#else	
-	if ((r = send(sock, record, recordlen, 0)) < 0) {
-#endif
+	if ((r = tls_socket_send(sock, record, recordlen, 0)) < 0) {
 		perror("tls_record_send");
 		error_print();
 		return -1;
@@ -1498,26 +1474,14 @@ int tls_record_send(const uint8_t *record, size_t recordlen, int sock)
 	return 1;
 }
 
-#ifdef WIN32
-int tls_record_do_recv(uint8_t *record, size_t *recordlen, SOCKET sock)
-#else
-int tls_record_do_recv(uint8_t *record, size_t *recordlen, int sock)
-#endif
+int tls_record_do_recv(uint8_t *record, size_t *recordlen, tls_socket_t sock)
 {
-#ifdef WIN32
-	int r;
-#else
-	ssize_t r;
-#endif
+	tls_ret_t r;
 	size_t len;
 
 	len = 5;
 	while (len) {
-#ifdef WIN32
-		if ((r = recv(sock, record + 5 - len, (int)len, 0)) < 0) {
-#else
-		if ((r = recv(sock, record + 5 - len, len, 0)) < 0) {
-#endif
+		if ((r = tls_socket_recv(sock, record + 5 - len, len, 0)) < 0) {
 			perror("tls_record_do_recv");
 			error_print();
 			return -1;
@@ -1546,11 +1510,7 @@ int tls_record_do_recv(uint8_t *record, size_t *recordlen, int sock)
 		return -1;
 	}
 	while (len) {
-#ifdef WIN32
-		if ((r = recv(sock, record + *recordlen - len, (int)len, 0)) < 0) {
-#else
 		if ((r = recv(sock, record + *recordlen - len, len, 0)) < 0) {
-#endif
 			perror("tls_record_do_recv");
 			error_print();
 			return -1;
@@ -1560,11 +1520,7 @@ int tls_record_do_recv(uint8_t *record, size_t *recordlen, int sock)
 	return 1;
 }
 
-#ifdef WIN32
-int tls_record_recv(uint8_t *record, size_t *recordlen, SOCKET sock)
-#else
-int tls_record_recv(uint8_t *record, size_t *recordlen, int sock)
-#endif
+int tls_record_recv(uint8_t *record, size_t *recordlen, tls_socket_t sock)
 {
 retry:
 	if (tls_record_do_recv(record, recordlen, sock) != 1) {
@@ -2303,11 +2259,7 @@ void tls_cleanup(TLS_CONNECT *conn)
 	gmssl_secure_clear(conn, sizeof(TLS_CONNECT));
 }
 
-#ifdef WIN32
-int tls_set_socket(TLS_CONNECT *conn, SOCKET sock)
-#else
-int tls_set_socket(TLS_CONNECT *conn, int sock)
-#endif
+int tls_set_socket(TLS_CONNECT *conn, tls_socket_t sock)
 {
 #if 0
 	int opts;
diff --git a/src/tls12.c b/src/tls12.c
index 27e8191f..81d7e24a 100644
--- a/src/tls12.c
+++ b/src/tls12.c
@@ -13,16 +13,6 @@
 #include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
-#ifdef WIN32
-#include <winsock2.h>
-#else
-#include <unistd.h>
-#include <fcntl.h>
-#include <sys/types.h>
-#include <arpa/inet.h>
-#include <sys/socket.h>
-#include <netinet/in.h>
-#endif
 #include <gmssl/rand.h>
 #include <gmssl/x509.h>
 #include <gmssl/error.h>
diff --git a/src/tls13.c b/src/tls13.c
index c33de356..aef4fbbc 100644
--- a/src/tls13.c
+++ b/src/tls13.c
@@ -13,15 +13,6 @@
 #include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
-#ifdef WIN32
-#include <winsock2.h>
-#else
-#include <unistd.h>
-#include <sys/types.h>
-#include <arpa/inet.h>
-#include <sys/socket.h>
-#include <netinet/in.h>
-#endif
 #include <gmssl/rand.h>
 #include <gmssl/x509.h>
 #include <gmssl/error.h>
diff --git a/src/tls_ext.c b/src/tls_ext.c
index 9f1ae0ac..cfb08c16 100644
--- a/src/tls_ext.c
+++ b/src/tls_ext.c
@@ -14,18 +14,6 @@
 #include <stdlib.h>
 #include <string.h>
 #include <fcntl.h>
-#ifdef WIN32
-
-#include <winsock2.h>		
-
-
-#else
-#include <unistd.h>
-#include <sys/types.h>
-#include <arpa/inet.h>
-#include <sys/socket.h>
-#include <netinet/in.h>
-#endif
 #include <gmssl/rand.h>
 #include <gmssl/x509.h>
 #include <gmssl/error.h>
diff --git a/tools/tlcp_client.c b/tools/tlcp_client.c
index a53cd22b..62cd577b 100644
--- a/tools/tlcp_client.c
+++ b/tools/tlcp_client.c
@@ -12,18 +12,6 @@
 #include <errno.h>
 #include <string.h>
 #include <stdlib.h>
-#ifdef WIN32
-// FIMXE: socket related headers should be moved to tls.h
-#include <winsock2.h>
-#else
-#include <unistd.h>
-#include <netdb.h>
-#include <sys/types.h>
-#include <arpa/inet.h>
-#include <sys/socket.h>
-#include <netinet/in.h>
-#endif
-
 #include <gmssl/tls.h>
 #include <gmssl/error.h>
 
@@ -49,11 +37,7 @@ int tlcp_client_main(int argc, char *argv[])
 	char *pass = NULL;
 	struct hostent *hp;
 	struct sockaddr_in server;
-#ifdef WIN32
-	SOCKET sock;
-#else	
-	int sock;
-#endif
+	tls_socket_t sock;
 	TLS_CTX ctx;
 	TLS_CONNECT conn;
 	char buf[1024] = {0};
@@ -206,11 +190,7 @@ bad:
 
 
 end:
-#ifdef WIN32
-	closesocket(sock);
-#else
-	close(sock);
-#endif
+	tls_socket_close(sock);
 	tls_ctx_cleanup(&ctx);
 	tls_cleanup(&conn);
 	return 0;
diff --git a/tools/tlcp_server.c b/tools/tlcp_server.c
index a0573fd9..d5d96a49 100644
--- a/tools/tlcp_server.c
+++ b/tools/tlcp_server.c
@@ -12,15 +12,6 @@
 #include <errno.h>
 #include <string.h>
 #include <stdlib.h>
-#ifdef WIN32
-#include <winsock2.h>
-#else
-#include <unistd.h>
-#include <sys/types.h>
-#include <arpa/inet.h>
-#include <sys/socket.h>
-#include <netinet/in.h>
-#endif
 #include <gmssl/mem.h>
 #include <gmssl/sm2.h>
 #include <gmssl/tls.h>
@@ -47,23 +38,11 @@ int tlcp_server_main(int argc , char **argv)
 	TLS_CONNECT conn;
 	char buf[1600] = {0};
 	size_t len = sizeof(buf);
-
-#ifdef WIN32
-	SOCKET sock;
-	SOCKET conn_sock;
-#else
-	int sock;
-	int conn_sock;	
-#endif
+	tls_socket_t sock;
+	tls_socket_t conn_sock;
 	struct sockaddr_in server_addr;
 	struct sockaddr_in client_addr;
-#ifdef WIN32
-	int client_addrlen;
-#else
-	socklen_t client_addrlen;
-#endif
-
-
+	tls_socklen_t client_addrlen;
 
 	argc--;
 	argv++;
@@ -194,7 +173,7 @@ restart:
 				if (rv < 0) fprintf(stderr, "%s: recv failure\n", prog);
 				else fprintf(stderr, "%s: Disconnected by remote\n", prog);
 
-				//close(conn.sock);
+				//tls_socket_close(conn.sock); // FIXME: 		
 				tls_cleanup(&conn);
 				goto restart;
 			}
@@ -202,11 +181,7 @@ restart:
 
 		if (tls_send(&conn, (uint8_t *)buf, len, &sentlen) != 1) {
 			fprintf(stderr, "%s: send failure, close connection\n", prog);
-#ifdef WIN32
-			closesocket(conn.sock);
-#else
-			close(conn.sock);
-#endif
+			tls_socket_close(conn.sock);
 			goto end;
 		}
 	}
diff --git a/tools/tls12_client.c b/tools/tls12_client.c
index 4ac0e16e..495f03ea 100644
--- a/tools/tls12_client.c
+++ b/tools/tls12_client.c
@@ -12,19 +12,6 @@
 #include <errno.h>
 #include <string.h>
 #include <stdlib.h>
-#ifdef WIN32
-#include <winsock2.h>
-#else
-
-#include <unistd.h>
-#include <netdb.h>
-#include <sys/types.h>
-#include <arpa/inet.h>
-#include <sys/socket.h>
-#include <netinet/in.h>
-#endif
-
-
 #include <gmssl/tls.h>
 #include <gmssl/error.h>
 
@@ -52,11 +39,7 @@ int tls12_client_main(int argc, char *argv[])
 	char *pass = NULL;
 	struct hostent *hp;
 	struct sockaddr_in server;
-#ifdef WIN32
-	SOCKET sock;
-#else	
-	int sock;
-#endif
+	tls_socket_t sock;
 	TLS_CTX ctx;
 	TLS_CONNECT conn;
 	char buf[1024] = {0};
@@ -202,11 +185,7 @@ bad:
 
 
 end:
-#ifdef WIN32
-	closesocket(sock);
-#else
-	close(sock);
-#endif
+	tls_socket_close(sock);
 	tls_ctx_cleanup(&ctx);
 	tls_cleanup(&conn);
 	return 0;
diff --git a/tools/tls12_server.c b/tools/tls12_server.c
index d983021a..ccd60e40 100644
--- a/tools/tls12_server.c
+++ b/tools/tls12_server.c
@@ -12,15 +12,6 @@
 #include <errno.h>
 #include <string.h>
 #include <stdlib.h>
-#ifdef WIN32
-#include <winsock2.h>
-#else
-#include <unistd.h>
-#include <sys/types.h>
-#include <arpa/inet.h>
-#include <sys/socket.h>
-#include <netinet/in.h>
-#endif
 #include <gmssl/mem.h>
 #include <gmssl/sm2.h>
 #include <gmssl/tls.h>
@@ -45,24 +36,11 @@ int tls12_server_main(int argc , char **argv)
 	TLS_CONNECT conn;
 	char buf[1600] = {0};
 	size_t len = sizeof(buf);
-
-#ifdef WIN32
-	SOCKET sock;
-	SOCKET conn_sock;
-#else
-	int sock;
-	int conn_sock;
-#endif
-
+	tls_socket_t sock;
+	tls_socket_t conn_sock;
 	struct sockaddr_in server_addr;
 	struct sockaddr_in client_addr;
-#ifdef WIN32
-	int client_addrlen;
-#else
-	socklen_t client_addrlen;
-#endif
-
-
+	tls_socklen_t client_addrlen;
 
 	argc--;
 	argv++;
@@ -179,7 +157,7 @@ restart:
 				if (rv < 0) fprintf(stderr, "%s: recv failure\n", prog);
 				else fprintf(stderr, "%s: Disconnected by remote\n", prog);
 
-				//close(conn.sock);
+				//tls_socket_close(conn.sock); // FIXME:		
 				tls_cleanup(&conn);
 				goto restart;
 			}
@@ -187,11 +165,7 @@ restart:
 
 		if (tls_send(&conn, (uint8_t *)buf, len, &sentlen) != 1) {
 			fprintf(stderr, "%s: send failure, close connection\n", prog);
-#ifdef WIN32
-			closesocket(conn.sock);
-#else
-			close(conn.sock);
-#endif
+			tls_socket_close(conn.sock);
 			goto end;
 		}
 	}
diff --git a/tools/tls13_client.c b/tools/tls13_client.c
index b6e9767e..3517fdbc 100644
--- a/tools/tls13_client.c
+++ b/tools/tls13_client.c
@@ -12,16 +12,6 @@
 #include <errno.h>
 #include <string.h>
 #include <stdlib.h>
-#ifdef WIN32
-#include <winsock2.h>
-#else
-#include <unistd.h>
-#include <netdb.h>
-#include <sys/types.h>
-#include <arpa/inet.h>
-#include <sys/socket.h>
-#include <netinet/in.h>
-#endif
 #include <gmssl/tls.h>
 #include <gmssl/error.h>
 
@@ -49,11 +39,7 @@ int tls13_client_main(int argc, char *argv[])
 	char *pass = NULL;
 	struct hostent *hp;
 	struct sockaddr_in server;
-#ifdef WIN32
-	SOCKET sock;
-#else	
-	int sock;
-#endif
+	tls_socket_t sock;
 	TLS_CTX ctx;
 	TLS_CONNECT conn;
 	char buf[1024] = {0};
@@ -198,11 +184,7 @@ bad:
 	}
 
 end:
-#ifdef WIN32
-	closesocket(sock);
-#else
-	close(sock);
-#endif
+	tls_socket_close(sock);
 	tls_ctx_cleanup(&ctx);
 	tls_cleanup(&conn);
 	return 0;
diff --git a/tools/tls13_server.c b/tools/tls13_server.c
index ca0ff93d..09308ac2 100644
--- a/tools/tls13_server.c
+++ b/tools/tls13_server.c
@@ -12,15 +12,6 @@
 #include <errno.h>
 #include <string.h>
 #include <stdlib.h>
-#ifdef WIN32
-#include <winsock2.h>
-#else
-#include <unistd.h>
-#include <sys/types.h>
-#include <arpa/inet.h>
-#include <sys/socket.h>
-#include <netinet/in.h>
-#endif
 #include <gmssl/mem.h>
 #include <gmssl/sm2.h>
 #include <gmssl/tls.h>
@@ -38,30 +29,16 @@ int tls13_server_main(int argc , char **argv)
 	char *keyfile = NULL;
 	char *pass = NULL;
 	char *cacertfile = NULL;
-
 	int server_ciphers[] = { TLS_cipher_sm4_gcm_sm3, };
-	
 	TLS_CTX ctx;
 	TLS_CONNECT conn;
 	char buf[1600] = {0};
 	size_t len = sizeof(buf);
-
-#ifdef WIN32
-	SOCKET sock;
-	SOCKET conn_sock;
-#else
-	int sock;
-	int conn_sock;	
-#endif
+	tls_socket_t sock;
+	tls_socket_t conn_sock;
 	struct sockaddr_in server_addr;
 	struct sockaddr_in client_addr;
-#ifdef WIN32
-	int client_addrlen;
-#else
-	socklen_t client_addrlen;
-#endif
-
-
+	tls_socklen_t client_addrlen;
 
 	argc--;
 	argv++;
@@ -186,11 +163,7 @@ restart:
 
 		if (tls13_send(&conn, (uint8_t *)buf, len, &sentlen) != 1) {
 			fprintf(stderr, "%s: send failure, close connection\n", prog);
-#ifdef WIN32
-			closesocket(conn.sock);
-#else
-			close(conn.sock);
-#endif
+			tls_socket_close(conn.sock);
 			goto end;
 		}
 	}