1 /*
2  * libwebsockets - small server side websockets and web server implementation
3  *
4  * Copyright (C) 2010 - 2019 Andy Green <andy@warmcat.com>
5  *
6  * Permission is hereby granted, free of charge, to any person obtaining a copy
7  * of this software and associated documentation files (the "Software"), to
8  * deal in the Software without restriction, including without limitation the
9  * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
10  * sell copies of the Software, and to permit persons to whom the Software is
11  * furnished to do so, subject to the following conditions:
12  *
13  * The above copyright notice and this permission notice shall be included in
14  * all copies or substantial portions of the Software.
15  *
16  * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
17  * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
18  * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
19  * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
20  * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
21  * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
22  * IN THE SOFTWARE.
23  */
24 
25 #include <string.h>
26 #include <esp_partition.h>
27 #include <esp_ota_ops.h>
28 #include <nvs.h>
29 
30 struct per_session_data__esplws_ota {
31 	struct lws_spa *spa;
32 	char filename[32];
33 	char result[LWS_PRE + 512];
34 	int result_len;
35 	int filename_length;
36 	esp_ota_handle_t otahandle;
37 	const esp_partition_t *part;
38 	long file_length;
39 	long last_rep;
40 	nvs_handle nvh;
41 	TimerHandle_t reboot_timer;
42 };
43 
44 struct per_vhost_data__esplws_ota {
45 	struct lws_context *context;
46 	struct lws_vhost *vhost;
47 	const struct lws_protocols *protocol;
48 };
49 
50 static const char * const ota_param_names[] = {
51 	"upload",
52 };
53 
54 enum enum_ota_param_names {
55 	EPN_UPLOAD,
56 };
57 
ota_reboot_timer_cb(TimerHandle_t t)58 static void ota_reboot_timer_cb(TimerHandle_t t)
59 {
60 	esp_restart();
61 }
62 
63 const esp_partition_t *
ota_choose_part(void)64 ota_choose_part(void)
65 {
66 	const esp_partition_t *bootpart, *part = NULL;
67 	esp_partition_iterator_t i;
68 
69 	bootpart = lws_esp_ota_get_boot_partition();
70 	i = esp_partition_find(ESP_PARTITION_TYPE_APP, ESP_PARTITION_SUBTYPE_ANY, NULL);
71 	while (i) {
72 		part = esp_partition_get(i);
73 
74 		/* cannot update ourselves */
75 		if (part == bootpart)
76 			goto next;
77 
78 		/* OTA Partition numbering is from _OTA_MIN to less than _OTA_MAX */
79 		if (part->subtype < ESP_PARTITION_SUBTYPE_APP_OTA_MIN ||
80 		    part->subtype >= ESP_PARTITION_SUBTYPE_APP_OTA_MAX)
81 			goto next;
82 
83 		break;
84 
85 next:
86 		i = esp_partition_next(i);
87 	}
88 
89 	if (!i) {
90 		lwsl_err("Can't find good OTA part\n");
91 		return NULL;
92 	}
93 	lwsl_notice("Directing OTA to part type %d/%d start 0x%x\n",
94 			part->type, part->subtype,
95 			(uint32_t)part->address);
96 
97 	return part;
98 }
99 
100 static int
ota_file_upload_cb(void * data,const char * name,const char * filename,char * buf,int len,enum lws_spa_fileupload_states state)101 ota_file_upload_cb(void *data, const char *name, const char *filename,
102 	       char *buf, int len, enum lws_spa_fileupload_states state)
103 {
104 	struct per_session_data__esplws_ota *pss =
105 			(struct per_session_data__esplws_ota *)data;
106 
107 	switch (state) {
108 	case LWS_UFS_OPEN:
109 		lwsl_notice("LWS_UFS_OPEN Filename %s\n", filename);
110 		lws_strncpy(pss->filename, filename, sizeof(pss->filename));
111 		if (strcmp(name, "ota"))
112 			return 1;
113 
114 		pss->part = ota_choose_part();
115 		if (!pss->part)
116 			return 1;
117 
118 		if (esp_ota_begin(pss->part, OTA_SIZE_UNKNOWN, &pss->otahandle) != ESP_OK) {
119 			lwsl_err("OTA: Failed to begin\n");
120 			return 1;
121 		}
122 
123 		pss->file_length = 0;
124 		pss->last_rep = -1;
125 		break;
126 
127 	case LWS_UFS_FINAL_CONTENT:
128 	case LWS_UFS_CONTENT:
129 		if (pss->file_length + len > pss->part->size) {
130 			lwsl_err("OTA: incoming file too large\n");
131 			return 1;
132 		}
133 
134 		if ((pss->file_length & ~0xffff) != (pss->last_rep & ~0xffff)) {
135 			lwsl_notice("writing 0x%lx...\n",
136 					pss->part->address + pss->file_length);
137 			pss->last_rep = pss->file_length;
138 		}
139 		if (esp_ota_write(pss->otahandle, buf, len) != ESP_OK) {
140 			lwsl_err("OTA: Failed to write\n");
141 			return 1;
142 		}
143 		pss->file_length += len;
144 
145 		if (state == LWS_UFS_CONTENT)
146 			break;
147 
148 		lwsl_notice("LWS_UFS_FINAL_CONTENT\n");
149 		if (esp_ota_end(pss->otahandle) != ESP_OK) {
150 			lwsl_err("OTA: end failed\n");
151 			return 1;
152 		}
153 
154 		if (esp_ota_set_boot_partition(pss->part) != ESP_OK) {
155 			lwsl_err("OTA: set boot part failed\n");
156 			return 1;
157 		}
158 
159 		pss->reboot_timer = xTimerCreate("x", pdMS_TO_TICKS(250), 0, NULL,
160 						 ota_reboot_timer_cb);
161 		xTimerStart(pss->reboot_timer, 0);
162 		break;
163 	}
164 
165 	return 0;
166 }
167 
168 static int
callback_esplws_ota(struct lws * wsi,enum lws_callback_reasons reason,void * user,void * in,size_t len)169 callback_esplws_ota(struct lws *wsi, enum lws_callback_reasons reason,
170 		    void *user, void *in, size_t len)
171 {
172 	struct per_session_data__esplws_ota *pss =
173 			(struct per_session_data__esplws_ota *)user;
174 	struct per_vhost_data__esplws_ota *vhd =
175 			(struct per_vhost_data__esplws_ota *)
176 			lws_protocol_vh_priv_get(lws_get_vhost(wsi),
177 					lws_get_protocol(wsi));
178 	unsigned char buf[LWS_PRE + 384], *start = buf + LWS_PRE - 1, *p = start,
179 	     *end = buf + sizeof(buf) - 1;
180 	int n;
181 
182 	switch (reason) {
183 
184 	case LWS_CALLBACK_PROTOCOL_INIT:
185 		vhd = lws_protocol_vh_priv_zalloc(lws_get_vhost(wsi),
186 				lws_get_protocol(wsi),
187 				sizeof(struct per_vhost_data__esplws_ota));
188 		vhd->context = lws_get_context(wsi);
189 		vhd->protocol = lws_get_protocol(wsi);
190 		vhd->vhost = lws_get_vhost(wsi);
191 		break;
192 
193 	case LWS_CALLBACK_PROTOCOL_DESTROY:
194 		if (!vhd)
195 			break;
196 		break;
197 
198 	/* OTA POST handling */
199 
200 	case LWS_CALLBACK_HTTP_BODY:
201 		/* create the POST argument parser if not already existing */
202 		// lwsl_notice("LWS_CALLBACK_HTTP_BODY (ota) %d %d %p\n", (int)pss->file_length, (int)len, pss->spa);
203 		lws_set_timeout(wsi, PENDING_TIMEOUT_HTTP_CONTENT, 30);
204 		if (!pss->spa) {
205 			pss->spa = lws_spa_create(wsi, ota_param_names,
206 					LWS_ARRAY_SIZE(ota_param_names), 4096,
207 					ota_file_upload_cb, pss);
208 			if (!pss->spa)
209 				return -1;
210 
211 			pss->filename[0] = '\0';
212 			pss->file_length = 0;
213 		}
214 		lws_esp32.upload = 1;
215 
216 		/* let it parse the POST data */
217 		if (lws_spa_process(pss->spa, in, len))
218 			return -1;
219 		break;
220 
221 	case LWS_CALLBACK_HTTP_BODY_COMPLETION:
222 		lwsl_notice("LWS_CALLBACK_HTTP_BODY_COMPLETION (ota)\n");
223 		/* call to inform no more payload data coming */
224 		lws_spa_finalize(pss->spa);
225 
226 		pss->result_len = snprintf(pss->result + LWS_PRE, sizeof(pss->result) - LWS_PRE - 1,
227 			"Rebooting after OTA update");
228 
229 		if (lws_add_http_header_status(wsi, HTTP_STATUS_OK, &p, end))
230 			goto bail;
231 
232 		if (lws_add_http_header_by_token(wsi, WSI_TOKEN_HTTP_CONTENT_TYPE,
233 				(unsigned char *)"text/html", 9, &p, end))
234 			goto bail;
235 		if (lws_add_http_header_content_length(wsi, pss->result_len, &p, end))
236 			goto bail;
237 		if (lws_finalize_http_header(wsi, &p, end))
238 			goto bail;
239 
240 		n = lws_write(wsi, start, p - start, LWS_WRITE_HTTP_HEADERS | LWS_WRITE_H2_STREAM_END);
241 		if (n < 0)
242 			goto bail;
243 
244 		lws_callback_on_writable(wsi);
245 		break;
246 
247 	case LWS_CALLBACK_HTTP_WRITEABLE:
248 		if (!pss->result_len)
249 			break;
250 		lwsl_debug("LWS_CALLBACK_HTTP_WRITEABLE: sending %d\n",
251 			   pss->result_len);
252 		n = lws_write(wsi, (unsigned char *)pss->result + LWS_PRE,
253 			      pss->result_len, LWS_WRITE_HTTP);
254 		if (n < 0)
255 			return 1;
256 
257 		if (lws_http_transaction_completed(wsi))
258 			return 1;
259 
260 		/* stop further service so we don't serve the probe GET to see if we rebooted */
261 		while (1);
262 
263 		break;
264 
265 	case LWS_CALLBACK_HTTP_DROP_PROTOCOL:
266 		/* called when our wsi user_space is going to be destroyed */
267 		if (pss->spa) {
268 			lws_spa_destroy(pss->spa);
269 			pss->spa = NULL;
270 		}
271 		lws_esp32.upload = 0;
272 		break;
273 
274 	default:
275 		break;
276 	}
277 
278 	return 0;
279 
280 bail:
281 	return 1;
282 }
283 
284 #define LWS_PLUGIN_PROTOCOL_ESPLWS_OTA \
285 	{ \
286 		"esplws-ota", \
287 		callback_esplws_ota, \
288 		sizeof(struct per_session_data__esplws_ota), \
289 		4096, 0, NULL, 900 \
290 	}
291 
292