From 9498a8df59f69acbf75f9aa69fef465350288bb8 Mon Sep 17 00:00:00 2001
From: zhangxiaoyu <zhangxiaoyu58@huawei.com>
Date: Mon, 1 Aug 2022 11:20:31 +0800
Subject: [PATCH 06/15] fix cri attach when stdout and stderr are false

Signed-off-by: zhangxiaoyu <zhangxiaoyu58@huawei.com>
---
 .../cri/websocket/service/attach_serve.cc     | 28 +++++++++++++------
 .../entry/cri/websocket/service/ws_server.cc  | 22 +++++++++++++++
 .../entry/cri/websocket/service/ws_server.h   |  2 ++
 3 files changed, 44 insertions(+), 8 deletions(-)

diff --git a/src/daemon/entry/cri/websocket/service/attach_serve.cc b/src/daemon/entry/cri/websocket/service/attach_serve.cc
index abb08363..57711180 100644
--- a/src/daemon/entry/cri/websocket/service/attach_serve.cc
+++ b/src/daemon/entry/cri/websocket/service/attach_serve.cc
@@ -101,31 +101,40 @@ int AttachServe::ExecuteStreamCommand(SessionData *lwsCtx, void *request)
         return -1;
     }
 
+    // stdout
     struct AttachContext stdoutContext = { 0 };
     stdoutContext.lwsCtx = lwsCtx;
     stdoutContext.sem = &attachSem;
-    stdoutContext.attachWriter = WsWriteStdoutToClient;
+    // write stdout to client if attach stdout is true
+    stdoutContext.attachWriter = m_request->attach_stdout ? WsWriteStdoutToClient : WsDoNotWriteStdoutToClient;
 
     struct io_write_wrapper stdoutstringWriter = { 0 };
     stdoutstringWriter.context = static_cast<void *>(&stdoutContext);
     stdoutstringWriter.write_func = AttachWriteToClient;
-    // the close function of StderrstringWriter is preferred unless StderrstringWriter is nullptr
-    stdoutstringWriter.close_func = m_request->attach_stderr ? nullptr : AttachConnectClosed;
+    stdoutstringWriter.close_func = AttachConnectClosed;
 
+    // stderr
     struct AttachContext stderrContext = { 0 };
     stderrContext.lwsCtx = lwsCtx;
-    stderrContext.sem = &attachSem;
-    stderrContext.attachWriter = WsWriteStderrToClient;
+    stderrContext.sem = nullptr;
+    // write stderr to client if attach stderr is true
+    stderrContext.attachWriter = m_request->attach_stderr ? WsWriteStderrToClient : WsDoNotWriteStderrToClient;
 
     struct io_write_wrapper stderrstringWriter = { 0 };
     stderrstringWriter.context = static_cast<void *>(&stderrContext);
     stderrstringWriter.write_func = AttachWriteToClient;
-    stderrstringWriter.close_func = m_request->attach_stderr ? AttachConnectClosed : nullptr;
+    stderrstringWriter.close_func = nullptr;
+
+    // Maybe attach stdout and stderr are both false.
+    // To make sure the close func sem_post, set attach stdout and stderr true.
+    bool record_attach_stdout = m_request->attach_stdout;
+    bool record_attach_stderr = m_request->attach_stderr;
+    m_request->attach_stdout=true;
+    m_request->attach_stderr=true;
 
     container_attach_response *m_response { nullptr };
     int ret = cb->container.attach(m_request, &m_response, m_request->attach_stdin ? lwsCtx->pipes.at(0) : -1,
-                                   m_request->attach_stdout ? &stdoutstringWriter : nullptr,
-                                   m_request->attach_stderr ? &stderrstringWriter : nullptr);
+                                   &stdoutstringWriter, &stderrstringWriter);
 
     if (ret != 0) {
         // join io copy thread in attach callback
@@ -139,11 +148,14 @@ int AttachServe::ExecuteStreamCommand(SessionData *lwsCtx, void *request)
         }
         WsWriteStdoutToClient(lwsCtx, message.c_str(), message.length());
     } else {
+        // wait io copy thread complete
         (void)sem_wait(&attachSem);
     }
 
     (void)sem_destroy(&attachSem);
     free_container_attach_response(m_response);
+    m_request->attach_stdout = record_attach_stdout;
+    m_request->attach_stderr = record_attach_stderr;
     return ret;
 }
 
diff --git a/src/daemon/entry/cri/websocket/service/ws_server.cc b/src/daemon/entry/cri/websocket/service/ws_server.cc
index 98c0fee0..08f2cff0 100644
--- a/src/daemon/entry/cri/websocket/service/ws_server.cc
+++ b/src/daemon/entry/cri/websocket/service/ws_server.cc
@@ -703,6 +703,28 @@ ssize_t WsWriteStderrToClient(void *context, const void *data, size_t len)
     return WsWriteToClient(context, data, len, STDERRCHANNEL);
 }
 
+ssize_t WsDoNotWriteStdoutToClient(void *context, const void *data, size_t len)
+{
+    if (context == nullptr) {
+        ERROR("websocket session context empty");
+        return -1;
+    }
+
+    TRACE("Ws do not write stdout to client");
+    return len;
+}
+
+ssize_t WsDoNotWriteStderrToClient(void *context, const void *data, size_t len)
+{
+    if (context == nullptr) {
+        ERROR("websocket session context empty");
+        return -1;
+    }
+
+    TRACE("Ws do not write stderr to client");
+    return len;
+}
+
 int closeWsConnect(void *context, char **err)
 {
     (void)err;
diff --git a/src/daemon/entry/cri/websocket/service/ws_server.h b/src/daemon/entry/cri/websocket/service/ws_server.h
index b871aabc..a2a180ec 100644
--- a/src/daemon/entry/cri/websocket/service/ws_server.h
+++ b/src/daemon/entry/cri/websocket/service/ws_server.h
@@ -113,6 +113,8 @@ private:
 
 ssize_t WsWriteStdoutToClient(void *context, const void *data, size_t len);
 ssize_t WsWriteStderrToClient(void *context, const void *data, size_t len);
+ssize_t WsDoNotWriteStdoutToClient(void *context, const void *data, size_t len);
+ssize_t WsDoNotWriteStderrToClient(void *context, const void *data, size_t len);
 int closeWsConnect(void *context, char **err);
 
 #endif // DAEMON_ENTRY_CRI_WEBSOCKET_SERVICE_WS_SERVER_H
-- 
2.25.1