Просмотр исходного кода

Proper shutdown, and ~session.

Steve Thielemann 3 лет назад
Родитель
Сommit
6744a0f28f
1 измененных файлов с 93 добавлено и 43 удалено
  1. 93 43
      twproxy.cpp

+ 93 - 43
twproxy.cpp

@@ -16,14 +16,14 @@
 
 #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}, host{hostname}, port{port} {}
+        resolver_{io_service}, server_{io_service}, host{hostname}, port{port} {
+  }
   void start(void) {
     std::cout << "session" << std::endl;
     auto self(shared_from_this());
@@ -32,7 +32,7 @@ public:
     do_read();
   }
 
-  ~session() { std::cout << "~session" << std::endl; }
+  ~session() { std::cout << "~session destructed" << std::endl; }
 
   void parse_auth(void) {
     // how many nulls should I be seeing?
@@ -48,17 +48,66 @@ public:
       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:
+      do_write("Failed to connect to server.\n\r");
+      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 << "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) {
-      for (boost::asio::ip::tcp::endpoint const &endpoint : results) {
-        std::cout << "GOT: " << endpoint << "\n";
-      }
+      // 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:
+      do_write("Unable to resolve?\n\r");
       std::cout << "Unable to resolve?" << std::endl;
+      std::cout << "SHUTDOWN ..." << std::endl;
+      socket_.shutdown(boost::asio::ip::tcp::socket::shutdown_both);
     }
   }
 
@@ -74,6 +123,7 @@ public:
             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));
@@ -83,37 +133,10 @@ public:
               do_write("\n\r");
 
               // Activate the connection to the server
-              // this works!
-              /*
-              boost::system::error_code ec;
-              boost::asio::ip::tcp::resolver::results_type results =
-                  self->resolver_.resolve("google.com", "80", ec);
-
-              for (boost::asio::ip::tcp::endpoint const &endpoint : results) {
-                std::cout << endpoint << "\n";
-              }
-*/
-              // resolver_.async_resolve(host, port, on_resolve);
-
-              // This works, but I can't std::bind to a callback?
-
-              /*
-                            resolver_.async_resolve(
-                                host, port,
-                                [&](boost::system::error_code ec,
-                                    boost::asio::ip::tcp::resolver::results_type
-                 results) { if (!ec) { for (const boost::asio::ip::tcp::endpoint
-                 &endpoint : results) { std::cout << endpoint << "\n";
-                                    }
-                                  } else {
-                                    std::cerr << "Something went wrong";
-                                  }
-                                });
-              */
-              /*
-                            resolver_.async_resolve(
-                                host, port,
-                            std::bind( &session::on_resolve, this, _1, _2)); */
+
+              /* // 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
@@ -125,15 +148,16 @@ public:
 
             } else if (length) {
               // std::cout << length << std::endl;
-              std::cout << read_buffer << std::endl;
-              std::string output = "Got [";
-              output += read_buffer;
-              output += "]\n";
-              do_write(output);
+              server_write(read_buffer);
+              std::cout << "C: " << read_buffer << std::endl;
+              // do_write(output);
             }
             do_read();
           } else {
             std::cout << "read_failed: connection 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));
           }
         });
   }
@@ -147,22 +171,48 @@ public:
 
           } 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;
 };
 
 /*