<commit_before>/**
 * FreeRDP: A Remote Desktop Protocol Implementation
 * FreeRDP Proxy Server
 *
 * Copyright 2019 <NAME> <<EMAIL>>
 * Copyright 2019 <NAME> <<EMAIL>>
 * Copyright 2019 <NAME> <<EMAIL>>
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include <winpr/crt.h>
#include <winpr/ssl.h>
#include <winpr/synch.h>
#include <winpr/string.h>
#include <winpr/winsock.h>
#include <winpr/thread.h>
#include <errno.h>

#include <freerdp/freerdp.h>
#include <freerdp/channels/wtsvc.h>
#include <freerdp/channels/channels.h>

#include "pf_server.h"
#include "pf_log.h"
#include "pf_config.h"
#include "pf_client.h"
#include "pf_context.h"
#include "pf_input.h"
#include "pf_update.h"
#include "pf_rdpgfx.h"
#include "pf_disp.h"
#include "pf_rail.h"
#include "pf_channels.h"
#include "pf_modules.h"

#define TAG PROXY_TAG("server")

static psPeerReceiveChannelData server_receive_channel_data_original = NULL;

static BOOL pf_server_parse_target_from_routing_token(rdpContext* context, char** target,
                                                      DWORD* port)
{
#define TARGET_MAX (100)
#define ROUTING_TOKEN_PREFIX "Cookie: msts="
	char* colon;
	size_t len;
	DWORD routing_token_length;
	const size_t prefix_len = strnlen(ROUTING_TOKEN_PREFIX, sizeof(ROUTING_TOKEN_PREFIX));
	const char* routing_token = freerdp_nego_get_routing_token(context, &routing_token_length);
	pServerContext* ps = (pServerContext*)context;

	if (!routing_token)
		return FALSE;

	if ((routing_token_length <= prefix_len) || (routing_token_length >= TARGET_MAX))
	{
		LOG_ERR(TAG, ps, "invalid routing token length: %" PRIu32 "", routing_token_length);
		return FALSE;
	}

	len = routing_token_length - prefix_len;
	*target = malloc(len + 1);

	if (!(*target))
		return FALSE;

	CopyMemory(*target, routing_token + prefix_len, len);
	*(*target + len) = '\0';
	colon = strchr(*target, ':');

	if (colon)
	{
		/* port is specified */
		unsigned long p = strtoul(colon + 1, NULL, 10);

		if (p > USHRT_MAX)
		{
			free(*target);
			return FALSE;
		}

		*port = (DWORD)p;
		*colon = '\0';
	}

	return TRUE;
}

static BOOL pf_server_get_target_info(rdpContext* context, rdpSettings* settings,
                                      proxyConfig* config)
{
	pServerContext* ps = (pServerContext*)context;

	LOG_INFO(TAG, ps, "fetching target from %s",
	         config->UseLoadBalanceInfo ? "load-balance-info" : "config");

	if (config->UseLoadBalanceInfo)
		return pf_server_parse_target_from_routing_token(context, &settings->ServerHostname,
		                                                 &settings->ServerPort);

	/* use hardcoded target info from configuration */
	if (!(settings->ServerHostname = _strdup(config->TargetHost)))
	{
		LOG_ERR(TAG, ps, "strdup failed!");
		return FALSE;
	}

	settings->ServerPort = config->TargetPort > 0 ? 3389 : settings->ServerPort;
	return TRUE;
}

/* Event callbacks */
/**
 * This callback is called when the entire connection sequence is done (as
 * described in MS-RDPBCGR section 1.3)
 *
 * The server may start sending graphics output and receiving keyboard/mouse
 * input after this callback returns.
 */
static BOOL pf_server_post_connect(freerdp_peer* peer)
{
	pServerContext* ps;
	pClientContext* pc;
	rdpSettings* client_settings;
	proxyData* pdata;
	ps = (pServerContext*)peer->context;
	pdata = ps->pdata;

	if (pdata->config->SessionCapture && !peer->settings->SupportGraphicsPipeline)
	{
		LOG_ERR(TAG, ps, "Session capture feature is enabled, only accepting connections with GFX");
		return FALSE;
	}

	pc = pf_context_create_client_context(peer->settings);
	if (pc == NULL)
	{
		LOG_ERR(TAG, ps, "failed to create client context!");
		return FALSE;
	}

	client_settings = pc->context.settings;

	/* keep both sides of the connection in pdata */
	proxy_data_set_client_context(pdata, pc);

	if (!pf_server_get_target_info(peer->context, client_settings, pdata->config))
	{
		LOG_INFO(TAG, ps, "pf_server_get_target_info failed!");
		return FALSE;
	}

	LOG_INFO(TAG, ps, "remote target is %s:%" PRIu16 "", client_settings->ServerHostname,
	         client_settings->ServerPort);

	if (!pf_server_channels_init(ps))
	{
		LOG_INFO(TAG, ps, "failed to initialize server's channels!");
		return FALSE;
	}

	/* Start a proxy's client in it's own thread */
	if (!(pdata->client_thread = CreateThread(NULL, 0, pf_client_start, pc, 0, NULL)))
	{
		LOG_ERR(TAG, ps, "failed to create client thread");
		return FALSE;
	}

	pf_server_register_input_callbacks(peer->input);
	pf_server_register_update_callbacks(peer->update);
	return pf_modules_run_hook(HOOK_TYPE_SERVER_POST_CONNECT, pdata);
}

static BOOL pf_server_activate(freerdp_peer* peer)
{
	peer->settings->CompressionLevel = PACKET_COMPR_TYPE_RDP8;
	return TRUE;
}

static BOOL pf_server_adjust_monitor_layout(freerdp_peer* peer)
{
	/* proxy as is, there's no need to do anything here */
	return TRUE;
}

static BOOL pf_server_receive_channel_data_hook(freerdp_peer* peer, UINT16 channelId,
                                                const BYTE* data, size_t size, UINT32 flags,
                                                size_t totalSize)
{
	pServerContext* ps = (pServerContext*)peer->context;
	pClientContext* pc = ps->pdata->pc;
	proxyData* pdata = ps->pdata;
	proxyConfig* config = pdata->config;
	size_t i;
	const char* channel_name = WTSChannelGetName(peer, channelId);

	/*<commit_msg>server: proxy: do not start client if post connect hook failed<commit_after>/**
 * FreeRDP: A Remote Desktop Protocol Implementation
 * FreeRDP Proxy Server
 *
 * Copyright 2019 <NAME> <<EMAIL>>
 * Copyright 2019 <NAME> <<EMAIL>>
 * Copyright 2019 <NAME> <<EMAIL>>
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include <winpr/crt.h>
#include <winpr/ssl.h>
#include <winpr/synch.h>
#include <winpr/string.h>
#include <winpr/winsock.h>
#include <winpr/thread.h>
#include <errno.h>

#include <freerdp/freerdp.h>
#include <freerdp/channels/wtsvc.h>
#include <freerdp/channels/channels.h>

#include "pf_server.h"
#include "pf_log.h"
#include "pf_config.h"
#include "pf_client.h"
#include "pf_context.h"
#include "pf_input.h"
#include "pf_update.h"
#include "pf_rdpgfx.h"
#include "pf_disp.h"
#include "pf_rail.h"
#include "pf_channels.h"
#include "pf_modules.h"

#define TAG PROXY_TAG("server")

static psPeerReceiveChannelData server_receive_channel_data_original = NULL;

static BOOL pf_server_parse_target_from_routing_token(rdpContext* context, char** target,
                                                      DWORD* port)
{
#define TARGET_MAX (100)
#define ROUTING_TOKEN_PREFIX "Cookie: msts="
	char* colon;
	size_t len;
	DWORD routing_token_length;
	const size_t prefix_len = strnlen(ROUTING_TOKEN_PREFIX, sizeof(ROUTING_TOKEN_PREFIX));
	const char* routing_token = freerdp_nego_get_routing_token(context, &routing_token_length);
	pServerContext* ps = (pServerContext*)context;

	if (!routing_token)
		return FALSE;

	if ((routing_token_length <= prefix_len) || (routing_token_length >= TARGET_MAX))
	{
		LOG_ERR(TAG, ps, "invalid routing token length: %" PRIu32 "", routing_token_length);
		return FALSE;
	}

	len = routing_token_length - prefix_len;
	*target = malloc(len + 1);

	if (!(*target))
		return FALSE;

	CopyMemory(*target, routing_token + prefix_len, len);
	*(*target + len) = '\0';
	colon = strchr(*target, ':');

	if (colon)
	{
		/* port is specified */
		unsigned long p = strtoul(colon + 1, NULL, 10);

		if (p > USHRT_MAX)
		{
			free(*target);
			return FALSE;
		}

		*port = (DWORD)p;
		*colon = '\0';
	}

	return TRUE;
}

static BOOL pf_server_get_target_info(rdpContext* context, rdpSettings* settings,
                                      proxyConfig* config)
{
	pServerContext* ps = (pServerContext*)context;

	LOG_INFO(TAG, ps, "fetching target from %s",
	         config->UseLoadBalanceInfo ? "load-balance-info" : "config");

	if (config->UseLoadBalanceInfo)
		return pf_server_parse_target_from_routing_token(context, &settings->ServerHostname,
		                                                 &settings->ServerPort);

	/* use hardcoded target info from configuration */
	if (!(settings->ServerHostname = _strdup(config->TargetHost)))
	{
		LOG_ERR(TAG, ps, "strdup failed!");
		return FALSE;
	}

	settings->ServerPort = config->TargetPort > 0 ? 3389 : settings->ServerPort;
	return TRUE;
}

/* Event callbacks */
/**
 * This callback is called when the entire connection sequence is done (as
 * described in MS-RDPBCGR section 1.3)
 *
 * The server may start sending graphics output and receiving keyboard/mouse
 * input after this callback returns.
 */
static BOOL pf_server_post_connect(freerdp_peer* peer)
{
	pServerContext* ps;
	pClientContext* pc;
	rdpSettings* client_settings;
	proxyData* pdata;
	ps = (pServerContext*)peer->context;
	pdata = ps->pdata;

	if (pdata->config->SessionCapture && !peer->settings->SupportGraphicsPipeline)
	{
		LOG_ERR(TAG, ps, "Session capture feature is enabled, only accepting connections with GFX");
		return FALSE;
	}

	pc = pf_context_create_client_context(peer->settings);
	if (pc == NULL)
	{
		LOG_ERR(TAG, ps, "failed to create client context!");
		return FALSE;
	}

	client_settings = pc->context.settings;

	/* keep both sides of the connection in pdata */
	proxy_data_set_client_context(pdata, pc);

	if (!pf_server_get_target_info(peer->context, client_settings, pdata->config))
	{
		LOG_INFO(TAG, ps, "pf_server_get_target_info failed!");
		return FALSE;
	}

	LOG_INFO(TAG, ps, "remote target is %s:%" PRIu16 "", client_settings->ServerHostname,
	         client_settings->ServerPort);

	if (!pf_server_channels_init(ps))
	{
		LOG_INFO(TAG, ps, "failed to initialize server's channels!");
		return FALSE;
	}

	if (!pf_modules_run_hook(HOOK_TYPE_SERVER_POST_CONNECT, pdata))
		return FALSE;

	/* Start a proxy's client in it's own thread */
	if (!(pdata->client_thread = CreateThread(NULL, 0, pf_client_start, pc, 0, NULL)))
	{
		LOG_ERR(TAG, ps, "failed to create client thread");
		return FALSE;
	}

	pf_server_register_input_callbacks(peer->input);
	pf_server_register_update_callbacks(peer->update);
	return TRUE;
}

static BOOL pf_server_activate(freerdp_peer* peer)
{
	peer->settings->CompressionLevel = PACKET_COMPR_TYPE_RDP8;
	return TRUE;
}

static BOOL pf_server_adjust_monitor_layout(freerdp_peer* peer)
{
	/* proxy as is, there's no need to do anything here */
	return TRUE;
}

static BOOL pf_server_receive_channel_data_hook(freerdp_peer* peer, UINT16 channelId,
                                                const BYTE* data, size_t size, UINT32 flags,
                                                size_t totalSize)
{
	pServerContext* ps = (pServerContext*)peer->context;
	pClientContext* pc = ps->pdata->pc;
	proxyData* pdata = ps->pdata;
	proxyConfig* config = pdata->config;
	size_t i;
	const char* channel_name = WTSChannelGetName(peer, channelId);

	/*
{'fraction_non_alphanumeric': 0.11616161616161616, 'fraction_numerical': 0.007239057239057239, 'mean_word_length': 4.049298767530812, 'pattern_counts': {'":': 0, '<': 43, '<?xml version=': 0, '>': 109, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 18, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3794702', 'n_tokens_mistral': 4081, 'n_tokens_neox': 3892, 'n_words': 1026}