twproxy2/session.cpp

#include <boost/bind.hpp>
#include <iostream>

#include <boost/format.hpp>
#include <boost/log/core.hpp>
#include <boost/log/trivial.hpp>

#include <regex>

#include "session.h"

#include <string>

// #include <boost/log/attributes/named_scope.hpp>

bool replace(std::string &str, const std::string &from, const std::string &to) {
  size_t start_pos = str.find(from);
  if (start_pos == std::string::npos)
    return false;
  do {
    str.replace(start_pos, from.length(), to);
  } while ((start_pos = str.find(from)) != std::string::npos);
  return true;
}

bool replace(std::string &str, const char *from, const char *to) {
  size_t start_pos = str.find(from);
  if (start_pos == std::string::npos)
    return false;
  do {
    str.replace(start_pos, strlen(from), to);
  } while ((start_pos = str.find(from)) != std::string::npos);
  return true;
}

void ansi_clean(std::string &str) {
  static std::regex ansi_cleaner("\x1b\[[0-9;]*[A-Zmh]",
                                 std::regex_constants::ECMAScript);
  str = std::regex_replace(str, ansi_cleaner, "");
}

void high_ascii(std::string &str) {
  static std::regex high_cleaner("[\x80-\xff]+",
                                 std::regex_constants::ECMAScript);
  str = std::regex_replace(str, high_cleaner, "#");
}

std::string clean_string(const std::string &source) {
  // BOOST_LOG_NAMED_SCOPE("clean_string");
  std::string clean = source;
  replace(clean, "\n", "\\n");
  replace(clean, "\r", "\\r");

  // ANSI too
  ansi_clean(clean);
  // BOOST_LOG_TRIVIAL(error) << "cleaned: " << clean;
  high_ascii(clean);

  replace(clean, "\x1b", "^");

  return clean;
}

session::session(boost::asio::ip::tcp::socket socket,
                 boost::asio::io_service &io_service, std::string hostname,
                 std::string port)
    : socket_(std::move(socket)),
      io_service_{io_service}, resolver_{io_service}, server_{io_service},
      timer_{io_service}, host{hostname}, port{port} {
  server_sent = 0;
  // BOOST_LOG_NAMED_SCOPE("session");
}

void session::start(void) {
  // BOOST_LOG_NAMED_SCOPE();

  // If I want the file and line number information, here's how to do it:
  // BUGZ_LOG(info) << boost::format("(%1%:%2%) ") % __FILE__ % __LINE__
  BOOST_LOG_TRIVIAL(info) << "session()";
  BUGZ_LOG(info) << "session::start()";
  auto self(shared_from_this());
  // read_buffer.reserve(1024);
  // do_write("Welcome!\n");
  do_read();
}

session::~session() {
  BOOST_LOG_TRIVIAL(info) << "~session";
  BUGZ_LOG(info) << "~session destructed";
}

void session::parse_auth(void) {
  // how many nulls should I be seeing?
  // \0user\0pass\0terminal/SPEED\0
  // If I don't have a proper rlogin value here, it isn't going
  // to work when I try to connect to the rlogin server.

  if (rlogin_auth.size() > 10)
    rlogin_name = rlogin_auth.c_str() + 1;
  else
    rlogin_name = "?";
}

void session::on_connect(const boost::system::error_code error) {
  // We've connected to the server!  WOOT WOOT!
  // BOOST_LOG_NAMED_SCOPE("session");

  if (!error) {
    BOOST_LOG_TRIVIAL(info) << "Connected to " << host;
    to_client("Connected...\n\r");
    connected = true;
    if (rlogin_auth[0] != 0) {
      // Ok, the rlogin information was junk --
      to_client("Let me make up some fake rlogin data for you...\n\r");
      char temp[] = "\0test\0test\0terminal/9600\0";
      std::string tmp(temp, sizeof(temp));
      to_server(tmp);
    } else {
      to_server(rlogin_auth);
    }

    server_read();
  } else {
    // TODO:
    std::string output =
        str(boost::format("Failed to connect: %1%:%2%\n\r") % host % port);
    to_client(output);
    BOOST_LOG_TRIVIAL(error) << "Failed to connect to " << host << ":" << port;
    BOOST_LOG_TRIVIAL(warning) << "socket.shutdown()";
    socket_.shutdown(boost::asio::ip::tcp::socket::shutdown_both);
  }
}

void session::dispatch_line(std::string line) {
  // Does this have \n\r still on it?  I don't want them.

  std::string temp = clean_string(line);
  BOOST_LOG_TRIVIAL(info) << "SL: " << temp; 
}

/*
Call this with whatever we just received.

That will allow me to send "just whatever I got"
this time around, rather then trying to figure out
what was just added to server_prompt.

What about \r, \b ?  Should that "reset" the server_prompt?

 */
void session::process_lines(std::string &received) {
  // break server_prompt into lines and send/process one by one.

  size_t pos, rpos;
  server_prompt.append(received);

  while ((pos = server_prompt.find('\n', 0)) != std::string::npos) {

    std::string line;

    // process "line" in received

    rpos = received.find('\n', 0);

    // get line to send to the client
    if (show_client) {
      // that is, if we're sending to the client!
      line = received.substr(0, rpos + 1);
      /*
      std::string clean = clean_string(line);
      BOOST_LOG_TRIVIAL(error) << "rpos/show_client:" << clean;
      */
      to_client(line);
    }
    received = received.substr(rpos + 1);

    // process "line" in server_prompt

    line = server_prompt.substr(0, pos + 1);
    server_prompt = server_prompt.substr(pos + 1);

    // Remove \n for dispatching
    std::string part = line.substr(0, pos);
    if (server_sent != 0) {
      line = line.substr(server_sent);
      server_sent = 0;
    };

    // display on?
    // to_client(line);

    // NOTE:  We get "TradeWars Game Server\n"  (Missing \r)
    // We add the \r with our injection line.

    // TODO(stevet): MOVE TO DEFAULT dispatcher
    // our first injection
    if (line.find("TradeWars Game Server") != std::string::npos) {
      to_client("\rTradeWars Proxy v2++ READY (~ to activate)\n\r");
    }

    // How should I handle \r in lines?  For now, remove it
    // but LOG that we did.
    if (replace(part, "\r", "")) {
      BOOST_LOG_TRIVIAL(warning) << "\\r removed from line";
    }

    dispatch_line(part);
  }

  // Ok, we have sent all of the \n lines.

  if (!received.empty())
    if (show_client) {
      to_client(received);
      std::string clean = clean_string(received);
      BOOST_LOG_TRIVIAL(error) << "show_client/leftovers:" << clean;
    }
}

void session::server_read(void) {
  auto self(shared_from_this());

  boost::asio::async_read(
      server_, boost::asio::buffer(server_buffer, sizeof(server_buffer) - 1),
      boost::asio::transfer_at_least(1),
      [this, self](boost::system::error_code ec, std::size_t length) {
        if (!ec) {
          server_buffer[length] = 0;

          // server_prompt.append(server_buffer, length);
          std::string received(server_buffer, length);
          process_lines(received);

          /*
          I don't believe I need to consume this,
          I'm not async_reading from a stream.
          */

          /*
            if (length) {
              // std::cout << length << std::endl;
              std::cout << "S: " << server_buffer << std::endl;
              do_write(server_buffer);
            }
            */

          server_read();
        } else {
          BOOST_LOG_TRIVIAL(warning) << "S: read_failed: socket.shutdown()";
          socket_.shutdown(boost::asio::ip::tcp::socket::shutdown_both);
        }
      });
}

void session::on_resolve(
    const boost::system::error_code error,
    const boost::asio::ip::tcp::resolver::results_type results) {
  //
  auto self(shared_from_this());

  if (!error) {
    // Take the first endpoint.
    boost::asio::ip::tcp::endpoint const &endpoint = *results;

    server_.async_connect(endpoint,
                          boost::bind(&session::on_connect, this,
                                      boost::asio::placeholders::error));

  } else {
    // TO DO:
    // BOOST_LOG_NAMED_SCOPE("session");
    BOOST_LOG_TRIVIAL(error) << "Unable to resolve: " << host;
    std::string output =
        str(boost::format("Unable to resolve: %1%\n\r") % host);
    to_client(output);
    BOOST_LOG_TRIVIAL(warning) << "socket.shutdown()";
    socket_.shutdown(boost::asio::ip::tcp::socket::shutdown_both);
  }
}

void session::do_read(void) {
  auto self(shared_from_this());

  boost::asio::async_read( // why can't I async_read_some here?
      socket_, boost::asio::buffer(read_buffer, sizeof(read_buffer) - 1),
      boost::asio::transfer_at_least(1),
      [this, self](boost::system::error_code ec, std::size_t length) {
        if (!ec) {
          read_buffer[length] = 0;
          if (rlogin_auth.empty()) {
            // first read should be rlogin information
            rlogin_auth.assign(read_buffer, length);

            // parse authentication information
            parse_auth();
            to_client(std::string(1, 0));

            to_client("Welcome, ");
            to_client(rlogin_name);
            to_client("\n\r");

            // Activate the connection to the server

            /* // this fails, and I'm not sure why.  I've used code like this
               before. resolver_.async_resolve( host, port, std::bind(
               &session::on_resolve, this, _1, _2)); */

            // This example shows using boost::bind, which WORKS.
            // https://stackoverflow.com/questions/6025471/bind-resolve-handler-to-resolver-async-resolve-using-boostasio
            resolver_.async_resolve(
                host, port,
                boost::bind(&session::on_resolve, this,
                            boost::asio::placeholders::error,
                            boost::asio::placeholders::iterator));

          } else if (length) {
            // Proxy Active?
            // BOOST_LOG_NAMED_SCOPE("session");

            if (talk_direct)
              to_server(read_buffer);
            BOOST_LOG_TRIVIAL(info) << "C: " << read_buffer;

            // do_write(output);
          }
          do_read();
        } else {
          BOOST_LOG_TRIVIAL(warning) << "C: read_failed";
          if (connected) {
            BOOST_LOG_TRIVIAL(warning) << "server.shutdown()";
            server_.shutdown(boost::asio::ip::tcp::socket::shutdown_both);
          }
        }
      });
}

void session::to_client(std::string message) {
  auto self(shared_from_this());
  // output the cleaned string (so I can see what we're sending in the
  // logs)

  std::string clean = clean_string(message);
  BOOST_LOG_TRIVIAL(trace) << "C: >>" << clean;

  boost::asio::async_write(
      socket_, boost::asio::buffer(message),
      [this, self](boost::system::error_code ec, std::size_t /*length*/) {
        if (!ec) {

        } else {
          BOOST_LOG_TRIVIAL(warning)
              << "C: write failed? closed? server.shutdown()";
          if (connected) {
            BOOST_LOG_TRIVIAL(warning) << "server.shutdown()";
            server_.shutdown(boost::asio::ip::tcp::socket::shutdown_both);
          }
        }
      });
}

void session::to_server(std::string message) {
  auto self(shared_from_this());
  boost::asio::async_write(
      server_, boost::asio::buffer(message),
      [this, self](boost::system::error_code ec, std::size_t /*length*/) {
        if (!ec) {

        } else {
          BOOST_LOG_TRIVIAL(warning)
              << "S: write failed? closed? socket.shutdown()";
          socket_.shutdown(boost::asio::ip::tcp::socket::shutdown_both);
        }
      });
}

/*
void session::on_shutdown(boost::system::error_code ec) {
  std::cout << "shutdown." << std::endl;
}
*/

server::server(boost::asio::io_service &io_service,
               const boost::asio::ip::tcp::endpoint &endpoint, std::string host,
               std::string port)
    : io_service_{io_service}, acceptor_{io_service_, endpoint}, host_{host},
      port_{port} {
  do_accept();
}

/**
 * setup async connect accept
 *
 * This creates a session for each connection.  Using make_shared allows the
 * session to automatically clean up when it is no longer active/has anything
 * running in the reactor.
 */
void server::do_accept(void) {
  acceptor_.async_accept([this](boost::system::error_code ec,
                                boost::asio::ip::tcp::socket socket) {
    if (!ec) {
      BOOST_LOG_TRIVIAL(info) << "server::do_accept()";
      std::make_shared<session>(std::move(socket), io_service_, host_, port_)
          ->start();
    }

    do_accept();
  });
}