Vermix-Gate/lib/data/users.cpp

#include "users.hpp"
#include <pb_encode.h>
#include <pb_decode.h>
#include <Arduino.h>
#include <cstring>
#include <functional>
#include <SHA256.h>
#include <logger.hpp>

// Definition of the global instance.
Users users;

// =================================================================
// == 1. NANOPB STREAM HELPERS for Arduino File Objects
// =================================================================

// A state structure to pass the Arduino File object to Nanopb's stream callbacks.
struct FileStreamState {
    File* file;
};

// Nanopb callback for writing data to a File.
static bool writeFileCallback(pb_ostream_t *stream, const uint8_t *buf, size_t count) {
    FileStreamState* state = (FileStreamState*)stream->state;
    return state->file->write(buf, count) == count;
}

// Nanopb callback for reading from a File.
static bool readFileCallback(pb_istream_t *stream, uint8_t *buf, size_t count) {
    FileStreamState* state = (FileStreamState*)stream->state;
    if (count == 0) return true;
    size_t len = state->file->read(buf, count);
    if (len < count) {
        stream->bytes_left = 0; // Signal End-Of-File
    }
    return len > 0 || count == 0;
}

// Helper to create a Nanopb output stream from an Arduino File object.
static pb_ostream_t pb_ostream_from_file(File &file) {
    FileStreamState* state = new FileStreamState{&file};
    pb_ostream_t stream = {&writeFileCallback, state, SIZE_MAX, 0};
    return stream;
}

// Helper to create a Nanopb input stream from an Arduino File object.
static pb_istream_t pb_istream_from_file(File &file) {
    FileStreamState* state = new FileStreamState{&file};
    pb_istream_t stream = {&readFileCallback, state, file.size()};
    return stream;
}

// This structure is used to pass data to the nanopb decode callbacks
// when decoding into a single, stack-allocated UserData object.
struct DecodeTarget {
    UserData* user_buffer;
};

// Nanopb callback to decode a username string into our temporary buffer.
static bool decodeUsernameCallback(pb_istream_t *stream, const pb_field_t *field, void **arg) {
    DecodeTarget* target = (DecodeTarget*)*arg;
    
    // 1. Store the length of the incoming string BEFORE reading from the stream.
    size_t len = stream->bytes_left;

    if (len >= sizeof(target->user_buffer->username)) {
        return false; // Input string is too long for our buffer.
    }
    
    // 2. Read 'len' bytes. This will modify stream->bytes_left.
    if (!pb_read(stream, (uint8_t*)target->user_buffer->username, len)) {
        return false;
    }
    
    // 3. Use the stored length 'len' to place the null terminator correctly.
    target->user_buffer->username[len] = '\0';
    
    return true;
}

// Nanopb callback to decode a password hash (bytes) into our temporary buffer.
static bool decodePasswordCallback(pb_istream_t *stream, const pb_field_t *field, void **arg) {
    DecodeTarget* target = (DecodeTarget*)*arg;
    if (stream->bytes_left != sizeof(target->user_buffer->password_hash)) {
        return false; // A SHA-256 hash MUST be 32 bytes.
    }
    return pb_read(stream, target->user_buffer->password_hash, stream->bytes_left);
}

// =================================================================
// == 2. USERS CLASS IMPLEMENTATION
// =================================================================

Users::Users() {}
Users::~Users() {}

void Users::hashPassword(const char* password, uint8_t* hash_out) const {
    SHA256 sha256;
    sha256.update(password, strlen(password));
    sha256.finalize(hash_out, 32);
}

// --- PRIVATE STATIC HELPERS (Moved inside the class scope) ---

// Generic helper function to iterate over all users in the file
// and execute a callback for each one. The callback can stop the iteration by returning false.
bool Users::_iterateUsers(const char* filename, const std::function<bool(UserData&)>& user_callback) {
    // check if file exists
    if(!LittleFS.exists(filename)){
        return true; // No file is not an error, just means no users to iterate.
    }

    File file = LittleFS.open(filename, "r");
    if (!file || file.size() == 0) {
        return true; // No file is not an error, just means no users to iterate.
    }

    pb_istream_t istream = pb_istream_from_file(file);
    
    UserData user_buffer; // Single user buffer on the stack
    DecodeTarget target = { &user_buffer };
    webui_WebUiUser user_msg = webui_WebUiUser_init_zero;
    user_msg.password_hash.funcs.decode = &decodePasswordCallback;
    user_msg.password_hash.arg = &target;

    bool continue_iterating = true;
    while (istream.bytes_left > 0) {
        memset(&user_buffer, 0, sizeof(user_buffer));

        if (pb_decode_delimited(&istream, webui_WebUiUser_fields, &user_msg)) {
            strcpy(user_buffer.username, user_msg.username);
            if (!user_callback(user_buffer)) {
                continue_iterating = false;
                break;
            }
        } else {
            // Decoding failed, stop iterating.
            break;
        }
    }


    delete (FileStreamState*)istream.state;
    file.close();
    return continue_iterating;
}

// Generic helper for streaming transformations (delete, update).
// The callback modifies a user buffer and returns true to keep it, false to delete it.
bool Users::_transformUsers(const char* src_filename, const std::function<bool(UserData&)>& transform_callback) {
    const char* tmp_filename = "/users.pb.tmp";
    File src_file = LittleFS.open(src_filename, "r");
    if (!src_file) return true; 
    
    File dst_file = LittleFS.open(tmp_filename, "w");
    if (!dst_file) {
        src_file.close();
        return false;
    }
    
    pb_istream_t istream = pb_istream_from_file(src_file);
    pb_ostream_t ostream = pb_ostream_from_file(dst_file);
    
    bool success = true;
    UserData user_buffer; // Our 1-user buffer

    // --- CORRECTION START ---

    // 1. A structure ONLY for decoding (reading)
    DecodeTarget decode_target = { &user_buffer };
    webui_WebUiUser decode_msg = webui_WebUiUser_init_zero;
    decode_msg.password_hash.funcs.decode = &decodePasswordCallback;
    decode_msg.password_hash.arg = &decode_target;

    // 2. A structure ONLY for encoding (writing)
    webui_WebUiUser encode_msg = webui_WebUiUser_init_zero;
    encode_msg.password_hash.funcs.encode = [](pb_ostream_t *stream, const pb_field_t *field, void * const *arg) {
        UserData* u = (UserData*)*arg;
        return pb_encode_tag_for_field(stream, field) && pb_encode_string(stream, u->password_hash, sizeof(u->password_hash));
    };

    // 3. Main loop: Decode with decode_msg, transform user_buffer, encode with encode_msg

    while (istream.bytes_left > 0 && pb_decode_delimited(&istream, webui_WebUiUser_fields, &decode_msg)) {
        strcpy(user_buffer.username, decode_msg.username);
        if (transform_callback(user_buffer)) {
            // Keep: Encode the content from the buffer with the encode_msg structure
            strcpy(encode_msg.username, user_buffer.username);
            encode_msg.password_hash.arg = &user_buffer;
            if (!pb_encode_delimited(&ostream, webui_WebUiUser_fields, &encode_msg)) {
                success = false;
                break;
            }
        }
    }

    // --- Rest of the function remains the same ---
    delete (FileStreamState*)istream.state;
    delete (FileStreamState*)ostream.state;
    src_file.close();
    dst_file.close();
    
    if (!success) {
        LittleFS.remove(tmp_filename);
        return false;
    }

    if (LittleFS.exists(src_filename)) {
        LittleFS.remove(src_filename);
    }
    return LittleFS.rename(tmp_filename, src_filename);
}

// --- PUBLIC READ-ONLY METHODS ---

size_t Users::getUserCount() {
    if(initialized_) {
        return userCount_;
    }

    size_t count = 0;
    _iterateUsers(filename_, [&](UserData& user) {
        count++;
        return true; // Continue iterating
    });
    userCount_ = count;
    initialized_ = true;
    return count;
}

bool Users::userExists(const char* username) const {
    bool found = false;
    _iterateUsers(filename_, [&](UserData& user) {
        if (strcmp(user.username, username) == 0) {
            found = true;
            return false; // Stop iterating
        }
        return true; // Continue iterating
    });
    return found;
}

bool Users::checkPassword(const char* username, const char* password) const {
    bool success = false;
    _iterateUsers(filename_, [&](UserData& user) {
        if (strcmp(user.username, username) == 0) {
            uint8_t provided_hash[32];
            hashPassword(password, provided_hash);
            if (memcmp(user.password_hash, provided_hash, 32) == 0) {
                success = true;
            }
            return false; // Stop iterating, user found.
        }
        return true; // Continue iterating
    });
    return success;
}

bool Users::loadUsernamesCallback(const std::function<bool(const char *)> &callback) const
{
    if (!callback) return false;
    return _iterateUsers(filename_, [&](UserData &user) {
        auto res = callback(user.username);
        if (!res) return false; // Stop iterating
        return true; // Continue iterating
    });
}

bool Users::addUser(const char* username, const char* password) {
    // --- Phase 1: Checking preconditions ---
    bool exists = userExists(username);
    size_t count = getUserCount();
    
    if (exists || count >= MAX_USERS) {
        if (exists) {
            LOG_ERROR("addUser: User already exists. Aborting.");
        } else {
            LOG_ERROR("addUser: User limit reached. Aborting.");
        }
        return false;
    }
    
    // --- Phase 2: Preparing the files ---
    const char* tmp_filename = "/users.pb.tmp";
    File dst_file = LittleFS.open(tmp_filename, "w");
    if (!dst_file) {
        LOG_ERROR("addUser: Failed to open temporary file. Check filesystem space and permissions.");
        return false;
    }

    File src_file = LittleFS.open(filename_, "r");
    if (src_file && src_file.size() > 0) {
        uint8_t buffer[128];
        while (src_file.available()) {
            size_t len = src_file.read(buffer, sizeof(buffer));
            dst_file.write(buffer, len);
        }
        src_file.close();
    } else {
        if (src_file) src_file.close();
    }

    // --- Phase 3: Appending the new user ---
    pb_ostream_t ostream = pb_ostream_from_file(dst_file);
    webui_WebUiUser new_user_msg = webui_WebUiUser_init_zero;
    
    // Prepare username for the encoder
    strcpy(new_user_msg.username, username);
    
    // Prepare password hash for the encoder
    uint8_t password_hash[32];
    hashPassword(password, password_hash); // Important: Hash is created here, not in the lambda
    new_user_msg.password_hash.arg = password_hash;
    new_user_msg.password_hash.funcs.encode = [](pb_ostream_t *stream, const pb_field_t *field, void * const *arg) {
        return pb_encode_tag_for_field(stream, field) && pb_encode_string(stream, (uint8_t*)*arg, 32);
    };

    bool status = pb_encode_delimited(&ostream, webui_WebUiUser_fields, &new_user_msg);
    
    // Important: Always release resources (stream state and file)
    delete (FileStreamState*)ostream.state;
    dst_file.close();

    if (!status) {
        LOG_ERROR("addUser: Protobuf encoding (pb_encode_delimited) failed!");
        LOG_DEBUG("addUser: Removing failed temporary file: %s\n", tmp_filename);
        LittleFS.remove(tmp_filename);
        return false;
    }

    // --- Phase 4: Atomic replacement of the old file ---
    if (LittleFS.exists(filename_)) {
        if (!LittleFS.remove(filename_)) {
            LOG_ERROR("addUser: Could not remove old file '%s'. Aborting.\n", filename_);
            LittleFS.remove(tmp_filename); // Cleanup temp file
            return false;
        }
    }
    
    bool rename_ok = LittleFS.rename(tmp_filename, filename_);
    
    if (rename_ok) {
    } else {
        LOG_ERROR("addUser: Final rename failed! Filesystem might be in an inconsistent state.");
        LittleFS.remove(tmp_filename); // Attempt to clean up
    }

    if (rename_ok) {
        userCount_++;
    }
    
    return rename_ok;
}

bool Users::deleteUser(const char* username) {
    bool user_was_found = false;

    bool success = _transformUsers(filename_, [&](UserData& user) {
        if (strcmp(user.username, username) == 0) {
            user_was_found = true;
            return false; // Delete this user
        }
        return true; // Keep this user
    });

    bool user_deleted = success && user_was_found;
    if (user_deleted) {
        userCount_--;
    }
    
    return user_deleted;
}

bool Users::updatePassword(const char* username, const char* newPassword) {
    bool user_was_found = false;
    bool success = _transformUsers(filename_, [&](UserData& user) {
        if (strcmp(user.username, username) == 0) {
            user_was_found = true;
            hashPassword(newPassword, user.password_hash);
        }
        return true; // Always keep the user, just maybe modified.
    });
    return success && user_was_found;
}

void Users::reset() {
    if (LittleFS.exists(filename_)) {
        LittleFS.remove(filename_);
    }
    userCount_ = 0;
    initialized_ = false;
}