#include "update_parser.h"

#include <cstring>

#include "version_compare.h"

namespace roro {

namespace {
uint16_t u16(const uint8_t* p) { return p[0] | (p[1] << 8); }
uint32_t u32(const uint8_t* p) { return p[0] | (p[1] << 8) | (p[2] << 16) | (uint32_t(p[3]) << 24); }
}  // namespace

void UpdateParser::fail(const std::string& why) {
    if (state_ == State::Image) sink_.abort();
    state_ = State::Failed;
    error_ = why;
}

void UpdateParser::parseHeader() {
    const uint8_t* h = header_;
    if (std::memcmp(h, update::kMagic, 8) != 0) return fail("not an update file");
    if (u16(h + 8) != update::kFormat || u16(h + 10) != update::kHeaderSize) return fail("unsupported update format");
    imageSize_ = u32(h + 12);
    if (imageSize_ == 0 || imageSize_ > maxImage_) return fail("image doesn't fit the update slot");
    std::memcpy(expectedHash_, h + 16, 32);
    version_.assign(reinterpret_cast<const char*>(h + 48), strnlen(reinterpret_cast<const char*>(h + 48), 32));

    size_t sigLen = u16(h + 80);
    if (sigLen == 0 || sigLen > update::kMaxSignature) return fail("missing signature");
    uint8_t digest[32];
    Sha256::hash(h, update::kSignedBytes, digest);
    if (!verifier_.verify(digest, h + 82, sigLen)) return fail("bad signature (wrong key)");

    downgrade_ = versionOlder(version_, installed_);
    if (!sink_.begin(imageSize_)) return fail("could not prepare the update slot");
    state_ = State::Image;
}

void UpdateParser::feed(const uint8_t* data, size_t len) {
    while (len > 0 && (state_ == State::Header || state_ == State::Image)) {
        if (state_ == State::Header) {
            size_t take = std::min(len, update::kHeaderSize - headerUsed_);
            std::memcpy(header_ + headerUsed_, data, take);
            headerUsed_ += take;
            data += take;
            len -= take;
            if (headerUsed_ == update::kHeaderSize) parseHeader();
        } else {
            size_t take = std::min(len, imageSize_ - received_);
            if (take == 0) return fail("data after the end of the image");
            hash_.update(data, take);
            if (!sink_.write(data, take)) return fail("writing the update failed");
            received_ += take;
            data += take;
            len -= take;
        }
    }
    if (len > 0 && state_ == State::Image) fail("data after the end of the image");
}

bool UpdateParser::end() {
    if (state_ == State::Done) return true;
    if (state_ != State::Image) {
        if (state_ == State::Header) fail("update file too short");
        return false;
    }
    if (received_ != imageSize_) {
        fail("update file too short");
        return false;
    }
    uint8_t got[32];
    hash_.finish(got);
    if (std::memcmp(got, expectedHash_, 32) != 0) {
        fail("image corrupted (hash mismatch)");
        return false;
    }
    if (!sink_.finish()) {
        state_ = State::Failed;
        error_ = "could not switch to the new image";
        return false;
    }
    state_ = State::Done;
    return true;
}

}  // namespace roro