Repository: NVIDIA/NVFlare

Issue:
Server memory leak when using AWS Fargate
Hi, the server seems to have a memory leak, we can see the memory rise over an hour before crashing and starting again. The final crash complains about `Too many open files`

The server had no application running on it during this time.

Another thing to note is the constant high CPU usage.

The server is running in an AWS Fargate container with 2CPU's and 4096Mb RAM

CPU/Memory usage 24h
![image](https://user-images.githubusercontent.com/24326299/156677980-442b3ebd-0ecf-488a-a6cf-5253627077d7.png)

Crash logs
```
E0303 10:35:57.812868828      25 tcp_server_posix.cc:213]    Failed accept4: Too many open files
--
2022-03-03 10:38:27,069 - FederatedServer - INFO - starting insecure server at 0.0.0.0:8002
deployed FL server trainer.
2022-03-03 10:38:27,075 - FedAdminServer - INFO - Starting Admin Server 0.0.0.0 on Port 8003
2022-03-03 10:38:27,075 - root - INFO - Server started
/app/start_service.sh: line 10:     8 Killed                  python -u -m nvflare.private.fed.app.server.server_train -m /workspace -s fed_server.json --set secure_train=false config_folder=config host=0.0.0.0
2022-03-03 10:43:29,115 - ClientManager - INFO - Re-activate dead client:ip-10-0-246-19.ap-southeast-2.compute.internal with token: 8de2503e-cb43-4cfd-91a9-14f0b4ae98be
2022-03-03 10:43:52,231 - ClientManager - INFO - Re-activate dead client:ip-10-0-128-176.ap-southeast-2.compute.internal with token: 77d51ef9-3ac3-4026-b26f-907eaa7cbc04
E0303 11:34:48.528427869      25 tcp_server_posix.cc:213]    Failed accept4: Too many open files
2022-03-03 11:37:56,203 - FederatedServer - INFO - starting insecure server at 0.0.0.0:8002
deployed FL server trainer.
2022-03-03 11:37:56,209 - FedAdminServer - INFO - Starting Admin Server 0.0.0.0 on Port 8003
2022-03-03 11:37:56,209 - root - INFO - Server started
/app/start_service.sh: line 10:     8 Killed                  python -u -m nvflare.private.fed.app.server.server_train -m /workspace -s fed_server.json --set secure_train=false config_folder=config host=0.0.0.0
2022-03-03 11:39:37,262 - ClientManager - INFO - Re-activate dead client:ip-10-0-128-176.ap-southeast-2.compute.internal with token: 77d51ef9-3ac3-4026-b26f-907eaa7cbc04
2022-03-03 11:39:40,542 - ClientManager - INFO - Re-activate dead client:ip-10-0-246-19.ap-southeast-2.compute.internal with token: 8de2503e-cb43-4cfd-91a9-14f0b4ae98be
E0303 12:34:14.240603486      24 tcp_server_posix.cc:213]    Failed accept4: Too many open files
2022-03-03 12:37:04,335 - FederatedServer - INFO - starting insecure server at 0.0.0.0:8002
deployed FL server trainer.
2022-03-03 12:37:04,341 - FedAdminServer - INFO - Starting Admin Server 0.0.0.0 on Port 8003
2022-03-03 12:37:04,368 - root - INFO - Server started
2022-03-03 12:42:02,992 - ClientManager - INFO - Re-activate dead client:ip-10-0-128-176.ap-southeast-2.compute.internal with token: 77d51ef9-3ac3-4026-b26f-907eaa7cbc04
2022-03-03 12:42:05,714 - ClientManager - INFO - Re-activate dead client:ip-10-0-246-19.ap-southeast-2.compute.internal with token: 8de2503e-cb43-4cfd-91a9-14f0b4ae98be
E0303 13:33:17.078870193      24 tcp_server_posix.cc:213]    Failed accept4: Too many open files
2022-03-03 13:36:20,660 - FederatedServer - INFO - starting insecure server at 0.0.0.0:8002
deployed FL server trainer.
2022-03-03 13:36:20,666 - FedAdminServer - INFO - Starting Admin Server 0.0.0.0 on Port 8003
2022-03-03 13:36:20,666 - root - INFO - Server started
2022-03-03 13:41:27,090 - ClientManager - INFO - Re-activate dead client:ip-10-0-128-176.ap-southeast-2.compute.internal with token: 77d51ef9-3ac3-4026-b26f-907eaa7cbc04
2022-03-03 13:41:34,427 - ClientManager - INFO - Re-activate dead client:ip-10-0-246-19.ap-southeast-2.compute.internal with token: 8de2503e-cb43-4cfd-91a9-14f0b4ae98be
E0303 14:32:29.195437902      27 tcp_server_posix.cc:213]    Failed accept4: Too many open files
2022-03-03 14:35:29,825 - FederatedServer - INFO - starting insecure server at 0.0.0.0:8002
deployed FL server trainer.
2022-03-03 14:35:29,831 - FedAdminServer - INFO - Starting Admin Server 0.0.0.0 on Port 8003
2022-03-03 14:35:29,831 - root - INFO - Server started
/app/start_service.sh: line 10:     8 Killed                  python -u -m nvflare.private.fed.app.server.server_train -m /workspace -s fed_server.json --set secure_train=false config_folder=config host=0.0.0.0
2022-03-03 14:37:50,369 - ClientManager - INFO - Re-activate dead client:ip-10-0-128-176.ap-southeast-2.compute.internal with token: 77d51ef9-3ac3-4026-b26f-907eaa7cbc04
2022-03-03 14:37:58,754 - ClientManager - INFO - Re-activate dead client:ip-10-0-246-19.ap-southeast-2.compute.internal with token: 8de2503e-cb43-4cfd-91a9-14f0b4ae98be
E0303 15:31:39.658074612      26 tcp_server_posix.cc:213]    Failed accept4: Too many open files
2022-03-03 15:34:46,429 - FederatedServer - INFO - starting insecure server at 0.0.0.0:8002
deployed FL server trainer.
2022-03-03 15:34:46,435 - FedAdminServer - INFO - Starting Admin Server 0.0.0.0 on Port 8003
2022-03-03 15:34:46,435 - root - INFO - Server started
/app/start_service.sh: line 10:     8 Killed                  python -u -m nvflare.private.fed.app.server.server_train -m /workspace -s fed_server.json --set secure_train=false config_folder=config host=0.0.0.0
2022-03-03 15:38:24,369 - ClientManager - INFO - Re-activate dead client:ip-10-0-246-19.ap-southeast-2.compute.internal with token: 8de2503e-cb43-4cfd-91a9-14f0b4ae98be
2022-03-03 15:39:12,769 - ClientManager - INFO - Re-activate dead client:ip-10-0-128-176.ap-southeast-2.compute.internal with token: 77d51ef9-3ac3-4026-b26f-907eaa7cbc04
E0303 16:30:55.183641326      26 tcp_server_posix.cc:213]    Failed accept4: Too many open files
2022-03-03 16:34:01,946 - FederatedServer - INFO - starting insecure server at 0.0.0.0:8002
deployed FL server trainer.
2022-03-03 16:34:01,952 - FedAdminServer - INFO - Starting Admin Server 0.0.0.0 on Port 8003
2022-03-03 16:34:01,952 - root - INFO - Server started
```


Discussion:
Hi @Nintorac ,

Can you elaborate on the steps of how you start the whole system?
So you just ran ./startup/start.sh on the server side and ran ./startup/start.sh on the client side?

Did you start any applications?
What other stuffs (other than NVFlare codes) did you run on these systems?
This is running as a container under AWS Fargate runtime, so nothing running alongside

Launching with the following script (couldn't get `start.sh` to work in docker compose, could probably return to using them with Fargate)

```
#!/bin/bash
export PYTHONPATH=/local/custom:$PYTHONPATH

if [[ -z "${SERVER}" ]]; then
  # Server env not set, running this as the server
  python -u -m nvflare.private.fed.app.server.server_train -m /workspace -s fed_server.json --set secure_train=false config_folder=config host=0.0.0.0
else
  python -u -m nvflare.private.fed.app.client.client_train -m /workspace -s fed_client.json --set secure_train=false uid=$(hostname) config_folder=config host=${SERVER}
fi
```

No applications running
Got it.
 To me this looks like an issue that is specific to running NVFlare on AWS Fargate.

I had run the server using AWS ECS and had no issues.
But I've never tried AWS Fargate though.
![image](https://user-images.githubusercontent.com/24326299/157133735-1e998191-9407-4822-8d6b-89d68375447c.png)

Don't think so. Here it is running on a c4.xlarge. It has gained 20% in 30 mins. The task definition still has the same requirements previously mentioned

There are no clients connected for this test. CPU usage is also still maxed
I've just started an c4.xlarge EC2 instance and running the server.
There are no clients connected.

Seeing from the HTOP command, the total CPU mem available is 7.28G.
The server process consume 0.9% MEM and shown 0.0% CPU (means a small amount).

I think the issue is AWS Fargate related.
Hmmm, I also cannot reproduce locally. 

It is happening with EC2 as well so definitely not Fargate related.

One potential difference I can think of is there will be constant stream of TCP health checks being run on both the admin and rpc endpoints...

After digging along that idea I have found that I can initiate the 100% CPU and RAM accumulation state by running this command which is effectively a TCP health check (AWS is performing these against the server automatically)

```
curl -v telnet://localhost:8003
```
For me it will hang, then when I press `ctrl + c`  the server goes into the degenerative state
I can't reproduce this locally or on EC2.

I see in your crash log `Too many open files`

Maybe you can try increase the `ulimit` inside your instance and see how that goes.
The issue is in [this function](https://github.com/NVIDIA/NVFlare/blob/6e6cb3657decdc4935c6d750c828d8b7ca18801a/nvflare/fuel/hci/conn.py#L35-L47).

When performing a TCP ping against the endpoint there is no data forthcoming and therefore no end of transmission packet. This results in an infinite loop, and the thread never dies. A seperate thread is launched for each TCP ping, none of them quit and the end result is that you fill up memory or open too many files.

The immediate fix for the infinite loop is something like
```

def is_socket_closed(sock: socket.socket) -> bool:
    """
    This solution thanks to Micheal Petrov

    https://stackoverflow.com/questions/48024720/python-how-to-check-if-socket-is-still-connected
    """
    try:
        # this will try to read bytes without blocking and also without removing them from buffer (peek only)
        data = sock.recv(16, socket.MSG_DONTWAIT | socket.MSG_PEEK)
        if len(data) == 0:
            return True
    except BlockingIOError:
        return False  # socket is open and reading from it would block
    except ConnectionResetError:
        return True  # socket was closed for some other reason
    except Exception as e:
        logger.exception("unexpected exception when checking if a socket is closed")
        return False
    return False

def receive_til_end(sock, end=ALL_END):
    total_data = []
    
    while True:
        data = str(sock.recv(1024), "utf-8")
        if end in data:
            total_data.append(data[: data.find(end)])
            break

        total_data.append(data)

    result = "".join(total_data)
    return result.replace(LINE_END, "")
```

Adding this in solves the infinite loop but results in some other exceptions that may need to be handled
Overall this function as is introduces several denial of service vulnerabilities 

a)  a set of adverserial clients can hold open as many connections as they want for as long as they want without closing which will eventually lead to `Too many open files`

b) an adverserial client is not limited by the amount of data it can send in a single request as such can overload the server quite easily, 
poc of this 
```
cat /dev/zero | base64 | telnet localhost 8003
```
explanation of POC
`cat /dev/zero` recieves an infinite stream of 0 bytes
`base64` base64 encodes the stream so that the server doesn't crash when utf-8 decoding
`telnet localhost 8003` pipes the infinite stream of base64 encoded 0's into the function `receive_til_end`

Warning don't run this. I get a hard system lockup and had to restart my machine

c) if an adversarial client does not send an end of transmission byte they can force the server to start an unbound number of threads without needing to maintain the connection (the issue I was experiencing) 

Hi @Nintorac ,

Thank you so much for discovering the root cause of this.
I guess my Amazon ECS server does not have the health checking / scanning on this port.

Thanks also for your POC steps.
Yes, this is a serious issue.

@nvidianz will fix this ASAP.

Resolving patch:
diff --git a/nvflare/fuel/hci/conn.py b/nvflare/fuel/hci/conn.py
--- a/nvflare/fuel/hci/conn.py
+++ b/nvflare/fuel/hci/conn.py
@@ -27,16 +27,20 @@
 
 LINE_END = "\x03"  # Indicates the end of the line (end of text)
 ALL_END = "\x04"  # Marks the end of a complete transmission (End of Transmission)
-
-
 MAX_MSG_SIZE = 1024
+MAX_DATA_SIZE = 512 * 1024 * 1024
+MAX_IDLE_TIME = 10
 
 
 def receive_til_end(sock, end=ALL_END):
     total_data = []
-
+    data_size = 0
+    sock.settimeout(MAX_IDLE_TIME)
     while True:
         data = str(sock.recv(1024), "utf-8")
+        data_size += len(data)
+        if data_size > MAX_DATA_SIZE:
+            raise BufferError(f"Data size exceeds limit ({MAX_DATA_SIZE} bytes)")
         if end in data:
             total_data.append(data[: data.find(end)])
             break
diff --git a/nvflare/fuel/hci/server/hci.py b/nvflare/fuel/hci/server/hci.py
--- a/nvflare/fuel/hci/server/hci.py
+++ b/nvflare/fuel/hci/server/hci.py
@@ -16,14 +16,14 @@
 import socketserver
 import ssl
 import threading
-import traceback
 
 from nvflare.fuel.hci.conn import Connection, receive_til_end
 from nvflare.fuel.hci.proto import validate_proto
 from nvflare.fuel.hci.security import get_certificate_common_name
-
 from .reg import ServerCommandRegister
 
+MAX_ADMIN_CONNECTIONS = 128
+
 
 class _MsgHandler(socketserver.BaseRequestHandler):
     """Message handler.
@@ -32,8 +32,23 @@ class _MsgHandler(socketserver.BaseRequestHandler):
     ServerCommandRegister.
     """
 
+    connections = 0
+    lock = threading.Lock()
+
+    def __init__(self, request, client_address, server):
+        # handle() is called in the constructor so logger must be initialized first
+        self.logger = logging.getLogger(self.__class__.__name__)
+        super().__init__(request, client_address, server)
+
     def handle(self):
         try:
+            with _MsgHandler.lock:
+                _MsgHandler.connections += 1
+
+            self.logger.debug(f"Concurrent admin connections: {_MsgHandler.connections}")
+            if _MsgHandler.connections > MAX_ADMIN_CONNECTIONS:
+                raise ConnectionRefusedError(f"Admin connection limit ({MAX_ADMIN_CONNECTIONS}) reached")
+
             conn = Connection(self.request, self.server)
 
             if self.server.use_ssl:
@@ -68,8 +83,13 @@ def handle(self):
 
             if not conn.ended:
                 conn.close()
-        except BaseException:
-            traceback.print_exc()
+        except BaseException as exc:
+            self.logger.error(f"Admin connection terminated due to exception: {str(exc)}")
+            if self.logger.getEffectiveLevel() <= logging.DEBUG:
+                self.logger.exception("Admin connection error")
+        finally:
+            with _MsgHandler.lock:
+                _MsgHandler.connections -= 1
 
 
 def initialize_hci():
@@ -121,7 +141,7 @@ def __init__(
             ctx.load_verify_locations(ca_cert)
             ctx.load_cert_chain(certfile=server_cert, keyfile=server_key)
 
-            # replace the socket with an ssl version of itself
+            # replace the socket with an SSL version of itself
             self.socket = ctx.wrap_socket(self.socket, server_side=True)
             self.use_ssl = True
 


Test patch: