boost_asio_test/tcp_connection.cpp at master · blackthunder0812/boost_asio_test · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
#include "tcp_connection.hpp"
#include "tcp_server.hpp"
#include <iostream>
#include <boost/endian/conversion.hpp>
#include <vector>
#include <functional>
#include "allocator_handler.hpp"

tcp_connection::tcp_connection(boost::asio::io_service &io_service, tcp_server *tcp_server_ptr) :
  socket_(io_service),
  tcp_server_ptr(tcp_server_ptr)
{

}

void tcp_connection::write(std::shared_ptr<std::vector<unsigned char>> message)
{
  boost::asio::async_write(socket_,
                           boost::asio::buffer(*message),
                           make_allocation_handler(write_handler_memory_,
                                                   std::bind(&tcp_connection::write_handler,
                                                             shared_from_this(),
                                                             std::placeholders::_1)));
}

void tcp_connection::close()
{
  tcp_server_ptr->get_connection_list().erase(this);
  socket_.close();
}

void tcp_connection::write_handler(const boost::system::error_code &err)
{
  if (err) {
    if (err != boost::asio::error::eof) {
      if (err == boost::asio::error::operation_aborted) {
        std::cerr << "Operation canceled" << std::endl;
      } else {
        std::cerr << "Error writing to client " << socket_.remote_endpoint().address().to_string() << ":" << socket_.remote_endpoint().port() << " : " << err.message() << std::endl;
      }
    } else {
      std::cout << "Connection to " << socket_.remote_endpoint().address().to_string() << ":" << socket_.remote_endpoint().port() << " closed" << std::endl;
    }
    tcp_server_ptr->get_connection_list().erase(this);
    socket_.close();
  }
}

void tcp_connection::read_header()
{
  boost::asio::async_read(socket_,
                          boost::asio::buffer(read_header_buffer),
                          make_allocation_handler(read_header_handler_memory_,
                                                  std::bind(&tcp_connection::read_header_handler,
                                                            shared_from_this(),
                                                            std::placeholders::_1)));
}

void tcp_connection::read_header_handler(const boost::system::error_code &err)
{
  if(err) {
    if (err != boost::asio::error::eof) {
      if (err == boost::asio::error::operation_aborted) {
        std::cerr << "Operation canceled" << std::endl;
      } else {
        std::cerr << "Error reading payload size from client " << socket_.remote_endpoint().address().to_string() << ":" << socket_.remote_endpoint().port() << " : " << err.message() << std::endl;
      }
    } else {
      std::cout << "Connection to " << socket_.remote_endpoint().address().to_string() << ":" << socket_.remote_endpoint().port() << " closed" << std::endl;
    }
    tcp_server_ptr->get_connection_list().erase(this);
    socket_.close();
  } else {
    if (read_header_buffer.data()[4] != 0x55 || read_header_buffer.data()[5] != 0xAA) {
      std::cerr << "Invalid signature: " << std::hex << read_header_buffer.data()[4] << read_header_buffer.data()[5] << std::endl;
      tcp_server_ptr->get_connection_list().erase(this);
      socket_.close();
    } else {
      unsigned int payload_length = boost::endian::big_to_native(*(unsigned int*)(read_header_buffer.data()));
      std::cout << "Payload length: " << payload_length << std::endl;
      if (payload_length == 0) { // heartbeat package: 0x00 0x00 0x00 0x00 0x55 0xAA 0x00 0x00
        read_header();
      } else {
        read_payload(payload_length);
      }
    }
  }
}

// TODO: Create custom handler allocator
void tcp_connection::read_payload(unsigned int payload_length)
{
  std::shared_ptr<std::vector<unsigned char>> payload_ptr(new std::vector<unsigned char>(payload_length));
  boost::asio::async_read(socket_,
                          boost::asio::buffer(payload_ptr->data(), payload_length),
                          make_allocation_handler(read_payload_handler_memory_,
                                                  std::bind(&tcp_connection::read_payload_handler,
                                                            shared_from_this(),
                                                            payload_ptr,
                                                            std::placeholders::_1)));
}

void tcp_connection::read_payload_handler(std::shared_ptr<std::vector<unsigned char>> payload_ptr, const boost::system::error_code &err)
{
  if(err) {
    if (err != boost::asio::error::eof) {
      if (err == boost::asio::error::operation_aborted) {
        std::cerr << "Operation canceled" << std::endl;
      } else {
        std::cerr << "Error reading from client " << socket_.remote_endpoint().address().to_string() << ":" << socket_.remote_endpoint().port() << " : " << err.message() << std::endl;
      }
    } else {
      std::cout << "Connection to " << socket_.remote_endpoint().address().to_string() << ":" << socket_.remote_endpoint().port() << " closed" << std::endl;
    }
    tcp_server_ptr->get_connection_list().erase(this);
    socket_.close();
  } else {
    process_message(payload_ptr);
    read_header();
  }
}

void tcp_connection::process_message(std::shared_ptr<std::vector<unsigned char> > payload_ptr)
{
  size_t payload_size = payload_ptr->size();
  for (size_t i = 0; i < payload_size; i++) {
    std::cout << (*payload_ptr)[i];
  }
  std::cout << std::endl;
}

std::shared_ptr<tcp_connection> tcp_connection::create_connection(boost::asio::io_service &io_service, tcp_server *tcp_server_ptr)
{
  return std::shared_ptr<tcp_connection>(new tcp_connection(io_service, tcp_server_ptr));
}

boost::asio::ip::tcp::socket &tcp_connection::socket()
{
  return socket_;
}

void tcp_connection::start()
{
  read_header();
}