1/*********************************************************************
2   PicoTCP. Copyright (c) 2012-2017 Altran Intelligent Systems. Some rights reserved.
3   See COPYING, LICENSE.GPLv2 and LICENSE.GPLv3 for usage.
4
5   .
6
7   Author: Daniele Lacamera
8 *********************************************************************/
9
10#include <pico_defines.h>
11#include <pico_stack.h>
12#include <pico_socket.h>
13#include <pico_tftp.h>
14#include <pico_strings.h>
15
16#ifdef DEBUG_TFTP
17#define tftp_dbg dbg
18#else
19#define tftp_dbg(...) do {} while(0)
20#endif
21
22/* a zero value means adaptative timeout! (2, 4, 8) */
23#define PICO_TFTP_TIMEOUT 2000U
24
25#define TFTP_MAX_RETRY 3
26
27#define TFTP_STATE_READ_REQUESTED   0
28#define TFTP_STATE_RX               1
29#define TFTP_STATE_LAST_ACK_SENT    2
30#define TFTP_STATE_WRITE_REQUESTED  3
31#define TFTP_STATE_TX               4
32#define TFTP_STATE_WAIT_OPT_CONFIRM 5
33#define TFTP_STATE_WAIT_LAST_ACK    6
34#define TFTP_STATE_CLOSING          7
35
36#define AUTOMA_STATES (TFTP_STATE_CLOSING + 1)
37
38/* MAX_OPTIONS_SIZE: "timeout" 255 "tsize" filesize =>  8 + 4 + 6 + 11 */
39#define MAX_OPTIONS_SIZE 29
40
41/* RRQ and WRQ packets (opcodes 1 and 2 respectively) */
42PACKED_STRUCT_DEF pico_tftp_hdr
43{
44    uint16_t opcode;
45};
46
47/* DATA or ACK (opcodes 3 and 4 respectively)*/
48PACKED_STRUCT_DEF pico_tftp_data_hdr
49{
50    uint16_t opcode;
51    uint16_t block;
52};
53
54/* ERROR (opcode 5) */
55PACKED_STRUCT_DEF pico_tftp_err_hdr
56{
57    uint16_t opcode;
58    uint16_t error_code;
59};
60
61#define PICO_TFTP_TOTAL_BLOCK_SIZE (PICO_TFTP_PAYLOAD_SIZE + (int32_t)sizeof(struct pico_tftp_data_hdr))
62#define tftp_payload(p) (((uint8_t *)(p)) + sizeof(struct pico_tftp_data_hdr))
63
64/* STATUS FLAGS */
65#define SESSION_STATUS_CLOSED       1
66#define SESSION_STATUS_APP_PENDING  2
67#define SESSION_STATUS_IN_CALLBACK  4
68#define SESSION_STATUS_APP_ACK     64
69
70struct pico_tftp_session {
71    int state;
72    int status;
73    int options;
74    int retry;
75    uint16_t packet_counter;
76    /* Current connection */
77    struct pico_socket *socket;
78    union pico_address remote_address;
79    uint16_t remote_port;
80    uint16_t localport;
81    pico_time wallclock_timeout;
82    pico_time bigger_wallclock;
83    struct pico_tftp_session *next;
84    uint32_t timer;
85    unsigned int active_timers;
86    void *argument;
87    int (*callback)(struct pico_tftp_session *session, uint16_t event, uint8_t *block, int32_t len, void *arg);
88    int32_t file_size;
89    int32_t len;
90    uint8_t option_timeout;
91    uint8_t tftp_block[PICO_TFTP_TOTAL_BLOCK_SIZE];
92    int32_t block_len;
93};
94
95struct server_t {
96    void (*listen_callback)(union pico_address *addr, uint16_t port, uint16_t opcode, char *filename, int32_t len);
97    struct pico_socket *listen_socket;
98    uint8_t tftp_block[PICO_TFTP_TOTAL_BLOCK_SIZE];
99};
100
101struct automa_events {
102    void (*ack)(struct pico_tftp_session *session, int32_t len, union pico_address *a, uint16_t port);
103    void (*data)(struct pico_tftp_session *session, int32_t len, union pico_address *a, uint16_t port);
104    void (*error)(struct pico_tftp_session *session, int32_t len, union pico_address *a, uint16_t port);
105    void (*oack)(struct pico_tftp_session *session, int32_t len, union pico_address *a, uint16_t port);
106    void (*timeout)(struct pico_tftp_session *session, pico_time t);
107};
108
109static struct server_t server;
110
111static struct pico_tftp_session *tftp_sessions = NULL;
112
113static inline void session_status_set(struct pico_tftp_session *session, int status)
114{
115    session->status |= status;
116}
117
118static inline void session_status_clear(struct pico_tftp_session *session, int status)
119{
120    session->status &= ~status;
121}
122
123static char *extract_arg_pointer(char *arg, char *end_arg, char **value)
124{
125    char *pos;
126
127    pos = get_string_terminator_position(arg, (size_t)(end_arg - arg));
128    if (!pos)
129        return NULL;
130
131    if (end_arg == ++pos)
132        return NULL;
133
134    arg = get_string_terminator_position(pos, (size_t)(end_arg - pos));
135
136    if (!arg)
137        return NULL;
138
139    *value = pos;
140    return arg + 1;
141}
142
143static int extract_value(char *str, uint32_t *value, uint32_t max)
144{
145    char *endptr;
146    unsigned long num;
147
148    num = strtoul(str, &endptr, 10);
149
150    if (endptr == str || *endptr || num > max)
151        return -1;
152
153    *value = (uint32_t)num;
154    return 0;
155}
156
157static int parse_optional_arguments(char *option_string, int32_t len, int *options, uint8_t *timeout, int32_t *filesize)
158{
159    char *pos;
160    char *end_args = option_string + len;
161    char *current_option;
162    int ret;
163    uint32_t value;
164
165    *options = 0;
166
167    while (option_string < end_args) {
168        current_option = option_string;
169        option_string = extract_arg_pointer(option_string, end_args, &pos);
170        if (!option_string)
171            return 0;
172
173        if (!pico_strncasecmp("timeout", current_option, (size_t)(pos - current_option))) {
174            ret = extract_value(pos, &value, PICO_TFTP_MAX_TIMEOUT);
175            if (ret)
176                return -1;
177
178            *timeout = (uint8_t)value;
179            *options |= PICO_TFTP_OPTION_TIME;
180        } else {
181            if (!pico_strncasecmp("tsize", current_option, (size_t)(pos - current_option))) {
182                ret = extract_value(pos, (uint32_t *)filesize, PICO_TFTP_MAX_FILESIZE);
183                if (ret)
184                    return -1;
185
186                if (*filesize < 0)
187                    return -1;
188
189                *options |= PICO_TFTP_OPTION_FILE;
190            }
191        }
192    }
193    return 0;
194}
195
196static inline struct pico_tftp_session *pico_tftp_session_create(struct pico_socket *sock, union pico_address *remote_addr)
197{
198    struct pico_tftp_session *session;
199
200    session = (struct pico_tftp_session *) PICO_ZALLOC(sizeof (struct pico_tftp_session));
201
202    if (!session)
203        pico_err = PICO_ERR_ENOMEM;
204    else {
205        session->state = 0;
206        session->status = 0;
207        session->options = 0;
208        session->packet_counter = 0u;
209        session->socket = sock;
210        session->wallclock_timeout = 0;
211        session->bigger_wallclock = 0;
212        session->active_timers = 0;
213        session->next = NULL;
214        session->localport = 0;
215        session->callback = NULL;
216        session->argument = NULL;
217        memcpy(&session->remote_address, remote_addr, sizeof(union pico_address));
218        session->remote_port = 0;
219        session->len = 0;
220    }
221
222    return session;
223}
224
225static struct pico_tftp_session *find_session_by_socket(struct pico_socket *tftp_socket)
226{
227    struct pico_tftp_session *pos = tftp_sessions;
228
229    for (; pos; pos = pos->next)
230        if (pos->socket == tftp_socket)
231            return pos;
232
233    return NULL;
234}
235
236/* **************** for future use...
237   static struct pico_tftp_session * find_session_by_localport(uint16_t localport)
238   {
239    struct pico_tftp_session *idx = tftp_sessions;
240
241    for (; idx; idx = idx->next)
242        if (idx->localport == localport)
243            return idx;
244
245    return NULL;
246   } *********************/
247
248static void add_session(struct pico_tftp_session *idx)
249{
250    struct pico_tftp_session *prev = NULL;
251    struct pico_tftp_session *pos;
252
253    for (pos = tftp_sessions; pos; prev = pos, pos = pos->next)
254        if (pos->localport > idx->localport)
255            break;
256
257    if (prev) {
258        idx->next = prev->next;
259        prev->next = idx;
260    } else {
261        idx->next = tftp_sessions;
262        tftp_sessions = idx;
263    }
264}
265
266/* Returns 0 if OK and -1 in case of errors */
267static int del_session(struct pico_tftp_session *idx)
268{
269    struct pico_tftp_session *prev = NULL;
270    struct pico_tftp_session *pos;
271
272    for (pos = tftp_sessions; pos; pos = pos->next) {
273        if (pos == idx) {
274            if (pos == tftp_sessions)
275                tftp_sessions = tftp_sessions->next;
276            else
277                prev->next = pos->next;
278
279            PICO_FREE(idx);
280            return 0;
281        }
282
283        prev = pos;
284    }
285    return -1;
286}
287
288static inline int do_callback(struct pico_tftp_session *session, uint16_t err, uint8_t *data, int32_t len)
289{
290    int ret;
291
292    session_status_set(session, SESSION_STATUS_IN_CALLBACK);
293    ret = session->callback(session, err, data, len, session->argument);
294    session_status_clear(session, SESSION_STATUS_IN_CALLBACK);
295
296    return ret;
297}
298
299static void timer_callback(pico_time now, void *arg);
300static void tftp_finish(struct pico_tftp_session *session);
301
302static void tftp_schedule_timeout(struct pico_tftp_session *session, pico_time interval)
303{
304    pico_time new_timeout = PICO_TIME_MS() + interval;
305
306    if (session->active_timers) {
307        if (session->bigger_wallclock > new_timeout) {
308            session->timer = pico_timer_add(interval + 1, timer_callback, session);
309            if (!session->timer) {
310                tftp_dbg("TFTP: Failed to start callback timer, deleting session\n");
311                tftp_finish(session);
312                return;
313            }
314            session->active_timers++;
315        }
316    } else {
317        session->timer = pico_timer_add(interval + 1, timer_callback, session);
318        if (!session->timer) {
319            tftp_dbg("TFTP: Failed to start callback timer, deleting session\n");
320            tftp_finish(session);
321            return;
322        }
323        session->active_timers++;
324        session->bigger_wallclock = new_timeout;
325    }
326
327    session->wallclock_timeout = new_timeout;
328}
329
330static void tftp_finish(struct pico_tftp_session *session)
331{
332    if (session->state != TFTP_STATE_CLOSING) {
333        pico_socket_close(session->socket);
334        session->state = TFTP_STATE_CLOSING;
335        if (session->active_timers) {
336            pico_timer_cancel(session->timer);
337            --session->active_timers;
338        }
339
340        session->wallclock_timeout = 0;
341        tftp_schedule_timeout(session, 5);
342    }
343}
344
345static void tftp_send(struct pico_tftp_session *session, int len)
346{
347    if (len)
348        session->len = len;
349    else
350        len = session->len;
351
352    pico_socket_sendto(session->socket, session->tftp_block, session->len, &session->remote_address, session->remote_port);
353}
354
355static void tftp_send_ack(struct pico_tftp_session *session)
356{
357    struct pico_tftp_data_hdr *dh;
358
359    dh = PICO_ZALLOC(sizeof(struct pico_tftp_data_hdr));
360    if (!dh)
361        return;
362
363    dh->opcode = short_be(PICO_TFTP_ACK);
364    dh->block = short_be(session->packet_counter);
365
366    if (session->socket) {
367        pico_socket_sendto(session->socket, dh, (int) sizeof(struct pico_tftp_err_hdr),
368                           &session->remote_address, session->remote_port);
369        tftp_schedule_timeout(session, PICO_TFTP_TIMEOUT);
370    }
371
372    PICO_FREE(dh);
373}
374
375static size_t prepare_options_string(struct pico_tftp_session *session, char *str_options, int32_t filesize)
376{
377    size_t len = 0;
378    int res;
379
380    if (session->options & PICO_TFTP_OPTION_TIME) {
381        strcpy(str_options, "timeout");
382        len += 8;
383        res = num2string(session->option_timeout, &str_options[len], 4);
384        if (res < 0)
385            return 0;
386
387        len += (size_t)res;
388    }
389
390    if (session->options & PICO_TFTP_OPTION_FILE) {
391        strcpy(&str_options[len], "tsize");
392        len += 6;
393        res = num2string(filesize, &str_options[len], 11);
394        if (res < 0)
395            return 0;
396
397        len += (size_t)res;
398    }
399
400    return len;
401}
402
403static void tftp_send_oack(struct pico_tftp_session *session)
404{
405    struct pico_tftp_hdr *hdr;
406    size_t options_size;
407    size_t options_pos = sizeof(struct pico_tftp_hdr);
408    uint8_t *buf;
409    char str_options[MAX_OPTIONS_SIZE] = {
410        0
411    };
412
413    options_size = prepare_options_string(session, str_options, session->file_size);
414
415    buf = PICO_ZALLOC(options_pos + options_size);
416    if (!buf) {
417        strcpy((char *)session->tftp_block, "Out of memory");
418        do_callback(session, PICO_TFTP_EV_ERR_LOCAL, session->tftp_block, 0);
419        tftp_finish(session);
420        return;
421    }
422
423    hdr = (struct pico_tftp_hdr *)buf;
424    hdr->opcode = short_be(PICO_TFTP_OACK);
425    memcpy(buf + options_pos, str_options, options_size);
426    (void)pico_socket_sendto(session->socket, buf, (int)(options_pos + options_size), &session->remote_address, session->remote_port);
427    PICO_FREE(buf);
428}
429
430static void tftp_send_req(struct pico_tftp_session *session, union pico_address *a, uint16_t port, const char *filename, uint16_t opcode)
431{
432#define OCTET_STRSIZ 7U
433    static const char octet[OCTET_STRSIZ] = {
434        0, 'o', 'c', 't', 'e', 't', 0
435    };
436    struct pico_tftp_hdr *hdr;
437    size_t len;
438    size_t options_size;
439    size_t options_pos;
440    uint8_t *buf;
441    char str_options[MAX_OPTIONS_SIZE] = {
442        0
443    };
444
445    if (!filename) {
446        return;
447    }
448
449    len = strlen(filename);
450
451    options_size = prepare_options_string(session, str_options, (opcode == PICO_TFTP_WRQ) ? (session->file_size) : (0));
452
453    options_pos = sizeof(struct pico_tftp_hdr) + OCTET_STRSIZ + len;
454    buf = PICO_ZALLOC(options_pos + options_size);
455    if (!buf) {
456        strcpy((char *)session->tftp_block, "Out of memory");
457        do_callback(session, PICO_TFTP_EV_ERR_LOCAL, session->tftp_block, 0);
458        tftp_finish(session);
459        return;
460    }
461
462    hdr = (struct pico_tftp_hdr *)buf;
463    hdr->opcode = short_be(opcode);
464    memcpy(buf + sizeof(struct pico_tftp_hdr), filename, len);
465    memcpy(buf + sizeof(struct pico_tftp_hdr) + len, octet, OCTET_STRSIZ);
466    memcpy(buf + options_pos, str_options, options_size);
467    (void)pico_socket_sendto(session->socket, buf, (int)(options_pos + options_size), a, port);
468    PICO_FREE(buf);
469}
470
471static void tftp_send_rx_req(struct pico_tftp_session *session, union pico_address *a, uint16_t port, const char *filename)
472{
473    tftp_send_req(session, a, port, filename, PICO_TFTP_RRQ);
474    session->state = TFTP_STATE_READ_REQUESTED;
475    tftp_schedule_timeout(session, PICO_TFTP_TIMEOUT);
476}
477
478static void tftp_send_tx_req(struct pico_tftp_session *session, union pico_address *a, uint16_t port, const char *filename)
479{
480    tftp_send_req(session, a, port, filename, PICO_TFTP_WRQ);
481    session->state = TFTP_STATE_WRITE_REQUESTED;
482    tftp_schedule_timeout(session, PICO_TFTP_TIMEOUT);
483}
484
485static int send_error(uint8_t *buf, struct pico_socket *sock, union pico_address *a, uint16_t port, uint16_t errcode, const char *errmsg)
486{
487    struct pico_tftp_err_hdr *eh;
488    int32_t len;
489    int32_t maxlen = PICO_TFTP_TOTAL_BLOCK_SIZE - sizeof(struct pico_tftp_err_hdr);
490
491    if (!errmsg)
492        len = 0;
493    else
494        len = (int32_t)strlen(errmsg);
495
496    eh = (struct pico_tftp_err_hdr *) buf;
497    eh->opcode = short_be(PICO_TFTP_ERROR);
498    eh->error_code = short_be(errcode);
499    if (len + 1 > maxlen)
500        len = maxlen;
501
502    if (len)
503        memcpy(tftp_payload(eh), errmsg, (size_t)len);
504
505    tftp_payload(eh)[len++] = (char)0;
506
507    return pico_socket_sendto(sock, eh, (int)(len + (int32_t)sizeof(struct pico_tftp_err_hdr)), a, port);
508}
509
510static void tftp_send_error(struct pico_tftp_session *session, union pico_address *a, uint16_t port, uint16_t errcode, const char *errmsg)
511{
512    struct pico_tftp_err_hdr *eh;
513    int32_t len;
514    int32_t maxlen = PICO_TFTP_TOTAL_BLOCK_SIZE - sizeof(struct pico_tftp_err_hdr);
515
516    if (!errmsg)
517        len = 0;
518    else
519        len = (int32_t)strlen(errmsg);
520
521    if (!a) {
522        a = &session->remote_address;
523        port = session->remote_port;
524    }
525
526    eh = (struct pico_tftp_err_hdr *) (session ? (session->tftp_block) : (server.tftp_block));
527    eh->opcode = short_be(PICO_TFTP_ERROR);
528    eh->error_code = short_be(errcode);
529    if (len + 1 > maxlen)
530        len = maxlen;
531
532    if (len)
533        memcpy(tftp_payload(eh), errmsg, (size_t)len);
534
535    tftp_payload(eh)[len++] = (char)0;
536    if (session) {
537        (void)pico_socket_sendto(session->socket, eh, (int) (len + (int32_t)sizeof(struct pico_tftp_err_hdr)), a, port);
538        tftp_finish(session);
539    } else
540        (void)pico_socket_sendto(server.listen_socket, eh, (int) (len + (int32_t)sizeof(struct pico_tftp_err_hdr)), a, port);
541}
542
543static void tftp_send_data(struct pico_tftp_session *session, const uint8_t *data, int32_t len)
544{
545    struct pico_tftp_data_hdr *dh;
546
547    dh = (struct pico_tftp_data_hdr *) session->tftp_block;
548    dh->opcode = short_be(PICO_TFTP_DATA);
549    dh->block = short_be(session->packet_counter++);
550
551    if (len < PICO_TFTP_PAYLOAD_SIZE)
552        session->state = TFTP_STATE_WAIT_LAST_ACK;
553    else
554        session->state = TFTP_STATE_TX;
555
556    memcpy(session->tftp_block + sizeof(struct pico_tftp_data_hdr), data, (size_t)len);
557    pico_socket_sendto(session->socket, session->tftp_block, (int)(len + (int32_t)sizeof(struct pico_tftp_data_hdr)),
558                       &session->remote_address, session->remote_port);
559    tftp_schedule_timeout(session, PICO_TFTP_TIMEOUT);
560}
561
562static inline void tftp_eval_finish(struct pico_tftp_session *session, int32_t len)
563{
564    if (len < PICO_TFTP_PAYLOAD_SIZE) {
565        pico_socket_close(session->socket);
566        session->state = TFTP_STATE_CLOSING;
567    }
568}
569
570static inline int tftp_data_prepare(struct pico_tftp_session *session, union pico_address *a, uint16_t port)
571{
572    if (!session->socket)
573        return -1;
574
575    if (pico_address_compare(a, &session->remote_address, session->socket->net->proto_number) != 0) {
576        tftp_send_error(session, a, port, TFTP_ERR_EXCEEDED, "TFTP busy, try again later.");
577        return -1;
578    }
579
580    return 0;
581}
582
583static void tftp_req(uint8_t *block, int32_t len, union pico_address *a, uint16_t port)
584{
585    struct pico_tftp_hdr *hdr = (struct pico_tftp_hdr *)block;
586    char *filename;
587    char *pos;
588    char *mode;
589    int ret;
590
591    switch (short_be(hdr->opcode)) {
592    case PICO_TFTP_RRQ:
593    case PICO_TFTP_WRQ:
594        filename = (char *)(block + sizeof(struct pico_tftp_hdr));
595        len -= (int32_t)sizeof(struct pico_tftp_hdr);
596
597        pos = extract_arg_pointer(filename, filename + len, &mode);
598        if (!pos) {
599            send_error(block, server.listen_socket, a, port, TFTP_ERR_EILL, "Invalid argument in request");
600            return;
601        }
602
603        ret = strcmp("octet", mode);
604        if (ret) {
605            send_error(block, server.listen_socket, a, port, TFTP_ERR_EILL, "Unsupported mode");
606            return;
607        }
608
609        /*ret = parse_optional_arguments((char *)(block + sizeof(struct pico_tftp_hdr)), len - sizeof(struct pico_tftp_hdr), &new_options, &new_timeout, &new_filesize);
610           if (ret) {
611            tftp_send_error(NULL, a, port, TFTP_ERR_EILL, "Bad request");
612            return;
613           } */
614
615        if (server.listen_callback) {
616            server.listen_callback(a, port, short_be(hdr->opcode), filename, len);
617        }
618
619        break;
620    default:
621        send_error(block, server.listen_socket, a, port, TFTP_ERR_EILL, "Illegal opcode");
622    }
623}
624
625static int event_ack_base(struct pico_tftp_session *session, int32_t len, union pico_address *a, uint16_t port)
626{
627    struct pico_tftp_data_hdr *dh;
628    uint16_t block_n;
629    const char *wrong_address = "Wrong address";
630    const char *wrong_block = "Wrong packet number";
631
632    (void)len;
633    if (pico_address_compare(a, &session->remote_address, session->socket->net->proto_number) != 0) {
634        strcpy((char *)session->tftp_block, wrong_address);
635        do_callback(session, PICO_TFTP_EV_ERR_PEER, session->tftp_block, len);
636        tftp_send_error(session, a, port, TFTP_ERR_EXCEEDED, wrong_address);
637        return -1;
638    }
639
640    dh = (struct pico_tftp_data_hdr *)session->tftp_block;
641    block_n = short_be(dh->block);
642    if (block_n != (session->packet_counter - 1U)) {
643        strcpy((char *)session->tftp_block, wrong_block);
644        do_callback(session, PICO_TFTP_EV_ERR_PEER, session->tftp_block, len);
645        tftp_send_error(session, a, port, TFTP_ERR_EILL, wrong_block);
646        return -1;
647    }
648
649    return 0;
650}
651
652static inline int event_ack0_check(struct pico_tftp_session *session, int32_t len, union pico_address *a, uint16_t port)
653{
654    struct pico_tftp_data_hdr *dh;
655    uint16_t block_n;
656
657    (void)len;
658    if (pico_address_compare(a, &session->remote_address, session->socket->net->proto_number) != 0) {
659        tftp_send_error(session, a, port, TFTP_ERR_EXCEEDED, "TFTP busy, try again later.");
660        return -1;
661    }
662
663    dh = (struct pico_tftp_data_hdr *)session->tftp_block;
664    block_n = short_be(dh->block);
665    if (block_n != 0) {
666        tftp_send_error(session, a, port, TFTP_ERR_EILL, "TFTP connection broken!");
667        return -1;
668    }
669
670    return 0;
671}
672
673static void event_ack0_wr(struct pico_tftp_session *session, int32_t len, union pico_address *a, uint16_t port)
674{
675    if (!event_ack0_check(session, len, a, port)) {
676        session->remote_port = port;
677        do_callback(session, PICO_TFTP_EV_OK, session->tftp_block, 0);
678    }
679}
680
681static void event_ack0_woc(struct pico_tftp_session *session, int32_t len, union pico_address *a, uint16_t port)
682{
683    if (!event_ack0_check(session, len, a, port))
684        do_callback(session, PICO_TFTP_EV_OPT, session->tftp_block, 0);
685}
686
687static void event_ack(struct pico_tftp_session *session, int32_t len, union pico_address *a, uint16_t port)
688{
689    if (!event_ack_base(session, len, a, port))
690        do_callback(session, PICO_TFTP_EV_OK, session->tftp_block, 0);
691}
692
693static void event_ack_last(struct pico_tftp_session *session, int32_t len, union pico_address *a, uint16_t port)
694{
695    if (!event_ack_base(session, len, a, port))
696        tftp_finish(session);
697}
698
699static void event_data(struct pico_tftp_session *session, int32_t len, union pico_address *a, uint16_t port)
700{
701    struct pico_tftp_data_hdr *dh;
702    int32_t payload_len = len - (int32_t)sizeof(struct pico_tftp_data_hdr);
703
704    if (tftp_data_prepare(session, a, port))
705        return;
706
707    dh = (struct pico_tftp_data_hdr *)session->tftp_block;
708    if (short_be(dh->block) > (session->packet_counter +  1U)) {
709        strcpy((char *)session->tftp_block, "Wrong/unexpected sequence number");
710        do_callback(session, PICO_TFTP_EV_ERR_LOCAL, session->tftp_block, 0);
711        tftp_send_error(session, a, port, TFTP_ERR_EILL, "TFTP connection broken!");
712        return;
713    }
714
715    if (short_be(dh->block) == (session->packet_counter + 1U)) {
716        session->packet_counter++;
717        if (do_callback(session, PICO_TFTP_EV_OK, tftp_payload(session->tftp_block), payload_len) >= 0) {
718            if (!(session->status & SESSION_STATUS_APP_ACK))
719                tftp_send_ack(session);
720        }
721
722        if (!(session->status & SESSION_STATUS_APP_ACK))
723            tftp_eval_finish(session, len);
724    }
725}
726
727static void event_data_rdr(struct pico_tftp_session *session, int32_t len, union pico_address *a, uint16_t port)
728{
729    if (tftp_data_prepare(session, a, port))
730        return;
731
732    session->remote_port = port;
733    session->state = TFTP_STATE_RX;
734    event_data(session, len, a, port);
735}
736
737static void event_data_rpl(struct pico_tftp_session *session, int32_t len, union pico_address *a, uint16_t port)
738{
739    struct pico_tftp_data_hdr *dh;
740
741    (void)len;
742    if (tftp_data_prepare(session, a, port))
743        return;
744
745    dh = (struct pico_tftp_data_hdr *)session->tftp_block;
746
747    if (short_be(dh->block) == session->packet_counter)
748        tftp_send_ack(session);
749}
750
751static void event_err(struct pico_tftp_session *session, int32_t len, union pico_address *a, uint16_t port)
752{
753    (void)a;
754    (void)port;
755    do_callback(session, PICO_TFTP_EV_ERR_PEER, session->tftp_block, len);
756    tftp_finish(session);
757}
758
759static inline void event_oack(struct pico_tftp_session *session, int32_t len, union pico_address *a, uint16_t port)
760{
761    char *option_string = (char *)session->tftp_block + sizeof(struct pico_tftp_hdr);
762    int ret;
763    int proposed_options = session->options;
764
765    (void)a;
766
767    session->remote_port = port;
768
769    ret = parse_optional_arguments(option_string, len - (int32_t)sizeof(struct pico_tftp_hdr), &session->options, &session->option_timeout, &session->file_size);
770    if (ret || (session->options & ~proposed_options)) {
771        do_callback(session, PICO_TFTP_EV_ERR_PEER, session->tftp_block, len);
772        tftp_send_error(session, a, port, TFTP_ERR_EOPT, "Invalid option");
773        return;
774    }
775
776    do_callback(session, PICO_TFTP_EV_OPT, session->tftp_block, len);
777}
778
779static void event_oack_rr(struct pico_tftp_session *session, int32_t len, union pico_address *a, uint16_t port)
780{
781    event_oack(session, len, a, port);
782    tftp_send_ack(session);
783    session->state = TFTP_STATE_RX;
784}
785
786static void event_oack_wr(struct pico_tftp_session *session, int32_t len, union pico_address *a, uint16_t port)
787{
788    event_oack(session, len, a, port);
789    session->state = TFTP_STATE_TX;
790}
791
792static void event_timeout(struct pico_tftp_session *session, pico_time t)
793{
794    pico_time new_timeout;
795    int factor;
796
797    (void)t;
798    if (++session->retry == TFTP_MAX_RETRY) {
799        strcpy((char *)session->tftp_block, "Network timeout");
800        do_callback(session, PICO_TFTP_EV_ERR_PEER, session->tftp_block, 0);
801        tftp_finish(session);
802        return;
803    }
804
805    tftp_send(session, 0);
806    if (session->options & PICO_TFTP_OPTION_TIME)
807        new_timeout = session->option_timeout * 1000U;
808    else {
809        new_timeout = PICO_TFTP_TIMEOUT;
810        for (factor = session->retry; factor; --factor)
811            new_timeout *= 2;
812    }
813
814    tftp_schedule_timeout(session, new_timeout);
815}
816
817static void event_timeout_closing(struct pico_tftp_session *session, pico_time t)
818{
819    (void)t;
820    if (session->active_timers == 0)
821        del_session(session);
822}
823
824static void event_timeout_final(struct pico_tftp_session *session, pico_time t)
825{
826    (void)t;
827
828    tftp_finish(session);
829}
830
831static void unexpected(struct pico_tftp_session *session, int32_t len, union pico_address *a, uint16_t port)
832{
833    (void)len;
834    tftp_send_error(session, a, port, TFTP_ERR_EILL, "Unexpected message");
835}
836
837static void null(struct pico_tftp_session *session, int32_t len, union pico_address *a, uint16_t port)
838{
839    (void)session;
840    (void)len;
841    (void)a;
842    (void)port;
843}
844
845static struct automa_events fsm[AUTOMA_STATES] = {
846    /*   STATE                       *     ACK          DATA            ERROR       OACK            TIMEOUT              */
847    /* ***************************************************************************************************************** */
848    { /* TFTP_STATE_READ_REQUESTED   */ unexpected,     event_data_rdr, event_err,  event_oack_rr,  event_timeout},
849    { /* TFTP_STATE_RX               */ unexpected,     event_data,     event_err,  unexpected,     event_timeout},
850    { /* TFTP_STATE_LAST_ACK_SENT    */ unexpected,     event_data_rpl, null,       unexpected,     event_timeout_final},
851    { /* TFTP_STATE_WRITE_REQUESTED  */ event_ack0_wr,  unexpected,     event_err,  event_oack_wr,  event_timeout},
852    { /* TFTP_STATE_TX               */ event_ack,      unexpected,     event_err,  unexpected,     event_timeout},
853    { /* TFTP_STATE_WAIT_OPT_CONFIRM */ event_ack0_woc, unexpected,     event_err,  unexpected,     event_timeout},
854    { /* TFTP_STATE_WAIT_LAST_ACK    */ event_ack_last, unexpected,     event_err,  unexpected,     event_timeout},
855    { /* TFTP_STATE_CLOSING          */ null,           null,           null,       null,           event_timeout_closing}
856};
857
858static void tftp_message_received(struct pico_tftp_session *session, uint8_t *block, int32_t len, union pico_address *a, uint16_t port)
859{
860    struct pico_tftp_hdr *th = (struct pico_tftp_hdr *) block;
861
862    if (!session->callback)
863        return;
864
865    session->wallclock_timeout = 0;
866
867    switch (short_be(th->opcode)) {
868    case PICO_TFTP_RRQ:
869    case PICO_TFTP_WRQ:
870        unexpected(session, len, a, port);
871        break;
872    case PICO_TFTP_DATA:
873        fsm[session->state].data(session, len, a, port);
874        break;
875    case PICO_TFTP_ACK:
876        fsm[session->state].ack(session, len, a, port);
877        break;
878    case PICO_TFTP_ERROR:
879        fsm[session->state].error(session, len, a, port);
880        break;
881    case PICO_TFTP_OACK:
882        fsm[session->state].oack(session, len, a, port);
883        break;
884    default:
885        tftp_send_error(session, NULL, 0, TFTP_ERR_EILL, "Illegal opcode");
886    }
887}
888
889static void tftp_cb(uint16_t ev, struct pico_socket *s)
890{
891    int r;
892    struct pico_tftp_session *session;
893    union pico_address ep;
894    uint16_t port = 0;
895
896    session = find_session_by_socket(s);
897    if (session) {
898        if (ev == PICO_SOCK_EV_ERR) {
899            strcpy((char *)session->tftp_block, "Socket Error");
900            do_callback(session, PICO_TFTP_EV_ERR_LOCAL, session->tftp_block, (int32_t)strlen((char *)session->tftp_block));
901            tftp_finish(session);
902            return;
903        }
904
905        r = pico_socket_recvfrom(s, session->tftp_block, PICO_TFTP_TOTAL_BLOCK_SIZE, &ep, &port);
906        if (r < (int)sizeof(struct pico_tftp_hdr))
907            return;
908
909        tftp_message_received(session, session->tftp_block, r, &ep, port);
910    } else {
911        if (!server.listen_socket || s != server.listen_socket) {
912            return;
913        }
914
915        r = pico_socket_recvfrom(s, server.tftp_block, PICO_TFTP_TOTAL_BLOCK_SIZE, &ep, &port);
916        if (r < (int)sizeof(struct pico_tftp_hdr))
917            return;
918
919        tftp_req(server.tftp_block, r, &ep, port);
920    }
921}
922
923static int application_rx_cb(struct pico_tftp_session *session, uint16_t event, uint8_t *block, int32_t len, void *arg)
924{
925    int *flag = (int *)arg;
926
927    (void)block;
928
929    switch (event) {
930    case PICO_TFTP_EV_ERR_PEER:
931    case PICO_TFTP_EV_ERR_LOCAL:
932        *flag = 0 - event;
933        break;
934    case PICO_TFTP_EV_OK:
935        session->len = len;
936        *flag = 1;
937        break;
938    case PICO_TFTP_EV_OPT:
939        break;
940    }
941    return 0;
942}
943
944static int application_tx_cb(struct pico_tftp_session *session, uint16_t event, uint8_t *block, int32_t len, void *arg)
945{
946    (void)session;
947    (void)block;
948    (void)len;
949
950    *(int*)arg = ((event == PICO_TFTP_EV_OK) || (event == PICO_TFTP_EV_OPT)) ? (1) : (0 - event);
951    return 0;
952}
953
954static void timer_callback(pico_time now, void *arg)
955{
956    struct pico_tftp_session *session = (struct pico_tftp_session *)arg;
957
958    --session->active_timers;
959    if (session->wallclock_timeout == 0) {
960        /* Timer is cancelled. */
961        return;
962    }
963
964    if (now >= session->wallclock_timeout) {
965        session->wallclock_timeout = 0ULL;
966        fsm[session->state].timeout(session, now);
967    } else {
968        tftp_schedule_timeout(session, session->wallclock_timeout - now);
969    }
970}
971
972static struct pico_socket *tftp_socket_open(uint16_t family, uint16_t localport)
973{
974    struct pico_socket *sock;
975    union pico_address local_address;
976
977    sock = pico_socket_open(family, PICO_PROTO_UDP, tftp_cb);
978    if (!sock)
979        return NULL;
980
981    localport = short_be(localport);
982
983    memset(&local_address, 0, sizeof(union pico_address));
984    if (pico_socket_bind(sock, &local_address, &localport) < 0) {
985        pico_socket_close(sock);
986        return NULL;
987    }
988
989    return sock;
990}
991
992static inline int tftp_start_check(struct pico_tftp_session *session, uint16_t port, const char *filename,
993                                   int (*user_cb)(struct pico_tftp_session *session, uint16_t err, uint8_t *block, int32_t len, void *arg))
994{
995    if (!session) {
996        pico_err = PICO_ERR_EINVAL;
997        return -1;
998    }
999
1000    if ((!server.listen_socket) && (port != short_be(PICO_TFTP_PORT))) {
1001        pico_err = PICO_ERR_EINVAL;
1002        return -1;
1003    }
1004
1005    if (!filename) {
1006        pico_err = PICO_ERR_EINVAL;
1007        return -1;
1008    }
1009
1010    if (!user_cb) {
1011        pico_err = PICO_ERR_EINVAL;
1012        return -1;
1013    }
1014
1015    return 0;
1016}
1017
1018/*   ***   EXPORTED FUNCTIONS   ***   */
1019
1020struct pico_tftp_session *pico_tftp_session_setup(union pico_address *a, uint16_t family)
1021{
1022    struct pico_socket *sock;
1023
1024    sock = tftp_socket_open(family, 0);
1025    if (!sock)
1026        return NULL;
1027
1028    return pico_tftp_session_create(sock, a);
1029}
1030
1031int pico_tftp_get_option(struct pico_tftp_session *session, uint8_t type, int32_t *value)
1032{
1033    if (!session) {
1034        pico_err = PICO_ERR_EINVAL;
1035        return -1;
1036    }
1037
1038    switch (type) {
1039    case PICO_TFTP_OPTION_FILE:
1040        if (session->options & PICO_TFTP_OPTION_FILE)
1041            *value = session->file_size;
1042        else {
1043            pico_err = PICO_ERR_ENOENT;
1044            return -1;
1045        }
1046
1047        break;
1048    case PICO_TFTP_OPTION_TIME:
1049        if (session->options & PICO_TFTP_OPTION_TIME)
1050            *value = session->option_timeout;
1051        else {
1052            pico_err = PICO_ERR_ENOENT;
1053            return -1;
1054        }
1055
1056        break;
1057    default:
1058        pico_err = PICO_ERR_EINVAL;
1059        return -1;
1060    }
1061
1062    return 0;
1063}
1064
1065int pico_tftp_set_option(struct pico_tftp_session *session, uint8_t type, int32_t value)
1066{
1067    if (!session) {
1068        pico_err = PICO_ERR_EINVAL;
1069        return -1;
1070    }
1071
1072    switch (type) {
1073    case PICO_TFTP_OPTION_FILE:
1074        if (value < 0) {
1075            pico_err = PICO_ERR_EINVAL;
1076            return -1;
1077        }
1078
1079        session->file_size = value;
1080        session->options |= PICO_TFTP_OPTION_FILE;
1081        break;
1082    case PICO_TFTP_OPTION_TIME:
1083        if (value > PICO_TFTP_MAX_TIMEOUT) {
1084            pico_err = PICO_ERR_EINVAL;
1085            return -1;
1086        }
1087
1088        session->option_timeout = (uint8_t)(value & 0xFF);
1089        if (value) {
1090            session->options |= PICO_TFTP_OPTION_TIME;
1091        } else {
1092            session->options &= ~PICO_TFTP_OPTION_TIME;
1093        }
1094
1095        break;
1096    default:
1097        pico_err = PICO_ERR_EINVAL;
1098        return -1;
1099    }
1100
1101    return 0;
1102}
1103
1104/* Active RX request from PicoTCP */
1105int pico_tftp_start_rx(struct pico_tftp_session *session, uint16_t port, const char *filename,
1106                       int (*user_cb)(struct pico_tftp_session *session, uint16_t event, uint8_t *block, int32_t len, void *arg), void *arg)
1107{
1108    if (tftp_start_check(session, port, filename, user_cb))
1109        return -1;
1110
1111    session->callback = user_cb;
1112    session->packet_counter = 0u;
1113    session->argument = arg;
1114
1115    add_session(session);
1116
1117    if (port != short_be(PICO_TFTP_PORT)) {
1118        session->remote_port = port;
1119        session->state = TFTP_STATE_RX;
1120        if (session->options & (PICO_TFTP_OPTION_FILE | PICO_TFTP_OPTION_TIME))
1121            tftp_send_oack(session);
1122        else
1123            tftp_send_ack(session);
1124    } else {
1125        tftp_send_rx_req(session, &session->remote_address, port, filename);
1126    }
1127
1128    return 0;
1129}
1130
1131int pico_tftp_start_tx(struct pico_tftp_session *session, uint16_t port, const char *filename,
1132                       int (*user_cb)(struct pico_tftp_session *session, uint16_t event, uint8_t *block, int32_t len, void *arg), void *arg)
1133{
1134    if (tftp_start_check(session, port, filename, user_cb))
1135        return -1;
1136
1137    session->callback = user_cb;
1138    session->packet_counter = 1u;
1139    session->argument = arg;
1140
1141    add_session(session);
1142
1143    if (port != short_be(PICO_TFTP_PORT)) {
1144        session->remote_port = port;
1145        if (session->options) {
1146            tftp_send_oack(session);
1147            session->state = TFTP_STATE_WAIT_OPT_CONFIRM;
1148        } else {
1149            do_callback(session, PICO_TFTP_EV_OK, NULL, 0);
1150        }
1151    } else
1152        tftp_send_tx_req(session, &session->remote_address, port, filename);
1153
1154    return 0;
1155}
1156
1157int pico_tftp_reject_request(union pico_address*addr, uint16_t port, uint16_t error_code, const char*error_message)
1158{
1159    return send_error(server.tftp_block, server.listen_socket, addr, port, error_code, error_message);
1160}
1161
1162int32_t pico_tftp_send(struct pico_tftp_session *session, const uint8_t *data, int32_t len)
1163{
1164    int32_t size;
1165
1166
1167    if (len < 0) {
1168        pico_err = PICO_ERR_EINVAL;
1169        return -1;
1170    }
1171
1172    size = len;
1173
1174    if (size > PICO_TFTP_PAYLOAD_SIZE) {
1175        pico_err = PICO_ERR_EINVAL;
1176        return -1;
1177    }
1178
1179    tftp_send_data(session, data, size);
1180
1181    return len;
1182}
1183
1184int pico_tftp_listen(uint16_t family, void (*cb)(union pico_address *addr, uint16_t port, uint16_t opcode, char *filename, int32_t len))
1185{
1186    struct pico_socket *sock;
1187
1188    if (server.listen_socket) {
1189        pico_err = PICO_ERR_EEXIST;
1190        return -1;
1191    }
1192
1193    sock = tftp_socket_open(family, PICO_TFTP_PORT);
1194    if (!sock)
1195        return -1;
1196
1197    server.listen_socket = sock;
1198    server.listen_callback = cb;
1199
1200    return 0;
1201}
1202
1203int pico_tftp_parse_request_args(char *args, int32_t len, int *options, uint8_t *timeout, int32_t *filesize)
1204{
1205    char *pos;
1206    char *end_args = args + len;
1207
1208    args = extract_arg_pointer(args, end_args, &pos);
1209
1210    return parse_optional_arguments(args, (int32_t)(end_args - args), options, timeout, filesize);
1211}
1212
1213int pico_tftp_abort(struct pico_tftp_session *session, uint16_t error, const char *reason)
1214{
1215    if (!session) {
1216        pico_err = PICO_ERR_EINVAL;
1217        return -1;
1218    }
1219
1220    if (!find_session_by_socket(session->socket)) {
1221        pico_err = PICO_ERR_EINVAL;
1222        return -1;
1223    }
1224
1225    tftp_send_error(session, NULL, 0, error, reason);
1226
1227    return 0;
1228}
1229
1230int pico_tftp_close_server(void)
1231{
1232    if (!server.listen_socket) {
1233        pico_err = PICO_ERR_EINVAL;
1234        return -1;
1235    }
1236
1237    pico_socket_close(server.listen_socket);
1238    server.listen_socket = NULL;
1239    return 0;
1240}
1241
1242int pico_tftp_get_file_size(struct pico_tftp_session *session, int32_t *file_size)
1243{
1244    return pico_tftp_get_option(session, PICO_TFTP_OPTION_FILE, file_size);
1245}
1246
1247struct pico_tftp_session *pico_tftp_app_setup(union pico_address *a, uint16_t port, uint16_t family, int *synchro)
1248{
1249    struct pico_tftp_session *session;
1250
1251    if (!synchro) {
1252        pico_err = PICO_ERR_EINVAL;
1253        return NULL;
1254    }
1255
1256    session = pico_tftp_session_setup(a, family);
1257    if (!session)
1258        return NULL;
1259
1260    session->remote_port = port;
1261    session->status |= SESSION_STATUS_APP_ACK;
1262    session->argument = synchro;
1263
1264    *synchro = 0;
1265
1266    return session;
1267}
1268
1269int pico_tftp_app_start_rx(struct pico_tftp_session *session, const char *filename)
1270{
1271    return pico_tftp_start_rx(session, session->remote_port, filename, application_rx_cb, session->argument);
1272}
1273
1274int pico_tftp_app_start_tx(struct pico_tftp_session *session, const char *filename)
1275{
1276    return pico_tftp_start_tx(session, session->remote_port, filename, application_tx_cb, session->argument);
1277}
1278
1279int32_t pico_tftp_get(struct pico_tftp_session *session, uint8_t *data, int32_t len)
1280{
1281    int synchro;
1282
1283    if (!session || len < session->len ) {
1284        pico_err = PICO_ERR_EINVAL;
1285        return -1;
1286    }
1287
1288    synchro = *(int*)session->argument;
1289    *(int*)session->argument = 0;
1290    if ((session->state != TFTP_STATE_RX) && (session->state != TFTP_STATE_READ_REQUESTED))
1291        return -1;
1292
1293    if (synchro < 0)
1294        return synchro;
1295
1296    memcpy(data, tftp_payload(session->tftp_block), (size_t)session->len);
1297    len = session->len;
1298
1299    tftp_send_ack(session);
1300    tftp_eval_finish(session, len);
1301    return len;
1302}
1303
1304int32_t pico_tftp_put(struct pico_tftp_session *session, uint8_t *data, int32_t len)
1305{
1306    int synchro;
1307
1308    if ((!session) || (!data) || (len < 0)) {
1309        pico_err = PICO_ERR_EINVAL;
1310        return -1;
1311    }
1312
1313    synchro = *(int*)session->argument;
1314    *(int*)session->argument = 0;
1315    if (synchro < 0)
1316        return synchro;
1317
1318    if (len > PICO_TFTP_PAYLOAD_SIZE)
1319        len = PICO_TFTP_PAYLOAD_SIZE;
1320
1321    pico_tftp_send(session, data, len);
1322    return len;
1323}
1324