From ab00d90f0342547e06083eb545a088752e35502d Mon Sep 17 00:00:00 2001
From: kubistika <kmizrachi18@gmail.com>
Date: Mon, 15 Jul 2019 10:55:57 +0300
Subject: [PATCH] server: proxy: implement NLA fallback

---
 server/proxy/pf_client.c  | 35 ++++++++++++++++++++++++++++++++---
 server/proxy/pf_context.h |  9 +++++++++
 2 files changed, 41 insertions(+), 3 deletions(-)

diff --git a/server/proxy/pf_client.c b/server/proxy/pf_client.c
index a313105..31ea8aa 100644
--- a/server/proxy/pf_client.c
+++ b/server/proxy/pf_client.c
@@ -208,7 +208,14 @@ static void pf_client_post_disconnect(freerdp* instance)
 	PubSub_UnsubscribeErrorInfo(instance->context->pubSub, pf_OnErrorInfo);
 	gdi_free(instance);
 
-	SetEvent(pdata->connectionClosed);
+	if (!context->during_connect_process)
+	{
+		/* proxy's client failed to connect, and now it's trying to connect without NLA, no need to shutdown
+		 * the connection between proxy's server and the original client.
+		 */
+		SetEvent(pdata->connectionClosed);
+	}
+
 	/* It's important to avoid calling `freerdp_peer_context_free` and `freerdp_peer_free` here,
 	 * in order to avoid double-free. Those objects will be freed by the server when needed.
 	 */
@@ -222,16 +229,38 @@ static void pf_client_post_disconnect(freerdp* instance)
 static DWORD WINAPI pf_client_thread_proc(LPVOID arg)
 {
 	freerdp* instance = (freerdp*)arg;
+	pClientContext* pc = (pClientContext*)instance->context;
 	DWORD nCount;
 	DWORD status;
 	HANDLE handles[64];
 
+	pc->during_connect_process = TRUE;
 	if (!freerdp_connect(instance))
 	{
-		WLog_ERR(TAG, "connection failure");
-		return 0;
+		if (instance->settings->NlaSecurity)
+		{
+			WLog_ERR(TAG, "freerdp_connect() failed, trying to connect without NLA");
+			/* disable NLA, enable TLS */
+			instance->settings->NlaSecurity = FALSE;
+			instance->settings->RdpSecurity = TRUE;
+			instance->settings->TlsSecurity = TRUE;
+
+			pc->during_connect_process = FALSE;
+			if (!freerdp_connect(instance))
+			{
+				WLog_ERR(TAG, "connection failure");
+				return 0;
+			}
+		}
+		else
+		{
+			WLog_ERR(TAG, "connection failure");
+			return 0;
+		}
 	}
 
+	pc->during_connect_process = FALSE;
+
 	while (!freerdp_shall_disconnect(instance))
 	{
 		nCount = freerdp_get_event_handles(instance->context, &handles[0], 64);
diff --git a/server/proxy/pf_context.h b/server/proxy/pf_context.h
index 2cd7fe6..b8ef9c8 100644
--- a/server/proxy/pf_context.h
+++ b/server/proxy/pf_context.h
@@ -68,6 +68,15 @@ struct p_client_context
 	RdpeiClientContext* rdpei;
 	RdpgfxClientContext* gfx;
 	DispClientContext* disp;
+
+	/*
+	 * Used for NLA fallback feature, to check if the server should close the connection.
+	 * When it is set to TRUE, proxy's client knows it shouldn't signal the server thread to closed
+	 * the connection when pf_client_post_disconnect is called, because it is trying to connect reconnect without NLA.
+	 * It must be set to TRUE before the first try, and to FALSE after the connection fully established,
+	 * to ensure graceful shutdown of the connection when it will be closed.
+	 */
+	BOOL during_connect_process;
 };
 typedef struct p_client_context pClientContext;
 
-- 
2.7.4