twproxy2/twproxy.cpp

//
// client.cpp
// ~~~~~~~~~~
//
// Copyright (c) 2003-2016 Christopher M. Kohlhoff (chris at kohlhoff dot com)
//
// Distributed under the Boost Software License, Version 1.0. (See accompanying
// file LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)
//

#include <boost/asio.hpp>
#include <boost/asio/ip/basic_resolver.hpp>
#include <boost/bind.hpp>
#include <cstdlib>
#include <iostream>

#include "config.h"

class session : public std::enable_shared_from_this<session> {
public:
  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}, host{hostname}, port{port} {
  }
  void start(void) {
    std::cout << "session" << std::endl;
    auto self(shared_from_this());
    // read_buffer.reserve(1024);
    // do_write("Welcome!\n");
    do_read();
  }

  ~session() { std::cout << "~session destructed" << std::endl; }

  void parse_auth(void) {
    // how many nulls should I be seeing?
    // \0user\0pass\0terminal/SPEED\0
    // Maybe in the future I'll care about parsing this out.  I don't right now.

    // Ok, yes I do!  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 on_connect(const boost::system::error_code error) {
    // We've connected to the server!  WOOT WOOT!
    if (!error) {
      std::cout << "Connected to server!" << std::endl;
      do_write("Connected...\n\r");
      connected = true;
      server_read();
    } else {
      // TODO:
      std::string output = "Failed to connect : ";
      output += host;
      output += " : ";
      output += port;
      output += "\n\r";
      do_write(output);
      
      std::cout << "Failed to connect to server." << std::endl;
      std::cout << "SHUTDOWN..." << std::endl;
      socket_.shutdown(boost::asio::ip::tcp::socket::shutdown_both);
    }
  }

  void 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;

            if (length) {
              // std::cout << length << std::endl;
              std::cout << "S: " << server_buffer << std::endl;
              do_write(server_buffer);
            }
            server_read();
          } else {
            std::cout << "S: read_failed: connection closed" << std::endl;
            socket_.shutdown(boost::asio::ip::tcp::socket::shutdown_both);
            // socket_.async_shutdown(boost::bind(&session::on_shutdown, this,
            // boost::asio::placeholders::error));
          }
        });
  }

  void 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:
      std::string output = "Unable to resolve: ";
      output += host;
      output += "\n\r";
      do_write(output); 
      std::cout << "Unable to resolve?" << std::endl;
      std::cout << "SHUTDOWN ..." << std::endl;
      socket_.shutdown(boost::asio::ip::tcp::socket::shutdown_both);
    }
  }

  void 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();
              do_write(std::string(1, 0));

              do_write("Welcome, ");
              do_write(rlogin_name);
              do_write("\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) {
              // std::cout << length << std::endl;
              server_write(read_buffer);
              std::cout << "C: " << read_buffer << std::endl;
              // do_write(output);
            }
            do_read();
          } else {
            std::cout << "C: read_failed: connection closed" << std::endl;
            if (connected)
              server_.shutdown(boost::asio::ip::tcp::socket::shutdown_both);
            // server_.async_shutdown(boost::bind(&session::on_shutdown, this,
            // boost::asio::placeholders::error));
          }
        });
  }

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

          } else {
            std::cout << "write failed? closed?" << std::endl;
            server_.shutdown(boost::asio::ip::tcp::socket::shutdown_both);
            // server_.async_shutdown(boost::bind(&session::on_shutdown, this,
            // boost::asio::placeholders::error));
          }
        });
  }

  void server_write(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 {
            std::cout << "write failed? closed?" << std::endl;
            socket_.shutdown(boost::asio::ip::tcp::socket::shutdown_both);
            // socket_.async_shutdown(boost::bind(&session::on_shutdown, this,
            // boost::asio::placeholders::error));
          }
        });
  }

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

private:
  boost::asio::ip::tcp::socket socket_;
  boost::asio::io_service &io_service_;
  boost::asio::ip::tcp::resolver resolver_;
  boost::asio::ip::tcp::socket server_;

  // std::string read_buffer;
  char read_buffer[1024];
  char server_buffer[1024];

  std::string rlogin_auth;
  std::string rlogin_name;
  std::string host;
  std::string port;
  bool connected = false;
};

/*
maybe move the resolver part to the server, so I don't need io_service?

I'm not sure what the socket connection part is going to need just yet,
so I probably won't move that just yet.  [NNY!]

*/

class server {

public:
  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();
  }

private:
  void do_accept() {
    acceptor_.async_accept([this](boost::system::error_code ec,
                                  boost::asio::ip::tcp::socket socket) {
      if (!ec) {
        std::make_shared<session>(std::move(socket), io_service_, host_, port_)
            ->start();
      }

      do_accept();
    });
  }
  boost::asio::io_service &io_service_;
  boost::asio::ip::tcp::acceptor acceptor_;
  std::string host_;
  std::string port_;
};

int main(int argc, char *argv[]) {
  // boost::json::json_value config;

  if (argc != 2) {
    std::cerr << "Usage: twproxy <filename>" << std::endl;
    return EXIT_FAILURE;
  }

  std::map<std::string, std::string> config = yaml_parse(argv[1]);

  /*
  try {
    // Parse the file as JSON
    config = yaml_parse( argv[1] );

  } catch (std::exception const &e) {
    std::cerr << "Caught exception: " << e.what() << std::endl;
    return EXIT_FAILURE;
  }
  */

  bool config_ok = true;

  // for (const char *key : {"server", "host", "port"}) {
  for (auto key : {"server", "host", "port"}) {
    auto pos = config.find(key);
    if (pos == config.end()) {
      config_ok = false;
      std::cout << "Config file missing: " << key << std::endl;
    }
  }
  if (!config_ok)
    return 2;

  int port = std::stoi(config["server"]);
  // int port = 9999; // 2002;

  try {

    boost::asio::io_service io_service;

    boost::asio::ip::tcp::endpoint endpoint(boost::asio::ip::tcp::v4(),
                                            port); // std::atoi(argv[i]));
    // connect to the BBS
    server s(io_service, endpoint, config["host"], config["port"]);
    //"127.0.0.1", "2023");

    io_service.run();
  } catch (std::exception &e) {
    std::cerr << "Exception: " << e.what() << "\n";
  }

  return EXIT_SUCCESS;
}