From e1bdd31e6f4ec0aa426cdc2e24b2c78fb7903cfe Mon Sep 17 00:00:00 2001
From: yangchen <yangchen145@huawei.com>
Date: Wed, 6 Dec 2023 09:26:45 +0800
Subject: [PATCH] wrap: support shutdown

---
 src/lstack/api/lstack_rtc_api.c            |  5 +++
 src/lstack/api/lstack_rtw_api.c            | 13 +++++++
 src/lstack/api/lstack_wrap.c               | 24 ++++++++++++
 src/lstack/core/lstack_dpdk.c              |  2 +-
 src/lstack/core/lstack_protocol_stack.c    | 45 ++++++++++++++++++++++
 src/lstack/core/lstack_thread_rpc.c        | 16 ++++++++
 src/lstack/include/lstack_protocol_stack.h |  3 ++
 src/lstack/include/lstack_rtc_api.h        |  1 +
 src/lstack/include/lstack_rtw_api.h        |  1 +
 src/lstack/include/lstack_thread_rpc.h     |  1 +
 10 files changed, 110 insertions(+), 1 deletion(-)

diff --git a/src/lstack/api/lstack_rtc_api.c b/src/lstack/api/lstack_rtc_api.c
index 5fad3e8..50d72bc 100644
--- a/src/lstack/api/lstack_rtc_api.c
+++ b/src/lstack/api/lstack_rtc_api.c
@@ -65,6 +65,11 @@ int rtc_close(int s)
     return lwip_close(s);
 }
 
+int rtc_shutdown(int fd, int how)
+{
+    return lwip_shutdown(fd, how);
+}
+
 int rtc_epoll_create(int flags)
 {
     if (stack_setup_app_thread() < 0) {
diff --git a/src/lstack/api/lstack_rtw_api.c b/src/lstack/api/lstack_rtw_api.c
index c524bf9..04944d5 100644
--- a/src/lstack/api/lstack_rtw_api.c
+++ b/src/lstack/api/lstack_rtw_api.c
@@ -22,8 +22,10 @@
 #include "lstack_protocol_stack.h"
 #include "lstack_cfg.h"
 #include "lstack_lwip.h"
+#include "gazelle_base_func.h"
 #include "lstack_rtw_api.h"
 
+
 int rtw_socket(int domain, int type, int protocol)
 {
     return rpc_call_socket(domain, type, protocol);
@@ -228,6 +230,16 @@ int rtw_close(int s)
     return stack_broadcast_close(s);
 }
 
+int rtw_shutdown(int fd, int how)
+{
+    struct lwip_sock *sock = get_socket_by_fd(fd);
+    if (sock && sock->wakeup && sock->wakeup->epollfd == fd) {
+	GAZELLE_RETURN(ENOTSOCK);
+    }
+
+    return stack_broadcast_shutdown(fd, how);
+}
+
 int rtw_epoll_ctl(int epfd, int op, int fd, struct epoll_event *event)
 {
     return lstack_rtw_epoll_ctl(epfd, op, fd, event);
@@ -242,3 +254,4 @@ int rtw_epoll_create(int flags)
 {
     return lstack_epoll_create(flags);
 }
+
diff --git a/src/lstack/api/lstack_wrap.c b/src/lstack/api/lstack_wrap.c
index abbf8a1..07d5f27 100644
--- a/src/lstack/api/lstack_wrap.c
+++ b/src/lstack/api/lstack_wrap.c
@@ -75,6 +75,7 @@ void wrap_api_init(void)
         g_wrap_api->epoll_wait_fn    = rtc_epoll_wait;
         g_wrap_api->poll_fn          = rtc_poll;
         g_wrap_api->close_fn         = rtc_close;
+        g_wrap_api->shutdown_fn      = rtc_shutdown;
         g_wrap_api->epoll_ctl_fn     = rtc_epoll_ctl;
         g_wrap_api->epoll_create1_fn = rtc_epoll_create1;
         g_wrap_api->epoll_create_fn  = rtc_epoll_create;
@@ -103,6 +104,7 @@ void wrap_api_init(void)
         g_wrap_api->epoll_wait_fn    = rtw_epoll_wait;
         g_wrap_api->poll_fn          = rtw_poll;
         g_wrap_api->close_fn         = rtw_close;
+        g_wrap_api->shutdown_fn      = rtw_shutdown;
         g_wrap_api->epoll_ctl_fn     = rtw_epoll_ctl;
         g_wrap_api->epoll_create1_fn = rtw_epoll_create1;
         g_wrap_api->epoll_create_fn  = rtw_epoll_create;
@@ -554,6 +556,20 @@ static inline int32_t do_close(int32_t s)
     return g_wrap_api->close_fn(s);
 }
 
+static int32_t do_shutdown(int fd, int how)
+{
+    struct lwip_sock *sock = NULL;
+    if (select_posix_path() == PATH_KERNEL || select_fd_posix_path(fd, &sock) == PATH_KERNEL) {
+        if (posix_api != NULL && !posix_api->ues_posix && g_wrap_api->shutdown_fn(fd, how) == 0) {
+            return 0;
+        } else {
+            return posix_api->shutdown_fn(fd, how);
+        }
+    }
+
+    return g_wrap_api->shutdown_fn(fd, how);
+}
+
 static int32_t do_poll(struct pollfd *fds, nfds_t nfds, int32_t timeout)
 {
     if ((select_posix_path() == PATH_KERNEL) || fds == NULL || nfds == 0) {
@@ -742,6 +758,10 @@ int32_t close(int32_t s)
 {
     return do_close(s);
 }
+int32_t shutdown(int fd, int how)
+{
+    return do_shutdown(fd, how);
+}
 int32_t poll(struct pollfd *fds, nfds_t nfds, int32_t timeout)
 {
     return do_poll(fds, nfds, timeout);
@@ -875,6 +895,10 @@ int32_t __wrap_close(int32_t s)
 {
     return do_close(s);
 }
+int32_t __wrap_shutdown(int fd, int how)
+{
+    return do_shutdown(fd, how);
+}
 int32_t __wrap_poll(struct pollfd *fds, nfds_t nfds, int32_t timeout)
 {
     return do_poll(fds, nfds, timeout);
diff --git a/src/lstack/core/lstack_dpdk.c b/src/lstack/core/lstack_dpdk.c
index a742335..1811a9d 100644
--- a/src/lstack/core/lstack_dpdk.c
+++ b/src/lstack/core/lstack_dpdk.c
@@ -294,7 +294,7 @@ int32_t dpdk_alloc_pktmbuf(struct rte_mempool *pool, struct rte_mbuf **mbufs, ui
 {
     int32_t ret = rte_pktmbuf_alloc_bulk(pool, mbufs, num);
     if (ret != 0) {
-        LSTACK_LOG(ERR, LSTACK, "rte_pktmbuf_alloc_bulk fail allocNum=%d, ret=%d, info:=%s \n",
+        LSTACK_LOG(ERR, LSTACK, "rte_pktmbuf_alloc_bulk fail allocNum=%d, ret=%d, info:%s \n",
                    num, ret, rte_strerror(-ret));
         return ret;
     }
diff --git a/src/lstack/core/lstack_protocol_stack.c b/src/lstack/core/lstack_protocol_stack.c
index 27eeafb..dc8f143 100644
--- a/src/lstack/core/lstack_protocol_stack.c
+++ b/src/lstack/core/lstack_protocol_stack.c
@@ -715,6 +715,27 @@ void stack_close(struct rpc_msg *msg)
     posix_api->close_fn(fd);
 }
 
+void stack_shutdown(struct rpc_msg *msg)
+{
+    int fd = msg->args[MSG_ARG_0].i;
+    int how = msg->args[MSG_ARG_1].i;
+    struct protocol_stack *stack = get_protocol_stack_by_fd(fd);
+    struct lwip_sock *sock = get_socket(fd);
+
+    if (sock && NETCONN_IS_DATAOUT(sock)) {
+        msg->recall_flag = 1;
+        rpc_call(&stack->rpc_queue, msg);
+        return;
+    }
+
+    msg->result = lwip_shutdown(fd, how);
+    if (msg->result != 0) {
+        LSTACK_LOG(ERR, LSTACK, "tid %ld, fd %d fail %ld\n", get_stack_tid(), fd, msg->result);
+    }
+
+    posix_api->shutdown_fn(fd, how);
+}
+
 void stack_bind(struct rpc_msg *msg)
 {
     msg->result = lwip_bind(msg->args[MSG_ARG_0].i, msg->args[MSG_ARG_1].cp, msg->args[MSG_ARG_2].socklen);
@@ -1041,6 +1062,29 @@ int32_t stack_broadcast_close(int32_t fd)
     return ret;
 }
 
+int stack_broadcast_shutdown(int fd, int how)
+{
+    int32_t ret = 0;
+    struct lwip_sock *sock = get_socket(fd);
+    if (sock == NULL) {
+        return -1;
+    }
+
+    do {
+        sock = sock->listen_next;
+        if (rpc_call_shutdown(fd, how)) {
+            ret = -1;
+        }
+
+        if (sock == NULL || sock->conn == NULL) {
+            break;
+        }
+        fd = sock->conn->socket;
+    } while (sock);
+
+    return ret;
+}
+
 /* choice one stack listen */
 int32_t stack_single_listen(int32_t fd, int32_t backlog)
 {
@@ -1204,3 +1248,4 @@ int32_t stack_broadcast_accept(int32_t fd, struct sockaddr *addr, socklen_t *add
 {
     return stack_broadcast_accept4(fd, addr, addrlen, 0);
 }
+
diff --git a/src/lstack/core/lstack_thread_rpc.c b/src/lstack/core/lstack_thread_rpc.c
index f23d935..4dc3da3 100644
--- a/src/lstack/core/lstack_thread_rpc.c
+++ b/src/lstack/core/lstack_thread_rpc.c
@@ -262,6 +262,21 @@ int32_t rpc_call_close(int fd)
     return rpc_sync_call(&stack->rpc_queue, msg);
 }
 
+int32_t rpc_call_shutdown(int fd, int how)
+{
+    struct protocol_stack *stack = get_protocol_stack_by_fd(fd);
+
+    struct rpc_msg *msg = rpc_msg_alloc(stack, stack_shutdown);
+    if (msg == NULL) {
+        return -1;
+    }
+
+    msg->args[MSG_ARG_0].i = fd;
+    msg->args[MSG_ARG_1].i = how;
+
+    return rpc_sync_call(&stack->rpc_queue, msg);
+}
+
 void rpc_call_clean_epoll(struct protocol_stack *stack, struct wakeup_poll *wakeup)
 {
     struct rpc_msg *msg = rpc_msg_alloc(stack, stack_clean_epoll);
@@ -465,3 +480,4 @@ int32_t rpc_call_send(int fd, const void *buf, size_t len, int flags)
 
     return 0;
 }
+
diff --git a/src/lstack/include/lstack_protocol_stack.h b/src/lstack/include/lstack_protocol_stack.h
index 2c581b3..e339b8d 100644
--- a/src/lstack/include/lstack_protocol_stack.h
+++ b/src/lstack/include/lstack_protocol_stack.h
@@ -126,6 +126,8 @@ void stack_broadcast_arp(struct rte_mbuf *mbuf, struct protocol_stack *cur_stack
 /* when fd is listenfd, listenfd of all protocol stack thread will be closed */
 int32_t stack_broadcast_close(int32_t fd);
 
+int stack_broadcast_shutdown(int fd, int how);
+
 /* listen sync to all protocol stack thread, so that any protocol stack thread can build connect */
 int32_t stack_broadcast_listen(int32_t fd, int backlog);
 int32_t stack_single_listen(int32_t fd, int32_t backlog);
@@ -152,6 +154,7 @@ void stack_clean_epoll(struct rpc_msg *msg);
 void stack_arp(struct rpc_msg *msg);
 void stack_socket(struct rpc_msg *msg);
 void stack_close(struct rpc_msg *msg);
+void stack_shutdown(struct rpc_msg *msg);
 void stack_bind(struct rpc_msg *msg);
 void stack_listen(struct rpc_msg *msg);
 void stack_accept(struct rpc_msg *msg);
diff --git a/src/lstack/include/lstack_rtc_api.h b/src/lstack/include/lstack_rtc_api.h
index dd90e59..3a41e6f 100644
--- a/src/lstack/include/lstack_rtc_api.h
+++ b/src/lstack/include/lstack_rtc_api.h
@@ -49,6 +49,7 @@ int rtc_poll(struct pollfd *fds, nfds_t nfds, int timeout);
 int rtc_epoll_wait(int epfd, struct epoll_event* events, int maxevents, int timeout);
 int rtc_socket(int domain, int type, int protocol);
 int rtc_close(int s);
+int rtc_shutdown(int fd, int how);
 int rtc_epoll_create(int flags);
 int rtc_epoll_create1(int flags);
 int rtc_epoll_ctl(int epfd, int op, int fd, struct epoll_event *event);
diff --git a/src/lstack/include/lstack_rtw_api.h b/src/lstack/include/lstack_rtw_api.h
index d0f77b7..a38b656 100644
--- a/src/lstack/include/lstack_rtw_api.h
+++ b/src/lstack/include/lstack_rtw_api.h
@@ -42,6 +42,7 @@ ssize_t rtw_sendto(int sockfd, const void *buf, size_t len, int flags,
 int rtw_epoll_wait(int epfd, struct epoll_event* events, int maxevents, int timeout);
 int rtw_poll(struct pollfd *fds, nfds_t nfds, int timeout);
 int rtw_close(int s);
+int rtw_shutdown(int fd, int how);
 int rtw_epoll_ctl(int epfd, int op, int fd, struct epoll_event *event);
 int rtw_epoll_create1(int flags);
 int rtw_epoll_create(int flags);
diff --git a/src/lstack/include/lstack_thread_rpc.h b/src/lstack/include/lstack_thread_rpc.h
index aa8c451..13f5ec2 100644
--- a/src/lstack/include/lstack_thread_rpc.h
+++ b/src/lstack/include/lstack_thread_rpc.h
@@ -74,6 +74,7 @@ int32_t rpc_call_connnum(struct protocol_stack *stack);
 int32_t rpc_call_arp(struct protocol_stack *stack, struct rte_mbuf *mbuf);
 int32_t rpc_call_socket(int32_t domain, int32_t type, int32_t protocol);
 int32_t rpc_call_close(int32_t fd);
+int32_t rpc_call_shutdown(int fd, int how);
 int32_t rpc_call_bind(int32_t fd, const struct sockaddr *addr, socklen_t addrlen);
 int32_t rpc_call_listen(int s, int backlog);
 int32_t rpc_call_accept(int fd, struct sockaddr *addr, socklen_t *addrlen, int flags);
-- 
2.27.0