/*

This file is part of cpp-ethereum.

cpp-ethereum is free software: you can redistribute it and/or modify

it under the terms of the GNU General Public License as published by

the Free Software Foundation, either version 3 of the License, or

(at your option) any later version.

cpp-ethereum is distributed in the hope that it will be useful,

but WITHOUT ANY WARRANTY; without even the implied warranty of

MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the

GNU General Public License for more details.

You should have received a copy of the GNU General Public License

along with cpp-ethereum. If not, see <http://www.gnu.org/licenses/>.

*/

/** @file NetConnection.cpp

* @author Alex Leverington <nessence@gmail.com>

* @author Gav Wood <i@gavwood.com>

* @date 2014

*/

#include <memory>

#include "NetConnection.h"

#include "NetMsg.h"

using namespace std;

using namespace dev;

NetConnection::NetConnection(boost::asio::io_service& _io_service):

m_socket(_io_service),

m_stopped(true),

m_originateConnection(false),

m_socketError()

{

m_stopped = true;

m_started = false;

}

NetConnection::NetConnection(boost::asio::io_service& _io_service, boost::asio::ip::tcp::endpoint _ep):

m_socket(_io_service),

m_endpoint(_ep),

m_stopped(true),

m_originateConnection(true),

m_socketError()

{

m_stopped = true;

m_started = false;

}

NetConnection::~NetConnection()

{

shutdown(true);

}

void NetConnection::setServiceMessageHandler(NetMsgServiceType _svc, messageHandler _h)

{

m_serviceMsgHandlers.insert(std::make_pair(_svc, _h));

}

void NetConnection::setDataMessageHandler(NetMsgServiceType _svc, messageHandler _h)

{

m_dataMsgHandlers.insert(std::make_pair(_svc, _h));

}

boost::asio::ip::tcp::socket* NetConnection::socket() {

return &m_socket;

}

void NetConnection::start()

{

assert(m_started.load() == false);

bool no = false;

if (!m_started.compare_exchange_strong(no, true))

return;

clog(RPCConnect) << (void*)this << (m_originateConnection ? "[egress]" : "[ingress]") << "start()";

if (!m_originateConnection)

handshake();

else

{

auto self(shared_from_this());

m_socket.async_connect(m_endpoint, [self, this](boost::system::error_code const& _ec)

{

if (_ec)

// todo: retry times 3 w/250ms sleep

shutdownWithError(_ec); // Connection failed

else

handshake();

});

}

}

void NetConnection::send(NetMsg const& _msg)

{

if (!m_started)

return;

lock_guard<mutex> l(x_writeQueue);

m_writeQueue.push_back(_msg.payload());

if (m_writeQueue.size() == 1 && !m_stopped)

doWrite();

}

void NetConnection::doWrite()

{

const bytes& bytes = m_writeQueue[0];

auto self(shared_from_this());

boost::asio::async_write(m_socket, boost::asio::buffer(bytes), [this, self](boost::system::error_code _ec, std::size_t /*length*/)

{

if (_ec)

return shutdownWithError(_ec);

else

{

lock_guard<mutex> l(x_writeQueue);

m_writeQueue.pop_front();

if (m_writeQueue.empty())

return;

}

doWrite();

});

}

void NetConnection::handshake(size_t _rlpLen)

{

// read header

if (m_recvdBytes == 0)

{

clog(RPCConnect) << (void*)this << (m_originateConnection ? "[egress]" : "[ingress]") << "handshake() started";

// write version packet

RLPStream s;

s.appendList(1);

s << "version";

shared_ptr<NetMsg> version(new NetMsg((NetMsgServiceType)0, 0, (NetMsgType)0, RLP(s.out())));

// as a special case, message is manually written as

// read/write isn't setup until connection is verified

auto self(shared_from_this());

boost::asio::async_write(m_socket, boost::asio::buffer(version->payload()), [this, self, version](boost::system::error_code _ec, size_t _len)

{

if (_ec)

return shutdownWithError(_ec); // handshake write failed

assert(_len);

// read length header

m_recvBuffer.resize(4);

boost::asio::async_read(m_socket, boost::asio::buffer(m_recvBuffer), [this, self](boost::system::error_code _ec, size_t _len)

{

if (_ec)

return shutdownWithError(_ec); // handshake read-header failed

if (_len == 0)

return shutdownWithError(boost::asio::error::connection_reset); // todo: reread

assert(_len == 4);

m_recvdBytes += _len;

size_t rlpLen = fromBigEndian<uint32_t>(bytesConstRef(m_recvBuffer.data(), 4));

if (rlpLen > 1024*64)

return shutdownWithError(boost::asio::error::connection_reset); // throw MessageTooLarge();

if (rlpLen < 3)

return shutdownWithError(boost::asio::error::connection_reset); // throw MessageTooSmall();

m_recvBuffer.resize(2 * (rlpLen + 4));

assert(rlpLen > 0);

assert(rlpLen < 14);

handshake(rlpLen);

});

});

return;

}

assert(_rlpLen > 0);

auto self(shared_from_this());

m_socket.async_read_some(boost::asio::buffer(boost::asio::buffer(m_recvBuffer) + m_recvdBytes), [this, self, _rlpLen](boost::system::error_code _ec, size_t _len)

{

if (_ec)

return shutdownWithError(_ec); // handshake read-data failed

m_recvdBytes += _len;

if (m_recvdBytes >= _rlpLen + 4)

{

try

{

NetMsg msg(bytesConstRef(m_recvBuffer.data(), _rlpLen + 4));

if (RLP(msg.rlp())[0].toString() == "version")

{

// It's possible for handshake to finish following shutdown

lock_guard<mutex> l(x_socketError);

if (m_started)

{

// start writes if pending messages were created during handshake

lock_guard<mutex> l(x_writeQueue);

if (m_writeQueue.size() > 0)

doWrite();

m_stopped = false;

}

else

return;

clog(RPCNote) << (m_originateConnection ? "[egress]" : "[ingress]") << "Version verified!";

}

else

clog(RPCNote) << (m_originateConnection ? "[egress]" : "[ingress]") << "Version Failed! ";

}

catch (std::exception const& _e)

{

shutdownWithError(boost::asio::error::connection_reset); // handshake failed negotiation

return;

}

m_recvdBytes = m_recvdBytes - _rlpLen - 4;

if (m_recvdBytes)

memmove(m_recvBuffer.data(), m_recvBuffer.data() + _rlpLen + 4, m_recvdBytes);

doRead();

return;

}

assert(_rlpLen > 0);

assert(_rlpLen < 14);

handshake(_rlpLen);

});

}

void NetConnection::doRead(size_t _rlpLen)

{

// Previously, shutdown wouldn't occur if _rlpLen was > 0, however with streaming protocols it's possible remote closes connection before writing a full response. Instead shutdown will defer until writes are complete (see send()).

if (m_stopped)

// if m_stopped is true it is certain shutdown already occurred

return closeSocket();

if (!_rlpLen && m_recvdBytes > 3)

{

// Read buffer if there's enough data for header

clog(RPCNote) << "NetConnection::doRead length header";

_rlpLen = fromBigEndian<uint32_t>(bytesConstRef(m_recvBuffer.data(), 4));

if (_rlpLen > 1024*1024*128)

// todo: managed via service messages and flow-control, requiring message-chunking

return shutdownWithError(boost::asio::error::connection_reset);

if (_rlpLen < 3)

return shutdownWithError(boost::asio::error::connection_reset);

if (4+2*_rlpLen > m_recvBuffer.size())

m_recvBuffer.resize(4+2*_rlpLen);

}

if (_rlpLen && m_recvdBytes >= _rlpLen + 4)

{

try

{

clog(RPCNote) << "NetConnection::doRead message";

NetMsg msg(bytesConstRef(m_recvBuffer.data(), _rlpLen + 4));

if (msg.service())

{

messageHandlers &hs = msg.type() ? m_dataMsgHandlers : m_serviceMsgHandlers;

// If handler is found, get pointer and call

auto hit = hs.find(msg.service());

if (hit!=hs.end())

hit->second(msg);

else

throw MessageServiceInvalid();

}

// else, control messages (ignored right now)

}

catch (std::exception const& _e)

{

// bad message

shutdownWithError(boost::asio::error::connection_reset); // MessageInvalid

return;

}

m_recvdBytes = m_recvdBytes - _rlpLen - 4;

if (m_recvdBytes)

memmove(m_recvBuffer.data(), m_recvBuffer.data() + _rlpLen + 4, m_recvdBytes);

doRead();

}

else

{

// otherwise read more data

auto self(shared_from_this());

m_socket.async_read_some(boost::asio::buffer(boost::asio::buffer(m_recvBuffer) + m_recvdBytes), [this, self, _rlpLen](boost::system::error_code _ec, size_t _len)

{

if (_ec)

return shutdownWithError(_ec);

assert(_len);

m_recvdBytes += _len;

doRead(_rlpLen);

});

}

}

bool NetConnection::connectionError() const

{

// Note: use of mutex here has issues

return m_socketError!=0;

}

bool NetConnection::connectionOpen() const

{

return m_stopped ? false : m_socket.is_open();

}

void NetConnection::shutdownWithError(boost::system::error_code _ec)

{

// If !started and already stopped, shutdown has already occured. (EOF or Operation canceled)

if (!m_started && m_stopped)

return;

assert(_ec);

{

lock_guard<mutex> l(x_socketError);

if (m_socketError != boost::system::error_code())

return;

m_socketError = _ec;

}

clog(RPCWarn) << (void*)this << "Socket shutdown due to error: " << _ec.message();

closeSocket();

shutdown(false);

}

void NetConnection::shutdown(bool _wait)

{

{

lock_guard<mutex> l(x_socketError);

bool yes = true;

if (!m_started.compare_exchange_strong(yes, false))

return;

// if socket never left stopped-state (pre-handshake)

if (m_stopped)

return closeSocket();

else

m_stopped = true;

// immediately close reads

// TODO: (if no error) wait for pending writes

closeSocket();

}

// when m_stopped is true, doRead will close socket

if (_wait)

while (m_socket.is_open() && !connectionError())

usleep(200);

m_recvBuffer.resize(0);

}

void NetConnection::closeSocket()

{

if (m_socket.is_open())

{

boost::system::error_code ec;

m_socket.shutdown(boost::asio::ip::tcp::socket::shutdown_both, ec);

m_socket.close();

}

}