372 lines
8.2 KiB
C++
372 lines
8.2 KiB
C++
#ifndef LWIP_OPEN_SRC
|
|
#define LWIP_OPEN_SRC
|
|
#endif
|
|
#include <functional>
|
|
#include <WiFiUdp.h>
|
|
#include "ArduinoOTA.h"
|
|
#include "MD5Builder.h"
|
|
#include "StreamString.h"
|
|
|
|
extern "C" {
|
|
#include "osapi.h"
|
|
#include "ets_sys.h"
|
|
#include "user_interface.h"
|
|
}
|
|
|
|
#include "lwip/opt.h"
|
|
#include "lwip/udp.h"
|
|
#include "lwip/inet.h"
|
|
#include "lwip/igmp.h"
|
|
#include "lwip/mem.h"
|
|
#include "include/UdpContext.h"
|
|
|
|
#if !defined(NO_GLOBAL_INSTANCES) && !defined(NO_GLOBAL_MDNS)
|
|
#include <ESP8266mDNS.h>
|
|
#endif
|
|
|
|
#ifdef DEBUG_ESP_OTA
|
|
#ifdef DEBUG_ESP_PORT
|
|
#define OTA_DEBUG DEBUG_ESP_PORT
|
|
#endif
|
|
#endif
|
|
|
|
ArduinoOTAClass::ArduinoOTAClass()
|
|
{
|
|
}
|
|
|
|
ArduinoOTAClass::~ArduinoOTAClass(){
|
|
if(_udp_ota){
|
|
_udp_ota->unref();
|
|
_udp_ota = 0;
|
|
}
|
|
}
|
|
|
|
void ArduinoOTAClass::onStart(THandlerFunction fn) {
|
|
_start_callback = fn;
|
|
}
|
|
|
|
void ArduinoOTAClass::onEnd(THandlerFunction fn) {
|
|
_end_callback = fn;
|
|
}
|
|
|
|
void ArduinoOTAClass::onProgress(THandlerFunction_Progress fn) {
|
|
_progress_callback = fn;
|
|
}
|
|
|
|
void ArduinoOTAClass::onError(THandlerFunction_Error fn) {
|
|
_error_callback = fn;
|
|
}
|
|
|
|
void ArduinoOTAClass::setPort(uint16_t port) {
|
|
if (!_initialized && !_port && port) {
|
|
_port = port;
|
|
}
|
|
}
|
|
|
|
void ArduinoOTAClass::setHostname(const char * hostname) {
|
|
if (!_initialized && !_hostname.length() && hostname) {
|
|
_hostname = hostname;
|
|
}
|
|
}
|
|
|
|
String ArduinoOTAClass::getHostname() {
|
|
return _hostname;
|
|
}
|
|
|
|
void ArduinoOTAClass::setPassword(const char * password) {
|
|
if (!_initialized && !_password.length() && password) {
|
|
MD5Builder passmd5;
|
|
passmd5.begin();
|
|
passmd5.add(password);
|
|
passmd5.calculate();
|
|
_password = passmd5.toString();
|
|
}
|
|
}
|
|
|
|
void ArduinoOTAClass::setPasswordHash(const char * password) {
|
|
if (!_initialized && !_password.length() && password) {
|
|
_password = password;
|
|
}
|
|
}
|
|
|
|
void ArduinoOTAClass::setRebootOnSuccess(bool reboot){
|
|
_rebootOnSuccess = reboot;
|
|
}
|
|
|
|
void ArduinoOTAClass::begin(bool useMDNS) {
|
|
if (_initialized)
|
|
return;
|
|
|
|
_useMDNS = useMDNS;
|
|
|
|
if (!_hostname.length()) {
|
|
char tmp[15];
|
|
sprintf(tmp, "esp8266-%06x", ESP.getChipId());
|
|
_hostname = tmp;
|
|
}
|
|
if (!_port) {
|
|
_port = 8266;
|
|
}
|
|
|
|
if(_udp_ota){
|
|
_udp_ota->unref();
|
|
_udp_ota = 0;
|
|
}
|
|
|
|
_udp_ota = new UdpContext;
|
|
_udp_ota->ref();
|
|
|
|
if(!_udp_ota->listen(IP_ADDR_ANY, _port))
|
|
return;
|
|
_udp_ota->onRx(std::bind(&ArduinoOTAClass::_onRx, this));
|
|
|
|
#if !defined(NO_GLOBAL_INSTANCES) && !defined(NO_GLOBAL_MDNS)
|
|
if(_useMDNS) {
|
|
MDNS.begin(_hostname.c_str());
|
|
|
|
if (_password.length()) {
|
|
MDNS.enableArduino(_port, true);
|
|
} else {
|
|
MDNS.enableArduino(_port);
|
|
}
|
|
}
|
|
#endif
|
|
_initialized = true;
|
|
_state = OTA_IDLE;
|
|
#ifdef OTA_DEBUG
|
|
OTA_DEBUG.printf("OTA server at: %s.local:%u\n", _hostname.c_str(), _port);
|
|
#endif
|
|
}
|
|
|
|
int ArduinoOTAClass::parseInt(){
|
|
char data[16];
|
|
uint8_t index;
|
|
char value;
|
|
while(_udp_ota->peek() == ' ') _udp_ota->read();
|
|
for(index = 0; index < sizeof(data); ++index){
|
|
value = _udp_ota->peek();
|
|
if(value < '0' || value > '9'){
|
|
data[index] = '\0';
|
|
return atoi(data);
|
|
}
|
|
data[index] = _udp_ota->read();
|
|
}
|
|
return 0;
|
|
}
|
|
|
|
String ArduinoOTAClass::readStringUntil(char end){
|
|
String res;
|
|
int value;
|
|
while(true){
|
|
value = _udp_ota->read();
|
|
if(value < 0 || value == '\0' || value == end){
|
|
return res;
|
|
}
|
|
res += static_cast<char>(value);
|
|
}
|
|
return res;
|
|
}
|
|
|
|
void ArduinoOTAClass::_onRx(){
|
|
if(!_udp_ota->next()) return;
|
|
IPAddress ota_ip;
|
|
|
|
if (_state == OTA_IDLE) {
|
|
int cmd = parseInt();
|
|
if (cmd != U_FLASH && cmd != U_FS)
|
|
return;
|
|
_ota_ip = _udp_ota->getRemoteAddress();
|
|
_cmd = cmd;
|
|
_ota_port = parseInt();
|
|
_ota_udp_port = _udp_ota->getRemotePort();
|
|
_size = parseInt();
|
|
_udp_ota->read();
|
|
_md5 = readStringUntil('\n');
|
|
_md5.trim();
|
|
if(_md5.length() != 32)
|
|
return;
|
|
|
|
ota_ip = _ota_ip;
|
|
|
|
if (_password.length()){
|
|
MD5Builder nonce_md5;
|
|
nonce_md5.begin();
|
|
nonce_md5.add(String(micros()));
|
|
nonce_md5.calculate();
|
|
_nonce = nonce_md5.toString();
|
|
|
|
char auth_req[38];
|
|
sprintf(auth_req, "AUTH %s", _nonce.c_str());
|
|
_udp_ota->append((const char *)auth_req, strlen(auth_req));
|
|
_udp_ota->send(ota_ip, _ota_udp_port);
|
|
_state = OTA_WAITAUTH;
|
|
return;
|
|
} else {
|
|
_state = OTA_RUNUPDATE;
|
|
}
|
|
} else if (_state == OTA_WAITAUTH) {
|
|
int cmd = parseInt();
|
|
if (cmd != U_AUTH) {
|
|
_state = OTA_IDLE;
|
|
return;
|
|
}
|
|
_udp_ota->read();
|
|
String cnonce = readStringUntil(' ');
|
|
String response = readStringUntil('\n');
|
|
if (cnonce.length() != 32 || response.length() != 32) {
|
|
_state = OTA_IDLE;
|
|
return;
|
|
}
|
|
|
|
String challenge = _password + ':' + String(_nonce) + ':' + cnonce;
|
|
MD5Builder _challengemd5;
|
|
_challengemd5.begin();
|
|
_challengemd5.add(challenge);
|
|
_challengemd5.calculate();
|
|
String result = _challengemd5.toString();
|
|
|
|
ota_ip = _ota_ip;
|
|
if(result.equalsConstantTime(response)) {
|
|
_state = OTA_RUNUPDATE;
|
|
} else {
|
|
_udp_ota->append("Authentication Failed", 21);
|
|
_udp_ota->send(ota_ip, _ota_udp_port);
|
|
if (_error_callback) _error_callback(OTA_AUTH_ERROR);
|
|
_state = OTA_IDLE;
|
|
}
|
|
}
|
|
|
|
while(_udp_ota->next()) _udp_ota->flush();
|
|
}
|
|
|
|
void ArduinoOTAClass::_runUpdate() {
|
|
IPAddress ota_ip = _ota_ip;
|
|
|
|
if (!Update.begin(_size, _cmd)) {
|
|
#ifdef OTA_DEBUG
|
|
OTA_DEBUG.println("Update Begin Error");
|
|
#endif
|
|
if (_error_callback) {
|
|
_error_callback(OTA_BEGIN_ERROR);
|
|
}
|
|
|
|
StreamString ss;
|
|
Update.printError(ss);
|
|
_udp_ota->append("ERR: ", 5);
|
|
_udp_ota->append(ss.c_str(), ss.length());
|
|
_udp_ota->send(ota_ip, _ota_udp_port);
|
|
delay(100);
|
|
_udp_ota->listen(IP_ADDR_ANY, _port);
|
|
_state = OTA_IDLE;
|
|
return;
|
|
}
|
|
_udp_ota->append("OK", 2);
|
|
_udp_ota->send(ota_ip, _ota_udp_port);
|
|
delay(100);
|
|
|
|
Update.setMD5(_md5.c_str());
|
|
WiFiUDP::stopAll();
|
|
WiFiClient::stopAll();
|
|
|
|
if (_start_callback) {
|
|
_start_callback();
|
|
}
|
|
if (_progress_callback) {
|
|
_progress_callback(0, _size);
|
|
}
|
|
|
|
WiFiClient client;
|
|
if (!client.connect(_ota_ip, _ota_port)) {
|
|
#ifdef OTA_DEBUG
|
|
OTA_DEBUG.printf("Connect Failed\n");
|
|
#endif
|
|
_udp_ota->listen(IP_ADDR_ANY, _port);
|
|
if (_error_callback) {
|
|
_error_callback(OTA_CONNECT_ERROR);
|
|
}
|
|
_state = OTA_IDLE;
|
|
}
|
|
// OTA sends little packets
|
|
client.setNoDelay(true);
|
|
|
|
uint32_t written, total = 0;
|
|
while (!Update.isFinished() && (client.connected() || client.available())) {
|
|
int waited = 1000;
|
|
while (!client.available() && waited--)
|
|
delay(1);
|
|
if (!waited){
|
|
#ifdef OTA_DEBUG
|
|
OTA_DEBUG.printf("Receive Failed\n");
|
|
#endif
|
|
_udp_ota->listen(IP_ADDR_ANY, _port);
|
|
if (_error_callback) {
|
|
_error_callback(OTA_RECEIVE_ERROR);
|
|
}
|
|
_state = OTA_IDLE;
|
|
}
|
|
written = Update.write(client);
|
|
if (written > 0) {
|
|
client.print(written, DEC);
|
|
total += written;
|
|
if(_progress_callback) {
|
|
_progress_callback(total, _size);
|
|
}
|
|
}
|
|
}
|
|
|
|
if (Update.end()) {
|
|
// Ensure last count packet has been sent out and not combined with the final OK
|
|
client.flush();
|
|
delay(1000);
|
|
client.print("OK");
|
|
client.flush();
|
|
delay(1000);
|
|
client.stop();
|
|
#ifdef OTA_DEBUG
|
|
OTA_DEBUG.printf("Update Success\n");
|
|
#endif
|
|
if (_end_callback) {
|
|
_end_callback();
|
|
}
|
|
if(_rebootOnSuccess){
|
|
#ifdef OTA_DEBUG
|
|
OTA_DEBUG.printf("Rebooting...\n");
|
|
#endif
|
|
//let serial/network finish tasks that might be given in _end_callback
|
|
delay(100);
|
|
ESP.restart();
|
|
}
|
|
} else {
|
|
_udp_ota->listen(IP_ADDR_ANY, _port);
|
|
if (_error_callback) {
|
|
_error_callback(OTA_END_ERROR);
|
|
}
|
|
Update.printError(client);
|
|
#ifdef OTA_DEBUG
|
|
Update.printError(OTA_DEBUG);
|
|
#endif
|
|
_state = OTA_IDLE;
|
|
}
|
|
}
|
|
|
|
//this needs to be called in the loop()
|
|
void ArduinoOTAClass::handle() {
|
|
if (_state == OTA_RUNUPDATE) {
|
|
_runUpdate();
|
|
_state = OTA_IDLE;
|
|
}
|
|
|
|
#if !defined(NO_GLOBAL_INSTANCES) && !defined(NO_GLOBAL_MDNS)
|
|
if(_useMDNS)
|
|
MDNS.update(); //handle MDNS update as well, given that ArduinoOTA relies on it anyways
|
|
#endif
|
|
}
|
|
|
|
int ArduinoOTAClass::getCommand() {
|
|
return _cmd;
|
|
}
|
|
|
|
#if !defined(NO_GLOBAL_INSTANCES) && !defined(NO_GLOBAL_ARDUINOOTA)
|
|
ArduinoOTAClass ArduinoOTA;
|
|
#endif
|