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, NULL, NULL);

   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