i2pd/WebSocks.cpp

#include "WebSocks.h"
#include "Log.h"
#include <string>

#ifdef WITH_EVENTS
#include "ClientContext.h"
#include "Identity.h"
#include "Destination.h"
#include "Streaming.h"
#include <functional>

#include <websocketpp/config/asio_no_tls.hpp>
#include <websocketpp/server.hpp>

#include <boost/property_tree/ini_parser.hpp>
#define GCC47_BOOST149 ((BOOST_VERSION == 104900) && (__GNUC__ == 4) && (__GNUC_MINOR__ >= 7))
#if !GCC47_BOOST149
#include <boost/property_tree/json_parser.hpp>
#endif

namespace i2p
{
namespace client
{
  typedef websocketpp::server<websocketpp::config::asio> WebSocksServerImpl;

  typedef std::function<void(std::shared_ptr<i2p::stream::Stream>)> StreamConnectFunc;


  struct IWebSocksConn
  {
    virtual void Close() = 0;
    virtual void GotMessage(const websocketpp::connection_hdl & conn, WebSocksServerImpl::message_ptr msg) = 0;
  };

  typedef std::shared_ptr<IWebSocksConn> WebSocksConn_ptr;

  WebSocksConn_ptr CreateWebSocksConn(const websocketpp::connection_hdl & conn, WebSocksImpl * parent);

	class WebSocksImpl
	{

    typedef std::mutex mutex_t;
    typedef std::unique_lock<mutex_t> lock_t;

    typedef std::shared_ptr<ClientDestination> Destination_t;
  public:

    typedef WebSocksServerImpl ServerImpl;
    typedef ServerImpl::message_ptr MessagePtr;

		WebSocksImpl(const std::string & addr, int port) :
      m_Run(false),
      m_Addr(addr),
      m_Port(port),
      m_Thread(nullptr)
		{
      m_Server.init_asio();
      m_Server.set_open_handler(std::bind(&WebSocksImpl::ConnOpened, this, std::placeholders::_1));
      i2p::data::PrivateKeys k = i2p::data::PrivateKeys::CreateRandomKeys(i2p::data::SIGNING_KEY_TYPE_EDDSA_SHA512_ED25519);
      m_Dest = std::make_shared<ClientDestination>(k, false);
		}

    ServerImpl::connection_ptr GetConn(const websocketpp::connection_hdl & conn)
    {
      return m_Server.get_con_from_hdl(conn);
    }

    void CloseConn(const websocketpp::connection_hdl & conn)
    {
      auto c = GetConn(conn);
      if(c) c->close(websocketpp::close::status::normal, "closed");
    }

    void CreateStreamTo(const std::string & addr, int port, StreamConnectFunc complete)
    {
      auto & addressbook = i2p::client::context.GetAddressBook();
      i2p::data::IdentHash ident;
      if(addressbook.GetIdentHash(addr, ident)) {
        // address found
        m_Dest->CreateStream(complete, ident, port);
      } else {
        // not found
        complete(nullptr);
      }
    }

    void ConnOpened(websocketpp::connection_hdl conn)
    {
      m_Conns.push_back(CreateWebSocksConn(conn, this));
    }

    void Start()
    {
      if(m_Run) return; // already started
      m_Server.listen(boost::asio::ip::address::from_string(m_Addr), m_Port);
      m_Server.start_accept();
      m_Run = true;
      m_Thread = new std::thread([&] (){
          while(m_Run) {
            try {
              m_Server.run();
            } catch( std::exception & ex) {
              LogPrint(eLogError, "Websocks runtime exception: ", ex.what());
            }
          }
      });
      m_Dest->Start();
    }

    void Stop()
    {
      for(const auto & conn : m_Conns)
        conn->Close();

      m_Dest->Stop();
      m_Run = false;
      m_Server.stop();
      if(m_Thread) {
        m_Thread->join();
        delete m_Thread;
      }
      m_Thread = nullptr;
    }

  private:
    std::vector<WebSocksConn_ptr> m_Conns;
    bool m_Run;
    ServerImpl m_Server;
    std::string m_Addr;
    int m_Port;
    std::thread * m_Thread;
    Destination_t m_Dest;
	};

  struct WebSocksConn : public IWebSocksConn
  {
    enum ConnState
    {
      eWSCInitial,
      eWSCTryConnect,
      eWSCFailConnect,
      eWSCOkayConnect,
      eWSCClose,
      eWSCEnd
    };

    typedef WebSocksServerImpl ServerImpl;
    typedef ServerImpl::message_ptr Message_t;
    typedef websocketpp::connection_hdl ServerConn;
    typedef std::shared_ptr<ClientDestination> Destination_t;
    typedef std::shared_ptr<i2p::stream::StreamingDestination> StreamDest_t;
    typedef std::shared_ptr<i2p::stream::Stream> Stream_t;

    ServerConn m_Conn;
    Stream_t m_Stream;
    ConnState m_State;
    WebSocksImpl * m_Parent;
    std::string m_RemoteAddr;
    int m_RemotePort;
    uint8_t m_RecvBuf[2048];

    WebSocksConn(const ServerConn & conn, WebSocksImpl * parent) :
      m_Conn(conn),
      m_Stream(nullptr),
      m_State(eWSCInitial),
      m_Parent(parent)
    {

    }

    ~WebSocksConn()
    {
      Close();
    }

    void EnterState(ConnState state)
    {
      LogPrint(eLogDebug, "websocks: state ", m_State, " -> ", state);
      switch(m_State)
      {
      case eWSCInitial:
        if (state == eWSCClose) {
          m_State = eWSCClose;
          // connection was opened but never used
          LogPrint(eLogInfo, "websocks: connection closed but never used");
          Close();
          return;
        } else if (state == eWSCTryConnect) {
          // we will try to connect
          m_State = eWSCTryConnect;
          m_Parent->CreateStreamTo(m_RemoteAddr, m_RemotePort, std::bind(&WebSocksConn::ConnectResult, this, std::placeholders::_1));
        } else {
          LogPrint(eLogWarning, "websocks: invalid state change ", m_State, " -> ", state);
        }
        return;
      case eWSCTryConnect:
        if(state == eWSCOkayConnect) {
          // we connected okay
          LogPrint(eLogDebug, "websocks: connected to ", m_RemoteAddr, ":", m_RemotePort);
          SendResponse("");
          m_State = eWSCOkayConnect;
        } else if(state == eWSCFailConnect) {
          // we did not connect okay
          LogPrint(eLogDebug, "websocks: failed to connect to ", m_RemoteAddr, ":", m_RemotePort);
          SendResponse("failed to connect");
          m_State = eWSCFailConnect;
          EnterState(eWSCInitial);
        } else if(state == eWSCClose) {
          // premature close
          LogPrint(eLogWarning, "websocks: websocket connection closed prematurely");
          m_State = eWSCClose;
        } else {
          LogPrint(eLogWarning, "websocks: invalid state change ", m_State, " -> ", state);
        }
        return;
      case eWSCFailConnect:
        if (state == eWSCInitial) {
          // reset to initial state so we can try connecting again
          m_RemoteAddr = "";
          m_RemotePort = 0;
          LogPrint(eLogDebug, "websocks: reset websocket conn to initial state");
          m_State = eWSCInitial;
        } else if (state == eWSCClose) {
          // we are going to close the connection
          m_State = eWSCClose;
          Close();
        } else {
          LogPrint(eLogWarning, "websocks: invalid state change ", m_State, " -> ", state);
        }
        return;
      case eWSCOkayConnect:
        if(state == eWSCClose) {
          // graceful close
          m_State = eWSCClose;
          Close();
        } else {
          LogPrint(eLogWarning, "websocks: invalid state change ", m_State, " -> ", state);
        }
      case eWSCClose:
        if(state == eWSCEnd) {
          LogPrint(eLogDebug, "websocks: socket ended");
        } else {
          LogPrint(eLogWarning, "websocks: invalid state change ", m_State, " -> ", state);
        }
        return;
      default:
        LogPrint(eLogError, "websocks: bad state ", m_State);
      }
    }

    void StartForwarding()
    {
      LogPrint(eLogDebug, "websocks: begin forwarding data");
			uint8_t b[1];
			m_Stream->Send(b, 0);
      AsyncRecv();
    }

    void HandleAsyncRecv(const boost::system::error_code &ec, std::size_t n)
    {
      if(ec) {
        // error
        LogPrint(eLogWarning, "websocks: connection error ", ec.message());
        EnterState(eWSCClose);
      } else {
        // forward data
        LogPrint(eLogDebug, "websocks recv ", n);
	  
				std::string str((char*)m_RecvBuf, n);
				auto conn = m_Parent->GetConn(m_Conn);
				if(!conn)  {
					LogPrint(eLogWarning, "websocks: connection is gone");
					EnterState(eWSCClose);
					return;
				}
				conn->send(str);
				AsyncRecv();
				
      }
    }

    void AsyncRecv()
    {
      m_Stream->AsyncReceive(
        boost::asio::buffer(m_RecvBuf, sizeof(m_RecvBuf)),
        std::bind(&WebSocksConn::HandleAsyncRecv, this, std::placeholders::_1, std::placeholders::_2), 60);
    }

    /** @brief send error message or empty string for success */
    void SendResponse(const std::string & errormsg)
    {
      boost::property_tree::ptree resp;
      if(errormsg.size()) {
        resp.put("error", errormsg);
        resp.put("success", 0);
      } else {
        resp.put("success", 1);
      }
      std::ostringstream ss;
      write_json(ss, resp);
      auto conn = m_Parent->GetConn(m_Conn);
      if(conn) conn->send(ss.str());
    }

    void ConnectResult(Stream_t stream)
    {
      m_Stream = stream;
      if(m_State == eWSCClose) {
        // premature close of websocket
        Close();
        return;
      }
      if(m_Stream) {
        // connect good
        EnterState(eWSCOkayConnect);
				StartForwarding();
      } else {
        // connect failed
        EnterState(eWSCFailConnect);
      }
    }
    
    virtual void GotMessage(const websocketpp::connection_hdl & conn, WebSocksServerImpl::message_ptr msg)
    {
      (void) conn;
      std::string payload = msg->get_payload();
      if(m_State == eWSCOkayConnect)
      {
        // forward to server
        LogPrint(eLogDebug, "websocks: forward ", payload.size());
        m_Stream->Send((uint8_t*)payload.c_str(), payload.size());
      } else if (m_State == eWSCInitial) {
        // recv connect request
        auto itr = payload.find(":");
        if(itr == std::string::npos) {
          // no port
          m_RemotePort = 0;
          m_RemoteAddr = payload;
        } else {
          // includes port
          m_RemotePort = std::stoi(payload.substr(itr+1));
          m_RemoteAddr = payload.substr(0, itr);
        }
        EnterState(eWSCTryConnect);
      } else {
        // wtf?
        LogPrint(eLogWarning, "websocks: got message in invalid state ", m_State);
      }
    }

    virtual void Close()
    {
      if(m_State == eWSCClose) {
        LogPrint(eLogDebug, "websocks: closing connection");
        if(m_Stream) m_Stream->Close();
        m_Parent->CloseConn(m_Conn);
        EnterState(eWSCEnd);
      } else {
        EnterState(eWSCClose);
      }
    }
  };

  WebSocksConn_ptr CreateWebSocksConn(const websocketpp::connection_hdl & conn, WebSocksImpl * parent)
  {
    auto ptr = std::make_shared<WebSocksConn>(conn, parent);
    auto c = parent->GetConn(conn);
    c->set_message_handler(std::bind(&WebSocksConn::GotMessage, ptr.get(), std::placeholders::_1, std::placeholders::_2));
    return ptr;
  }

}
}
#else

// no websocket support

namespace i2p
{
namespace client
{
	class WebSocksImpl
	{
	public:
		WebSocksImpl(const std::string & addr, int port)
		{
		}

		~WebSocksImpl()
		{
		}

		void Start()
		{
			LogPrint(eLogInfo, "WebSockets not enabled on compile time");
		}

		void Stop()
		{

		}

	};
}
}

#endif
namespace i2p
{
namespace client
{
	WebSocks::WebSocks(const std::string & addr, int port) : m_Impl(new WebSocksImpl(addr, port)) {}
	WebSocks::~WebSocks() { delete m_Impl; }

	void WebSocks::Start()
	{
		m_Impl->Start();
	}

	void WebSocks::Stop()
	{
		m_Impl->Stop();
	}
}
}