1// Copyright 2017 The Fuchsia Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5#include <fbl/atomic.h>
6#include <tftp/tftp.h>
7#include <unittest/unittest.h>
8
9#include <limits.h>
10#include <pthread.h>
11#include <stdio.h>
12#include <string.h>
13#include <unistd.h>
14
15// This test simulates a tftp file transfer by running two threads. Both the
16// file and transport interfaces are implemented in memory buffers.
17
18typedef enum { DIR_SEND, DIR_RECEIVE } xfer_dir_t;
19
20struct test_params {
21    xfer_dir_t direction;
22    uint32_t filesz;
23    uint16_t winsz;
24    uint16_t blksz;
25};
26
27static uint8_t *src_file;
28static uint8_t *dst_file;
29
30/* FAUX FILES INTERFACE */
31
32typedef struct {
33    uint8_t* buf;
34    char filename[PATH_MAX + 1];
35    size_t filesz;
36} file_info_t;
37
38// Allocate our src and dst buffers, filling both with random values.
39int initialize_files(struct test_params* tp) {
40    src_file = reinterpret_cast<uint8_t*>(malloc(tp->filesz));
41    dst_file = reinterpret_cast<uint8_t*>(malloc(tp->filesz));
42
43    if (!src_file || !dst_file) {
44        return 1;
45    }
46
47    int* src_as_ints = (int*)src_file;
48    int* dst_as_ints = (int*)dst_file;
49
50    size_t ndx;
51    srand(0);
52    for (ndx = 0; ndx < tp->filesz / sizeof(int); ndx++) {
53        src_as_ints[ndx] = rand();
54        dst_as_ints[ndx] = rand();
55    }
56    for (ndx = (tp->filesz / sizeof(int)) * sizeof(int);
57         ndx < tp->filesz;
58         ndx++) {
59        src_file[ndx] = static_cast<uint8_t>(rand());
60        dst_file[ndx] = static_cast<uint8_t>(rand());
61    }
62
63    return 0;
64}
65
66int compare_files(size_t filesz) {
67    return memcmp(src_file, dst_file, filesz);
68}
69
70const char* file_get_filename(file_info_t* file_info) {
71    return file_info->filename;
72}
73
74ssize_t file_open_read(const char* filename, void* file_cookie) {
75    auto* file_info = reinterpret_cast<file_info_t*>(file_cookie);
76    file_info->buf = src_file;
77    strncpy(file_info->filename, filename, PATH_MAX);
78    file_info->filename[PATH_MAX] = '\0';
79    return file_info->filesz;
80}
81
82tftp_status file_open_write(const char* filename,
83                            size_t size,
84                            void* file_cookie) {
85    auto* file_info = reinterpret_cast<file_info_t*>(file_cookie);
86    file_info->buf = dst_file;
87    file_info->filename[PATH_MAX] = '\0';
88    strncpy(file_info->filename, filename, PATH_MAX);
89    return TFTP_NO_ERROR;
90}
91
92// Every SHORT_READ_FREQ writes we will read a smaller amount instead, to verify behavior when
93// a read operation returns a length less than the requested amount.
94#define SHORT_READ_FREQ 10
95
96tftp_status file_read(void* data, size_t* length, off_t offset,
97                      void* file_cookie) {
98    auto* file_info = reinterpret_cast<file_info_t*>(file_cookie);
99    if ((size_t)offset > file_info->filesz) {
100        // Something has gone wrong in libtftp
101        return TFTP_ERR_INTERNAL;
102    }
103    if ((offset + *length) > file_info->filesz) {
104        *length = file_info->filesz - offset;
105    }
106    static size_t read_count = 0;
107    if (read_count++ % SHORT_READ_FREQ == 0) {
108        *length /= 2;
109    }
110    memcpy(data, &file_info->buf[offset], *length);
111    return TFTP_NO_ERROR;
112}
113
114// Every SHORT_WRITE_FREQ writes we will write a smaller amount instead, to verify behavior when
115// a write operation returns a length less than the requested amount.
116#define SHORT_WRITE_FREQ 10
117
118tftp_status file_write(const void* data, size_t* length, off_t offset,
119                       void* file_cookie) {
120    auto* file_info = reinterpret_cast<file_info_t*>(file_cookie);
121    if (((size_t)offset > file_info->filesz) || ((offset + *length) > file_info->filesz)) {
122        // Something has gone wrong in libtftp
123        return TFTP_ERR_INTERNAL;
124    }
125    static size_t write_count = 0;
126    if (write_count++ % SHORT_WRITE_FREQ == 0) {
127        *length /= 2;
128    }
129    memcpy(&file_info->buf[offset], data, *length);
130    return TFTP_NO_ERROR;
131}
132
133void file_close(void* file_cookie) {
134}
135
136/* FAUX SOCKET INTERFACE */
137
138#define FAKE_SOCK_BUF_SZ 65536
139typedef struct {
140    uint8_t buf[FAKE_SOCK_BUF_SZ];
141    size_t size = FAKE_SOCK_BUF_SZ;
142    fbl::atomic<size_t> read_ndx;
143    fbl::atomic<size_t> write_ndx;
144} fake_socket_t;
145static fake_socket_t client_out_socket;
146static fake_socket_t server_out_socket;
147
148typedef struct {
149    fake_socket_t* in_sock;
150    fake_socket_t* out_sock;
151} transport_info_t;
152
153void clear_sockets(void) {
154    client_out_socket.read_ndx.store(0);
155    client_out_socket.write_ndx.store(0);
156    server_out_socket.read_ndx.store(0);
157    server_out_socket.write_ndx.store(0);
158}
159
160// Initialize "sockets" for either client or server.
161void transport_init(transport_info_t* transport_info, bool is_server) {
162    if (is_server) {
163        transport_info->in_sock = &client_out_socket;
164        transport_info->out_sock = &server_out_socket;
165    } else {
166        transport_info->in_sock = &server_out_socket;
167        transport_info->out_sock = &client_out_socket;
168    }
169}
170
171// Write to our circular message buffer.
172void write_to_buf(fake_socket_t* sock, void* data, size_t size) {
173    uint8_t* in_buf = reinterpret_cast<uint8_t*>(data);
174    uint8_t* out_buf = sock->buf;
175    size_t curr_offset = sock->write_ndx.load() % sock->size;
176    if (curr_offset + size <= sock->size) {
177        memcpy(&out_buf[curr_offset], in_buf, size);
178    } else {
179        size_t first_size = sock->size - curr_offset;
180        size_t second_size = size - first_size;
181        memcpy(out_buf + curr_offset, in_buf, first_size);
182        memcpy(out_buf, in_buf + first_size, second_size);
183    }
184    sock->write_ndx.fetch_add(size);
185}
186
187// Send a message. Note that the buffer's read_ndx and write_ndx don't wrap,
188// which makes it easier to recognize underflow.
189tftp_status transport_send(void* data, size_t len, void* transport_cookie) {
190    auto* transport_info = reinterpret_cast<transport_info_t*>(transport_cookie);
191    fake_socket_t* sock = transport_info->out_sock;
192    while ((sock->write_ndx.load() + sizeof(len) + len - sock->read_ndx.load())
193           > sock->size) {
194        // Wait for the other thread to catch up
195        usleep(10);
196    }
197    write_to_buf(sock, &len, sizeof(len));
198    write_to_buf(sock, data, len);
199    return TFTP_NO_ERROR;
200}
201
202// Read from our circular message buffer. If |move_ptr| is false, just peeks at
203// the data (reads without updating the read pointer).
204void read_from_buf(fake_socket_t* sock, void* data, size_t size,
205                   bool move_ptr) {
206    uint8_t* in_buf = sock->buf;
207    uint8_t* out_buf = reinterpret_cast<uint8_t*>(data);
208    size_t curr_offset = sock->read_ndx.load() % sock->size;
209    if (curr_offset + size <= sock->size) {
210        memcpy(out_buf, &in_buf[curr_offset], size);
211    } else {
212        size_t first_size = sock->size - curr_offset;
213        size_t second_size = size - first_size;
214        memcpy(out_buf, in_buf + curr_offset, first_size);
215        memcpy(out_buf + first_size, in_buf, second_size);
216    }
217    if (move_ptr) {
218        sock->read_ndx.fetch_add(size);
219    }
220}
221
222// Receive a message. Note that the buffer's read_ndx and write_ndx don't
223// wrap, which makes it easier to recognize underflow.
224int transport_recv(void* data, size_t len, bool block, void* transport_cookie) {
225    auto* transport_info = reinterpret_cast<transport_info_t*>(transport_cookie);
226    if (block) {
227        while ((transport_info->in_sock->read_ndx.load() + sizeof(size_t)) >=
228               transport_info->in_sock->write_ndx.load()) {
229            usleep(10);
230        }
231    } else if ((transport_info->in_sock->read_ndx.load() + sizeof(size_t)) >=
232               transport_info->in_sock->write_ndx.load()) {
233        return TFTP_ERR_TIMED_OUT;
234    }
235    size_t block_len;
236    read_from_buf(transport_info->in_sock, &block_len, sizeof(block_len),
237                  false);
238    if (block_len > len) {
239        return TFTP_ERR_BUFFER_TOO_SMALL;
240    }
241    transport_info->in_sock->read_ndx.fetch_add(sizeof(block_len));
242    read_from_buf(transport_info->in_sock, data, block_len, true);
243    return static_cast<int>(block_len);
244}
245
246int transport_timeout_set(uint32_t timeout_ms, void* transport_cookie) {
247    return 0;
248}
249
250/// SEND THREAD
251
252bool run_client_test(struct test_params* tp) {
253    BEGIN_HELPER;
254
255    // Configure TFTP session
256    tftp_session* session;
257    size_t session_size = tftp_sizeof_session();
258    void* session_buf = malloc(session_size);
259    ASSERT_NONNULL(session_buf, "memory allocation failed");
260
261    tftp_status status = tftp_init(&session, session_buf, session_size);
262    ASSERT_EQ(status, TFTP_NO_ERROR, "unable to initialize a tftp session");
263
264    // Configure file interface
265    file_info_t file_info;
266    file_info.filesz = tp->filesz;
267    tftp_file_interface file_callbacks = { file_open_read,
268                                           file_open_write,
269                                           file_read,
270                                           file_write,
271                                           file_close };
272    status = tftp_session_set_file_interface(session, &file_callbacks);
273    ASSERT_EQ(status, TFTP_NO_ERROR, "could not set file interface");
274
275    // Configure transport interface
276    transport_info_t transport_info;
277    transport_init(&transport_info, false);
278
279    tftp_transport_interface transport_callbacks = { transport_send,
280                                                     transport_recv,
281                                                     transport_timeout_set };
282    status = tftp_session_set_transport_interface(session,
283                                                  &transport_callbacks);
284    ASSERT_EQ(status, TFTP_NO_ERROR, "could not set transport interface");
285
286    // Allocate intermediate buffers
287    size_t buf_sz = tp->blksz > PATH_MAX ?
288                    tp->blksz + 2 : PATH_MAX + 2;
289    char* msg_in_buf = reinterpret_cast<char*>(malloc(buf_sz));
290    ASSERT_NONNULL(msg_in_buf, "memory allocation failure");
291    char* msg_out_buf = reinterpret_cast<char*>(malloc(buf_sz));
292    ASSERT_NONNULL(msg_out_buf, "memory allocation failure");
293
294    char err_msg_buf[128];
295
296    // Set our preferred transport options
297    tftp_set_options(session, &tp->blksz, NULL, &tp->winsz);
298
299    tftp_request_opts opts = {};
300    opts.inbuf = msg_in_buf;
301    opts.inbuf_sz = buf_sz;
302    opts.outbuf = msg_out_buf;
303    opts.outbuf_sz = buf_sz;
304    opts.err_msg = err_msg_buf;
305    opts.err_msg_sz = sizeof(err_msg_buf);
306
307    if (tp->direction == DIR_SEND) {
308        status = tftp_push_file(session, &transport_info, &file_info, "abc.txt",
309                                "xyz.txt", &opts);
310        EXPECT_GE(status, 0, "failed to send file");
311    } else {
312        status = tftp_pull_file(session, &transport_info, &file_info, "abc.txt",
313                                "xyz.txt", &opts);
314        EXPECT_GE(status, 0, "failed to receive file");
315    }
316
317    free(session);
318    END_HELPER;
319}
320
321void* tftp_client_main(void* arg) {
322    auto* tp = reinterpret_cast<test_params*>(arg);
323    run_client_test(tp);
324    pthread_exit(NULL);
325}
326
327/// RECV THREAD
328
329bool run_server_test(struct test_params* tp) {
330    BEGIN_HELPER;
331
332    // Configure TFTP session
333    tftp_session* session;
334    size_t session_size = tftp_sizeof_session();
335    void* session_buf = malloc(session_size);
336    ASSERT_NONNULL(session_buf, "memory allocation failed");
337
338    tftp_status status = tftp_init(&session, session_buf, session_size);
339    ASSERT_EQ(status, TFTP_NO_ERROR, "unable to initiate a tftp session");
340
341    // Configure file interface
342    file_info_t file_info;
343    file_info.filesz = tp->filesz;
344    tftp_file_interface file_callbacks = { file_open_read,
345                                           file_open_write,
346                                           file_read,
347                                           file_write,
348                                           file_close };
349    status = tftp_session_set_file_interface(session, &file_callbacks);
350    ASSERT_EQ(status, TFTP_NO_ERROR, "could not set file interface");
351
352    // Configure transport interface
353    transport_info_t transport_info;
354    transport_init(&transport_info, true);
355    tftp_transport_interface transport_callbacks = { transport_send,
356                                                     transport_recv,
357                                                     transport_timeout_set };
358    status = tftp_session_set_transport_interface(session,
359                                                  &transport_callbacks);
360    ASSERT_EQ(status, TFTP_NO_ERROR, "could not set transport interface");
361
362    // Allocate intermediate buffers
363    size_t buf_sz = tp->blksz > PATH_MAX ?
364                    tp->blksz + 2 : PATH_MAX + 2;
365    char* msg_in_buf = reinterpret_cast<char*>(malloc(buf_sz));
366    ASSERT_NONNULL(msg_in_buf, "memory allocation failure");
367    char* msg_out_buf = reinterpret_cast<char*>(malloc(buf_sz));
368    ASSERT_NONNULL(msg_out_buf, "memory allocation failure");
369
370    char err_msg_buf[128];
371    tftp_handler_opts opts = { .inbuf = msg_in_buf,
372                               .inbuf_sz = buf_sz,
373                               .outbuf = msg_out_buf,
374                               .outbuf_sz = &buf_sz,
375                               .err_msg = err_msg_buf,
376                               .err_msg_sz = sizeof(err_msg_buf) };
377    do {
378        status = tftp_service_request(session, &transport_info, &file_info,
379                                      &opts);
380    } while (status == TFTP_NO_ERROR);
381    EXPECT_EQ(status, TFTP_TRANSFER_COMPLETED, "failed to receive file");
382    free(session);
383    END_HELPER;
384}
385
386void* tftp_server_main(void* arg) {
387    auto* tp = reinterpret_cast<test_params*>(arg);
388    run_server_test(tp);
389    pthread_exit(NULL);
390}
391
392bool run_one_test(struct test_params* tp) {
393    BEGIN_TEST;
394    int init_result = initialize_files(tp);
395    ASSERT_EQ(init_result, 0, "failure to initialize state");
396
397    clear_sockets();
398
399    pthread_t client_thread, server_thread;
400    pthread_create(&client_thread, NULL, tftp_client_main, tp);
401    pthread_create(&server_thread, NULL, tftp_server_main, tp);
402
403    pthread_join(client_thread, NULL);
404    pthread_join(server_thread, NULL);
405
406    int compare_result = compare_files(tp->filesz);
407    EXPECT_EQ(compare_result, 0, "output file mismatch");
408    END_TEST;
409}
410
411bool test_tftp_send_file(void) {
412    struct test_params tp = {.direction = DIR_SEND, .filesz = 1000000, .winsz = 20, .blksz = 1000};
413    return run_one_test(&tp);
414}
415
416bool test_tftp_send_file_wrapping_block_count(void) {
417    // Wraps block count 4 times
418    struct test_params tp = {.direction = DIR_SEND, .filesz = 2100000, .winsz = 9999, .blksz = 8};
419    return run_one_test(&tp);
420}
421
422bool test_tftp_send_file_lg_window(void) {
423    // Make sure that a window size > 255 works properly
424    struct test_params tp = {.direction = DIR_SEND, .filesz = 1000000, .winsz = 1024,
425                             .blksz = 1024};
426    return run_one_test(&tp);
427}
428
429bool test_tftp_receive_file(void) {
430    struct test_params tp = {.direction = DIR_RECEIVE, .filesz = 1000000, .winsz = 20,
431                             .blksz = 1000};
432    return run_one_test(&tp);
433}
434
435bool test_tftp_receive_file_wrapping_block_count(void) {
436    // Wraps block count 4 times
437    struct test_params tp = {.direction = DIR_RECEIVE, .filesz = 2100000, .winsz = 8192,
438                             .blksz = 8};
439    return run_one_test(&tp);
440}
441
442bool test_tftp_receive_file_lg_window(void) {
443    // Make sure that a window size > 255 works properly
444    struct test_params tp = {.direction = DIR_RECEIVE, .filesz = 1000000, .winsz = 1024,
445                             .blksz = 1024};
446    return run_one_test(&tp);
447}
448
449BEGIN_TEST_CASE(tftp_transfer_file)
450RUN_TEST(test_tftp_send_file)
451RUN_TEST(test_tftp_send_file_wrapping_block_count)
452RUN_TEST(test_tftp_send_file_lg_window)
453RUN_TEST(test_tftp_receive_file)
454RUN_TEST(test_tftp_receive_file_wrapping_block_count)
455RUN_TEST(test_tftp_receive_file_lg_window)
456END_TEST_CASE(tftp_transfer_file)
457
458