1

#include "sock_test.h"

     2

     3

#include <stdlib.h>

     4

#include <string.h>

     5

#include <assert.h>

     6

     7

/**

     8

 * Grow buf->vecs as needed, to ensure that buf->write_vec is valid

     9

 */

    10

static err_t sock_test_grow_buf (struct io_buf *buf)

    11

{

    12

    size_t read_vec_offset = buf->read_vec ? (buf->read_vec - buf->vecs) : 0;

    13

    size_t write_vec_offset = buf->write_vec ? (buf->write_vec - buf->vecs) : 0;

    14

    struct io_vec *v;

    15

    struct io_vec *vecs_tmp = buf->vecs;

    16

    17

    // don't grow if not full

    18

    if (buf->vecs && buf->write_vec < buf->vecs + buf->count)

    19

        return SUCCESS;

    20

    21

    // new size

    22

    buf->count = buf->count * 2 + 1;

    23

    24

    // grow

    25

    if ((buf->vecs = realloc(buf->vecs, buf->count * sizeof(struct io_vec))) == NULL) {

    26

        // restore old value

    27

        buf->vecs = vecs_tmp;

    28

    29

        return ERR_CALLOC;

    30

    }

    31

    32

    // set vec

    33

    buf->write_vec = buf->vecs + write_vec_offset;

    34

    buf->read_vec = buf->vecs + read_vec_offset;

    35

    36

    // zero new vecs

    37

    for (v = buf->write_vec; v < buf->vecs + buf->count; v++) {

    38

        v->buf = NULL;

    39

        v->len = 0;

    40

    }

    41

    42

    // ok

    43

    return SUCCESS;

    44

}

    45

    46

static err_t sock_test_read (struct sock_stream *base_sock, void *buf_ptr, size_t *len, struct error_info *err)

    47

{

    48

    struct sock_test *sock = SOCK_FROM_BASE(base_sock, struct sock_test);

    49

    struct io_buf *buf = &sock->recv_buf;

    50

    struct io_vec *vec = buf->read_vec;

    51

    52

    // EOF/nonblock if we're past the end of the last vector

    53

    if (!vec || vec == buf->vecs + buf->count || buf->off >= vec->len) {

    54

        if (sock->nonblocking && !sock->eof) {

    55

            // wait for more to be fed in

    56

            *len = 0;

    57

            return SUCCESS;

    58

    59

        } else {

    60

            // EOF!

    61

            return SET_ERROR(err, ERR_READ_EOF);

    62

        }

    63

    }

    64

    65

    // amount of data available in this iovec

    66

    size_t available = vec->len - buf->off;

    67

    68

    // amount to read

    69

    size_t to_read = *len;

    70

    71

    // trim down?

    72

    if (to_read > available)

    73

        to_read = available;

    74

    75

    // copy

    76

    memcpy(buf_ptr, vec->buf + buf->off, to_read);

    77

    78

    // consumed the whole vec?

    79

    if (to_read < available) {

    80

        // move offset

    81

        buf->off += to_read;

    82

    83

    } else {

    84

        // next vector

    85

        buf->read_vec++;

    86

    }

    87

    88

    // update len

    89

    *len = to_read;

    90

    91

    // ok

    92

    return SUCCESS;

    93

}

    94

    95

static err_t sock_test_write (struct sock_stream *base_sock, const void *buf_ptr, size_t *len, struct error_info *err)

    96

{

    97

    struct sock_test *sock = SOCK_FROM_BASE(base_sock, struct sock_test);

    98

    struct io_buf *buf = &sock->send_buf;

    99

   100

    // ensure there's room

   101

    if ((ERROR_CODE(err) = sock_test_grow_buf(buf)))

   102

        goto error;

   103

   104

    // the next buffer

   105

    struct io_vec *vec = buf->write_vec;

   106

   107

    // store

   108

    vec->len = *len;

   109

    assert((vec->buf = malloc(vec->len)));

   110

    memcpy(vec->buf, buf_ptr, vec->len);

   111

   112

    // move vec onwards

   113

    buf->write_vec++;

   114

   115

    // ok

   116

    return SUCCESS;

   117

   118

error:

   119

    return ERROR_CODE(err);

   120

}

   121

   122

static err_t sock_test_event_init (struct sock_stream *base_sock)

   123

{

   124

    struct sock_test *sock = SOCK_FROM_BASE(base_sock, struct sock_test);

   125

   126

    // set the nonblocking flag

   127

    sock->nonblocking = true;

   128

   129

    return SUCCESS;

   130

}

   131

   132

static err_t sock_test_event_enable (struct sock_stream *base_sock, short mask)

   133

{

   134

    struct sock_test *sock = SOCK_FROM_BASE(base_sock, struct sock_test);

   135

   136

    // store mask

   137

    sock->ev_mask = mask;

   138

   139

    return SUCCESS;

   140

}

   141

   142

static void sock_test_release (struct sock_stream *base_sock)

   143

{

   144

    struct sock_test *sock = SOCK_FROM_BASE(base_sock, struct sock_test);

   145

   146

    sock_test_destroy(sock);

   147

}

   148

   149

/*

   150

 * Our sock_stream_type

   151

 */

   152

static struct sock_stream_type sock_test_type = {

   153

    .methods                = {

   154

        .read               = &sock_test_read,

   155

        .write              = &sock_test_write,

   156

        .event_init         = &sock_test_event_init,

   157

        .event_enable       = &sock_test_event_enable,

   158

        .release            = &sock_test_release,

   159

    },

   160

};

   161

   162

struct sock_test* sock_test_create (void)

   163

{

   164

    struct sock_test *sock;

   165

   166

    // allocate

   167

    assert((sock = calloc(1, sizeof(*sock))));

   168

   169

    // initialize base with our sock_stream_type

   170

    sock_stream_init(SOCK_TEST_BASE(sock), &sock_test_type);

   171

   172

    // ok

   173

    return sock;

   174

}

   175

   176

void sock_test_destroy (struct sock_test *sock)

   177

{

   178

    size_t i;

   179

    struct io_buf *sbuf = &sock->send_buf, *rbuf = &sock->recv_buf;

   180

   181

    // free the send buffer

   182

    for (i = 0; i < sbuf->count; i++) {

   183

        free(sbuf->vecs[i].buf);

   184

    }

   185

   186

    // free the buffer vector lists

   187

    free(sbuf->vecs);

   188

    free(rbuf->vecs);

   189

   190

    // free the sock itself

   191

    free(sock);

   192

}

   193

   194

void sock_test_set_recv_buffer (struct sock_test *sock, struct io_vec *vecs, size_t count, bool eof)

   195

{

   196

    struct io_buf *buf = &sock->recv_buf;

   197

   198

    // allocate + copy

   199

    assert((buf->vecs = calloc(count, sizeof(struct io_vec))));

   200

    memcpy(buf->vecs, vecs, count * sizeof(struct io_vec));

   201

   202

    // set

   203

    buf->count = count;

   204

    buf->read_vec = buf->vecs;

   205

    buf->write_vec = buf->vecs + count;

   206

    buf->off = 0;

   207

   208

    // set EOF flag?

   209

    if (eof)

   210

        sock->eof = true;

   211

}

   212

   213

void sock_test_notify_events (struct sock_test *sock)

   214

{

   215

    // notify if events are enabled

   216

    if (sock->ev_mask) {

   217

        // zero mask

   218

        int mask = sock->ev_mask;

   219

        sock->ev_mask = 0;

   220

   221

        sock_stream_invoke_callbacks(SOCK_TEST_BASE(sock), mask);

   222

    }

   223

}

   224

   225

void sock_test_add_recv_vec (struct sock_test *sock, struct io_vec new_vec)

   226

{

   227

    struct io_buf *buf = &sock->recv_buf;

   228

   229

    // ensure there's room

   230

    assert(sock_test_grow_buf(buf) == SUCCESS);

   231

   232

    // copy    

   233

    *(buf->write_vec++) = new_vec;

   234

   235

    // notify

   236

    sock_test_notify_events(sock);

   237

}

   238

   239

void sock_test_add_recv_str (struct sock_test *sock, const char *str)

   240

{

   241

    struct io_vec vec = {

   242

        (char*) str, strlen(str)

   243

    };

   244

   245

    sock_test_add_recv_vec(sock, vec);

   246

}

   247

   248

void sock_test_set_recv_eof (struct sock_test *sock)

   249

{

   250

    sock->eof = true;

   251

   252

    sock_test_notify_events(sock);

   253

}

   254

   255

void sock_test_get_send_data (struct sock_test *sock, char **buf_ptr, size_t *len_ptr)

   256

{

   257

    struct io_buf *buf = &sock->send_buf;

   258

    size_t len = 0, i, off = 0;

   259

    char *out;

   260

   261

    // calculate total size

   262

    for (i = 0; i < buf->count; i++) {

   263

        len += buf->vecs[i].len;

   264

    }

   265

   266

    // alloc

   267

    assert((out = malloc(len)));

   268

   269

    // copy

   270

    for (i = 0; i < buf->count; i++) {

   271

        struct io_vec *vec = buf->vecs + i;

   272

   273

        memcpy(out + off, vec->buf, vec->len);

   274

        off += vec->len;

   275

   276

        // zero

   277

        free(vec->buf); vec->buf = NULL;

   278

        vec->len = 0;

   279

    }

   280

   281

    // update return

   282

    *buf_ptr = out;

   283

    *len_ptr = len;

   284

   285

    // update write_vec

   286

    buf->write_vec = buf->vecs;

   287

}

   288