shortener-bot/main.cpp

#include <algorithm>
#include <chrono>
#include <cstdlib>
#include <iostream>
#include <memory>
#include <stdexcept>
#include <string>
#include <sys/select.h>
#include <termios.h>
#include <unistd.h>

#include <spdlog/spdlog.h>
#include <uv.h>
#include "common.h"
#include "commands.h"
#include "curl/multi.h"
#include "td/telegram/td_api.h"
#include "td/telegram/td_api.hpp"

const char *APP_GREETING = "--------------------\n"\
                           "| shortener-bot v0 |\n"\
                           "--------------------\n\n";

void async_startup(uv_idle_t *h);
void configure_blocked_signals();
void configure_logging();
std::array<std::unique_ptr<uv_signal_t, void(*)(uv_signal_t*)>, 3> configure_shutdown_signals(uv_loop_t *loop, context *ctx);
CommandLineParams parse_command_line(int argc, char **argv);
void print_greeting();
void print_usage(const char *programName);
void shutdown_signal_handler(uv_signal_t *h, int signum);
int curl_socket_cb(CURL*, curl_socket_t, int, context*, void*);
int curl_timeout_cb(CURLM*, long, context*);
void uv_socket_cb(uv_poll_t *handle, int status, int events);
void uv_curl_timeout_cb(uv_timer_t *handle);
void check_curl_multi_info(context *ctx);
void gc_pending_queries(uv_timer_t *handle);

int main(int argc, char **argv) {
  CommandLineParams cmdLineParams;
  try {
    cmdLineParams = parse_command_line(argc, argv);
  } catch (CommandLineParamError e) {
    switch (e) {
      case CommandLineParamError::NoApiCreds:
        print_usage(argv[0]);
        return 1;
    }
  }
  configure_logging();
  configure_blocked_signals();
  print_greeting();


  uv_loop_t *defaultLoop = uv_default_loop();
  context ctx;
  ctx.apiCreds = cmdLineParams.apiCreds;

  spdlog::info("Initializing CURL");
  curl_global_init(CURL_GLOBAL_DEFAULT);
  CURLM *curl = curl_multi_init();
  ctx.curl = curl;
  curl_multi_setopt(curl, CURLMOPT_SOCKETFUNCTION, &curl_socket_cb);
  curl_multi_setopt(curl, CURLMOPT_SOCKETDATA, &ctx);
  curl_multi_setopt(curl, CURLMOPT_TIMERFUNCTION, &curl_timeout_cb);
  curl_multi_setopt(curl, CURLMOPT_TIMERDATA, &ctx);
  uv_timer_init(defaultLoop, &ctx.curlTimer);
  ctx.curlTimer.data = &ctx;

  spdlog::info("Creating Telegram client");
  TelegramClient tg(defaultLoop);
  ctx.tg = &tg;

  uv_idle_t startupHandle;
  startupHandle.data = (void*)&ctx;
  uv_idle_init(defaultLoop, &startupHandle);
  uv_idle_start(&startupHandle, async_startup);

  uv_timer_t pendingQueryGcTimer;
  uv_timer_init(defaultLoop, &pendingQueryGcTimer);
  pendingQueryGcTimer.data = &ctx;
  uv_timer_start(&pendingQueryGcTimer, gc_pending_queries, 60000, 60000);

  auto signalHandles = configure_shutdown_signals(defaultLoop, &ctx);

  termios termSettings;
  tcgetattr(1, &termSettings);
  termSettings.c_lflag &= ~(ECHO | ICANON);
  tcsetattr(1, 0, &termSettings);

  spdlog::info("Entering event loop");
  uv_run(defaultLoop, UV_RUN_DEFAULT);
  spdlog::info("Leaving event loop");

  termSettings.c_lflag |= ECHO | ICANON;
  tcsetattr(1, 0, &termSettings);

  spdlog::info("Cleaning up");
  curl_multi_cleanup(curl);
  curl_global_cleanup();
  return 0;
}

void async_startup(uv_idle_t *h) {
  context *ctx = (context*)h->data;
  spdlog::info("Starting telegram client (using api id {}, hash {})", ctx->apiCreds.id, ctx->apiCreds.hash);

  ctx->tg->add_update_handler([ctx](void*, td_api::Object &update){
    td_api::downcast_call(update, overloaded(
      [ctx](td_api::updateAuthorizationState &upd) {
        td_api::downcast_call(*upd.authorization_state_, overloaded(
          [ctx](td_api::authorizationStateWaitTdlibParameters &state) {
            auto request = td_api::make_object<td_api::setTdlibParameters>();
            request->database_directory_ = "tdata";
            request->use_message_database_ = false;
            request->use_secret_chats_ = false;
            request->api_id_ = ctx->apiCreds.id;
            request->api_hash_ = ctx->apiCreds.hash;
            request->system_language_code_ = "en";
            request->device_model_ = "server";
            request->application_version_ = "1.0.0";
            ctx->tg->send_query(std::move(request), {});
          },
          [ctx](td_api::authorizationStateWaitPhoneNumber &state) {
            ctx->tg->send_query(td_api::make_object<td_api::checkAuthenticationBotToken>(ctx->apiCreds.botToken), {}); // TODO add error handler
          },
          [ctx](td_api::authorizationStateReady &state) {
            spdlog::info("Telegram ready");
          },
          [ctx,&upd](td_api::Object &obj){
            spdlog::debug("unknown authorization state ID {}", upd.authorization_state_->get_id());
          }
        ));
      },
      [ctx](td_api::updateNewMessage &upd) {
        if (upd.message_->is_outgoing_) return;
        cmd::handle_regular_message(ctx, *upd.message_);
      },
      [ctx](td_api::updateNewInlineQuery &upd) {
        cmd::handle_inline_query(ctx, upd);
      },
      [ctx](td_api::updateNewInlineCallbackQuery &upd) {
        cmd::handle_callback_query(ctx, upd);
      },
      [](td_api::Object &obj){}
    ));
  }, nullptr);
  ctx->tg->start();

  uv_close((uv_handle_t*)h, nullptr);
}

void configure_logging() {
  auto loggingSink = std::make_shared<spdlog::sinks::ansicolor_sink<spdlog::details::console_mutex>>(stdout, spdlog::color_mode::automatic);
  loggingSink->set_color(spdlog::level::debug, "\033[38;5;61m");
  auto logger = std::make_shared<spdlog::logger>("main", loggingSink);
  spdlog::set_default_logger(logger);
  spdlog::set_level(spdlog::level::debug);
}

void configure_blocked_signals() {
  sigset_t signalSet;
  sigemptyset(&signalSet);
  sigaddset(&signalSet, SIGUSR1);
  sigaddset(&signalSet, SIGUSR2);
  sigaddset(&signalSet, SIGHUP);
  sigprocmask(SIG_BLOCK, &signalSet, nullptr);
}

std::unique_ptr<uv_signal_t, void(*)(uv_signal_t*)> register_signal_handle(uv_loop_t *loop, context *ctx, int signum, void(*signal_handler)(uv_signal_t*, int)) {
  spdlog::debug("register signal handler for {}", signum);
  auto deleter = [](uv_signal_t *handle) {
    uv_close((uv_handle_t*)handle, [](uv_handle_t *h){ delete h; });
  };
  std::unique_ptr<uv_signal_t, void(*)(uv_signal_t*)> handle(new uv_signal_t, deleter);
  uv_signal_init(loop, handle.get());
  uv_signal_start(handle.get(), signal_handler, signum);
  return handle;
}

std::array<std::unique_ptr<uv_signal_t, void(*)(uv_signal_t*)>, 3> configure_shutdown_signals(uv_loop_t *loop, context *ctx) {
  return {register_signal_handle(loop, ctx, SIGINT, shutdown_signal_handler),
          register_signal_handle(loop, ctx, SIGQUIT, shutdown_signal_handler),
          register_signal_handle(loop, ctx, SIGTERM, shutdown_signal_handler)};
}

void shutdown_signal_handler(uv_signal_t *h, int signum) {
  spdlog::info("Stopping");
  uv_stop(uv_default_loop());
}

CommandLineParams parse_command_line(int argc, char **argv) {
  if (argc < 4) throw CommandLineParamError::NoApiCreds;
  char *apiIdEnd = nullptr;
  int apiId = strtol(argv[1], &apiIdEnd, 10);
  if (apiIdEnd == argv[1]) throw CommandLineParamError::NoApiCreds;
  return CommandLineParams {apiId, std::string(argv[2]), std::string(argv[3])};
}

void print_greeting() {
  std::cout << APP_GREETING << std::flush;
}

void print_usage(const char *programName) {
  std::cout << "usage: " << programName << " <api_id> <api_hash> <token>\n";
}

struct curl_context {
  uv_poll_t *pollHandle;
  curl_socket_t socketFd;
  context *ctx;

  curl_context(uv_loop_t *loop, context *ctx, int fd) {
    spdlog::debug("creating new curl context for socket {}", fd);
    pollHandle = new uv_poll_t;
    uv_poll_init(loop, pollHandle, fd);
    pollHandle->data = this;
    socketFd = fd;
    this->ctx = ctx;
  }
  ~curl_context() {
    spdlog::debug("destroying curl context for socket {}", socketFd);
    uv_close((uv_handle_t*)pollHandle, [](uv_handle_t *h){ delete h; });
  }
};

int curl_socket_cb(CURL *curl, curl_socket_t s, int action, context* ctx, void* socketPtr) {
  (void)curl;
  spdlog::debug("curl socket callback: {}", action);
  curl_context *curlCtx;
  switch (action) {
    case CURL_POLL_IN:
    case CURL_POLL_INOUT:
    case CURL_POLL_OUT: {
      curlCtx = socketPtr ? (curl_context*)socketPtr : new curl_context(uv_default_loop(), ctx, s);
      curl_multi_assign(ctx->curl, s, curlCtx);
      int events = 0;
      if (action != CURL_POLL_IN) events |= UV_WRITABLE;
      if (action != CURL_POLL_OUT) events |= UV_READABLE;
      uv_poll_start(curlCtx->pollHandle, events, uv_socket_cb);
      break;
    }
    case CURL_POLL_REMOVE:
      if (socketPtr) {
        uv_poll_stop(((curl_context*)socketPtr)->pollHandle);
        curl_multi_assign(ctx->curl, s, NULL);
        delete (curl_context*)socketPtr;
      }
  }
  return 0;
}

int curl_timeout_cb(CURLM* curl, long timeout, context* ctx) {
  (void)curl;
  spdlog::debug("curl timeout change: {} ms", timeout);
  if (timeout < 0) uv_timer_stop(&ctx->curlTimer);
  else {
    if (timeout == 0) timeout = 1;
    uv_timer_start(&ctx->curlTimer, uv_curl_timeout_cb, timeout, 0);
  }
  return 0;
}

void uv_socket_cb(uv_poll_t *handle, int status, int events) {
  spdlog::debug("libuv: socket {}", events & UV_READABLE ? (events & UV_WRITABLE ? "readable+writable" : "readable") : "writable");
  (void)status;
  auto ctx = reinterpret_cast<curl_context*>(handle->data);
  int flags = 0;
  if (events & UV_READABLE) flags |= CURL_CSELECT_IN;
  if (events & UV_WRITABLE) flags |= CURL_CSELECT_OUT;
  int runningHandles;
  curl_multi_socket_action(ctx->ctx->curl, ctx->socketFd, flags, &runningHandles);
  check_curl_multi_info(ctx->ctx);
}

void uv_curl_timeout_cb(uv_timer_t *handle) {
  spdlog::debug("libuv: curl timeout ended");
  auto ctx = reinterpret_cast<context*>(handle->data);
  if (ctx) {
    int runningHandles;
    curl_multi_socket_action(ctx->curl, CURL_SOCKET_TIMEOUT, 0, &runningHandles);
    check_curl_multi_info(ctx);
  }
}


void check_curl_multi_info(context *ctx) {
  CURLMsg *message;
  int pending;
  while ((message = curl_multi_info_read(ctx->curl, &pending))) {
    switch (message->msg) {
      case CURLMSG_DONE: {
        spdlog::debug("curl transfer done");
        curl_multi_remove_handle(ctx->curl, message->easy_handle);
        curl_easy_cleanup(message->easy_handle);
        try {
          active_request &req = ctx->requests.at(message->easy_handle);
          req.doneCallback(req);
        } catch (std::out_of_range) {}
      }
      default:
        break;
    }
  }
}

size_t curl_receive_cb(char *ptr, size_t size, size_t nmemb, CURL *curl) {
  spdlog::debug("received {} bytes from server: '{}'", nmemb, std::string(ptr, nmemb));
  
  context *ctx = nullptr;
  curl_easy_getinfo(curl, CURLINFO_PRIVATE, &ctx);
  spdlog::debug("ctx pointer: {}", (size_t)ctx);
  try {
    active_request &req = ctx->requests.at(curl);
    req.receivedData.insert(req.receivedData.end(), ptr, ptr + nmemb);
  } catch (std::out_of_range) {
    spdlog::error("No associated request data for handle {}", (size_t)curl);
  }

  return nmemb;
}

void gc_pending_queries(uv_timer_t *handle) {
  context *ctx = reinterpret_cast<context*>(handle->data);

  std::map<uint64_t, pending_inline_query> oldQueries(ctx->inlineQueries);
  std::map<uint64_t, pending_inline_query> newQueries;

  auto now = std::chrono::steady_clock::now();
  for (auto i = oldQueries.begin(); i != oldQueries.end(); i++) {
    if (i->second.ttl > now) {
      newQueries.insert(*i);
    } else {
      spdlog::debug("deleting a pending query {} that expired {} seconds ago", i->first, std::chrono::duration_cast<std::chrono::seconds>(now - i->second.ttl).count());
    }
  }

  ctx->inlineQueries = newQueries;
}