1#include "utils.h"
2#include <pico_stack.h>
3#include <pico_tftp.h>
4#include <pico_ipv4.h>
5#include <pico_ipv6.h>
6#include <sys/types.h>
7#include <sys/stat.h>
8#include <unistd.h>
9#include <inttypes.h>
10
11/* Let's use linux fs */
12#include <fcntl.h>
13
14#include <ctype.h>
15
16/*** START TFTP ***/
17#ifdef PICO_SUPPORT_TFTP
18#define TFTP_MODE_SRV 0
19#define TFTP_MODE_CLI 1
20#define TFTP_MODE_PSH 2
21#define TFTP_TX_COUNT 2000
22#define TFTP_PAYLOAD_SIZE 512
23unsigned char tftp_txbuf[TFTP_PAYLOAD_SIZE];
24static uint16_t family;
25
26struct command_t {
27    char operation;
28    char *filename;
29    union pico_address server_address;
30    struct command_t *next;
31};
32
33struct note_t {
34    char *filename;
35    int fd;
36    char direction;
37    int32_t filesize;
38    struct note_t *next;
39};
40
41struct note_t *clipboard = NULL;
42
43struct note_t *add_note(const char *filename, int fd, char direction)
44{
45    struct note_t *note = PICO_ZALLOC(sizeof(struct note_t));
46
47    note->filename = strdup(filename);
48    note->fd = fd;
49    note->direction = direction;
50    note->filesize = 0;
51    note->next = clipboard;
52    clipboard = note;
53    return note;
54}
55
56void del_note(struct note_t *note)
57{
58    struct note_t *prev;
59
60    if (note == clipboard)
61    {
62        clipboard = clipboard->next;
63        if (note->filename)
64            free (note->filename);
65
66        PICO_FREE(note);
67    } else {
68        for (prev = clipboard; prev->next; prev = prev->next)
69            if (prev->next == note) {
70                prev->next = note->next;
71                if (note->filename)
72                    free (note->filename);
73
74                PICO_FREE(note);
75                break;
76            }
77
78    }
79}
80
81struct command_t *add_command(struct command_t *commands, char operation,
82                              char *filename, union pico_address *server_address)
83{
84    struct command_t *command = PICO_ZALLOC(sizeof(struct command_t));
85
86    command->operation = operation;
87    command->filename = filename;
88    memcpy(&command->server_address, server_address, sizeof(union pico_address));
89    command->next = commands;
90    return command;
91}
92
93int32_t get_filesize(const char *filename)
94{
95    int ret;
96    struct stat buf;
97
98    ret = stat(filename, &buf);
99    if (ret)
100        return -1;
101
102    return buf.st_size;
103}
104
105struct note_t *setup_transfer(char operation, const char *filename)
106{
107    int fd;
108
109    printf("operation %c\n", operation);
110    fd = open(filename, (toupper(operation) == 'T') ? O_RDONLY : O_WRONLY | O_EXCL | O_CREAT, 0666);
111    if (fd < 0) {
112        perror("open");
113        fprintf(stderr, "Unable to handle file %s\n", filename);
114        return NULL;
115    }
116
117    return add_note(filename, fd, operation);
118}
119
120int cb_tftp_tx(struct pico_tftp_session *session, uint16_t event, uint8_t *block, int32_t len, void *arg)
121{
122    struct note_t *note = (struct note_t *) arg;
123
124    if (event != PICO_TFTP_EV_OK) {
125        fprintf(stderr, "TFTP: Error %" PRIu16 ": %s\n", event, block);
126        exit(1);
127    }
128
129    len = read(note->fd, tftp_txbuf, PICO_TFTP_PAYLOAD_SIZE);
130
131    if (len >= 0) {
132        note->filesize += len;
133        pico_tftp_send(session, tftp_txbuf, len);
134        if (len < PICO_TFTP_PAYLOAD_SIZE) {
135            printf("TFTP: file %s (%" PRId32 " bytes) TX transfer complete!\n", note->filename, note->filesize);
136            close(note->fd);
137            del_note(note);
138        }
139    } else {
140        perror("read");
141        fprintf(stderr, "Filesystem error reading file %s, cancelling current transfer\n", note->filename);
142        pico_tftp_abort(session, TFTP_ERR_EACC, "Error on read");
143        del_note(note);
144    }
145
146    if (!clipboard) {
147        if (!pico_timer_add(3000, deferred_exit, NULL)) {
148            printf("Failed to start exit timer, exiting now\n");
149            exit(1);
150        }
151    }
152
153    return len;
154}
155
156int cb_tftp_tx_opt(struct pico_tftp_session *session, uint16_t event, uint8_t *block, int32_t len, void *arg)
157{
158    int ret;
159    int32_t filesize;
160
161    if (event == PICO_TFTP_EV_OPT) {
162        ret = pico_tftp_get_option(session, PICO_TFTP_OPTION_FILE, &filesize);
163        if (ret)
164            printf("TFTP: Option filesize is not used\n");
165        else
166            printf("TFTP: We expect to transmit %" PRId32 " bytes\n", filesize);
167
168        event = PICO_TFTP_EV_OK;
169    }
170
171    return cb_tftp_tx(session, event, block, len, arg);
172}
173
174int cb_tftp_rx(struct pico_tftp_session *session, uint16_t event, uint8_t *block, int32_t len, void *arg)
175{
176    struct note_t *note = (struct note_t *) arg;
177    int ret;
178
179    if (event != PICO_TFTP_EV_OK) {
180        fprintf(stderr, "TFTP: Error %" PRIu16 ": %s\n", event, block);
181        exit(1);
182    }
183
184    if (!note)
185        return 0;
186
187    note->filesize += len;
188    if (write(note->fd, block, len) < 0) {
189        perror("write");
190        fprintf(stderr, "Filesystem error writing file %s, cancelling current transfer\n", note->filename);
191        pico_tftp_abort(session, TFTP_ERR_EACC, "Error on write");
192        del_note(note);
193    } else {
194        if (len != PICO_TFTP_PAYLOAD_SIZE) {
195            printf("TFTP: file %s (%" PRId32 " bytes) RX transfer complete!\n", note->filename, note->filesize);
196            close(note->fd);
197            del_note(note);
198        }
199    }
200
201    if (!clipboard) {
202        if (!pico_timer_add(3000, deferred_exit, NULL)) {
203            printf("Failed to start exit timer, exiting now\n");
204            exit(1);
205        }
206    }
207
208    return len;
209}
210
211int cb_tftp_rx_opt(struct pico_tftp_session *session, uint16_t event, uint8_t *block, int32_t len, void *arg)
212{
213    int ret;
214    int32_t filesize;
215
216    if (event == PICO_TFTP_EV_OPT) {
217        ret = pico_tftp_get_option(session, PICO_TFTP_OPTION_FILE, &filesize);
218        if (ret)
219            printf("TFTP: Option filesize is not used\n");
220        else
221            printf("TFTP: We expect to receive %" PRId32 " bytes\n", filesize);
222
223        return 0;
224    }
225
226    return cb_tftp_rx(session, event, block, len, arg);
227}
228
229struct pico_tftp_session *make_session_or_die(union pico_address *addr, uint16_t family)
230{
231    struct pico_tftp_session *session;
232
233    session = pico_tftp_session_setup(addr, family);
234    if (!session) {
235        fprintf(stderr, "TFTP: Error in session setup\n");
236        exit(3);
237    }
238
239    return session;
240}
241
242struct note_t *transfer_prepare(struct pico_tftp_session **psession, char operation, const char *filename, union pico_address *addr, uint16_t family)
243{
244    struct note_t *note;
245
246    note = setup_transfer(operation, filename);
247    *psession = make_session_or_die(addr, family);
248    return note;
249}
250
251void start_rx(struct pico_tftp_session *session, const char *filename, uint16_t port,
252              int (*rx_callback)(struct pico_tftp_session *session, uint16_t err, uint8_t *block, int32_t len, void *arg),
253              struct note_t *note)
254{
255    if (pico_tftp_start_rx(session, port, filename, rx_callback, note)) {
256        fprintf(stderr, "TFTP: Error in initialization\n");
257        exit(1);
258    }
259}
260
261void start_tx(struct pico_tftp_session *session, const char *filename, uint16_t port,
262              int (*tx_callback)(struct pico_tftp_session *session, uint16_t err, uint8_t *block, int32_t len, void *arg),
263              struct note_t *note)
264{
265    if (pico_tftp_start_tx(session, port, filename, tx_callback, note)) {
266        fprintf(stderr, "TFTP: Error in initialization\n");
267        exit(1);
268    }
269}
270
271void tftp_listen_cb(union pico_address *addr, uint16_t port, uint16_t opcode, char *filename, int32_t len)
272{
273    struct note_t *note;
274    struct pico_tftp_session *session;
275
276    printf("TFTP listen callback (BASIC) from remote port %" PRIu16 ".\n", short_be(port));
277    if (opcode == PICO_TFTP_RRQ) {
278        printf("Received TFTP get request for %s\n", filename);
279        note = transfer_prepare(&session, 't', filename, addr, family);
280        start_tx(session, filename, port, cb_tftp_tx, note);
281    } else if (opcode == PICO_TFTP_WRQ) {
282        printf("Received TFTP put request for %s\n", filename);
283        note = transfer_prepare(&session, 'r', filename, addr, family);
284        start_rx(session, filename, port, cb_tftp_rx, note);
285    }
286}
287
288void tftp_listen_cb_opt(union pico_address *addr, uint16_t port, uint16_t opcode, char *filename, int32_t len)
289{
290    struct note_t *note;
291    struct pico_tftp_session *session;
292    int options;
293    uint8_t timeout;
294    int32_t filesize;
295    int ret;
296
297    printf("TFTP listen callback (OPTIONS) from remote port %" PRIu16 ".\n", short_be(port));
298    /* declare the options we want to support */
299    ret = pico_tftp_parse_request_args(filename, len, &options, &timeout, &filesize);
300    if (ret)
301        pico_tftp_reject_request(addr, port, TFTP_ERR_EOPT, "Malformed request");
302
303    if (opcode == PICO_TFTP_RRQ) {
304        printf("Received TFTP get request for %s\n", filename);
305        note = transfer_prepare(&session, 'T', filename, addr, family);
306
307        if (options & PICO_TFTP_OPTION_TIME)
308            pico_tftp_set_option(session, PICO_TFTP_OPTION_TIME, timeout);
309
310        if (options & PICO_TFTP_OPTION_FILE) {
311            ret = get_filesize(filename);
312            if (ret < 0) {
313                pico_tftp_reject_request(addr, port, TFTP_ERR_ENOENT, "File not found");
314                return;
315            }
316
317            pico_tftp_set_option(session, PICO_TFTP_OPTION_FILE, ret);
318        }
319
320        start_tx(session, filename, port, cb_tftp_tx_opt, note);
321    } else { /* opcode == PICO_TFTP_WRQ */
322        printf("Received TFTP put request for %s\n", filename);
323
324        note = transfer_prepare(&session, 'R', filename, addr, family);
325        if (options & PICO_TFTP_OPTION_TIME)
326            pico_tftp_set_option(session, PICO_TFTP_OPTION_TIME, timeout);
327
328        if (options & PICO_TFTP_OPTION_FILE)
329            pico_tftp_set_option(session, PICO_TFTP_OPTION_FILE, filesize);
330
331        start_rx(session, filename, port, cb_tftp_rx_opt, note);
332    }
333}
334
335void print_usage(int exit_code)
336{
337    printf("\nUsage: tftp:OPTION:[OPTION]...\n"
338           "\nOtions can be repeated. Every option may be one of the following:\n"
339           "\ts\t\t\t starts the basic server (RFC1350)\n"
340           "\tS\t\t\t starts the server with option handling capability\n"
341           "\tt:file:ip\t\t PUT request (without options) for file to server ip\n"
342           "\tT:file:ip\t\t PUT request for file to server ip\n"
343           "\tr:file:ip\t\t GET request (without options) for file to server ip\n"
344           "\tR:file:ip\t\t GET request for file to server ip\n"
345           "Example:\n"
346           "\t\t tftp:S:T:firstFile:10.40.0.2:R:another.file:10.40.0.5:T:secondFile:10.40.0.2\n\n");
347    exit(exit_code);
348}
349
350struct command_t *parse_arguments_recursive(struct command_t *commands, char *arg)
351{
352    char *next;
353    char *operation;
354    char *filename;
355    char *address;
356    static union pico_address remote_address;
357    int ret;
358    struct command_t *new_cmd = NULL;
359
360    if (!arg)
361        return commands;
362
363    next = cpy_arg(&operation, arg);
364    switch (*operation) {
365    case 'S':
366    case 's':
367        filename = address = NULL;
368        break;
369    case 'T':
370    case 'R':
371    case 't':
372    case 'r':
373        if (!next) {
374            fprintf(stderr, "Incomplete client command %s (filename componet is missing)\n", arg);
375            return NULL;
376        }
377
378        next = cpy_arg(&filename, next);
379        if (!next) {
380            fprintf(stderr, "Incomplete client command %s (address component is missing)\n", arg);
381            return NULL;
382        }
383
384        next = cpy_arg(&address, next);
385        if (!IPV6_MODE)
386            ret = pico_string_to_ipv4(address, &remote_address.ip4.addr);
387        else
388            ret = pico_string_to_ipv6(address, remote_address.ip6.addr);
389
390        if (ret < 0) {
391            fprintf(stderr, "Invalid IP address %s\n", address);
392            print_usage(2);
393        }
394
395        if (address)
396            free(address);
397
398        break;
399    default:
400        fprintf(stderr, "Invalid command %s\n", operation);
401        return NULL;
402    };
403
404    new_cmd = add_command(commands, *operation, filename, &remote_address);
405    free(operation);
406    return parse_arguments_recursive(new_cmd, next);
407}
408
409struct command_t *parse_arguments(char *arg)
410{
411    struct command_t *reversed = parse_arguments_recursive(NULL, arg);
412    struct command_t *commands = NULL;
413    struct command_t *current;
414
415    if (!reversed) {
416        fprintf(stderr, "Wrong command line!\n");
417        print_usage(1);
418    }
419
420    while (reversed) {
421        current = reversed;
422        reversed = reversed->next;
423        current->next = commands;
424        commands = current;
425    }
426    return commands;
427}
428
429void app_tftp(char *arg)
430{
431    struct command_t *commands, *old_cmd;
432    struct note_t *note;
433    struct pico_tftp_session *session;
434    int is_server_enabled = 0;
435    int filesize;
436
437    family = IPV6_MODE ? PICO_PROTO_IPV6 : PICO_PROTO_IPV4;
438
439    commands = parse_arguments(arg);
440    while (commands) {
441
442        if (toupper(commands->operation) != 'S')
443            note = transfer_prepare(&session, commands->operation, commands->filename, &commands->server_address, family);
444
445        switch (commands->operation) {
446        case 'S':
447        case 's':
448            if (!is_server_enabled) {
449                pico_tftp_listen(PICO_PROTO_IPV4, (commands->operation == 'S') ? tftp_listen_cb_opt : tftp_listen_cb);
450                is_server_enabled = 1;
451            }
452
453            break;
454        case 'T':
455            filesize = get_filesize(commands->filename);
456            if (filesize < 0) {
457                fprintf(stderr, "TFTP: unable to read size of file %s\n", commands->filename);
458                exit(3);
459            }
460
461            pico_tftp_set_option(session, PICO_TFTP_OPTION_FILE, filesize);
462            start_tx(session, commands->filename, short_be(PICO_TFTP_PORT), cb_tftp_tx_opt, note);
463            break;
464        case 't':
465            start_tx(session, commands->filename, short_be(PICO_TFTP_PORT), cb_tftp_tx, note);
466            break;
467        case 'R':
468            pico_tftp_set_option(session, PICO_TFTP_OPTION_FILE, 0);
469            start_rx(session, commands->filename, short_be(PICO_TFTP_PORT), cb_tftp_rx_opt, note);
470            break;
471        case 'r':
472            start_rx(session, commands->filename, short_be(PICO_TFTP_PORT), cb_tftp_rx, note);
473        }
474        old_cmd = commands;
475        commands = commands->next;
476        if (old_cmd->filename)
477            free(old_cmd->filename);
478
479        /* commands are allocated using PICO_ZALLOC, so use PICO_FREE */
480        PICO_FREE(old_cmd);
481    }
482}
483
484#endif
485/* END TFTP */
486