More tftp improvements
[oweals/nmrpflash.git] / tftp.c
1 /**
2  * nmrpflash - Netgear Unbrick Utility
3  * Copyright (C) 2016 Joseph Lehner <joseph.c.lehner@gmail.com>
4  *
5  * nmrpflash is free software: you can redistribute it and/or modify
6  * it under the terms of the GNU General Public License as published by
7  * the Free Software Foundation, either version 3 of the License, or
8  * (at your option) any later version.
9  *
10  * nmrpflash is distributed in the hope that it will be useful,
11  * but WITHOUT ANY WARRANTY; without even the implied warranty of
12  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
13  * GNU General Public License for more details.
14  *
15  * You should have received a copy of the GNU General Public License
16  * along with nmrpflash.  If not, see <http://www.gnu.org/licenses/>.
17  *
18  */
19
20 #include <string.h>
21 #include <unistd.h>
22 #include <stdlib.h>
23 #include <stdio.h>
24 #include <errno.h>
25 #include <fcntl.h>
26 #include <ctype.h>
27 #include "nmrpd.h"
28
29 #ifndef O_BINARY
30 #define O_BINARY 0
31 #endif
32
33 #define TFTP_PKT_SIZE 516
34
35 static const char *opcode_names[] = {
36         "RRQ", "WRQ", "DATA", "ACK", "ERR", "OACK"
37 };
38
39 enum tftp_opcode {
40         RRQ  = 1,
41         WRQ  = 2,
42         DATA = 3,
43         ACK  = 4,
44         ERR  = 5,
45         OACK = 6
46 };
47
48 static bool is_netascii(const char *str)
49 {
50         uint8_t *p = (uint8_t*)str;
51
52         for (; *p; ++p) {
53                 if (*p < 0x20 || *p > 0x7f) {
54                         return false;
55                 }
56         }
57
58         return true;
59 }
60
61 static inline char *pkt_mknum(char *pkt, uint16_t n)
62 {
63         *(uint16_t*)pkt = htons(n);
64         return pkt + 2;
65 }
66
67 static inline uint16_t pkt_num(char *pkt)
68 {
69         return ntohs(*(uint16_t*)pkt);
70 }
71
72 static char *pkt_mkopt(char *pkt, const char *opt, const char* val)
73 {
74         strcpy(pkt, opt);
75         pkt += strlen(pkt) + 1;
76         strcpy(pkt, val);
77         pkt += strlen(val) + 1;
78         return pkt;
79 }
80
81 static bool pkt_nextstr(char **pkt, char **str, size_t *rem)
82 {
83         size_t len;
84
85         if (!isprint(**pkt) || !(len = strnlen(*pkt, *rem))) {
86                 return false;
87         } else if (str) {
88                 *str = *pkt;
89         }
90
91         *pkt += len + 1;
92
93         if (*rem > 1) {
94                 *rem -= len + 1;
95         } else {
96                 *rem = 0;
97         }
98
99         return true;
100 }
101
102 static bool pkt_nextopt(char **pkt, char **opt, char **val, size_t *rem)
103 {
104         return pkt_nextstr(pkt, opt, rem) && pkt_nextstr(pkt, val, rem);
105 }
106
107 static char *pkt_optval(char* pkt, const char* name)
108 {
109         size_t rem = 512;
110         char *opt, *val;
111         pkt += 2;
112
113         while (pkt_nextopt(&pkt, &opt, &val, &rem)) {
114                 if (!strcmp(name, opt)) {
115                         return val;
116                 }
117         }
118
119         return NULL;
120 }
121
122 static size_t pkt_xrqlen(char *pkt)
123 {
124         size_t rem = 512;
125
126         pkt += 2;
127         while (pkt_nextopt(&pkt, NULL, NULL, &rem)) {
128                 ;
129         }
130
131         return 514 - rem;
132 }
133
134 static void pkt_mkwrq(char *pkt, const char *filename, unsigned blksize)
135 {
136         size_t len = 2;
137
138         filename = leafname(filename);
139         if (!tftp_is_valid_filename(filename)) {
140                 fprintf(stderr, "Overlong/illegal filename; using 'firmware'.\n");
141                 filename = "firmware";
142         } else if (!strcmp(filename, "-")) {
143                 filename = "firmware";
144         }
145
146         pkt = pkt_mknum(pkt, WRQ);
147         pkt = pkt_mkopt(pkt, filename, "octet");
148
149         if (blksize && blksize != 512) {
150                 pkt = pkt_mkopt(pkt, "blksize", lltostr(blksize, 10));
151         }
152 }
153
154 static inline void pkt_print(char *pkt, FILE *fp)
155 {
156         uint16_t opcode = pkt_num(pkt);
157         size_t rem;
158         char *opt, *val;
159
160         if (!opcode || opcode > OACK) {
161                 fprintf(fp, "(%d)", opcode);
162         } else {
163                 fprintf(fp, "%s", opcode_names[opcode - 1]);
164                 if (opcode == ACK || opcode == DATA) {
165                         fprintf(fp, "(%d)", pkt_num(pkt + 2));
166                 } else if (opcode == WRQ || opcode == RRQ) {
167                         fprintf(fp, "(%s, %s)", pkt + 2, pkt + 2 + strlen(pkt + 2) + 1);
168                 } else if (opcode == OACK) {
169                                 fprintf(fp, "(");
170                                 rem = 512;
171                                 pkt += 2;
172                                 while (pkt_nextopt(&pkt, &opt, &val, &rem)) {
173                                         fprintf(fp, " %s=%s ", opt, val);
174                                 }
175                                 fprintf(fp, ")");
176                 }
177         }
178 }
179
180 static ssize_t tftp_recvfrom(int sock, char *pkt, uint16_t* port,
181                 unsigned timeout, size_t pktlen)
182 {
183         ssize_t len;
184         struct sockaddr_in src;
185 #ifndef NMRPFLASH_WINDOWS
186         socklen_t alen;
187 #else
188         int alen;
189 #endif
190
191         len = select_fd(sock, timeout);
192         if (len < 0) {
193                 return -1;
194         } else if (!len) {
195                 return 0;
196         }
197
198         alen = sizeof(src);
199         len = recvfrom(sock, pkt, pktlen, 0, (struct sockaddr*)&src, &alen);
200         if (len < 0) {
201                 sock_perror("recvfrom");
202                 return -1;
203         }
204
205         *port = ntohs(src.sin_port);
206
207         uint16_t opcode = pkt_num(pkt);
208
209         if (opcode == ERR) {
210                 fprintf(stderr, "Error (%d): %.511s\n", pkt_num(pkt + 2), pkt + 4);
211                 return -1;
212         } else if (isprint(pkt[0])) {
213                 /* In case of a firmware checksum error, the EX2700 I've tested this
214                  * on sends a raw UDP packet containing just an error message starting
215                  * at offset 0. The limit of 32 chars is arbitrary.
216                  */
217                 fprintf(stderr, "Error: %.32s\n", pkt);
218                 return -2;
219         } else if (!opcode || opcode > OACK) {
220                 fprintf(stderr, "Received invalid packet: ");
221                 pkt_print(pkt, stderr);
222                 fprintf(stderr, ".\n");
223                 return -1;
224         }
225
226         if (verbosity > 2) {
227                 printf(">> ");
228                 pkt_print(pkt, stdout);
229                 printf("\n");
230         }
231
232         return len;
233 }
234
235 static ssize_t tftp_sendto(int sock, char *pkt, size_t len,
236                 struct sockaddr_in *dst)
237 {
238         ssize_t sent;
239
240         switch (pkt_num(pkt)) {
241                 case RRQ:
242                 case WRQ:
243                 case OACK:
244                         len = pkt_xrqlen(pkt);
245                         break;
246                 case DATA:
247                         len += 4;
248                         break;
249                 case ACK:
250                         len = 4;
251                         break;
252                 case ERR:
253                         len = 4 + strlen(pkt + 4);
254                         break;
255                 default:
256                         fprintf(stderr, "Attempted to send invalid packet ");
257                         pkt_print(pkt, stderr);
258                         fprintf(stderr, "; this is a bug!\n");
259                         return -1;
260         }
261
262         if (verbosity > 2) {
263                 printf("<< ");
264                 pkt_print(pkt, stdout);
265                 printf("\n");
266         }
267
268         sent = sendto(sock, pkt, len, 0, (struct sockaddr*)dst, sizeof(*dst));
269         if (sent < 0) {
270                 sock_perror("sendto");
271         }
272
273         return sent;
274 }
275
276 const char *leafname(const char *path)
277 {
278         const char *slash, *bslash;
279
280         slash = strrchr(path, '/');
281         bslash = strrchr(path, '\\');
282
283         if (slash && bslash) {
284                 path = 1 + (slash > bslash ? slash : bslash);
285         } else if (slash) {
286                 path = 1 + slash;
287         } else if (bslash) {
288                 path = 1 + bslash;
289         }
290
291         return path;
292 }
293
294 #ifdef NMRPFLASH_WINDOWS
295 void sock_perror(const char *msg)
296 {
297         win_perror2(msg, WSAGetLastError());
298 }
299 #endif
300
301 inline bool tftp_is_valid_filename(const char *filename)
302 {
303         return strlen(filename) <= 500 && is_netascii(filename);
304 }
305
306 int tftp_put(struct nmrpd_args *args)
307 {
308         struct sockaddr_in addr;
309         uint16_t block, port, op, blksize;
310         ssize_t len, last_len;
311         int fd, sock, ret, timeout, errors, ackblock;
312         char rx[2048], tx[2048];
313         const char *file_remote = args->file_remote;
314         char *val, *end;
315
316         sock = -1;
317         ret = -1;
318         fd = -1;
319
320         if (g_interrupted) {
321                 goto cleanup;
322         }
323
324         if (!strcmp(args->file_local, "-")) {
325                 fd = STDIN_FILENO;
326                 if (!file_remote) {
327                         file_remote = "firmware";
328                 }
329         } else {
330                 fd = open(args->file_local, O_RDONLY | O_BINARY);
331                 if (fd < 0) {
332                         xperror("open");
333                         ret = fd;
334                         goto cleanup;
335                 } else if (!file_remote) {
336                         file_remote = args->file_local;
337                 }
338         }
339
340         sock = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
341         if (sock < 0) {
342                 sock_perror("socket");
343                 ret = sock;
344                 goto cleanup;
345         }
346
347         memset(&addr, 0, sizeof(addr));
348
349         addr.sin_family = AF_INET;
350
351         if (args->ipaddr_intf) {
352                 if ((addr.sin_addr.s_addr = inet_addr(args->ipaddr_intf)) == INADDR_NONE) {
353                         xperror("inet_addr");
354                         goto cleanup;
355                 }
356
357                 if (bind(sock, (struct sockaddr*)&addr, sizeof(addr)) != 0) {
358                         sock_perror("bind");
359                         goto cleanup;
360                 }
361         }
362
363         if ((addr.sin_addr.s_addr = inet_addr(args->ipaddr)) == INADDR_NONE) {
364                 xperror("inet_addr");
365                 goto cleanup;
366         }
367         addr.sin_port = htons(args->port);
368
369         blksize = 512;
370         block = 0;
371         last_len = -1;
372         len = 0;
373         errors = 0;
374         /* Not really, but this way the loop sends our WRQ before receiving */
375         timeout = 1;
376
377         pkt_mkwrq(tx, file_remote, 1456);
378
379         while (!g_interrupted) {
380                 ackblock = -1;
381                 op = pkt_num(rx);
382
383                 if (!timeout) {
384                         if (op == ACK) {
385                                 ackblock = pkt_num(rx + 2);
386                         } else if (op == OACK) {
387                                 ackblock = 0;
388                                 if ((val = pkt_optval(rx, "blksize"))) {
389                                         blksize = strtol(val, &end, 10);
390                                         if (!blksize || (*end != '\0')) {
391                                                 fprintf(stderr, "Error: invalid blksize in OACK: %s\n", val);
392                                                 ret = -1;
393                                                 goto cleanup;
394                                         }
395                                 }
396                         }
397                 }
398
399                 if (timeout || ackblock == block) {
400                         if (!timeout) {
401                                 ++block;
402                                 pkt_mknum(tx, DATA);
403                                 pkt_mknum(tx + 2, block);
404                                 len = read(fd, tx + 4, blksize);
405                                 if (len < 0) {
406                                         xperror("read");
407                                         ret = len;
408                                         goto cleanup;
409                                 } else if (!len) {
410                                         if (last_len != blksize && last_len != -1) {
411                                                 break;
412                                         }
413                                 }
414
415                                 last_len = len;
416                         }
417
418                         ret = tftp_sendto(sock, tx, len, &addr);
419                         if (ret < 0) {
420                                 goto cleanup;
421                         }
422                 } else if ((op != OACK && op != ACK) || ackblock > block) {
423                         if (verbosity) {
424                                 fprintf(stderr, "Expected ACK(%d), got ", block);
425                                 pkt_print(rx, stderr);
426                                 fprintf(stderr, ".\n");
427                         }
428
429                         if (ackblock != -1 && ++errors > 5) {
430                                 fprintf(stderr, "Protocol error; bailing out.\n");
431                                 ret = -1;
432                                 goto cleanup;
433                         }
434                 }
435
436                 ret = tftp_recvfrom(sock, rx, &port, args->rx_timeout, blksize + 2);
437                 if (ret < 0) {
438                         goto cleanup;
439                 } else if (!ret) {
440                         if (++timeout < 5 || (!block && timeout < 10)) {
441                                 continue;
442                         } else if (block) {
443                                 fprintf(stderr, "Timeout while waiting for ACK(%d).\n", block);
444                         } else {
445                                 fprintf(stderr, "Timeout while waiting for initial reply.\n");
446                         }
447                         ret = -1;
448                         goto cleanup;
449                 } else {
450                         timeout = 0;
451                         ret = 0;
452
453                         if (!block && port != args->port) {
454                                 if (verbosity > 1) {
455                                         printf("Switching to port %d\n", port);
456                                 }
457                                 addr.sin_port = htons(port);
458                         }
459                 }
460         }
461
462         ret = !g_interrupted ? 0 : -1;
463
464 cleanup:
465         if (fd >= 0) {
466                 close(fd);
467         }
468
469         if (sock >= 0) {
470 #ifndef NMRPFLASH_WINDOWS
471                 shutdown(sock, SHUT_RDWR);
472                 close(sock);
473 #else
474                 shutdown(sock, SD_BOTH);
475                 closesocket(sock);
476 #endif
477         }
478
479         return ret;
480 }