From 9816a45d35ffa187b5830a417147afa9695dcead Mon Sep 17 00:00:00 2001
From: liqiang <liqiang64@huawei.com>
Date: Sat, 3 Jun 2023 11:01:23 +0800
Subject: uds add vsock mode

Signed-off-by: liqiang <liqiang64@huawei.com>
---
 qtfs/ipc/Makefile   |  7 +++++++
 qtfs/ipc/uds_main.c | 43 ++++++++++++++++++++++++++++++++++++++-----
 qtfs/ipc/uds_main.h | 12 +++++++++---
 3 files changed, 54 insertions(+), 8 deletions(-)

diff --git a/qtfs/ipc/Makefile b/qtfs/ipc/Makefile
index 2136f4f..3b851f3 100644
--- a/qtfs/ipc/Makefile
+++ b/qtfs/ipc/Makefile
@@ -3,6 +3,9 @@ CFLAGS += -g -O2
 CFLAGS += -fstack-protector-strong
 CFLAGS += -fPIE -pie -fPIC
 CFLAGS += -D_FORTIFY_SOURCE=2
+ifdef UDS_TEST_MODE
+CFLAGS += -DUDS_TEST_MODE
+endif
 LDFLAGS += -s
 LDFLAGS += -Wl,-z,now
 LDFLAGS += -Wl,-z,noexecstack
@@ -12,6 +15,10 @@ all: udsproxyd libudsproxy.so
 
 udsproxyd: uds_event.o uds_main.o
 	gcc $(LDFLAGS) -o udsproxyd $^ -I../ $(DEPGLIB)
+	@test -z $(UDS_TEST_MODE) || echo "Important risk warning: The test mode is turned on,\
+	and udsproxyd will expose the network port, which will bring security risks and is only for\
+	testing! If you do not understand the risks, please don't use or compile again without\
+	UDS_TEST_MODE."
 
 uds_event.o:
 	cc $(CFLAGS) -c -o uds_event.o uds_event.c $(DEPGLIB)
diff --git a/qtfs/ipc/uds_main.c b/qtfs/ipc/uds_main.c
index 2af4f06..adf0936 100644
--- a/qtfs/ipc/uds_main.c
+++ b/qtfs/ipc/uds_main.c
@@ -37,6 +37,7 @@
 #include <sys/resource.h>
 #include <sys/prctl.h>
 #include <sys/file.h>
+#include <linux/vm_sockets.h>
 
 #include "comm.h"
 #include "uds_main.h"
@@ -179,15 +180,24 @@ void uds_main_loop(int efd, struct uds_thread_arg *arg)
 #define UDS_MAX_LISTEN_NUM 64
 int uds_build_tcp_connection(struct uds_conn_arg *arg)
 {
+	int family = AF_VSOCK;
 	if (arg->cs > UDS_SOCKET_SERVER) {
 		uds_err("cs type %d is error.", arg->cs);
 		return -1;
 	}
-	struct sockaddr_in sock_addr = {
-		.sin_family = AF_INET,
-	};
-	int sock_fd = socket(AF_INET, SOCK_STREAM, 0);
+#ifdef UDS_TEST_MODE
+	family = AF_INET;
+	struct sockaddr_in sock_addr;
+	memset(&sock_addr, 0, sizeof(sock_addr));
+	sock_addr.sin_family = AF_INET;
+#else
+	family = AF_VSOCK;
+	struct sockaddr_vm sock_addr;
+	memset(&sock_addr, 0, sizeof(sock_addr));
+	sock_addr.svm_family = AF_VSOCK;
+#endif
 
+	int sock_fd = socket(family, SOCK_STREAM, 0);
 	if (sock_fd < 0) {
 		uds_err("As %s failed, socket fd: %d, errno:%d.",
 				(arg->cs == UDS_SOCKET_CLIENT) ? "client" : "server",
@@ -197,8 +207,13 @@ int uds_build_tcp_connection(struct uds_conn_arg *arg)
 	arg->sockfd = sock_fd;
 
 	if (arg->cs == UDS_SOCKET_SERVER) {
+#ifdef UDS_TEST_MODE
 		sock_addr.sin_port = htons(p_uds_var->tcp.port);
 		sock_addr.sin_addr.s_addr = inet_addr(p_uds_var->tcp.addr);
+#else
+		sock_addr.svm_port = p_uds_var->vsock.port;
+		sock_addr.svm_cid = p_uds_var->vsock.cid;
+#endif
 		if (bind(sock_fd, (struct sockaddr *)&sock_addr, sizeof(sock_addr)) < 0) {
 			uds_err("As tcp server failed, bind error, errno:%d.",
 					errno);
@@ -209,13 +224,22 @@ int uds_build_tcp_connection(struct uds_conn_arg *arg)
 			goto close_and_return;
 		}
 	} else {
+#ifdef UDS_TEST_MODE
 		sock_addr.sin_port = htons(p_uds_var->tcp.peerport);
 		sock_addr.sin_addr.s_addr = inet_addr(p_uds_var->tcp.peeraddr);
-		if (connect(arg->sockfd, (struct sockaddr *)&sock_addr, sizeof(struct sockaddr_in)) < 0) {
+#else
+		sock_addr.svm_port = p_uds_var->vsock.peerport;
+		sock_addr.svm_cid = p_uds_var->vsock.peercid;
+#endif	
+		if (connect(arg->sockfd, (struct sockaddr *)&sock_addr, sizeof(sock_addr)) < 0) {
 			goto close_and_return;
 		}
 		arg->connfd = sock_fd;
+#ifdef UDS_TEST_MODE
 		uds_log("Connect to tcp server successed, ip:%s port:%u", p_uds_var->tcp.peeraddr, p_uds_var->tcp.peerport);
+#else
+		uds_log("Connect to vsock server successed, cid:%u port:%u", p_uds_var->vsock.peercid, p_uds_var->vsock.peerport);
+#endif
 	}
 
 	return 0;
@@ -599,6 +623,7 @@ static int uds_glob_var_init(char *argv[])
 		uds_err("work thread var malloc failed.");
 		return -1;
 	}
+#ifdef UDS_TEST_MODE
 	p_uds_var->tcp.port = atoi(argv[3]);
 	strncpy(p_uds_var->tcp.addr, argv[2], sizeof(p_uds_var->tcp.addr) - 1);
 	p_uds_var->tcp.peerport = atoi(argv[5]);
@@ -607,6 +632,14 @@ static int uds_glob_var_init(char *argv[])
 	uds_log("uds proxy param thread num:%d ip:%s port:%u peerip:%s port:%u",
 			p_uds_var->work_thread_num, p_uds_var->tcp.addr, p_uds_var->tcp.port,
 			 p_uds_var->tcp.peeraddr, p_uds_var->tcp.peerport);
+#else
+	// vsock param: <thread num> <local cid> <local port> <peer cid> <peer port>
+	// port and peerport is checked before
+	p_uds_var->vsock.cid = atoi(argv[2]);
+	p_uds_var->vsock.port = myport;
+	p_uds_var->vsock.peercid = atoi(argv[4]);
+	p_uds_var->vsock.peerport = peerport;
+#endif
 	g_event_var = (struct uds_event_global_var *)malloc(sizeof(struct uds_event_global_var) * p_uds_var->work_thread_num);
 	if (g_event_var == NULL) {
 		free(p_uds_var->efd);
diff --git a/qtfs/ipc/uds_main.h b/qtfs/ipc/uds_main.h
index 3903ec8..de400f8 100644
--- a/qtfs/ipc/uds_main.h
+++ b/qtfs/ipc/uds_main.h
@@ -111,15 +111,21 @@ struct uds_global_var {
 	struct uds_thread_arg *work_thread;
 	int loglevel;
 	char **logstr;
+#ifdef UDS_TEST_MODE
 	struct _tcp {
 		char addr[20];
 		unsigned short port;
 		char peeraddr[20];
 		unsigned short peerport;
 	} tcp;
-	struct _uds {
-		char sun_path[UDS_SUN_PATH_LEN];
-	} uds;
+#else
+	struct _vsock {
+		unsigned int cid;
+		unsigned int port;
+		unsigned int peercid;
+		unsigned int peerport;
+	} vsock;
+#endif
 };
 enum uds_cs {
 	UDS_SOCKET_CLIENT = 1,
-- 
2.33.0