Avoid leaking file descriptors for terminated sessions. fix #392

pyrdp/mitm/TCPMITM.py

@@ -99,6 +99,7 @@ def onClientDisconnection(self, reason):
         if self.recorder.recordFilename:
             self.statCounter.logReport(self.log, {"replayFilename":
                                                   self.recorder.recordFilename})
+            self.recorder.finalize()
         else:
             self.statCounter.logReport(self.log)
 

pyrdp/recording/recorder.py

@@ -35,6 +35,7 @@ def __init__(self, transports: List[LayerChainItem]):
         }
 
         self.topLayers = []
+        self.openFileLayers = []
         self.recordFilename = None
 
         for transport in transports:
@@ -52,6 +53,9 @@ def addTransport(self, transportLayer: LayerChainItem):
         player.setPrevious(transportLayer)
         self.topLayers.append(player)
 
+        if isinstance(transportLayer, FileLayer):
+            self.openFileLayers.append(transportLayer)
+
     def setParser(self, messageType: PlayerPDUType, parser: Parser):
         """
         Set the parser to use for a given message type.
@@ -81,6 +85,11 @@ def record(self, pdu: Optional[PDU], messageType: PlayerPDUType):
     def getCurrentTimeStamp(self) -> int:
         return PlayerLayer.timeStampFunction()
 
+    def finalize(self):
+        """When the recording is finished"""
+        for layer in self.openFileLayers:
+            layer.close()
+
 
 class FileLayer(LayerChainItem):
     """
@@ -114,3 +123,7 @@ def sendBytes(self, data: bytes):
             self.fd.write(data)
         else:
             log.error("Recording file handle closed, cannot write message: %(message)s", {"message": data})
+
+    def close(self):
+        if self.fd:
+            self.fd.close()