[core,gateway] fix various issues with arm transport

* fix possible leaks
* split big functions
* fix missing NULL checks

libfreerdp/core/gateway/arm.c

@@ -305,89 +305,81 @@ BOOL arm_resolve_endpoint(rdpContext* context, DWORD timeout)
 	}
 
 	rdpArm* arm = (rdpArm*)calloc(1, sizeof(rdpArm));
-	char* message = NULL;
-
-	if (arm)
-	{
-		arm->context = context;
-		arm->settings = arm->context->settings;
-		if (!arm->settings)
-			goto arm_error;
-
-		arm->tls = freerdp_tls_new(arm->settings);
-		HttpResponse* response = NULL;
-		long StatusCode;
-
-		if (!arm->tls)
-			goto arm_error;
-
-		arm->http = http_context_new();
-
-		if (!arm->http)
-			goto arm_error;
-
-		if (!http_context_set_uri(arm->http, "/api/arm/v2/connections/") ||
-		    !http_context_set_accept(arm->http, "application/json") ||
-		    !http_context_set_cache_control(arm->http, "no-cache") ||
-		    !http_context_set_pragma(arm->http, "no-cache") ||
-		    !http_context_set_connection(arm->http, "Keep-Alive") ||
-		    !http_context_set_user_agent(arm->http, "FreeRDP/3.0") ||
-		    !http_context_set_x_ms_user_agent(arm->http, "FreeRDP/3.0") ||
-		    !http_context_set_host(
-		        arm->http, freerdp_settings_get_string(arm->settings, FreeRDP_GatewayHostname)))
-			goto arm_error;
-
-		if (!arm_tls_connect(arm, arm->tls, timeout))
-			goto arm_error;
-
-		message = arm_create_request_json(arm);
-		if (!message)
-			goto arm_error;
-
-		if (!arm_send_http_request(arm, arm->tls, "POST", "application/json", message,
-		                           strlen(message)))
-			goto arm_error;
-
-		free(message);
-		message = NULL;
-
-		response = http_response_recv(arm->tls, TRUE);
-		if (!response)
-			goto arm_error;
-
-		StatusCode = http_response_get_status_code(response);
-		if (StatusCode == HTTP_STATUS_OK)
-		{
-			message = calloc(http_response_get_body_length(response) + 1, sizeof(char));
-			if (!message)
-				goto arm_error;
-
-			strncpy(message, (const char*)http_response_get_body(response),
-			        http_response_get_body_length(response));
-
-			WLog_DBG(TAG, "Got HTTP Response data: %s", message);
-			if (!arm_fill_gateway_parameters(arm, message))
-				goto arm_error;
-
-			free(message);
-			message = NULL;
-		}
-		else
-		{
-			char buffer[64] = { 0 };
-			WLog_ERR(TAG, "Unexpected HTTP status: %s",
-			         http_status_string_format(StatusCode, buffer, ARRAYSIZE(buffer)));
-			goto arm_error;
-		}
-	}
-	else
+	if (!arm)
 		return FALSE;
 
-	arm_free(arm);
-	return TRUE;
+	char* message = NULL;
+	BOOL rc = FALSE;
+
+	arm->context = context;
+	arm->settings = arm->context->settings;
+	if (!arm->settings)
+		goto arm_error;
+
+	arm->tls = freerdp_tls_new(arm->settings);
+	HttpResponse* response = NULL;
+	long StatusCode;
+
+	if (!arm->tls)
+		goto arm_error;
+
+	arm->http = http_context_new();
+
+	if (!arm->http)
+		goto arm_error;
+
+	if (!http_context_set_uri(arm->http, "/api/arm/v2/connections/") ||
+	    !http_context_set_accept(arm->http, "application/json") ||
+	    !http_context_set_cache_control(arm->http, "no-cache") ||
+	    !http_context_set_pragma(arm->http, "no-cache") ||
+	    !http_context_set_connection(arm->http, "Keep-Alive") ||
+	    !http_context_set_user_agent(arm->http, "FreeRDP/3.0") ||
+	    !http_context_set_x_ms_user_agent(arm->http, "FreeRDP/3.0") ||
+	    !http_context_set_host(arm->http,
+	                           freerdp_settings_get_string(arm->settings, FreeRDP_GatewayHostname)))
+		goto arm_error;
+
+	if (!arm_tls_connect(arm, arm->tls, timeout))
+		goto arm_error;
+
+	message = arm_create_request_json(arm);
+	if (!message)
+		goto arm_error;
+
+	if (!arm_send_http_request(arm, arm->tls, "POST", "application/json", message, strlen(message)))
+		goto arm_error;
+
+	response = http_response_recv(arm->tls, TRUE);
+	if (!response)
+		goto arm_error;
+
+	StatusCode = http_response_get_status_code(response);
+	if (StatusCode == HTTP_STATUS_OK)
+	{
+		char* msg = calloc(http_response_get_body_length(response) + 1, sizeof(char));
+		if (!msg)
+			goto arm_error;
+
+		memcpy(msg, http_response_get_body(response), http_response_get_body_length(response));
+
+		WLog_DBG(TAG, "Got HTTP Response data: %s", msg);
+		const BOOL res = arm_fill_gateway_parameters(arm, msg);
+		free(msg);
+		if (!res)
+			goto arm_error;
+	}
+	else
+	{
+		char buffer[64] = { 0 };
+		WLog_ERR(TAG, "Unexpected HTTP status: %s",
+		         http_status_string_format(StatusCode, buffer, ARRAYSIZE(buffer)));
+		goto arm_error;
+	}
+
+	rc = TRUE;
 arm_error:
 	arm_free(arm);
 	free(message);
-	return FALSE;
+	return rc;
 #endif
 }

libfreerdp/core/gateway/arm.h

@@ -20,6 +20,8 @@
 #ifndef FREERDP_LIB_CORE_GATEWAY_ARM_H
 #define FREERDP_LIB_CORE_GATEWAY_ARM_H
 
+#include <freerdp/api.h>
+
 FREERDP_LOCAL BOOL arm_resolve_endpoint(rdpContext* context, DWORD timeout);
 
 #endif /* FREERDP_LIB_CORE_GATEWAY_ARM_H */

libfreerdp/core/gateway/http.c

@@ -142,15 +142,28 @@ static void* copy_string(const void* ptr)
 HttpContext* http_context_new(void)
 {
 	HttpContext* context = (HttpContext*)calloc(1, sizeof(HttpContext));
-	if (context)
-	{
-		context->cookies = ListDictionary_New(FALSE);
-		ListDictionary_KeyObject(context->cookies)->fnObjectFree = free;
-		ListDictionary_ValueObject(context->cookies)->fnObjectFree = free;
-		ListDictionary_KeyObject(context->cookies)->fnObjectNew = copy_string;
-		ListDictionary_ValueObject(context->cookies)->fnObjectNew = copy_string;
-	}
+	if (!context)
+		return NULL;
+
+	context->cookies = ListDictionary_New(FALSE);
+	if (!context->cookies)
+		goto fail;
+
+	wObject* key = ListDictionary_KeyObject(context->cookies);
+	wObject* value = ListDictionary_ValueObject(context->cookies);
+	if (!key || !value)
+		goto fail;
+
+	key->fnObjectFree = free;
+	key->fnObjectNew = copy_string;
+	value->fnObjectFree = free;
+	value->fnObjectNew = copy_string;
+
 	return context;
+
+fail:
+	http_context_free(context);
+	return NULL;
 }
 
 BOOL http_context_set_method(HttpContext* context, const char* Method)
@@ -1037,7 +1050,7 @@ int http_chuncked_read(BIO* bio, BYTE* pBuffer, size_t size,
 			break;
 			case ChunkStateFooter:
 			{
-				char _dummy[2];
+				char _dummy[2] = { 0 };
 				WINPR_ASSERT(encodingContext->nextOffset == 0);
 				WINPR_ASSERT(encodingContext->headerFooterPos < 2);
 				ERR_clear_error();
@@ -1339,6 +1352,18 @@ const BYTE* http_response_get_body(HttpResponse* response)
 	return response->BodyContent;
 }
 
+static BOOL set_compare(wListDictionary* dict)
+{
+	WINPR_ASSERT(dict);
+	wObject* key = ListDictionary_KeyObject(dict);
+	wObject* value = ListDictionary_KeyObject(dict);
+	if (!key || !value)
+		return FALSE;
+	key->fnObjectEquals = strings_equals_nocase;
+	value->fnObjectEquals = strings_equals_nocase;
+	return TRUE;
+}
+
 HttpResponse* http_response_new(void)
 {
 	HttpResponse* response = (HttpResponse*)calloc(1, sizeof(HttpResponse));
@@ -1351,20 +1376,22 @@ HttpResponse* http_response_new(void)
 	if (!response->Authenticates)
 		goto fail;
 
+	if (!set_compare(response->Authenticates))
+		goto fail;
+
 	response->SetCookie = ListDictionary_New(FALSE);
 
 	if (!response->SetCookie)
 		goto fail;
 
+	if (!set_compare(response->SetCookie))
+		goto fail;
+
 	response->data = Stream_New(NULL, 2048);
 
 	if (!response->data)
 		goto fail;
 
-	ListDictionary_KeyObject(response->Authenticates)->fnObjectEquals = strings_equals_nocase;
-	ListDictionary_KeyObject(response->SetCookie)->fnObjectEquals = strings_equals_nocase;
-	ListDictionary_ValueObject(response->Authenticates)->fnObjectEquals = strings_equals_nocase;
-
 	response->TransferEncoding = TransferEncodingIdentity;
 	return response;
 fail:

libfreerdp/core/gateway/wst.c

@@ -101,10 +101,6 @@ static BOOL wst_auth_init(rdpWst* wst, rdpTls* tls, TCHAR* authPkg)
 	SEC_WINNT_AUTH_IDENTITY identity = { 0 };
 	int rc;
 
-	wst->auth = credssp_auth_new(context);
-	if (!wst->auth)
-		return FALSE;
-
 	if (!credssp_auth_init(wst->auth, authPkg, tls->Bindings))
 		return FALSE;
 
@@ -351,6 +347,102 @@ static BOOL wst_send_http_request(rdpWst* wst, rdpTls* tls)
 	return (status >= 0);
 }
 
+static BOOL wst_handle_ok_or_forbidden(rdpWst* wst, HttpResponse** ppresponse, DWORD timeout,
+                                       long* pStatusCode)
+{
+	WINPR_ASSERT(wst);
+	WINPR_ASSERT(ppresponse);
+	WINPR_ASSERT(*ppresponse);
+	WINPR_ASSERT(pStatusCode);
+
+	/* AVD returns a 403 response with a ARRAffinity cookie set. retry with that cookie */
+	const char* affinity = http_response_get_setcookie(*ppresponse, "ARRAffinity");
+	if (affinity && freerdp_settings_get_bool(wst->settings, FreeRDP_GatewayArmTransport))
+	{
+		WLog_DBG(TAG, "Got Affinity cookie %s", affinity);
+		http_context_set_cookie(wst->http, "ARRAffinity", affinity);
+		http_response_free(*ppresponse);
+		/* Terminate this connection and make a new one with the Loadbalancing Cookie */
+		int fd = BIO_get_fd(wst->tls->bio, NULL);
+		if (fd >= 0)
+			closesocket((SOCKET)fd);
+		freerdp_tls_free(wst->tls);
+
+		wst->tls = freerdp_tls_new(wst->settings);
+		if (!wst_tls_connect(wst, wst->tls, timeout))
+			return FALSE;
+
+		if (freerdp_settings_get_string(wst->settings, FreeRDP_GatewayHttpExtAuthBearer) &&
+		    freerdp_settings_get_bool(wst->settings, FreeRDP_GatewayArmTransport))
+		{
+			char* urlWithAuth = NULL;
+			size_t urlLen = 0;
+			char firstParam = (strchr(wst->gwpath, '?') > 0 ? '&' : '?');
+			winpr_asprintf(
+			    &urlWithAuth, &urlLen, arm_query_param, wst->gwpath, firstParam,
+			    freerdp_settings_get_string(wst->settings, FreeRDP_GatewayHttpExtAuthBearer));
+			if (!urlWithAuth)
+				return FALSE;
+			free(wst->gwpath);
+			wst->gwpath = urlWithAuth;
+			http_context_set_uri(wst->http, wst->gwpath);
+			http_context_enable_websocket_upgrade(wst->http, TRUE);
+		}
+
+		if (!wst_send_http_request(wst, wst->tls))
+			return FALSE;
+		*ppresponse = http_response_recv(wst->tls, TRUE);
+		if (!*ppresponse)
+			return FALSE;
+
+		*pStatusCode = http_response_get_status_code(*ppresponse);
+	}
+
+	return TRUE;
+}
+
+static BOOL wst_handle_denied(rdpWst* wst, HttpResponse** ppresponse, long* pStatusCode)
+{
+	WINPR_ASSERT(wst);
+	WINPR_ASSERT(ppresponse);
+	WINPR_ASSERT(*ppresponse);
+	WINPR_ASSERT(pStatusCode);
+
+	BOOL success = FALSE;
+
+	if (freerdp_settings_get_string(wst->settings, FreeRDP_GatewayHttpExtAuthBearer))
+		return FALSE;
+
+	if (!wst_auth_init(wst, wst->tls, AUTH_PKG))
+		return FALSE;
+	if (!wst_send_http_request(wst, wst->tls))
+		return FALSE;
+
+	http_response_free(*ppresponse);
+	*ppresponse = http_response_recv(wst->tls, TRUE);
+	if (!*ppresponse)
+		return FALSE;
+
+	while (!credssp_auth_is_complete(wst->auth))
+	{
+		if (!wst_recv_auth_token(wst->auth, *ppresponse))
+			return FALSE;
+
+		if (credssp_auth_have_output_token(wst->auth))
+		{
+			if (!wst_send_http_request(wst, wst->tls))
+				return FALSE;
+
+			http_response_free(*ppresponse);
+			*ppresponse = http_response_recv(wst->tls, TRUE);
+			if (!*ppresponse)
+				return FALSE;
+		}
+	}
+	*pStatusCode = http_response_get_status_code(*ppresponse);
+	return success;
+}
+
 BOOL wst_connect(rdpWst* wst, DWORD timeout)
 {
 	HttpResponse* response = NULL;
@@ -378,106 +470,22 @@ BOOL wst_connect(rdpWst* wst, DWORD timeout)
 	}
 
 	StatusCode = http_response_get_status_code(response);
-
+	BOOL success = TRUE;
 	switch (StatusCode)
 	{
 		case HTTP_STATUS_FORBIDDEN:
 		case HTTP_STATUS_OK:
-		{
-			/* AVD returns a 403 response with a ARRAffinity cookie set. retry with that cookie */
-			const char* affinity = http_response_get_setcookie(response, "ARRAffinity");
-			if (affinity && freerdp_settings_get_bool(wst->settings, FreeRDP_GatewayArmTransport))
-			{
-				WLog_DBG(TAG, "Got Affinity cookie %s", affinity);
-				http_context_set_cookie(wst->http, "ARRAffinity", affinity);
-				http_response_free(response);
-				/* Terminate this connection and make a new one with the Loadbalancing Cookie */
-				int fd = BIO_get_fd(wst->tls->bio, NULL);
-				if (fd >= 0)
-					closesocket((SOCKET)fd);
-				freerdp_tls_free(wst->tls);
-
-				wst->tls = freerdp_tls_new(wst->settings);
-				if (!wst_tls_connect(wst, wst->tls, timeout))
-					return FALSE;
-
-				if (freerdp_settings_get_string(wst->settings, FreeRDP_GatewayHttpExtAuthBearer) &&
-				    freerdp_settings_get_bool(wst->settings, FreeRDP_GatewayArmTransport))
-				{
-					char* urlWithAuth = NULL;
-					size_t urlLen = 0;
-					char firstParam = (strchr(wst->gwpath, '?') > 0 ? '&' : '?');
-					winpr_asprintf(&urlWithAuth, &urlLen, arm_query_param, wst->gwpath, firstParam,
-					               freerdp_settings_get_string(wst->settings,
-					                                           FreeRDP_GatewayHttpExtAuthBearer));
-					if (!urlWithAuth)
-						return FALSE;
-					free(wst->gwpath);
-					wst->gwpath = urlWithAuth;
-					http_context_set_uri(wst->http, wst->gwpath);
-					http_context_enable_websocket_upgrade(wst->http, TRUE);
-				}
-
-				if (!wst_send_http_request(wst, wst->tls))
-					return FALSE;
-				response = http_response_recv(wst->tls, TRUE);
-				if (!response)
-				{
-					return FALSE;
-				}
-				StatusCode = http_response_get_status_code(response);
-			}
+			success = wst_handle_ok_or_forbidden(wst, &response, timeout, &StatusCode);
 			break;
-		}
+
 		case HTTP_STATUS_DENIED:
-		{
-			/* if the response is HTTP_STATUS_DENIED and no bearer authentication was used, try with
-			 * credssp */
-			http_response_free(response);
-			if (freerdp_settings_get_string(wst->settings, FreeRDP_GatewayHttpExtAuthBearer))
-				return FALSE;
-
-			if (!wst_auth_init(wst, wst->tls, AUTH_PKG))
-				return FALSE;
-			if (!wst_send_http_request(wst, wst->tls))
-				return FALSE;
-
-			response = http_response_recv(wst->tls, TRUE);
-			if (!response)
-			{
-				return FALSE;
-			}
-			while (!credssp_auth_is_complete(wst->auth))
-			{
-				if (!wst_recv_auth_token(wst->auth, response))
-				{
-					http_response_free(response);
-					return FALSE;
-				}
-
-				if (credssp_auth_have_output_token(wst->auth))
-				{
-					http_response_free(response);
-
-					if (!wst_send_http_request(wst, wst->tls))
-						return FALSE;
-
-					response = http_response_recv(wst->tls, TRUE);
-					if (!response)
-					{
-						return FALSE;
-					}
-				}
-			}
-			StatusCode = http_response_get_status_code(response);
-			credssp_auth_free(wst->auth);
-			wst->auth = NULL;
-			break;
-		}
+			success = wst_handle_denied(wst, &response, &StatusCode);
 		default:
 			break;
 	}
 	http_response_free(response);
+	if (!success)
+		return FALSE;
 
 	if (StatusCode == HTTP_STATUS_SWITCH_PROTOCOLS)
 	{
@@ -815,7 +823,9 @@ rdpWst* wst_new(rdpContext* context)
 
 		BIO_set_data(wst->frontBio, wst);
 		InitializeCriticalSection(&wst->writeSection);
-		wst->auth = NULL;
+		wst->auth = credssp_auth_new(context);
+		if (!wst->auth)
+			goto wst_alloc_error;
 	}
 
 	return wst;

libfreerdp/core/transport.c

@@ -474,7 +474,7 @@ BOOL transport_connect(rdpTransport* transport, const char* hostname, UINT16 por
 		if (freerdp_settings_get_bool(settings, FreeRDP_GatewayArmTransport))
 		{
 			if (!arm_resolve_endpoint(context, timeout))
-				WLog_ERR(TAG, "ARM Endpoint resolve failed");
+				return FALSE;
 		}
 		if (settings->GatewayUrl)
 		{