Further simplifications
[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                         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 = 512;
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, 1456);
377
378         while (!g_interrupted) {
379                 ackblock = -1;
380                 op = pkt_num(rx);
381
382                 if (!timeout) {
383                         if (op == ACK) {
384                                 ackblock = pkt_num(rx + 2);
385                         } else if (op == OACK) {
386                                 ackblock = 0;
387                                 if ((p = pkt_optval(rx, "blksize"))) {
388                                         blksize = atoi(p);
389                                 }
390                         }
391                 }
392
393                 if (timeout || ackblock == block) {
394                         if (!timeout) {
395                                 ++block;
396                                 pkt_mknum(tx, DATA);
397                                 pkt_mknum(tx + 2, block);
398                                 len = read(fd, tx + 4, blksize);
399                                 if (len < 0) {
400                                         xperror("read");
401                                         ret = len;
402                                         goto cleanup;
403                                 } else if (!len) {
404                                         if (last_len != blksize && last_len != -1) {
405                                                 break;
406                                         }
407                                 }
408
409                                 last_len = len;
410                         }
411
412                         ret = tftp_sendto(sock, tx, len, &addr);
413                         if (ret < 0) {
414                                 goto cleanup;
415                         }
416                 } else if ((op != OACK && op != ACK) || ackblock > block) {
417                         if (verbosity) {
418                                 fprintf(stderr, "Expected ACK(%d), got ", block);
419                                 pkt_print(rx, stderr);
420                                 fprintf(stderr, ".\n");
421                         }
422
423                         if (ackblock != -1 && ++errors > 5) {
424                                 fprintf(stderr, "Protocol error; bailing out.\n");
425                                 ret = -1;
426                                 goto cleanup;
427                         }
428                 }
429
430                 ret = tftp_recvfrom(sock, rx, &port, args->rx_timeout, blksize + 2);
431                 if (ret < 0) {
432                         goto cleanup;
433                 } else if (!ret) {
434                         if (++timeout < 5 || (!block && timeout < 10)) {
435                                 continue;
436                         } else if (block) {
437                                 fprintf(stderr, "Timeout while waiting for ACK(%d).\n", block);
438                         } else {
439                                 fprintf(stderr, "Timeout while waiting for initial reply.\n");
440                         }
441                         ret = -1;
442                         goto cleanup;
443                 } else {
444                         timeout = 0;
445                         ret = 0;
446
447                         if (!block && port != args->port) {
448                                 if (verbosity > 1) {
449                                         printf("Switching to port %d\n", port);
450                                 }
451                                 addr.sin_port = htons(port);
452                         }
453                 }
454         }
455
456         ret = !g_interrupted ? 0 : -1;
457
458 cleanup:
459         if (fd >= 0) {
460                 close(fd);
461         }
462
463         if (sock >= 0) {
464 #ifndef NMRPFLASH_WINDOWS
465                 shutdown(sock, SHUT_RDWR);
466                 close(sock);
467 #else
468                 shutdown(sock, SD_BOTH);
469                 closesocket(sock);
470 #endif
471         }
472
473         return ret;
474 }