1

     2

#include "NetworkTCP.hh"

     3

#include "Engine.hh"

     4

     5

#include <cstdlib>

     6

#include <cassert>

     7

     8

NetworkBuffer::NetworkBuffer (NetworkSocket &socket, size_t size_hint) :

     9

    socket(socket), buf(0), size(0), offset(0) {

    10

    11

    // allocate initial buffer

    12

    if ((buf = (char *) malloc(size_hint)) == NULL)

    13

       throw NetworkBufferError("malloc failed");

    14

    15

    // remember size

    16

    size = size_hint;

    17

}

    18

    19

NetworkBuffer::~NetworkBuffer (void) {

    20

    free(buf);

    21

}

    22

    23

void NetworkBuffer::resize (size_t item_size) {

    24

    size_t new_size = size;

    25

    26

    // grow new_size until item_size fits

    27

    while (offset + item_size > new_size)

    28

        new_size *= 2;

    29

    30

    // grow if needed

    31

    if (new_size != size) {

    32

        // realloc buffer

    33

        if ((buf = (char *) realloc((void *) buf, new_size)) == NULL)

    34

            throw NetworkBufferError("realloc failed");

    35

    36

        // update size

    37

        size = new_size;

    38

    39

    } else if (new_size > (offset + item_size) * 4) {

    40

        // XXX: shrink?

    41

    }

    42

}

    43

    44

void NetworkBuffer::trim (size_t prefix_size) {

    45

    // update offset

    46

    offset -= prefix_size;

    47

    48

    // shift the buffer forwards from (buf + prefix) -> (buf), copying (old_offset - prefix) bytes

    49

    memmove(buf, buf + prefix_size, offset);

    50

}

    51

    52

bool NetworkBuffer::try_read (size_t item_size) {

    53

    int ret;

    54

    size_t to_read = item_size;

    55

    56

    // keept reads at at least NETWORK_CHUNK_SIZE bytes

    57

    if (to_read < NETWORK_TCP_CHUNK_SIZE)

    58

        to_read = NETWORK_TCP_CHUNK_SIZE;

    59

    60

    // resize buffer if needed

    61

    resize(to_read);

    62

    63

    // read once

    64

    try {

    65

        ret = socket.recv(buf + offset, to_read);

    66

    67

    } catch (CL_Error &e) {

    68

        if (errno == EAGAIN)

    69

            return false;

    70

    71

        else

    72

            throw NetworkSocketOSError(socket, "recv");

    73

    }

    74

    75

    // handle EOF

    76

    if (ret == 0)

    77

        throw NetworkSocketEOFError(socket, "recv");

    78

    79

    assert(ret >= 0);

    80

    81

    // update offset

    82

    offset += ret;

    83

    84

    // did we get enough?

    85

    if ((unsigned int) ret < item_size)

    86

        return false;

    87

    else

    88

        return true;

    89

} 

    90

    91

bool NetworkBuffer::peek_prefix (uint16_t &ref) {

    92

    if (offset < sizeof(uint16_t))

    93

        return false;

    94

    95

    ref = ntohs(*((uint16_t *) (buf)));

    96

    97

    return true;

    98

}

    99

   100

bool NetworkBuffer::peek_prefix (uint32_t &ref) {

   101

    if (offset < sizeof(uint32_t))

   102

        return false;

   103

   104

    ref = ntohl(*((uint32_t *) (buf)));

   105

   106

    return true;

   107

}

   108

   109

template <typename PrefixType> PrefixType NetworkBuffer::read_prefix (char *buf_ptr, size_t buf_max) {

   110

    PrefixType prefix = 0;

   111

    size_t missing = 0;

   112

   113

    do {    

   114

        // do we have the prefix?

   115

        if (peek_prefix(prefix)) {

   116

            // do we already have the payload?

   117

            if (offset >= sizeof(PrefixType) + prefix) {

   118

                break;

   119

   120

            } else {

   121

                missing = (sizeof(PrefixType) + prefix) - offset;

   122

            }

   123

   124

        } else {

   125

            missing = sizeof(PrefixType);

   126

        }

   127

   128

        // sanity-check

   129

        assert(missing);

   130

   131

        // try and read the missing data

   132

        if (try_read(missing) == false) {

   133

            // if unable to read what we need, return zero.

   134

            return 0;

   135

        }

   136

   137

        // assess the situation again

   138

    } while (true);

   139

   140

    // copy the data over, unless it's too large

   141

    if (prefix <= buf_max) {

   144

   145

        // trim the bytes out

   146

        trim(sizeof(PrefixType) + prefix);

   147

   148

        // return

   149

        return prefix;

   150

   151

    } else {

   152

        // trim the bytes out

   153

        trim(sizeof(PrefixType) + prefix);

   154

   155

        throw NetworkBufferError("recv prefix overflow");   

   156

    }

   157

}

   158

   159

void NetworkBuffer::push_write (char *buf_ptr, size_t buf_size) {

   160

    int ret;

   161

   162

    // try and short-circuit writes unless we have already buffered data

   163

    if (offset == 0) {

   164

        try {

   165

            // attempt to send something

   166

            ret = socket.send(buf_ptr, buf_size);

   167

   168

        } catch (CL_Error &e) {

   169

            // ignore EAGAIN, detect this by setting ret to -1

   170

            if (errno != EAGAIN)

   171

                throw NetworkSocketOSError(socket, "send");

   172

   173

            ret = -1;

   174

        }

   175

   176

        // if we managed to send something, adjust buf/size and buffer

   177

        if (ret > 0) {

   178

            // sanity-check

   179

            assert(buf_size >= (unsigned int) ret);

   180

   181

            buf_ptr += ret;

   182

            buf_size -= ret;

   183

   184

            // if that was all, we're done

   185

            if (buf_size == 0)

   186

                return;

   187

        }

   188

    }

   189

   190

    // resize to fit buf_size more bytes

   191

    resize(buf_size);

   192

   193

    // copy into our internal buffer

   194

    memcpy(buf + offset, buf_ptr, buf_size);

   195

}

   196

   197

void NetworkBuffer::flush_write (void) {

   198

    int ret;

   199

   200

    // ignore if we don't have any data buffered

   201

    if (offset == 0)

   202

        return;

   203

   204

    // attempt to write as much as possible

   205

    try {

   206

        ret = socket.send(buf, offset);

   207

   208

    } catch (CL_Error &e) {

   209

        // ignore EAGAIN and just return

   210

        if (errno == EAGAIN)

   211

            return;

   212

   213

        else

   214

            throw NetworkSocketOSError(socket, "send");

   215

    }

   216

   217

    // trim the buffer

   218

    trim(ret);

   219

}

   220

   221

void NetworkBuffer::write_prefix (char *buf, uint16_t prefix) {

   222

    uint16_t nval = htons(prefix);

   223

   224

    push_write((char*) &nval, sizeof(uint16_t)); 

   225

    push_write(buf, prefix);

   226

}

   227

   228

void NetworkBuffer::write_prefix (char *buf, uint32_t prefix) {

   229

    uint32_t nval = htonl(prefix);

   230

   231

    push_write((char*) &nval, sizeof(uint32_t)); 

   232

    push_write(buf, prefix);

   233

}

   234

   235

NetworkTCPTransport::NetworkTCPTransport (NetworkSocket socket) :

   236

    socket(socket), in(socket, NETWORK_TCP_INITIAL_IN_BUF), out(socket, NETWORK_TCP_INITIAL_OUT_BUF) {

   237

   238

    // connect signals

   239

    slots.connect(socket.sig_read_triggered(), this, &NetworkTCPTransport::on_read);

   240

    slots.connect(socket.sig_write_triggered(), this, &NetworkTCPTransport::on_write);

   241

    slots.connect(socket.sig_disconnected(), this, &NetworkTCPTransport::on_disconnected);

   242

}

   243

   244

   245

void NetworkTCPTransport::on_read (void) {

   246

    uint16_t prefix;

   247

    NetworkPacket packet;

   248

   249

    // let the in stream read length-prefixed packets and pass them on to handle_packet

   250

    while ((prefix = in.read_prefix<uint16_t>(packet.get_buf(), packet.get_buf_size())) > 0) {

   251

        packet.set_data_size(prefix);

   252

        _sig_packet(packet);

   253

    }

   254

}

   255

   256

void NetworkTCPTransport::on_write (void) {

   257

    // just flush the output buffer

   258

    out.flush_write();

   259

}

   260

   261

void NetworkTCPTransport::on_disconnected (void) {

   262

    // pass right through

   263

    _sig_disconnect();

   264

}

   265

   266

void NetworkTCPTransport::write_packet (const NetworkPacket &packet) {

   267

    uint16_t prefix = packet.get_data_size();

   268

   269

    if (prefix != packet.get_data_size())

   270

        throw CL_Error("send prefix overflow");

   271

   272

    try {

   273

        // just write to the output buffer

   274

        out.write_prefix((char *) packet.get_buf(), prefix);

   275

   276

    } catch (Error &e) {

   277

        const char *err = e.what();

   278

   279

        Engine::log(ERROR, "tcp.write_packet") << err;

   280

   281

        throw;    

   282

    }

   283

}

   284

   285

NetworkTCPServer::NetworkTCPServer (const NetworkAddress &listen_addr) :

   286

    socket(CL_Socket::tcp, CL_Socket::ipv4) {

   287

   288

    // bind

   289

    socket.bind(listen_addr);

   290

   291

    // assign slots

   292

    slots.connect(socket.sig_read_triggered(), this, &NetworkTCPServer::on_accept);

   293

   294

    // listen

   295

    socket.listen(NETWORK_LISTEN_BACKLOG);

   296

   297

    // use nonblocking sockets

   298

    socket.set_nonblocking(true);

   299

}

   300

   301

   302

void NetworkTCPServer::on_accept (void) {

   303

    // accept a new socket

   304

    NetworkSocket client_sock = socket.accept();

   305

   306

    // create a new NetworkTCPTransport

   307

    NetworkTCPTransport *client = buildTransport(client_sock);

   308

   309

    // let our user handle it

   310

    _sig_client(client);

   311

}

   312

   313

NetworkTCPTransport* NetworkTCPServer::buildTransport (CL_Socket &socket) {

   314

    return new NetworkTCPTransport(socket);

   315

}

   316

   317

NetworkTCPClient::NetworkTCPClient (const NetworkAddress &connect_addr) :

   318

    NetworkTCPTransport(NetworkSocket(CL_Socket::tcp, CL_Socket::ipv4)) {

   319

   320

    // connect

   321

    socket.connect(connect_addr);

   322

   323

    // use nonblocking sockets

   324

    socket.set_nonblocking(true);

   325

}