Implement TFTP blksize option
[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 void pkt_mknum(char *pkt, uint16_t n)
62 {
63         *(uint16_t*)pkt = htons(n);
64 }
65
66 static inline uint16_t pkt_num(char *pkt)
67 {
68         return ntohs(*(uint16_t*)pkt);
69 }
70
71 static char *pkt_mkopt(char *pkt, const char *opt, const char* val)
72 {
73         strcpy(pkt, opt);
74         pkt += strlen(pkt) + 1;
75         strcpy(pkt, val);
76         pkt += strlen(val) + 1;
77         return pkt;
78 }
79
80 static bool pkt_nextstr(char **pkt, char **str, size_t *rem)
81 {
82         size_t len;
83
84         if (!isprint(**pkt) || !(len = strnlen(*pkt, *rem))) {
85                 return false;
86         }
87
88         *str = *pkt;
89         *pkt += len + 1;
90
91         if (*rem > 1) {
92                 *rem -= len + 1;
93         } else {
94                 *rem = 0;
95         }
96
97         return true;
98 }
99
100 static bool pkt_nextopt(char **pkt, char **opt, char **val, size_t *rem)
101 {
102         return pkt_nextstr(pkt, opt, rem) && pkt_nextstr(pkt, val, rem);
103 }
104
105 static char *pkt_optval(char* pkt, const char* name)
106 {
107         size_t rem = 512;
108         char *opt, *val;
109         pkt += 2;
110
111         while (pkt_nextopt(&pkt, &opt, &val, &rem)) {
112                 if (!strcmp(name, opt)) {
113                         return val;
114                 }
115         }
116
117         return NULL;
118 }
119
120 static size_t pkt_xrqlen(char *pkt)
121 {
122         size_t len = 2, rem = 512;
123         char *opt, *val;
124
125         pkt += 2;
126
127         while (pkt_nextopt(&pkt, &opt, &val, &rem)) {
128                 len += strlen(opt) + strlen(val) + 2;
129         }
130
131         return len;
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_mknum(pkt, WRQ);
147         pkt = pkt_mkopt(pkt + 2, 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 blksize)
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, blksize, 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                         len = pkt_xrqlen(pkt);
244                         break;
245                 case DATA:
246                         len += 4;
247                         break;
248                 case ACK:
249                         len = 4;
250                         break;
251                 case ERR:
252                         len = 4 + strlen(pkt + 4);
253                         break;
254                 default:
255                         fprintf(stderr, "Attempted to send invalid packet ");
256                         pkt_print(pkt, stderr);
257                         fprintf(stderr, "; this is a bug!\n");
258                         return -1;
259         }
260
261         if (verbosity > 2) {
262                 printf("<< ");
263                 pkt_print(pkt, stdout);
264                 printf("\n");
265         }
266
267         sent = sendto(sock, pkt, len, 0, (struct sockaddr*)dst, sizeof(*dst));
268         if (sent < 0) {
269                 sock_perror("sendto");
270         }
271
272         return sent;
273 }
274
275 const char *leafname(const char *path)
276 {
277         const char *slash, *bslash;
278
279         slash = strrchr(path, '/');
280         bslash = strrchr(path, '\\');
281
282         if (slash && bslash) {
283                 path = 1 + (slash > bslash ? slash : bslash);
284         } else if (slash) {
285                 path = 1 + slash;
286         } else if (bslash) {
287                 path = 1 + bslash;
288         }
289
290         return path;
291 }
292
293 #ifdef NMRPFLASH_WINDOWS
294 void sock_perror(const char *msg)
295 {
296         win_perror2(msg, WSAGetLastError());
297 }
298 #endif
299
300 inline bool tftp_is_valid_filename(const char *filename)
301 {
302         return strlen(filename) <= 500 && is_netascii(filename);
303 }
304
305 int tftp_put(struct nmrpd_args *args)
306 {
307         struct sockaddr_in addr;
308         uint16_t block, port, op, blksize;
309         ssize_t len, last_len;
310         int fd, sock, ret, timeout, errors, ackblock;
311         char rx[2048], tx[2048];
312         const char *file_remote = args->file_remote;
313         char *p;
314
315         sock = -1;
316         ret = -1;
317         fd = -1;
318
319         if (g_interrupted) {
320                 goto cleanup;
321         }
322
323         if (!strcmp(args->file_local, "-")) {
324                 fd = STDIN_FILENO;
325                 if (!file_remote) {
326                         file_remote = "firmware";
327                 }
328         } else {
329                 fd = open(args->file_local, O_RDONLY | O_BINARY);
330                 if (fd < 0) {
331                         xperror("open");
332                         ret = fd;
333                         goto cleanup;
334                 } else if (!file_remote) {
335                         file_remote = args->file_local;
336                 }
337         }
338
339         sock = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
340         if (sock < 0) {
341                 sock_perror("socket");
342                 ret = sock;
343                 goto cleanup;
344         }
345
346         memset(&addr, 0, sizeof(addr));
347
348         addr.sin_family = AF_INET;
349
350         if (args->ipaddr_intf) {
351                 if ((addr.sin_addr.s_addr = inet_addr(args->ipaddr_intf)) == INADDR_NONE) {
352                         xperror("inet_addr");
353                         goto cleanup;
354                 }
355
356                 if (bind(sock, (struct sockaddr*)&addr, sizeof(addr)) != 0) {
357                         sock_perror("bind");
358                         goto cleanup;
359                 }
360         }
361
362         if ((addr.sin_addr.s_addr = inet_addr(args->ipaddr)) == INADDR_NONE) {
363                 xperror("inet_addr");
364                 goto cleanup;
365         }
366         addr.sin_port = htons(args->port);
367
368         blksize = 1468;
369         block = 0;
370         last_len = -1;
371         len = 0;
372         errors = 0;
373         /* Not really, but this way the loop sends our WRQ before receiving */
374         timeout = 1;
375
376         pkt_mkwrq(tx, file_remote, blksize);
377
378         while (!g_interrupted) {
379                 ackblock = -1;
380                 op = pkt_num(rx);
381
382                 if (!timeout) {
383                         if (pkt_num(rx) == ACK) {
384                                 ackblock = pkt_num(rx + 2);
385                                 if (!ackblock) {
386                                         blksize = 512;
387                                 }
388                         } else if (pkt_num(rx) == OACK) {
389                                 op = ACK;
390                                 ackblock = 0;
391                                 if ((p = pkt_optval(rx, "blksize"))) {
392                                         blksize = atoi(p);
393                                 } else {
394                                         blksize = 512;
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 != 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 }