Merge pull request #31 from ejohnstown/echoserver

Echoserver Enhancement

Author: dgarske

examples/echoserver/echoserver.c

@@ -90,7 +90,9 @@ typedef int SOCKET_T;
 #endif
 
 typedef struct {
-    SOCKET_T clientFd;
+    WOLFSSH* ssh;
+    SOCKET_T fd;
+    uint32_t id;
 } thread_ctx_t;
 
 
@@ -269,27 +271,49 @@ static INLINE void tcp_bind(SOCKET_T* sockFd, uint16_t port, int useAnyAddr)
 }
 
 
-static int find_char(uint8_t ch, const uint8_t* buf, uint32_t bufSz)
+static uint8_t find_char(const uint8_t* str, const uint8_t* buf, uint32_t bufSz)
 {
+    const uint8_t* cur;
+
     while (bufSz) {
-        if (ch == *buf)
-            return 1;
-        else {
-            buf++;
-            bufSz--;
+        cur = str;
+        while (*cur != '\0') {
+            if (*cur == *buf)
+                return *cur;
+            cur++;
         }
+        buf++;
+        bufSz--;
     }
 
     return 0;
 }
 
 
+static int dump_stats(thread_ctx_t* ctx)
+{
+    char stats[1024];
+    uint32_t statsSz;
+    uint32_t txCount, rxCount, seq, peerSeq;
+
+    wolfSSH_GetStats(ctx->ssh, &txCount, &rxCount, &seq, &peerSeq);
+
+    sprintf(stats,
+            "Statistics for Thread #%u:\r\n"
+            "  txCount = %u\r\n  rxCount = %u\r\n"
+            "  seq = %u\r\n  peerSeq = %u\r\n",
+            ctx->id, txCount, rxCount, seq, peerSeq);
+    statsSz = (uint32_t)strlen(stats);
+
+    fprintf(stderr, "%s", stats);
+    return wolfSSH_stream_send(ctx->ssh, (uint8_t*)stats, statsSz);
+}
+
 static THREAD_RETURN CYASSL_THREAD server_worker(void* vArgs)
 {
-    WOLFSSH* ssh = (WOLFSSH*)vArgs;
-    SOCKET_T clientFd = wolfSSH_get_fd(ssh);
+    thread_ctx_t* threadCtx = (thread_ctx_t*)vArgs;
 
-    if (wolfSSH_accept(ssh) == WS_SUCCESS) {
+    if (wolfSSH_accept(threadCtx->ssh) == WS_SUCCESS) {
         uint8_t* buf = NULL;
         uint8_t* tmpBuf;
         int bufSz, backlogSz = 0, rxSz, txSz, stop = 0, txSum;
@@ -304,7 +328,7 @@ static THREAD_RETURN CYASSL_THREAD server_worker(void* vArgs)
                 buf = tmpBuf;
 
             if (!stop) {
-                rxSz = wolfSSH_stream_read(ssh,
+                rxSz = wolfSSH_stream_read(threadCtx->ssh,
                                            buf + backlogSz,
                                            EXAMPLE_BUFFER_SZ);
                 if (rxSz > 0) {
@@ -313,15 +337,25 @@ static THREAD_RETURN CYASSL_THREAD server_worker(void* vArgs)
                     txSz = 0;
 
                     while (backlogSz != txSum && txSz >= 0 && !stop) {
-                        txSz = wolfSSH_stream_send(ssh,
+                        txSz = wolfSSH_stream_send(threadCtx->ssh,
                                                    buf + txSum,
                                                    backlogSz - txSum);
 
                         if (txSz > 0) {
-                            if (find_char(0x03, buf + txSum, txSz))
-                                stop = 1;
-                            else
-                                txSum += txSz;
+                            uint8_t c;
+                            const uint8_t matches[] = { 0x03, 0x04, 0x05, 0x00 };
+
+                            c = find_char(matches, buf + txSum, txSz);
+                            switch (c) {
+                                case 0x03:
+                                    stop = 1;
+                                    break;
+                                case 0x05:
+                                    if (dump_stats(threadCtx) <= 0)
+                                        stop = 1;
+                                default:
+                                    txSum += txSz;
+                            }
                         }
                         else if (txSz != WS_REKEYING)
                             stop = 1;
@@ -338,8 +372,9 @@ static THREAD_RETURN CYASSL_THREAD server_worker(void* vArgs)
 
         free(buf);
     }
-    close(clientFd);
-    wolfSSH_free(ssh);
+    close(threadCtx->fd);
+    wolfSSH_free(threadCtx->ssh);
+    free(threadCtx);
 
     return 0;
 }
@@ -634,6 +669,7 @@ int main(void)
     PwMapList pwMapList;
     SOCKET_T listenFd = 0;
     uint32_t defaultHighwater = EXAMPLE_HIGHWATER_MARK;
+    uint32_t threadCount = 0;
 
     #ifdef DEBUG_WOLFSSH
         wolfSSH_Debugging_ON();
@@ -682,12 +718,22 @@ int main(void)
 
     tcp_bind(&listenFd, SERVER_PORT_NUMBER, 0);
 
+    if (listen(listenFd, 5) != 0)
+        err_sys("tcp listen failed");
+
     for (;;) {
         SOCKET_T      clientFd = 0;
         SOCKADDR_IN_T clientAddr;
         SOCKLEN_T     clientAddrSz = sizeof(clientAddr);
         THREAD_TYPE   thread;
         WOLFSSH*      ssh;
+        thread_ctx_t* threadCtx;
+
+        threadCtx = (thread_ctx_t*)malloc(sizeof(thread_ctx_t));
+        if (threadCtx == NULL) {
+            fprintf(stderr, "Couldn't allocate thread context data.\n");
+            exit(EXIT_FAILURE);
+        }
 
         ssh = wolfSSH_new(ctx);
         if (ssh == NULL) {
@@ -701,17 +747,18 @@ int main(void)
             wolfSSH_SetHighwater(ssh, defaultHighwater);
         }
 
-        if (listen(listenFd, 5) != 0)
-            err_sys("tcp listen failed");
-
         clientFd = accept(listenFd, (struct sockaddr*)&clientAddr,
                                                                  &clientAddrSz);
         if (clientFd == -1)
             err_sys("tcp accept failed");
 
         wolfSSH_set_fd(ssh, clientFd);
 
-        pthread_create(&thread, 0, server_worker, ssh);
+        threadCtx->ssh = ssh;
+        threadCtx->fd = clientFd;
+        threadCtx->id = threadCount++;
+
+        pthread_create(&thread, 0, server_worker, threadCtx);
         pthread_detach(thread);
     }
 

src/ssh.c

@@ -557,6 +557,32 @@ int wolfSSH_CTX_UsePrivateKey_buffer(WOLFSSH_CTX* ctx,
 }
 
 
+void wolfSSH_GetStats(WOLFSSH* ssh, uint32_t* txCount, uint32_t* rxCount,
+                      uint32_t* seq, uint32_t* peerSeq)
+{
+    uint32_t rTxCount = 0;
+    uint32_t rRxCount = 0;
+    uint32_t rSeq = 0;
+    uint32_t rPeerSeq = 0;
+
+    if (ssh != NULL) {
+        rTxCount = ssh->txCount;
+        rRxCount = ssh->rxCount;
+        rSeq = ssh->seq;
+        rPeerSeq = ssh->peerSeq;
+    }
+
+    if (txCount != NULL)
+        *txCount = rTxCount;
+    if (rxCount != NULL)
+        *rxCount = rRxCount;
+    if (seq != NULL)
+        *seq = rSeq;
+    if (peerSeq != NULL)
+        *peerSeq = rPeerSeq;
+}
+
+
 int wolfSSH_KDF(uint8_t hashId, uint8_t keyId,
                 uint8_t* key, uint32_t keySz,
                 const uint8_t* k, uint32_t kSz,

wolfssh/ssh.h

@@ -136,6 +136,9 @@ WOLFSSH_API int wolfSSH_channel_send(WOLFSSH_CHANNEL*, uint8_t*, uint32_t);
 WOLFSSH_API int wolfSSH_worker(WOLFSSH*);
 WOLFSSH_API int wolfSSH_TriggerKeyExchange(WOLFSSH*);
 
+WOLFSSH_API void wolfSSH_GetStats(WOLFSSH*,
+                                  uint32_t*, uint32_t*, uint32_t*, uint32_t*);
+
 WOLFSSH_API int wolfSSH_KDF(uint8_t, uint8_t, uint8_t*, uint32_t,
                 const uint8_t*, uint32_t, const uint8_t*, uint32_t,
                 const uint8_t*, uint32_t);