/**

 * FreeRDP: A Remote Desktop Protocol Implementation

 * SSH Agent Virtual Channel Extension

 *

 * Copyright 2013 Christian Hofstaedtler

 * Copyright 2015 Thincast Technologies GmbH

 * Copyright 2015 DI (FH) Martin Haimberger <martin.haimberger@thincast.com>

 * Copyright 2017 Ben Cohen

 *

 * 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.

 */

/*

 * sshagent_main.c: DVC plugin to forward queries from RDP to the ssh-agent

 *

 * This relays data to and from an ssh-agent program equivalent running on the

 * RDP server to an ssh-agent running locally.  Unlike the normal ssh-agent,

 * which sends data over an SSH channel, the data is send over an RDP dynamic

 * virtual channel.

 *

 * protocol specification:

 *     Forward data verbatim over RDP dynamic virtual channel named "sshagent"

 *     between a ssh client on the xrdp server and the real ssh-agent where

 *     the RDP client is running.  Each connection by a separate client to

 *     xrdp-ssh-agent gets a separate DVC invocation.

 */

#include <freerdp/config.h>

#include <stdio.h>

#include <stdlib.h>

#include <sys/types.h>

#include <sys/socket.h>

#include <sys/un.h>

#include <pwd.h>

#include <unistd.h>

#include <errno.h>

#include <winpr/crt.h>

#include <winpr/assert.h>

#include <winpr/synch.h>

#include <winpr/thread.h>

#include <winpr/stream.h>

#include "sshagent_main.h"

#include <freerdp/freerdp.h>

#include <freerdp/client/channels.h>

#include <freerdp/channels/log.h>

#define TAG CHANNELS_TAG("sshagent.client")

typedef struct

{

  IWTSListenerCallback iface;

  IWTSPlugin* plugin;

  IWTSVirtualChannelManager* channel_mgr;

  rdpContext* rdpcontext;

  const char* agent_uds_path;

} SSHAGENT_LISTENER_CALLBACK;

typedef struct

{

  GENERIC_CHANNEL_CALLBACK generic;

  rdpContext* rdpcontext;

  int agent_fd;

  HANDLE thread;

  CRITICAL_SECTION lock;

} SSHAGENT_CHANNEL_CALLBACK;

typedef struct

{

  IWTSPlugin iface;

  SSHAGENT_LISTENER_CALLBACK* listener_callback;

  rdpContext* rdpcontext;

  BOOL connected;

} SSHAGENT_PLUGIN;

/**

 * Function to open the connection to the sshagent

 *

 * @return The fd on success, otherwise -1

 */

static int connect_to_sshagent(const char* udspath)

{

  WINPR_ASSERT(udspath);

  int agent_fd = socket(AF_UNIX, SOCK_STREAM, 0);

  if (agent_fd == -1)

  {

    WLog_ERR(TAG, "Can't open Unix domain socket!");

    return -1;

  }

  struct sockaddr_un addr = WINPR_C_ARRAY_INIT;

  addr.sun_family = AF_UNIX;

  strncpy(addr.sun_path, udspath, sizeof(addr.sun_path) - 1);

  int rc = connect(agent_fd, (struct sockaddr*)&addr, sizeof(addr));

  if (rc != 0)

  {

    WLog_ERR(TAG, "Can't connect to Unix domain socket \"%s\"!", udspath);

    close(agent_fd);

    return -1;

  }

  return agent_fd;

}

/**

 * Entry point for thread to read from the ssh-agent socket and forward

 * the data to RDP

 *

 * @return 0

 */

static DWORD WINAPI sshagent_read_thread(LPVOID data)

{

  SSHAGENT_CHANNEL_CALLBACK* callback = (SSHAGENT_CHANNEL_CALLBACK*)data;

  WINPR_ASSERT(callback);

  BYTE buffer[4096] = WINPR_C_ARRAY_INIT;

  int going = 1;

  UINT status = CHANNEL_RC_OK;

  while (going)

  {

    const ssize_t bytes_read = read(callback->agent_fd, buffer, sizeof(buffer));

    if (bytes_read == 0)

    {

      /* Socket closed cleanly at other end */

      going = 0;

    }

    else if (bytes_read < 0)

    {

      if (errno != EINTR)

      {

        WLog_ERR(TAG, "Error reading from sshagent, errno=%d", errno);

        status = ERROR_READ_FAULT;

        going = 0;

      }

    }

    else if ((size_t)bytes_read > ULONG_MAX)

    {

      status = ERROR_READ_FAULT;

      going = 0;

    }

    else

    {

      /* Something read: forward to virtual channel */

      IWTSVirtualChannel* channel = callback->generic.channel;

      status = channel->Write(channel, (ULONG)bytes_read, buffer, nullptr);

      if (status != CHANNEL_RC_OK)

      {

        going = 0;

      }

    }

  }

  close(callback->agent_fd);

  if (status != CHANNEL_RC_OK)

    setChannelError(callback->rdpcontext, status, "sshagent_read_thread reported an error");

  ExitThread(status);

  return status;

}

/**

 * Callback for data received from the RDP server; forward this to ssh-agent

 *

 * @return 0 on success, otherwise a Win32 error code

 */

static UINT sshagent_on_data_received(IWTSVirtualChannelCallback* pChannelCallback, wStream* data)

{

  SSHAGENT_CHANNEL_CALLBACK* callback = (SSHAGENT_CHANNEL_CALLBACK*)pChannelCallback;

  WINPR_ASSERT(callback);

  BYTE* pBuffer = Stream_Pointer(data);

  size_t cbSize = Stream_GetRemainingLength(data);

  BYTE* pos = pBuffer;

  /* Forward what we have received to the ssh agent */

  size_t bytes_to_write = cbSize;

  errno = 0;

  while (bytes_to_write > 0)

  {

    const ssize_t bytes_written = write(callback->agent_fd, pos, bytes_to_write);

    if (bytes_written < 0)

    {

      if (errno != EINTR)

      {

        WLog_ERR(TAG, "Error writing to sshagent, errno=%d", errno);

        return ERROR_WRITE_FAULT;

      }

    }

    else

    {

      bytes_to_write -= WINPR_ASSERTING_INT_CAST(size_t, bytes_written);

      pos += bytes_written;

    }

  }

  /* Consume stream */

  Stream_Seek(data, cbSize);

  return CHANNEL_RC_OK;

}

/**

 * Callback for when the virtual channel is closed

 *

 * @return 0 on success, otherwise a Win32 error code

 */

static UINT sshagent_on_close(IWTSVirtualChannelCallback* pChannelCallback)

{

  SSHAGENT_CHANNEL_CALLBACK* callback = (SSHAGENT_CHANNEL_CALLBACK*)pChannelCallback;

  WINPR_ASSERT(callback);

  /* Call shutdown() to wake up the read() in sshagent_read_thread(). */

  shutdown(callback->agent_fd, SHUT_RDWR);

  EnterCriticalSection(&callback->lock);

  if (WaitForSingleObject(callback->thread, INFINITE) == WAIT_FAILED)

  {

    UINT error = GetLastError();

    WLog_ERR(TAG, "WaitForSingleObject failed with error %" PRIu32 "!", error);

    return error;

  }

  (void)CloseHandle(callback->thread);

  LeaveCriticalSection(&callback->lock);

  DeleteCriticalSection(&callback->lock);

  SSHAGENT_PLUGIN* sshagent = (SSHAGENT_PLUGIN*)callback->generic.plugin;

  WINPR_ASSERT(sshagent);

  sshagent->connected = FALSE;

  free(callback);

  return CHANNEL_RC_OK;

}

/**

 * Callback for when a new virtual channel is opened

 *

 * @return 0 on success, otherwise a Win32 error code

 */

// NOLINTBEGIN(readability-non-const-parameter)

static UINT sshagent_on_new_channel_connection(IWTSListenerCallback* pListenerCallback,

                                               IWTSVirtualChannel* pChannel, BYTE* Data,

                                               BOOL* pbAccept,

                                               IWTSVirtualChannelCallback** ppCallback)

// NOLINTEND(readability-non-const-parameter)

{

  SSHAGENT_LISTENER_CALLBACK* listener_callback = (SSHAGENT_LISTENER_CALLBACK*)pListenerCallback;

  WINPR_UNUSED(Data);

  WINPR_UNUSED(pbAccept);

  SSHAGENT_PLUGIN* sshagent = (SSHAGENT_PLUGIN*)listener_callback->plugin;

  WINPR_ASSERT(sshagent);

  if (sshagent->connected)

  {

    WLog_ERR(TAG, "Channel already connected, terminating.");

    return ERROR_DEVICE_ALREADY_ATTACHED;

  }

  SSHAGENT_CHANNEL_CALLBACK* callback =

      (SSHAGENT_CHANNEL_CALLBACK*)calloc(1, sizeof(SSHAGENT_CHANNEL_CALLBACK));

  if (!callback)

  {

    WLog_ERR(TAG, "calloc failed!");

    return CHANNEL_RC_NO_MEMORY;

  }

  /* Now open a connection to the local ssh-agent.  Do this for each

   * connection to the plugin in case we mess up the agent session. */

  callback->agent_fd = connect_to_sshagent(listener_callback->agent_uds_path);

  if (callback->agent_fd == -1)

  {

    free(callback);

    return CHANNEL_RC_INITIALIZATION_ERROR;

  }

  InitializeCriticalSection(&callback->lock);

  GENERIC_CHANNEL_CALLBACK* generic = &callback->generic;

  generic->iface.OnDataReceived = sshagent_on_data_received;

  generic->iface.OnClose = sshagent_on_close;

  generic->plugin = listener_callback->plugin;

  generic->channel_mgr = listener_callback->channel_mgr;

  generic->channel = pChannel;

  callback->rdpcontext = listener_callback->rdpcontext;

  callback->thread = CreateThread(nullptr, 0, sshagent_read_thread, (void*)callback, 0, nullptr);

  if (!callback->thread)

  {

    WLog_ERR(TAG, "CreateThread failed!");

    DeleteCriticalSection(&callback->lock);

    free(callback);

    return CHANNEL_RC_INITIALIZATION_ERROR;

  }

  *ppCallback = &callback->generic.iface;

  sshagent->connected = TRUE;

  return CHANNEL_RC_OK;

}

/**

 * Callback for when the plugin is initialised

 *

 * @return 0 on success, otherwise a Win32 error code

 */

static UINT sshagent_plugin_initialize(IWTSPlugin* pPlugin, IWTSVirtualChannelManager* pChannelMgr)

{

  SSHAGENT_PLUGIN* sshagent = (SSHAGENT_PLUGIN*)pPlugin;

  WINPR_ASSERT(sshagent);

  WINPR_ASSERT(pChannelMgr);

  sshagent->listener_callback =

      (SSHAGENT_LISTENER_CALLBACK*)calloc(1, sizeof(SSHAGENT_LISTENER_CALLBACK));

  if (!sshagent->listener_callback)

  {

    WLog_ERR(TAG, "calloc failed!");

    return CHANNEL_RC_NO_MEMORY;

  }

  sshagent->listener_callback->rdpcontext = sshagent->rdpcontext;

  sshagent->listener_callback->iface.OnNewChannelConnection = sshagent_on_new_channel_connection;

  sshagent->listener_callback->plugin = pPlugin;

  sshagent->listener_callback->channel_mgr = pChannelMgr;

  // NOLINTNEXTLINE(concurrency-mt-unsafe)

  sshagent->listener_callback->agent_uds_path = getenv("SSH_AUTH_SOCK");

  if (sshagent->listener_callback->agent_uds_path == nullptr)

  {

    WLog_ERR(TAG, "Environment variable $SSH_AUTH_SOCK undefined!");

    free(sshagent->listener_callback);

    sshagent->listener_callback = nullptr;

    return CHANNEL_RC_INITIALIZATION_ERROR;

  }

  return pChannelMgr->CreateListener(pChannelMgr, "SSHAGENT", 0,

                                     (IWTSListenerCallback*)sshagent->listener_callback, nullptr);

}

/**

 * Callback for when the plugin is terminated

 *

 * @return 0 on success, otherwise a Win32 error code

 */

static UINT sshagent_plugin_terminated(IWTSPlugin* pPlugin)

{

  SSHAGENT_PLUGIN* sshagent = (SSHAGENT_PLUGIN*)pPlugin;

  free(sshagent);

  return CHANNEL_RC_OK;

}

/**

 * Main entry point for sshagent DVC plugin

 *

 * @return 0 on success, otherwise a Win32 error code

 */

FREERDP_ENTRY_POINT(UINT VCAPITYPE sshagent_DVCPluginEntry(IDRDYNVC_ENTRY_POINTS* pEntryPoints))

{

  UINT status = CHANNEL_RC_OK;

  WINPR_ASSERT(pEntryPoints);

  SSHAGENT_PLUGIN* sshagent = (SSHAGENT_PLUGIN*)pEntryPoints->GetPlugin(pEntryPoints, "sshagent");

  if (!sshagent)

  {

    sshagent = (SSHAGENT_PLUGIN*)calloc(1, sizeof(SSHAGENT_PLUGIN));

    if (!sshagent)

    {

      WLog_ERR(TAG, "calloc failed!");

      return CHANNEL_RC_NO_MEMORY;

    }

    sshagent->iface.Initialize = sshagent_plugin_initialize;

    sshagent->iface.Connected = nullptr;

    sshagent->iface.Disconnected = nullptr;

    sshagent->iface.Terminated = sshagent_plugin_terminated;

    sshagent->rdpcontext = pEntryPoints->GetRdpContext(pEntryPoints);

    status = pEntryPoints->RegisterPlugin(pEntryPoints, "sshagent", &sshagent->iface);

  }

  return status;

}