simple-object-server/src/update_handler.cpp

#include "update_handler.hpp"
#include "validation.hpp"
#include <nlohmann/json.hpp>
#include <drogon/MultiPart.h>
#include <iostream>

namespace simple_object_storage {

UpdateHandler::UpdateHandler(Server& server) : server_(server) {}

void UpdateHandler::handle_update_object(const drogon::HttpRequestPtr& req, std::function<void(const drogon::HttpResponsePtr &)>&& callback) {
    auto resp = drogon::HttpResponse::newHttpResponse();
    std::map<std::string, std::string> params;
    
    // Validate authentication and rate limit (no required query params, just auth)
    if (!server_.validate_write_request(req, resp, {}, params)) {
        server_.add_security_headers(resp);
        callback(resp);
        return;
    }

    // Parse request based on content type
    nlohmann::json body;
    const auto& contentType = req->getHeader("content-type");
    
    if (contentType.find("application/json") != std::string::npos) {
        // Handle JSON content
        try {
            body = nlohmann::json::parse(req->getBody());
        } catch (const nlohmann::json::parse_error& e) {
            resp->setStatusCode(drogon::k400BadRequest);
            nlohmann::json response = {{"result", "error"}, {"error", "Invalid JSON body"}};
            resp->setBody(response.dump());
            resp->setContentTypeCode(drogon::CT_APPLICATION_JSON);
            server_.add_security_headers(resp);
            callback(resp);
            return;
        }
    } else if (contentType.find("multipart/form-data") != std::string::npos) {
        // Handle multipart form data
        body = nlohmann::json::object();
        
        // Parse the multipart form data
        drogon::MultiPartParser formParser;
        if (formParser.parse(req) != 0) {
            resp->setStatusCode(drogon::k400BadRequest);
            nlohmann::json response = {{"result", "error"}, {"error", "Failed to parse multipart data"}};
            resp->setBody(response.dump());
            resp->setContentTypeCode(drogon::CT_APPLICATION_JSON);
            server_.add_security_headers(resp);
            callback(resp);
            return;
        }
        
        auto &parameters = formParser.getParameters();
        
        // Get hash from form data
        auto hash_it = parameters.find("hash");
        if (hash_it != parameters.end()) {
            body["hash"] = hash_it->second;
        }
        
        // Parse metadata if present
        auto meta_it = parameters.find("metadata");
        if (meta_it != parameters.end()) {
            try {
                body["metadata"] = nlohmann::json::parse(meta_it->second);
            } catch (const nlohmann::json::parse_error& e) {
                // If parsing as JSON fails, treat it as a string
                body["metadata"] = meta_it->second;
            }
        }
    } else if (contentType.find("application/x-www-form-urlencoded") != std::string::npos) {
        // Handle URL-encoded form data
        body = nlohmann::json::object();
        const auto& form = req->getParameters();
        
        // Get hash from form data
        auto hash_it = form.find("hash");
        if (hash_it != form.end()) {
            body["hash"] = hash_it->second;
        }
        
        // Parse metadata if present
        auto meta_it = form.find("metadata");
        if (meta_it != form.end()) {
            try {
                body["metadata"] = nlohmann::json::parse(meta_it->second);
            } catch (const nlohmann::json::parse_error& e) {
                // If parsing as JSON fails, treat it as a string
                body["metadata"] = meta_it->second;
            }
        }
    } else {
        resp->setStatusCode(drogon::k400BadRequest);
        nlohmann::json response = {{"result", "error"}, {"error", "Unsupported content type. Use application/json or form data."}};
        resp->setBody(response.dump());
        resp->setContentTypeCode(drogon::CT_APPLICATION_JSON);
        server_.add_security_headers(resp);
        callback(resp);
        return;
    }

    // Check for required fields
    if (!body.contains("hash") || !body.contains("metadata")) {
        resp->setStatusCode(drogon::k400BadRequest);
        nlohmann::json response = {{"result", "error"}, {"error", "Missing 'hash' or 'metadata' field in request body"}};
        resp->setBody(response.dump());
        resp->setContentTypeCode(drogon::CT_APPLICATION_JSON);
        server_.add_security_headers(resp);
        callback(resp);
        return;
    }

    std::string hash = body["hash"].get<std::string>();
    nlohmann::json new_metadata = body["metadata"];

    // Validate hash format
    auto hashValidation = InputValidator::validateHash(hash);
    if (!hashValidation.valid) {
        resp->setStatusCode(drogon::k400BadRequest);
        nlohmann::json response = {{"result", "error"}, {"error", "Invalid hash: " + hashValidation.error}};
        resp->setBody(response.dump());
        resp->setContentTypeCode(drogon::CT_APPLICATION_JSON);
        server_.add_security_headers(resp);
        callback(resp);
        return;
    }

    // Validate metadata
    auto metadataValidation = InputValidator::validateMetadata(new_metadata);
    if (!metadataValidation.valid) {
        resp->setStatusCode(drogon::k400BadRequest);
        nlohmann::json response = {{"result", "error"}, {"error", "Invalid metadata: " + metadataValidation.error}};
        resp->setBody(response.dump());
        resp->setContentTypeCode(drogon::CT_APPLICATION_JSON);
        server_.add_security_headers(resp);
        callback(resp);
        return;
    }

    // If metadata contains labeltags, validate them
    if (new_metadata.contains("labeltags") && new_metadata["labeltags"].is_array()) {
        std::vector<std::string> labeltags;
        for (const auto& lt : new_metadata["labeltags"]) {
            if (!lt.is_string()) {
                resp->setStatusCode(drogon::k400BadRequest);
                nlohmann::json response = {{"result", "error"}, {"error", "Invalid label:tag format in metadata"}};
                resp->setBody(response.dump());
                resp->setContentTypeCode(drogon::CT_APPLICATION_JSON);
                server_.add_security_headers(resp);
                callback(resp);
                return;
            }
            labeltags.push_back(lt.get<std::string>());
        }
        
        if (!labeltags.empty()) {
            auto labeltagsValidation = InputValidator::validateLabelTags(labeltags);
            if (!labeltagsValidation.valid) {
                resp->setStatusCode(drogon::k400BadRequest);
                nlohmann::json response = {{"result", "error"}, {"error", labeltagsValidation.error}};
                resp->setBody(response.dump());
                resp->setContentTypeCode(drogon::CT_APPLICATION_JSON);
                server_.add_security_headers(resp);
                callback(resp);
                return;
            }
        }
    }

    // Get the object entry
    dbEntry entry;
    if (!server_.db_->get(hash, entry)) {
        resp->setStatusCode(drogon::k404NotFound);
        nlohmann::json response = {{"result", "error"}, {"error", "Object not found for hash: " + hash}};
        resp->setBody(response.dump());
        resp->setContentTypeCode(drogon::CT_APPLICATION_JSON);
        server_.add_security_headers(resp);
        callback(resp);
        return;
    }

    // Prepare new entry for database merge
    dbEntry new_entry;
    new_entry.hash = hash;
    new_entry.metadata = new_metadata;
    
    // Extract labeltags if provided in the new metadata
    if (new_metadata.contains("labeltags") && new_metadata["labeltags"].is_array()) {
        new_entry.labeltags = new_metadata["labeltags"].get<std::vector<std::string>>();
    }
    // If no labeltags provided, new_entry.labeltags will be empty and merge will preserve existing ones

    if (!server_.db_->update_or_insert(new_entry)) {
        resp->setStatusCode(drogon::k500InternalServerError);
        nlohmann::json response = {{"result", "error"}, {"error", "Failed to update metadata for hash: " + hash}};
        resp->setBody(response.dump());
        resp->setContentTypeCode(drogon::CT_APPLICATION_JSON);
        server_.add_security_headers(resp);
        callback(resp);
        return;
    }

    nlohmann::json response = {{"result", "success"}, {"hash", hash}};
    resp->setBody(response.dump());
    resp->setContentTypeCode(drogon::CT_APPLICATION_JSON);
    server_.add_security_headers(resp);
    callback(resp);
}

} // namespace simple_object_storage