Эх сурвалжийг харах

ssh: use select() to disconnect the client on disabling sshd

similar to the telnet logic
Sergey Alirzaev 5 жил өмнө
parent
commit
8939751fda

+ 22 - 10
modules/SSH_Server/server.c

@@ -25,6 +25,7 @@
 #include <stdbool.h>
 #include <stdbool.h>
 #include "FreeRTOS.h"
 #include "FreeRTOS.h"
 #include "task.h"
 #include "task.h"
+#include "semphr.h"
 #include "wolfssh_test.h"
 #include "wolfssh_test.h"
 #include "cli.h"
 #include "cli.h"
 #include "CLI_Commands.h"
 #include "CLI_Commands.h"
@@ -55,9 +56,9 @@ typedef struct {
     WOLFSSH* ssh;
     WOLFSSH* ssh;
     SOCKET_T fd;
     SOCKET_T fd;
     word32 id;
     word32 id;
-    char nonBlock;
 } thread_ctx_t;
 } thread_ctx_t;
 
 
+static const bool nonBlock = true;
 
 
 #ifndef EXAMPLE_HIGHWATER_MARK
 #ifndef EXAMPLE_HIGHWATER_MARK
     #define EXAMPLE_HIGHWATER_MARK 0x3FFF8000 /* 1GB - 32kB */
     #define EXAMPLE_HIGHWATER_MARK 0x3FFF8000 /* 1GB - 32kB */
@@ -138,7 +139,7 @@ static void *server_worker(void* vArgs)
     user_level_t user_id;
     user_level_t user_id;
     wolfSSH_SetUserAuthCtx(threadCtx->ssh, &user_id);
     wolfSSH_SetUserAuthCtx(threadCtx->ssh, &user_id);
 
 
-    if (!threadCtx->nonBlock)
+    if (!nonBlock)
         ret = wolfSSH_accept(threadCtx->ssh);
         ret = wolfSSH_accept(threadCtx->ssh);
     else
     else
         ret = NonBlockSSH_accept(threadCtx->ssh);
         ret = NonBlockSSH_accept(threadCtx->ssh);
@@ -159,11 +160,29 @@ static void *server_worker(void* vArgs)
                 uint8_t buf[EXAMPLE_BUFFER_SZ];
                 uint8_t buf[EXAMPLE_BUFFER_SZ];
                 int rxSz = 0;
                 int rxSz = 0;
                 do {
                 do {
+                    if (nonBlock) {
+                        SOCKET_T sockfd;
+                        int select_ret = 0;
+
+                        sockfd = (SOCKET_T)wolfSSH_get_fd(threadCtx->ssh);
+
+                        select_ret = tcp_select(sockfd, 1);
+                        if (!sSettings.sSSH.SSHEnable ||
+                            (select_ret != WS_SELECT_RECV_READY &&
+                            select_ret != WS_SELECT_ERROR_READY &&
+                            select_ret != WS_SELECT_TIMEOUT)) {
+                            break;
+                        }
+                    }
                     rxSz = wolfSSH_stream_read(threadCtx->ssh, buf, sizeof(buf));
                     rxSz = wolfSSH_stream_read(threadCtx->ssh, buf, sizeof(buf));
                     if (rxSz <= 0) {
                     if (rxSz <= 0) {
                         rxSz = wolfSSH_get_error(threadCtx->ssh);
                         rxSz = wolfSSH_get_error(threadCtx->ssh);
                     }
                     }
-                } while (rxSz == WS_WANT_READ || rxSz == WS_WANT_WRITE);
+                } while ((rxSz == WS_WANT_READ || rxSz == WS_WANT_WRITE) && sSettings.sSSH.SSHEnable);
+                if (!sSettings.sSSH.SSHEnable) {
+                    stop = true;
+                    break;
+                }
 
 
                 if (rxSz > 0) {
                 if (rxSz > 0) {
                     cli_getchar(cli_state, buf[0], true);    // TODO handle rxSz > 1
                     cli_getchar(cli_state, buf[0], true);    // TODO handle rxSz > 1
@@ -300,7 +319,6 @@ static void ssh_server(void *arg)
     word32 threadCount = 0;
     word32 threadCount = 0;
     const char multipleConnections = 0;
     const char multipleConnections = 0;
     char useEcc = 1;
     char useEcc = 1;
-    char nonBlock = 0;
 
 
     if (wolfSSH_Init() != WS_SUCCESS) {
     if (wolfSSH_Init() != WS_SUCCESS) {
         printf("Couldn't initialize wolfSSH.\n");
         printf("Couldn't initialize wolfSSH.\n");
@@ -376,7 +394,6 @@ static void ssh_server(void *arg)
             threadCtx->ssh = ssh;
             threadCtx->ssh = ssh;
             threadCtx->fd = clientFd;
             threadCtx->fd = clientFd;
             threadCtx->id = threadCount++;
             threadCtx->id = threadCount++;
-            threadCtx->nonBlock = nonBlock;
 
 
 #ifndef SINGLE_THREADED
 #ifndef SINGLE_THREADED
             ThreadStart(server_worker, threadCtx, &thread);
             ThreadStart(server_worker, threadCtx, &thread);
@@ -413,9 +430,4 @@ void ssh_server_init(void)
 
 
 void ssh_server_restart(void)
 void ssh_server_restart(void)
 {
 {
-    for (unsigned i = 0; i < array_len(cli_states); ++i) {
-        if (cli_states[i].state != STATE_UNUSED && cli_states[i].send == cli_send) {
-            wolfSSH_shutdown(cli_states[i].num_connect);
-        }
-    }
 }
 }