From bd5a88e95a9e34c55841ebd2e27802eb688816ca Mon Sep 17 00:00:00 2001
From: David Fort <contact@hardening-consulting.com>
Date: Tue, 22 Oct 2019 23:35:18 +0200
Subject: [PATCH] rdp: count in/out bytes and packets

---
 include/freerdp/freerdp.h |  1 +
 libfreerdp/core/rdp.c     | 22 ++++++++++++++++++++++
 libfreerdp/core/rdp.h     |  4 ++++
 3 files changed, 27 insertions(+)

diff --git a/include/freerdp/freerdp.h b/include/freerdp/freerdp.h
index 8f81d0791..6c43e2260 100644
--- a/include/freerdp/freerdp.h
+++ b/include/freerdp/freerdp.h
@@ -454,6 +454,7 @@ FREERDP_API int freerdp_message_queue_process_pending_messages(
 FREERDP_API UINT32 freerdp_error_info(freerdp* instance);
 FREERDP_API void freerdp_set_error_info(rdpRdp* rdp, UINT32 error);
 FREERDP_API BOOL freerdp_send_error_info(rdpRdp* rdp);
+FREERDP_API BOOL freerdp_get_stats(rdpRdp* rdp, UINT64 *inBytes, UINT64 *outBytes, UINT64 *inPackets, UINT64 *outPackets);
 
 FREERDP_API void freerdp_get_version(int* major, int* minor, int* revision);
 FREERDP_API const char* freerdp_get_version_string(void);
diff --git a/libfreerdp/core/rdp.c b/libfreerdp/core/rdp.c
index f0a47fc6f..f0ad8294a 100644
--- a/libfreerdp/core/rdp.c
+++ b/libfreerdp/core/rdp.c
@@ -654,6 +654,8 @@ BOOL rdp_send_data_pdu(rdpRdp* rdp, wStream* s, BYTE type, UINT16 channel_id)
 	WLog_DBG(TAG, "%s: sending data (type=0x%x size=%"PRIuz" channelId=%"PRIu16")", __FUNCTION__,
 	         type, Stream_Length(s), channel_id);
 
+	rdp->outBytes += Stream_GetRemainingLength(s);
+	rdp->outPackets++;
 	if (transport_write(rdp->transport, s) < 0)
 		goto fail;
 
@@ -1194,6 +1196,7 @@ static int rdp_recv_tpkt_pdu(rdpRdp* rdp, wStream* s)
 		return -1;
 	}
 
+	rdp->inPackets++;
 	if (freerdp_shall_disconnect(rdp->instance))
 		return 0;
 
@@ -1354,6 +1357,8 @@ int rdp_recv_callback(rdpTransport* transport, wStream* s, void* extra)
 	int status = 0;
 	rdpRdp* rdp = (rdpRdp*) extra;
 
+	rdp->inBytes += Stream_GetRemainingLength(s);
+
 	/*
 	 * At any point in the connection sequence between when all
 	 * MCS channels have been joined and when the RDP connection
@@ -1586,6 +1591,23 @@ int rdp_check_fds(rdpRdp* rdp)
 	return status;
 }
 
+BOOL freerdp_get_stats(rdpRdp* rdp, UINT64 *inBytes, UINT64 *outBytes, UINT64 *inPackets, UINT64 *outPackets)
+{
+	if (!rdp)
+		return FALSE;
+
+	if (inBytes)
+		*inBytes = rdp->inBytes;
+	if (outBytes)
+		*outBytes = rdp->outBytes;
+	if (inPackets)
+		*inPackets = rdp->inPackets;
+	if (outPackets)
+		*outPackets = rdp->outPackets;
+
+	return TRUE;
+}
+
 /**
  * Instantiate new RDP module.
  * @return new RDP module
diff --git a/libfreerdp/core/rdp.h b/libfreerdp/core/rdp.h
index 31011f410..88103f033 100644
--- a/libfreerdp/core/rdp.h
+++ b/libfreerdp/core/rdp.h
@@ -175,6 +175,10 @@ struct rdp_rdp
 	BOOL resendFocus;
 	BOOL deactivation_reactivation;
 	BOOL AwaitCapabilities;
+	UINT64 inBytes;
+	UINT64 inPackets;
+	UINT64 outBytes;
+	UINT64 outPackets;
 };
 
 FREERDP_LOCAL BOOL rdp_read_security_header(wStream* s, UINT16* flags, UINT16* length);