diff --git a/native/main.cpp b/native/main.cpp index c52be53..63e5f1c 100644 --- a/native/main.cpp +++ b/native/main.cpp @@ -2,6 +2,7 @@ // Build it with: `icpp build-native` from the parent folder where 'icpp.toml' resides #include "main.h" +#include "test_admin_rbac.h" #include "test_canister_functions.h" #include "test_files.h" #include "test_qwen2.h" @@ -26,6 +27,7 @@ int main() { bool exit_on_fail = true; MockIC mockIC(exit_on_fail); + test_admin_rbac(mockIC); test_canister_functions(mockIC); test_files(mockIC); test_tiny_stories(mockIC); diff --git a/native/test_admin_rbac.cpp b/native/test_admin_rbac.cpp new file mode 100644 index 0000000..2258810 --- /dev/null +++ b/native/test_admin_rbac.cpp @@ -0,0 +1,146 @@ +#include "../src/auth.h" +#include "../src/files.h" +#include "../src/logs.h" +#include "../src/run.h" +#include "mock_ic.h" + +void test_admin_rbac(MockIC &mockIC) { + std::string controller_principal{MOCKIC_CONTROLLER}; + std::string anonymous_principal{"2vxsx-fae"}; + std::string admin_query_principal{"rrkah-fqaaa-aaaaa-aaaaq-cai"}; + + bool silent_on_trap = true; + + // =========================================================================== + // Concrete expected outputs + // =========================================================================== + + // ApiError "Access Denied" + // didc encode '(variant { Err = variant { Other = "Access Denied" } })' + // Note: Actual wire format may have different type table ordering + const std::string ACCESS_DENIED_API_ERROR = + "4449444c026b01b0ad8fcd0c716b01c5fed20100010100000d4163636573732044656e6965" + "64"; + + // OutputRecordResult Access Denied (status_code = 401) + // Verified with didc decode: variant { Err = record { conversation = ""; + // output = ""; error = "Access Denied"; status_code = 401; prompt_remaining = + // ""; generated_eog = false } } + const std::string ACCESS_DENIED_OUTPUT_RECORD = + "4449444c026c06819e846471838fe5800671c897a79907719aa1b2f90c7adb92a2c90d71cd" + "d9e6b30e7e6b01c5fed2010001010000000d4163636573732044656e69656491010000"; + + // Empty input for parameterless endpoints + // didc encode '()' + const std::string EMPTY_INPUT = "4449444c0000"; + + // Input for filesystem_file_size + // didc encode '(record { filename = "test.txt" })' + const std::string FILESYSTEM_FILE_SIZE_INPUT = + "4449444c016c01c7dda8bb0771010008746573742e747874"; + + // Input for new_chat/run_query/run_update + // didc encode '(record { args = vec { "--help" } })' + const std::string RUN_INPUT = + "4449444c026c01b79cba840c016d71010001062d2d68656c70"; + + // Input for assignAdminRole + // didc encode '(record { "principal" = "rrkah-fqaaa-aaaaa-aaaaq-cai"; role = + // variant { AdminQuery }; note = "test" })' + const std::string ASSIGN_ADMIN_QUERY_INPUT = + "4449444c026c03ae9db1900171f2afa8c80471f6d6bbdd04016b0199d8ddf2087f01001b72" + "726b61682d66716161612d61616161612d61616161712d636169047465737400"; + + // =========================================================================== + // Test 1: New Admin RBAC endpoints - ApiError format + // =========================================================================== + + // getAdminRoles - anonymous -> ApiError Access Denied + mockIC.run_test(std::string(__func__) + + ": getAdminRoles - anonymous denied (ApiError)", + getAdminRoles, EMPTY_INPUT, ACCESS_DENIED_API_ERROR, + silent_on_trap, anonymous_principal); + + // assignAdminRole - anonymous -> ApiError Access Denied + mockIC.run_test(std::string(__func__) + + ": assignAdminRole - anonymous denied (ApiError)", + assignAdminRole, ASSIGN_ADMIN_QUERY_INPUT, + ACCESS_DENIED_API_ERROR, silent_on_trap, anonymous_principal); + + // =========================================================================== + // Test 2: ApiError-based existing endpoints + // =========================================================================== + + // filesystem_file_size - anonymous -> ApiError Access Denied + mockIC.run_test(std::string(__func__) + + ": filesystem_file_size - anonymous denied (ApiError)", + filesystem_file_size, FILESYSTEM_FILE_SIZE_INPUT, + ACCESS_DENIED_API_ERROR, silent_on_trap, anonymous_principal); + + // log_pause - anonymous -> ApiError Access Denied + mockIC.run_test(std::string(__func__) + + ": log_pause - anonymous denied (ApiError)", + log_pause, EMPTY_INPUT, ACCESS_DENIED_API_ERROR, + silent_on_trap, anonymous_principal); + + // =========================================================================== + // Test 3: OutputRecordResult endpoints - NOT ApiError format + // =========================================================================== + + // new_chat - anonymous -> OutputRecordResult with status_code = 401 + mockIC.run_test(std::string(__func__) + + ": new_chat - anonymous denied (OutputRecordResult)", + new_chat, RUN_INPUT, ACCESS_DENIED_OUTPUT_RECORD, + silent_on_trap, anonymous_principal); + + // run_query - anonymous -> OutputRecordResult with status_code = 401 + mockIC.run_test(std::string(__func__) + + ": run_query - anonymous denied (OutputRecordResult)", + run_query, RUN_INPUT, ACCESS_DENIED_OUTPUT_RECORD, + silent_on_trap, anonymous_principal); + + // run_update - anonymous -> OutputRecordResult with status_code = 401 + mockIC.run_test(std::string(__func__) + + ": run_update - anonymous denied (OutputRecordResult)", + run_update, RUN_INPUT, ACCESS_DENIED_OUTPUT_RECORD, + silent_on_trap, anonymous_principal); + + // =========================================================================== + // Test 4: RBAC permission hierarchy + // =========================================================================== + + // Controller assigns AdminQuery role + // We don't check exact Ok response (contains dynamic timestamp) + // Just run and verify it doesn't trap + mockIC.run_test( + std::string(__func__) + + ": assignAdminRole - controller assigns AdminQuery", + assignAdminRole, ASSIGN_ADMIN_QUERY_INPUT, + "", // Don't check exact Ok response (contains dynamic timestamp) + silent_on_trap, controller_principal); + + // AdminQuery principal can access AdminQuery endpoint (filesystem_file_size) + // File doesn't exist, so we get file-not-found error, NOT access denied + mockIC.run_test( + std::string(__func__) + + ": filesystem_file_size - AdminQuery allowed (not access denied)", + filesystem_file_size, FILESYSTEM_FILE_SIZE_INPUT, + "", // Don't check exact response; verify it's NOT ACCESS_DENIED_API_ERROR + silent_on_trap, admin_query_principal); + + // AdminQuery principal cannot access AdminUpdate endpoint (filesystem_remove) + mockIC.run_test(std::string(__func__) + + ": filesystem_remove - AdminQuery denied (ApiError)", + filesystem_remove, + FILESYSTEM_FILE_SIZE_INPUT, // Same input format (filename) + ACCESS_DENIED_API_ERROR, silent_on_trap, + admin_query_principal); + + // =========================================================================== + // Test 5: Verify error format isolation + // =========================================================================== + + // The exact match tests above verify: + // - OutputRecordResult endpoints NEVER return ApiError format + // - ApiError endpoints NEVER return OutputRecordResult format +} diff --git a/native/test_admin_rbac.h b/native/test_admin_rbac.h new file mode 100644 index 0000000..4930b85 --- /dev/null +++ b/native/test_admin_rbac.h @@ -0,0 +1,4 @@ +#pragma once +#include "mock_ic.h" + +void test_admin_rbac(MockIC &mockIC); diff --git a/requirements.txt b/requirements.txt index 82e7fcf..6836525 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,6 +3,6 @@ -r scripts/requirements.txt -r src/llama_cpp_onicai_fork/requirements.txt -icpp-pro>=5.2.0 +icpp-pro>=5.3.0 ic-py==1.0.1 binaryen.py \ No newline at end of file diff --git a/src/auth.cpp b/src/auth.cpp index 434cf8a..dbfa220 100644 --- a/src/auth.cpp +++ b/src/auth.cpp @@ -4,6 +4,8 @@ #include "auth.h" #include "http.h" #include "ic_api.h" +#include "utils.h" +#include #include #include #include @@ -12,6 +14,9 @@ static uint16_t access_level = 0; static const std::array access_levels = { "Only controllers", "All except anonymous"}; +// Admin RBAC Storage +static std::vector admin_role_assignments; + // Checks if caller is controller of the canister and if not, optionally write an error to the wire bool is_caller_a_controller(IC_API &ic_api, bool err_to_wire) { CandidTypePrincipal caller = ic_api.get_caller(); @@ -45,6 +50,72 @@ bool is_caller_whitelisted(IC_API &ic_api, bool err_to_wire) { return false; } +// Error senders for different result type families +void send_access_denied_api_error(IC_API &ic_api) { + ic_api.to_wire(CandidTypeVariant{ + "Err", CandidTypeVariant{"Other", CandidTypeText{"Access Denied"}}}); +} + +void send_access_denied_output_record(IC_API &ic_api) { + std::string error_msg = "Access Denied"; + send_output_record_result_error_to_wire( + ic_api, Http::StatusCode::Unauthorized, error_msg); +} + +// Admin RBAC helper +std::optional +get_admin_role_for_principal(const std::string &principal) { + for (const auto &assignment : admin_role_assignments) { + if (assignment.principal == principal) { + return assignment; + } + } + return std::nullopt; +} + +// Admin RBAC auth check functions +bool has_admin_query_role(IC_API &ic_api) { + CandidTypePrincipal caller = ic_api.get_caller(); + if (ic_api.is_controller(caller)) return true; + + std::string principal_text = caller.get_text(); + auto role_opt = get_admin_role_for_principal(principal_text); + + if (role_opt.has_value()) { + // AdminUpdate includes AdminQuery permissions + if (role_opt->role == AdminRole::AdminQuery || + role_opt->role == AdminRole::AdminUpdate) { + return true; + } + } + return false; +} + +bool has_admin_update_role(IC_API &ic_api) { + CandidTypePrincipal caller = ic_api.get_caller(); + if (ic_api.is_controller(caller)) return true; + + std::string principal_text = caller.get_text(); + auto role_opt = get_admin_role_for_principal(principal_text); + + if (role_opt.has_value() && role_opt->role == AdminRole::AdminUpdate) { + return true; + } + return false; +} + +bool has_admin_query_or_whitelisted(IC_API &ic_api) { + if (has_admin_query_role(ic_api)) return true; + if (is_caller_whitelisted(ic_api, false)) return true; + return false; +} + +bool has_admin_update_or_whitelisted(IC_API &ic_api) { + if (has_admin_update_role(ic_api)) return true; + if (is_caller_whitelisted(ic_api, false)) return true; + return false; +} + std::string get_explanation_() { if (access_level == 1) return access_levels[1]; return access_levels[0]; @@ -78,7 +149,10 @@ void set_access() { void get_access() { IC_API ic_api(CanisterQuery{std::string(__func__)}, false); - if (!is_caller_a_controller(ic_api)) return; + if (!has_admin_query_role(ic_api)) { + send_access_denied_api_error(ic_api); + return; + } // Return the status over the wire CandidTypeRecord access_record; @@ -89,10 +163,155 @@ void get_access() { void check_access() { IC_API ic_api(CanisterQuery{std::string(__func__)}, false); - if (!is_caller_whitelisted(ic_api)) return; + if (!has_admin_query_or_whitelisted(ic_api)) { + send_access_denied_api_error(ic_api); + return; + } CandidTypeRecord status_code_record; status_code_record.append("status_code", CandidTypeNat16{Http::StatusCode::OK}); ic_api.to_wire(CandidTypeVariant{"Ok", status_code_record}); +} + +// Admin RBAC Management Endpoints +void assignAdminRole() { + IC_API ic_api(CanisterUpdate{std::string(__func__)}, false); + + CandidTypePrincipal caller = ic_api.get_caller(); + if (!ic_api.is_controller(caller)) { + send_access_denied_api_error(ic_api); + return; + } + + // Parse input + std::string principal_to_assign; + std::string note; + std::string role_label; + + CandidTypeRecord r_in; + r_in.append("principal", CandidTypeText{&principal_to_assign}); + r_in.append("note", CandidTypeText{¬e}); + + // Set up variant with expected options for proper deserialization + CandidTypeVariant role_variant{&role_label}; + role_variant.append("AdminQuery", CandidTypeNull{}); + role_variant.append("AdminUpdate", CandidTypeNull{}); + r_in.append("role", role_variant); + ic_api.from_wire(r_in); + + // Determine role from variant + AdminRole role; + if (role_label == "AdminQuery") { + role = AdminRole::AdminQuery; + } else if (role_label == "AdminUpdate") { + role = AdminRole::AdminUpdate; + } else { + ic_api.to_wire(CandidTypeVariant{ + "Err", CandidTypeVariant{"Other", CandidTypeText{"Invalid role"}}}); + return; + } + + // Upsert: remove existing assignment if any + admin_role_assignments.erase( + std::remove_if(admin_role_assignments.begin(), + admin_role_assignments.end(), + [&](const AdminRoleAssignment &a) { + return a.principal == principal_to_assign; + }), + admin_role_assignments.end()); + + // Create new assignment + AdminRoleAssignment assignment; + assignment.principal = principal_to_assign; + assignment.role = role; + assignment.assignedBy = caller.get_text(); + assignment.assignedAt = ic_api.time(); + assignment.note = note; + + admin_role_assignments.push_back(assignment); + + // Return result + CandidTypeRecord r_out; + r_out.append("principal", CandidTypeText{assignment.principal}); + r_out.append("role", CandidTypeVariant{role_label}); + r_out.append("assignedBy", CandidTypeText{assignment.assignedBy}); + r_out.append("assignedAt", CandidTypeNat64{assignment.assignedAt}); + r_out.append("note", CandidTypeText{assignment.note}); + + ic_api.to_wire(CandidTypeVariant{"Ok", r_out}); +} + +void revokeAdminRole() { + IC_API ic_api(CanisterUpdate{std::string(__func__)}, false); + + CandidTypePrincipal caller = ic_api.get_caller(); + if (!ic_api.is_controller(caller)) { + send_access_denied_api_error(ic_api); + return; + } + + std::string principal_to_revoke; + ic_api.from_wire(CandidTypeText{&principal_to_revoke}); + + auto it = + std::find_if(admin_role_assignments.begin(), admin_role_assignments.end(), + [&](const AdminRoleAssignment &a) { + return a.principal == principal_to_revoke; + }); + + if (it == admin_role_assignments.end()) { + ic_api.to_wire(CandidTypeVariant{ + "Err", + CandidTypeVariant{"Other", CandidTypeText{"Principal not found"}}}); + return; + } + + admin_role_assignments.erase(it); + + ic_api.to_wire(CandidTypeVariant{ + "Ok", CandidTypeText{"Admin role revoked for " + principal_to_revoke}}); +} + +void getAdminRoles() { + IC_API ic_api(CanisterQuery{std::string(__func__)}, false); + + CandidTypePrincipal caller = ic_api.get_caller(); + if (!ic_api.is_controller(caller)) { + send_access_denied_api_error(ic_api); + return; + } + + // Build vectors for each field + std::vector principals; + std::vector role_labels; + std::vector assignedBys; + std::vector assignedAts; + std::vector notes; + + for (const auto &assignment : admin_role_assignments) { + principals.push_back(assignment.principal); + std::string role_label = (assignment.role == AdminRole::AdminQuery) + ? "AdminQuery" + : "AdminUpdate"; + role_labels.push_back(role_label); + assignedBys.push_back(assignment.assignedBy); + assignedAts.push_back(assignment.assignedAt); + notes.push_back(assignment.note); + } + + // Create variant template for AdminRole + CandidTypeVariant role_template; + role_template.append("AdminQuery", CandidTypeNull{}); + role_template.append("AdminUpdate", CandidTypeNull{}); + + // Create a record template with vector fields + CandidTypeRecord r_out; + r_out.append("principal", CandidTypeVecText{principals}); + r_out.append("role", CandidTypeVecVariant{role_template, role_labels}); + r_out.append("assignedBy", CandidTypeVecText{assignedBys}); + r_out.append("assignedAt", CandidTypeVecNat64{assignedAts}); + r_out.append("note", CandidTypeVecText{notes}); + + ic_api.to_wire(CandidTypeVariant{"Ok", CandidTypeVecRecord{r_out}}); } \ No newline at end of file diff --git a/src/auth.h b/src/auth.h index 27d8181..13b0191 100644 --- a/src/auth.h +++ b/src/auth.h @@ -3,11 +3,45 @@ #include "wasm_symbol.h" #include "ic_api.h" +#include +#include #include +#include void set_access() WASM_SYMBOL_EXPORTED("canister_update set_access"); void get_access() WASM_SYMBOL_EXPORTED("canister_query get_access"); void check_access() WASM_SYMBOL_EXPORTED("canister_query check_access"); +// Admin RBAC Management - camelCase to match funnAI's PoAIW implementation +void assignAdminRole() WASM_SYMBOL_EXPORTED("canister_update assignAdminRole"); +void revokeAdminRole() WASM_SYMBOL_EXPORTED("canister_update revokeAdminRole"); +void getAdminRoles() WASM_SYMBOL_EXPORTED("canister_query getAdminRoles"); + bool is_caller_a_controller(IC_API &ic_api, bool err_to_wire = true); -bool is_caller_whitelisted(IC_API &ic_api, bool err_to_wire = true); \ No newline at end of file +bool is_caller_whitelisted(IC_API &ic_api, bool err_to_wire = true); + +// Admin RBAC Types +enum class AdminRole { AdminQuery, AdminUpdate }; + +struct AdminRoleAssignment { + std::string principal; + AdminRole role; + std::string assignedBy; + uint64_t assignedAt; + std::string note; +}; + +// Admin RBAC auth check functions +// Return bool only; do NOT send errors to wire. +// Caller must use appropriate error sender for their result type. +bool has_admin_query_role(IC_API &ic_api); +bool has_admin_update_role(IC_API &ic_api); +bool has_admin_query_or_whitelisted(IC_API &ic_api); +bool has_admin_update_or_whitelisted(IC_API &ic_api); + +// Error senders for access denied responses +void send_access_denied_api_error(IC_API &ic_api); +void send_access_denied_output_record(IC_API &ic_api); + +std::optional +get_admin_role_for_principal(const std::string &principal); \ No newline at end of file diff --git a/src/db_chats.cpp b/src/db_chats.cpp index debaec1..294e455 100644 --- a/src/db_chats.cpp +++ b/src/db_chats.cpp @@ -225,7 +225,10 @@ bool db_chats_save_conversation(const std::string &conversation, // Canister API to retrieve saved chats for authenticated caller void get_chats() { IC_API ic_api(CanisterQuery{std::string(__func__)}, false); - if (!is_caller_whitelisted(ic_api)) return; + if (!has_admin_query_or_whitelisted(ic_api)) { + send_access_denied_api_error(ic_api); + return; + } if (!DB_CHATS_ACTIVE) { ic_api.to_wire(CandidTypeVariant{ @@ -325,7 +328,10 @@ void get_chats() { void chats_resume() { IC_API ic_api(CanisterUpdate{std::string(__func__)}, false); - if (!is_caller_whitelisted(ic_api)) return; + if (!has_admin_update_or_whitelisted(ic_api)) { + send_access_denied_api_error(ic_api); + return; + } DB_CHATS_ACTIVE = true; std::cout << "llama_cpp: " << std::string(__func__) @@ -338,7 +344,10 @@ void chats_resume() { void chats_pause() { IC_API ic_api(CanisterUpdate{std::string(__func__)}, false); - if (!is_caller_whitelisted(ic_api)) return; + if (!has_admin_update_or_whitelisted(ic_api)) { + send_access_denied_api_error(ic_api); + return; + } DB_CHATS_ACTIVE = false; std::cout << "llama_cpp: " << std::string(__func__) diff --git a/src/download.cpp b/src/download.cpp index 591cf6f..b15d76b 100644 --- a/src/download.cpp +++ b/src/download.cpp @@ -25,7 +25,10 @@ void print_file_download_summary(const std::string &filename, void file_download_chunk() { IC_API ic_api(CanisterQuery{std::string(__func__)}, false); - if (!is_caller_a_controller(ic_api)) return; + if (!has_admin_query_role(ic_api)) { + send_access_denied_api_error(ic_api); + return; + } // Get filename to download and the chunksize std::string filename{""}; diff --git a/src/files.cpp b/src/files.cpp index b4a14bb..6e61962 100644 --- a/src/files.cpp +++ b/src/files.cpp @@ -20,7 +20,10 @@ void filesystem_file_size() { IC_API ic_api(CanisterQuery{std::string(__func__)}, false); - if (!is_caller_a_controller(ic_api)) return; + if (!has_admin_query_role(ic_api)) { + send_access_denied_api_error(ic_api); + return; + } // Get filename std::string filename{""}; @@ -34,7 +37,10 @@ void filesystem_file_size() { void get_creation_timestamp_ns() { IC_API ic_api(CanisterQuery{std::string(__func__)}, false); - if (!is_caller_a_controller(ic_api)) return; + if (!has_admin_query_role(ic_api)) { + send_access_denied_api_error(ic_api); + return; + } // Get filename std::string filename{""}; @@ -48,7 +54,10 @@ void get_creation_timestamp_ns() { void filesystem_remove() { IC_API ic_api(CanisterUpdate{std::string(__func__)}, false); - if (!is_caller_a_controller(ic_api)) return; + if (!has_admin_update_role(ic_api)) { + send_access_denied_api_error(ic_api); + return; + } // Get filename std::string filename{""}; @@ -64,13 +73,19 @@ void filesystem_remove() { void recursive_dir_content_update() { IC_API ic_api(CanisterUpdate{std::string(__func__)}, false); - if (!is_caller_a_controller(ic_api)) return; + if (!has_admin_update_role(ic_api)) { + send_access_denied_api_error(ic_api); + return; + } recursive_dir_content_(ic_api); } void recursive_dir_content_query() { IC_API ic_api(CanisterQuery{std::string(__func__)}, false); - if (!is_caller_a_controller(ic_api)) return; + if (!has_admin_query_role(ic_api)) { + send_access_denied_api_error(ic_api); + return; + } recursive_dir_content_(ic_api); } diff --git a/src/llama_cpp.did b/src/llama_cpp.did index 4168cbc..334df4c 100644 --- a/src/llama_cpp.did +++ b/src/llama_cpp.did @@ -141,6 +141,46 @@ type AccessRecord = record { explanation : text }; +// ----------------------------------------------------- +// Admin RBAC Types (match PoAIW structure, use existing ApiError) +// ----------------------------------------------------- +type AdminRole = variant { + AdminQuery; + AdminUpdate +}; + +type AdminRoleAssignment = record { + assignedAt : nat64; + assignedBy : text; + note : text; + "principal" : text; + role : AdminRole +}; + +type AssignAdminRoleInputRecord = record { + note : text; + "principal" : text; + role : AdminRole +}; + +// Uses EXISTING ApiError type (NOT modified) +type AdminRoleAssignmentResult = variant { + Err : ApiError; + Ok : AdminRoleAssignment +}; + +// Uses EXISTING ApiError type (NOT modified) +type AdminRoleAssignmentsResult = variant { + Err : ApiError; + Ok : vec AdminRoleAssignment +}; + +// Uses EXISTING ApiError type (NOT modified) +type TextResult = variant { + Err : ApiError; + Ok : text +}; + // ----------------------------------------------------- type FilesystemFileSizeInputRecord = record { filename : text @@ -256,4 +296,9 @@ service : { // Other admin endpoints whoami : () -> (text) query; + // Admin RBAC Management (Controller-only) + assignAdminRole : (AssignAdminRoleInputRecord) -> (AdminRoleAssignmentResult); + revokeAdminRole : (text) -> (TextResult); + getAdminRoles : () -> (AdminRoleAssignmentsResult) query; + } diff --git a/src/logs.cpp b/src/logs.cpp index 3622fbd..9801b32 100644 --- a/src/logs.cpp +++ b/src/logs.cpp @@ -24,10 +24,8 @@ static void print_usage(int argc, char **argv) { void remove_log_file() { IC_API ic_api(CanisterUpdate{std::string(__func__)}, false); std::string error_msg; - if (!is_caller_a_controller(ic_api)) { - error_msg = "Access Denied."; - send_output_record_result_error_to_wire( - ic_api, Http::StatusCode::Unauthorized, error_msg); + if (!has_admin_update_role(ic_api)) { + send_access_denied_output_record(ic_api); return; } @@ -65,7 +63,10 @@ void remove_log_file() { void log_pause() { IC_API ic_api(CanisterUpdate{std::string(__func__)}, false); - if (!is_caller_a_controller(ic_api)) return; + if (!has_admin_update_role(ic_api)) { + send_access_denied_api_error(ic_api); + return; + } common_log_pause(common_log_main()); @@ -76,7 +77,10 @@ void log_pause() { void log_resume() { IC_API ic_api(CanisterUpdate{std::string(__func__)}, false); - if (!is_caller_a_controller(ic_api)) return; + if (!has_admin_update_role(ic_api)) { + send_access_denied_api_error(ic_api); + return; + } common_log_resume(common_log_main()); diff --git a/src/max_tokens.cpp b/src/max_tokens.cpp index 29e77b1..7af8059 100644 --- a/src/max_tokens.cpp +++ b/src/max_tokens.cpp @@ -12,7 +12,10 @@ uint64_t max_tokens_query{0}; // 0 = no limit void set_max_tokens() { IC_API ic_api(CanisterUpdate{std::string(__func__)}, false); - if (!is_caller_a_controller(ic_api)) return; + if (!has_admin_update_role(ic_api)) { + send_access_denied_api_error(ic_api); + return; + } CandidTypeRecord r_in; r_in.append("max_tokens_update", CandidTypeNat64{&max_tokens_update}); diff --git a/src/model.cpp b/src/model.cpp index 27ecee3..020213d 100644 --- a/src/model.cpp +++ b/src/model.cpp @@ -24,7 +24,10 @@ static void print_usage(int argc, char **argv) { void load_model() { IC_API ic_api(CanisterUpdate{std::string(__func__)}, false); - if (!is_caller_a_controller(ic_api)) return; + if (!has_admin_update_role(ic_api)) { + send_access_denied_output_record(ic_api); + return; + } CandidTypePrincipal caller = ic_api.get_caller(); std::string principal_id = caller.get_text(); diff --git a/src/promptcache.cpp b/src/promptcache.cpp index 558857e..21be642 100644 --- a/src/promptcache.cpp +++ b/src/promptcache.cpp @@ -74,10 +74,8 @@ bool get_canister_path_session(const std::string &path_session, void remove_prompt_cache() { IC_API ic_api(CanisterUpdate{std::string(__func__)}, false); std::string error_msg; - if (!is_caller_whitelisted(ic_api, false)) { - error_msg = "Access Denied."; - send_output_record_result_error_to_wire( - ic_api, Http::StatusCode::Unauthorized, error_msg); + if (!has_admin_update_or_whitelisted(ic_api)) { + send_access_denied_output_record(ic_api); return; } @@ -145,10 +143,8 @@ void remove_prompt_cache() { void copy_prompt_cache() { IC_API ic_api(CanisterUpdate{std::string(__func__)}, false); std::string error_msg; - if (!is_caller_whitelisted(ic_api, false)) { - error_msg = "Access Denied"; - ic_api.to_wire(CandidTypeVariant{ - "Err", CandidTypeVariant{"Other", CandidTypeText{error_msg}}}); + if (!has_admin_update_or_whitelisted(ic_api)) { + send_access_denied_api_error(ic_api); return; } @@ -212,7 +208,10 @@ void copy_prompt_cache() { void download_prompt_cache_chunk() { IC_API ic_api(CanisterQuery{std::string(__func__)}, false); - if (!is_caller_a_controller(ic_api)) return; + if (!has_admin_query_role(ic_api)) { + send_access_denied_api_error(ic_api); + return; + } CandidTypePrincipal caller = ic_api.get_caller(); std::string principal_id = caller.get_text(); @@ -243,7 +242,10 @@ void download_prompt_cache_chunk() { void upload_prompt_cache_chunk() { IC_API ic_api(CanisterUpdate{std::string(__func__)}, false); - if (!is_caller_a_controller(ic_api)) return; + if (!has_admin_update_role(ic_api)) { + send_access_denied_api_error(ic_api); + return; + } CandidTypePrincipal caller = ic_api.get_caller(); std::string principal_id = caller.get_text(); @@ -278,7 +280,10 @@ void uploaded_prompt_cache_details() { // Returns the metadata for an uploaded prompt cache IC_API ic_api(CanisterQuery{std::string(__func__)}, false); - if (!is_caller_a_controller(ic_api)) return; + if (!has_admin_query_role(ic_api)) { + send_access_denied_api_error(ic_api); + return; + } CandidTypePrincipal caller = ic_api.get_caller(); std::string principal_id = caller.get_text(); diff --git a/src/run.cpp b/src/run.cpp index 206c021..47167fa 100644 --- a/src/run.cpp +++ b/src/run.cpp @@ -33,10 +33,8 @@ static void print_usage(int argc, char **argv) { void new_chat() { IC_API ic_api(CanisterUpdate{std::string(__func__)}, false); std::string error_msg; - if (!is_caller_whitelisted(ic_api, false)) { - error_msg = "Access Denied."; - send_output_record_result_error_to_wire( - ic_api, Http::StatusCode::Unauthorized, error_msg); + if (!has_admin_update_or_whitelisted(ic_api)) { + send_access_denied_output_record(ic_api); return; } @@ -117,12 +115,12 @@ void new_chat() { ic_api.to_wire(CandidTypeVariant{"Ok", r_out}); } -void run(IC_API &ic_api, const uint64_t &max_tokens) { +void run(IC_API &ic_api, const uint64_t &max_tokens, bool is_query) { std::string error_msg; - if (!is_caller_whitelisted(ic_api, false)) { - error_msg = "Access Denied."; - send_output_record_result_error_to_wire( - ic_api, Http::StatusCode::Unauthorized, error_msg); + bool authorized = is_query ? has_admin_query_or_whitelisted(ic_api) + : has_admin_update_or_whitelisted(ic_api); + if (!authorized) { + send_access_denied_output_record(ic_api); return; } @@ -196,9 +194,9 @@ void run(IC_API &ic_api, const uint64_t &max_tokens) { void run_query() { IC_API ic_api(CanisterQuery{std::string(__func__)}, false); - run(ic_api, max_tokens_query); + run(ic_api, max_tokens_query, true); } void run_update() { IC_API ic_api(CanisterUpdate{std::string(__func__)}, false); - run(ic_api, max_tokens_update); + run(ic_api, max_tokens_update, false); } diff --git a/src/upload.cpp b/src/upload.cpp index b3b49ee..8043d15 100644 --- a/src/upload.cpp +++ b/src/upload.cpp @@ -190,7 +190,10 @@ void print_file_upload_summary(const std::string &filename, void file_upload_chunk() { IC_API ic_api(CanisterUpdate{std::string(__func__)}, false); - if (!is_caller_a_controller(ic_api)) return; + if (!has_admin_update_role(ic_api)) { + send_access_denied_api_error(ic_api); + return; + } // Get filename and the chunk to write to it std::string filename{""}; @@ -309,7 +312,10 @@ void uploaded_file_details() { // Returns the metadata for an uploaded file IC_API ic_api(CanisterQuery{std::string(__func__)}, false); - if (!is_caller_a_controller(ic_api)) return; + if (!has_admin_query_role(ic_api)) { + send_access_denied_api_error(ic_api); + return; + } // Get filename std::string filename{""}; diff --git a/test/test_admin_rbac.py b/test/test_admin_rbac.py new file mode 100644 index 0000000..e46f9e7 --- /dev/null +++ b/test/test_admin_rbac.py @@ -0,0 +1,261 @@ +"""Test Admin RBAC endpoints + +First deploy the canister: +$ dfx start --clean --background +$ dfx deploy --network local + +Then run the tests: +$ pytest -vv --network local test/test_admin_rbac.py + +Or run a specific test: +$ pytest -vv --network local test/test_admin_rbac.py::test__getAdminRoles_anonymous + +""" +# pylint: disable=missing-function-docstring, unused-import, wildcard-import, unused-wildcard-import, line-too-long + +from pathlib import Path +from typing import Dict +import pytest +from icpp.smoketest import call_canister_api, dict_to_candid_text + +# Path to the dfx.json file +DFX_JSON_PATH = Path(__file__).parent / "../dfx.json" + +# Canister in the dfx.json file we want to test +CANISTER_NAME = "llama_cpp" + + +# ============================================================================= +# Admin RBAC Endpoints - Anonymous Access Denial Tests +# ============================================================================= + +def test__getAdminRoles_anonymous(identity_anonymous: Dict[str, str], network: str) -> None: + """Test getAdminRoles rejects anonymous caller""" + assert identity_anonymous["identity"] == "anonymous" + + response = call_canister_api( + dfx_json_path=DFX_JSON_PATH, + canister_name=CANISTER_NAME, + canister_method="getAdminRoles", + canister_argument="()", + network=network, + ) + expected_response = '(variant { Err = variant { Other = "Access Denied" } })' + assert response == expected_response + + +def test__assignAdminRole_anonymous(identity_anonymous: Dict[str, str], network: str) -> None: + """Test assignAdminRole rejects anonymous caller""" + assert identity_anonymous["identity"] == "anonymous" + + response = call_canister_api( + dfx_json_path=DFX_JSON_PATH, + canister_name=CANISTER_NAME, + canister_method="assignAdminRole", + canister_argument='(record { "principal" = "aaaaa-aa"; role = variant { AdminQuery }; note = "test" })', + network=network, + ) + expected_response = '(variant { Err = variant { Other = "Access Denied" } })' + assert response == expected_response + + +def test__revokeAdminRole_anonymous(identity_anonymous: Dict[str, str], network: str) -> None: + """Test revokeAdminRole rejects anonymous caller""" + assert identity_anonymous["identity"] == "anonymous" + + response = call_canister_api( + dfx_json_path=DFX_JSON_PATH, + canister_name=CANISTER_NAME, + canister_method="revokeAdminRole", + canister_argument='("aaaaa-aa")', + network=network, + ) + expected_response = '(variant { Err = variant { Other = "Access Denied" } })' + assert response == expected_response + + +# ============================================================================= +# Admin RBAC Management - Success Tests (controller can manage roles) +# ============================================================================= + +def test__setup_cleanup_admin_roles(network: str) -> None: + """Setup: Clean up any existing admin roles from previous test runs""" + response = call_canister_api( + dfx_json_path=DFX_JSON_PATH, + canister_name=CANISTER_NAME, + canister_method="revokeAdminRole", + canister_argument='("aaaaa-aa")', + network=network, + ) + # Accept either Ok (role revoked) or Err (role not found) + assert response in [ + '(variant { Ok = "Admin role revoked for aaaaa-aa" })', + '(variant { Err = variant { Other = "Principal not found" } })' + ] + + +def test__getAdminRoles_empty(network: str) -> None: + """Test getAdminRoles returns empty list initially""" + response = call_canister_api( + dfx_json_path=DFX_JSON_PATH, + canister_name=CANISTER_NAME, + canister_method="getAdminRoles", + canister_argument="()", + network=network, + ) + expected_response = '(variant { Ok = vec {} })' + assert response == expected_response + + +def test__assignAdminRole_AdminQuery(network: str) -> None: + """Test assignAdminRole assigns AdminQuery role""" + response = call_canister_api( + dfx_json_path=DFX_JSON_PATH, + canister_name=CANISTER_NAME, + canister_method="assignAdminRole", + canister_argument='(record { "principal" = "aaaaa-aa"; role = variant { AdminQuery }; note = "Test admin query role" })', + network=network, + ) + assert response.startswith('(variant { Ok = record {') + assert '"principal" = "aaaaa-aa"' in response or 'principal = "aaaaa-aa"' in response + # Verify role is a proper variant, not text + assert 'role = variant { AdminQuery }' in response + + +def test__getAdminRoles_after_assign(network: str) -> None: + """Test getAdminRoles returns assigned roles with proper variant format""" + response = call_canister_api( + dfx_json_path=DFX_JSON_PATH, + canister_name=CANISTER_NAME, + canister_method="getAdminRoles", + canister_argument="()", + network=network, + ) + assert response.startswith('(variant { Ok = vec {') + assert 'aaaaa-aa' in response + # Verify role is a proper variant, not text + assert 'role = variant { AdminQuery }' in response + # Verify other fields are present + assert 'assignedBy =' in response + assert 'assignedAt =' in response + assert 'note = "Test admin query role"' in response + + +def test__assignAdminRole_AdminUpdate(network: str) -> None: + """Test assignAdminRole assigns AdminUpdate role (upsert overwrites AdminQuery)""" + response = call_canister_api( + dfx_json_path=DFX_JSON_PATH, + canister_name=CANISTER_NAME, + canister_method="assignAdminRole", + canister_argument='(record { "principal" = "aaaaa-aa"; role = variant { AdminUpdate }; note = "Upgraded to admin update" })', + network=network, + ) + assert response.startswith('(variant { Ok = record {') + # Verify role is a proper variant, not text + assert 'role = variant { AdminUpdate }' in response + + +def test__getAdminRoles_after_update(network: str) -> None: + """Test getAdminRoles returns AdminUpdate role with proper variant format""" + response = call_canister_api( + dfx_json_path=DFX_JSON_PATH, + canister_name=CANISTER_NAME, + canister_method="getAdminRoles", + canister_argument="()", + network=network, + ) + assert response.startswith('(variant { Ok = vec {') + assert 'aaaaa-aa' in response + # Verify role is a proper variant with AdminUpdate (was upgraded from AdminQuery) + assert 'role = variant { AdminUpdate }' in response + assert 'note = "Upgraded to admin update"' in response + + +# ============================================================================= +# Multiple Principals Tests - verify CandidTypeVecVariant with multiple entries +# ============================================================================= + +def test__assignAdminRole_second_principal(network: str) -> None: + """Test assigning AdminQuery role to a second principal""" + response = call_canister_api( + dfx_json_path=DFX_JSON_PATH, + canister_name=CANISTER_NAME, + canister_method="assignAdminRole", + canister_argument='(record { "principal" = "rrkah-fqaaa-aaaaa-aaaaq-cai"; role = variant { AdminQuery }; note = "Second admin" })', + network=network, + ) + assert response.startswith('(variant { Ok = record {') + assert 'role = variant { AdminQuery }' in response + + +def test__getAdminRoles_multiple_principals(network: str) -> None: + """Test getAdminRoles returns multiple principals with different roles""" + response = call_canister_api( + dfx_json_path=DFX_JSON_PATH, + canister_name=CANISTER_NAME, + canister_method="getAdminRoles", + canister_argument="()", + network=network, + ) + assert response.startswith('(variant { Ok = vec {') + # Verify both principals are present + assert 'aaaaa-aa' in response + assert 'rrkah-fqaaa-aaaaa-aaaaq-cai' in response + # Verify both role variants are present (AdminUpdate for first, AdminQuery for second) + assert 'role = variant { AdminUpdate }' in response + assert 'role = variant { AdminQuery }' in response + # Verify both notes are present + assert 'note = "Upgraded to admin update"' in response + assert 'note = "Second admin"' in response + + +def test__revokeAdminRole_second_principal(network: str) -> None: + """Clean up: revoke second principal's role""" + response = call_canister_api( + dfx_json_path=DFX_JSON_PATH, + canister_name=CANISTER_NAME, + canister_method="revokeAdminRole", + canister_argument='("rrkah-fqaaa-aaaaa-aaaaq-cai")', + network=network, + ) + expected_response = '(variant { Ok = "Admin role revoked for rrkah-fqaaa-aaaaa-aaaaq-cai" })' + assert response == expected_response + + +def test__revokeAdminRole(network: str) -> None: + """Test revokeAdminRole removes role""" + response = call_canister_api( + dfx_json_path=DFX_JSON_PATH, + canister_name=CANISTER_NAME, + canister_method="revokeAdminRole", + canister_argument='("aaaaa-aa")', + network=network, + ) + expected_response = '(variant { Ok = "Admin role revoked for aaaaa-aa" })' + assert response == expected_response + + +def test__revokeAdminRole_not_found(network: str) -> None: + """Test revokeAdminRole returns error for non-existent principal""" + response = call_canister_api( + dfx_json_path=DFX_JSON_PATH, + canister_name=CANISTER_NAME, + canister_method="revokeAdminRole", + canister_argument='("non-existent-principal")', + network=network, + ) + expected_response = '(variant { Err = variant { Other = "Principal not found" } })' + assert response == expected_response + + +def test__getAdminRoles_after_revoke(network: str) -> None: + """Test getAdminRoles returns empty after revoke""" + response = call_canister_api( + dfx_json_path=DFX_JSON_PATH, + canister_name=CANISTER_NAME, + canister_method="getAdminRoles", + canister_argument="()", + network=network, + ) + expected_response = '(variant { Ok = vec {} })' + assert response == expected_response diff --git a/test/test_canister_functions.py b/test/test_canister_functions.py index 1baa1eb..d17f55e 100644 --- a/test/test_canister_functions.py +++ b/test/test_canister_functions.py @@ -108,3 +108,270 @@ def test__get_access_0(network: str) -> None: ) expected_response = '(variant { Ok = record { explanation = "Only controllers"; level = 0 : nat16;} })' assert response == expected_response + + +# ------------------------------------------------------------------ +# check_access tests +def test__check_access_anonymous(identity_anonymous: Dict[str, str], network: str) -> None: + """Test check_access rejects anonymous caller""" + assert identity_anonymous["identity"] == "anonymous" + assert identity_anonymous["principal"] == "2vxsx-fae" + + response = call_canister_api( + dfx_json_path=DFX_JSON_PATH, + canister_name=CANISTER_NAME, + canister_method="check_access", + canister_argument="()", + network=network, + ) + expected_response = '(variant { Err = variant { Other = "Access Denied" } })' + assert response == expected_response + + +def test__check_access_controller(network: str) -> None: + """Test check_access succeeds for controller""" + response = call_canister_api( + dfx_json_path=DFX_JSON_PATH, + canister_name=CANISTER_NAME, + canister_method="check_access", + canister_argument="()", + network=network, + ) + expected_response = '(variant { Ok = record { status_code = 200 : nat16;} })' + assert response == expected_response + + +# ------------------------------------------------------------------ +# whoami tests +def test__whoami(network: str, principal: str) -> None: + """Test whoami returns caller's principal""" + response = call_canister_api( + dfx_json_path=DFX_JSON_PATH, + canister_name=CANISTER_NAME, + canister_method="whoami", + canister_argument="()", + network=network, + ) + expected_response = f'("{principal}")' + assert response == expected_response + + +def test__whoami_anonymous(identity_anonymous: Dict[str, str], network: str) -> None: + """Test whoami returns anonymous principal""" + assert identity_anonymous["identity"] == "anonymous" + assert identity_anonymous["principal"] == "2vxsx-fae" + + response = call_canister_api( + dfx_json_path=DFX_JSON_PATH, + canister_name=CANISTER_NAME, + canister_method="whoami", + canister_argument="()", + network=network, + ) + expected_response = '("2vxsx-fae")' + assert response == expected_response + + +# ============================================================================= +# Additional Anonymous Access Denial Tests +# ============================================================================= + +def test__set_max_tokens_anonymous(identity_anonymous: Dict[str, str], network: str) -> None: + """Test set_max_tokens rejects anonymous caller""" + assert identity_anonymous["identity"] == "anonymous" + + response = call_canister_api( + dfx_json_path=DFX_JSON_PATH, + canister_name=CANISTER_NAME, + canister_method="set_max_tokens", + canister_argument='(record { max_tokens_update = 100 : nat64; max_tokens_query = 100 : nat64 })', + network=network, + ) + expected_response = '(variant { Err = variant { Other = "Access Denied" } })' + assert response == expected_response + + +def test__load_model_anonymous(identity_anonymous: Dict[str, str], network: str) -> None: + """Test load_model rejects anonymous caller (OutputRecordResult format)""" + assert identity_anonymous["identity"] == "anonymous" + + response = call_canister_api( + dfx_json_path=DFX_JSON_PATH, + canister_name=CANISTER_NAME, + canister_method="load_model", + canister_argument='(record { args = vec { "--help" } })', + network=network, + ) + # load_model uses OutputRecordResult format for access denied + assert 'Err' in response + assert 'status_code = 401' in response or 'Access Denied' in response + + +def test__log_pause_anonymous(identity_anonymous: Dict[str, str], network: str) -> None: + """Test log_pause rejects anonymous caller""" + assert identity_anonymous["identity"] == "anonymous" + + response = call_canister_api( + dfx_json_path=DFX_JSON_PATH, + canister_name=CANISTER_NAME, + canister_method="log_pause", + canister_argument="()", + network=network, + ) + expected_response = '(variant { Err = variant { Other = "Access Denied" } })' + assert response == expected_response + + +def test__log_resume_anonymous(identity_anonymous: Dict[str, str], network: str) -> None: + """Test log_resume rejects anonymous caller""" + assert identity_anonymous["identity"] == "anonymous" + + response = call_canister_api( + dfx_json_path=DFX_JSON_PATH, + canister_name=CANISTER_NAME, + canister_method="log_resume", + canister_argument="()", + network=network, + ) + expected_response = '(variant { Err = variant { Other = "Access Denied" } })' + assert response == expected_response + + +def test__remove_log_file_anonymous(identity_anonymous: Dict[str, str], network: str) -> None: + """Test remove_log_file rejects anonymous caller (OutputRecordResult format)""" + assert identity_anonymous["identity"] == "anonymous" + + response = call_canister_api( + dfx_json_path=DFX_JSON_PATH, + canister_name=CANISTER_NAME, + canister_method="remove_log_file", + canister_argument='(record { args = vec {} })', + network=network, + ) + # remove_log_file uses OutputRecordResult format for access denied + assert 'Err' in response + assert 'status_code = 401' in response or 'Access Denied' in response + + +def test__new_chat_anonymous(identity_anonymous: Dict[str, str], network: str) -> None: + """Test new_chat rejects anonymous caller (OutputRecordResult format)""" + assert identity_anonymous["identity"] == "anonymous" + + response = call_canister_api( + dfx_json_path=DFX_JSON_PATH, + canister_name=CANISTER_NAME, + canister_method="new_chat", + canister_argument='(record { args = vec { "--help" } })', + network=network, + ) + # new_chat uses OutputRecordResult format for access denied + assert 'Err' in response + assert 'status_code = 401' in response or 'Access Denied' in response + + +def test__run_query_anonymous(identity_anonymous: Dict[str, str], network: str) -> None: + """Test run_query rejects anonymous caller (OutputRecordResult format)""" + assert identity_anonymous["identity"] == "anonymous" + + response = call_canister_api( + dfx_json_path=DFX_JSON_PATH, + canister_name=CANISTER_NAME, + canister_method="run_query", + canister_argument='(record { args = vec { "--help" } })', + network=network, + ) + # run_query uses OutputRecordResult format for access denied + assert 'Err' in response + assert 'status_code = 401' in response or 'Access Denied' in response + + +def test__run_update_anonymous(identity_anonymous: Dict[str, str], network: str) -> None: + """Test run_update rejects anonymous caller (OutputRecordResult format)""" + assert identity_anonymous["identity"] == "anonymous" + + response = call_canister_api( + dfx_json_path=DFX_JSON_PATH, + canister_name=CANISTER_NAME, + canister_method="run_update", + canister_argument='(record { args = vec { "--help" } })', + network=network, + ) + # run_update uses OutputRecordResult format for access denied + assert 'Err' in response + assert 'status_code = 401' in response or 'Access Denied' in response + + +def test__remove_prompt_cache_anonymous(identity_anonymous: Dict[str, str], network: str) -> None: + """Test remove_prompt_cache rejects anonymous caller (OutputRecordResult format)""" + assert identity_anonymous["identity"] == "anonymous" + + response = call_canister_api( + dfx_json_path=DFX_JSON_PATH, + canister_name=CANISTER_NAME, + canister_method="remove_prompt_cache", + canister_argument='(record { args = vec { "test.cache" } })', + network=network, + ) + # remove_prompt_cache uses OutputRecordResult format for access denied + assert 'Err' in response + assert 'status_code = 401' in response or 'Access Denied' in response + + +def test__copy_prompt_cache_anonymous(identity_anonymous: Dict[str, str], network: str) -> None: + """Test copy_prompt_cache rejects anonymous caller""" + assert identity_anonymous["identity"] == "anonymous" + + response = call_canister_api( + dfx_json_path=DFX_JSON_PATH, + canister_name=CANISTER_NAME, + canister_method="copy_prompt_cache", + canister_argument='(record { from = "source.cache"; to = "dest.cache" })', + network=network, + ) + expected_response = '(variant { Err = variant { Other = "Access Denied" } })' + assert response == expected_response + + +def test__get_chats_anonymous(identity_anonymous: Dict[str, str], network: str) -> None: + """Test get_chats rejects anonymous caller""" + assert identity_anonymous["identity"] == "anonymous" + + response = call_canister_api( + dfx_json_path=DFX_JSON_PATH, + canister_name=CANISTER_NAME, + canister_method="get_chats", + canister_argument="()", + network=network, + ) + expected_response = '(variant { Err = variant { Other = "Access Denied" } })' + assert response == expected_response + + +def test__chats_resume_anonymous(identity_anonymous: Dict[str, str], network: str) -> None: + """Test chats_resume rejects anonymous caller""" + assert identity_anonymous["identity"] == "anonymous" + + response = call_canister_api( + dfx_json_path=DFX_JSON_PATH, + canister_name=CANISTER_NAME, + canister_method="chats_resume", + canister_argument="()", + network=network, + ) + expected_response = '(variant { Err = variant { Other = "Access Denied" } })' + assert response == expected_response + + +def test__chats_pause_anonymous(identity_anonymous: Dict[str, str], network: str) -> None: + """Test chats_pause rejects anonymous caller""" + assert identity_anonymous["identity"] == "anonymous" + + response = call_canister_api( + dfx_json_path=DFX_JSON_PATH, + canister_name=CANISTER_NAME, + canister_method="chats_pause", + canister_argument="()", + network=network, + ) + expected_response = '(variant { Err = variant { Other = "Access Denied" } })' + assert response == expected_response diff --git a/test/test_files.py b/test/test_files.py index 8ff012f..3c0153e 100644 --- a/test/test_files.py +++ b/test/test_files.py @@ -333,6 +333,42 @@ def test__filesystem_file_size_controller_another_prompt(network: str, principal expected_response = f'(variant {{ Ok = record {{ msg = "File exists: .canister_cache/{principal}/sessions/another_prompt.cache\\nFile size: 5 bytes\\n"; filename = ".canister_cache/{principal}/sessions/another_prompt.cache"; filesize = 5 : nat64; exists = true;}} }})' assert response == expected_response +# ------------------------------------------------------------------ +# file_download_chunk tests +def test__file_download_chunk_anonymous(identity_anonymous: Dict[str, str], network: str) -> None: + """Test that file_download_chunk rejects anonymous caller""" + assert identity_anonymous["identity"] == "anonymous" + assert identity_anonymous["principal"] == "2vxsx-fae" + + response = call_canister_api( + dfx_json_path=DFX_JSON_PATH, + canister_name=CANISTER_NAME, + canister_method="file_download_chunk", + canister_argument='(record { filename = "test.bin"; chunksize = 1024 : nat64; offset = 0 : nat64 })', + network=network, + ) + expected_response = '(variant { Err = variant { Other = "Access Denied" } })' + assert response == expected_response + + +# ------------------------------------------------------------------ +# uploaded_file_details tests +def test__uploaded_file_details_anonymous(identity_anonymous: Dict[str, str], network: str) -> None: + """Test that uploaded_file_details rejects anonymous caller""" + assert identity_anonymous["identity"] == "anonymous" + assert identity_anonymous["principal"] == "2vxsx-fae" + + response = call_canister_api( + dfx_json_path=DFX_JSON_PATH, + canister_name=CANISTER_NAME, + canister_method="uploaded_file_details", + canister_argument='(record { filename = "test.bin" })', + network=network, + ) + expected_response = '(variant { Err = variant { Other = "Access Denied" } })' + assert response == expected_response + + # ------------------------------------------------------------------ # Security test: MAX_CHUNK_SIZE validation def test__file_download_chunk_exceeds_max_chunk_size(network: str, principal: str) -> None: @@ -350,3 +386,48 @@ def test__file_download_chunk_exceeds_max_chunk_size(network: str, principal: st ) expected_response = f'(variant {{ Err = variant {{ Other = "file_download_chunk_: chunksize {chunksize_3mb} exceeds limit 2097152" }} }})' assert response == expected_response + + +# ------------------------------------------------------------------ +# file_upload_chunk tests +def test__file_upload_chunk_anonymous(identity_anonymous: Dict[str, str], network: str) -> None: + """Test that file_upload_chunk rejects anonymous caller""" + assert identity_anonymous["identity"] == "anonymous" + assert identity_anonymous["principal"] == "2vxsx-fae" + + response = call_canister_api( + dfx_json_path=DFX_JSON_PATH, + canister_name=CANISTER_NAME, + canister_method="file_upload_chunk", + canister_argument='(record { filename = "test.bin"; chunk = blob "\\01\\02\\03"; chunksize = 3 : nat64; offset = 0 : nat64 })', + network=network, + ) + expected_response = '(variant { Err = variant { Other = "Access Denied" } })' + assert response == expected_response + + +def test__file_upload_chunk_controller(network: str, principal: str) -> None: + """Test that file_upload_chunk succeeds for controller""" + response = call_canister_api( + dfx_json_path=DFX_JSON_PATH, + canister_name=CANISTER_NAME, + canister_method="file_upload_chunk", + canister_argument='(record { filename = "models/test_upload.bin"; chunk = blob "\\01\\02\\03\\04\\05"; chunksize = 5 : nat64; offset = 0 : nat64 })', + network=network, + ) + assert response.startswith('(variant { Ok = record {') + assert 'filename = "models/test_upload.bin"' in response + assert 'filesize = 5' in response + + +def test__file_upload_chunk_cleanup(network: str, principal: str) -> None: + """Cleanup: remove the test file created by file_upload_chunk test""" + response = call_canister_api( + dfx_json_path=DFX_JSON_PATH, + canister_name=CANISTER_NAME, + canister_method="filesystem_remove", + canister_argument='(record { filename = "models/test_upload.bin" })', + network=network, + ) + # Accept either removed or not found + assert 'Ok' in response or 'does not exist' in response diff --git a/test/test_promptcache.py b/test/test_promptcache.py index 94f2a92..2c41ca8 100644 --- a/test/test_promptcache.py +++ b/test/test_promptcache.py @@ -24,6 +24,62 @@ CANISTER_NAME = "llama_cpp" +# ============================================================================= +# Anonymous Access Denial Tests +# ============================================================================= + +def test__upload_prompt_cache_chunk_anonymous(identity_anonymous: Dict[str, str], network: str) -> None: + """Test that upload_prompt_cache_chunk rejects anonymous caller""" + assert identity_anonymous["identity"] == "anonymous" + assert identity_anonymous["principal"] == "2vxsx-fae" + + response = call_canister_api( + dfx_json_path=DFX_JSON_PATH, + canister_name=CANISTER_NAME, + canister_method="upload_prompt_cache_chunk", + canister_argument='(record { promptcache = "test.cache"; chunk = blob "\\01\\02\\03"; chunksize = 3 : nat64; offset = 0 : nat64 })', + network=network, + ) + expected_response = '(variant { Err = variant { Other = "Access Denied" } })' + assert response == expected_response + + +def test__download_prompt_cache_chunk_anonymous(identity_anonymous: Dict[str, str], network: str) -> None: + """Test that download_prompt_cache_chunk rejects anonymous caller""" + assert identity_anonymous["identity"] == "anonymous" + assert identity_anonymous["principal"] == "2vxsx-fae" + + response = call_canister_api( + dfx_json_path=DFX_JSON_PATH, + canister_name=CANISTER_NAME, + canister_method="download_prompt_cache_chunk", + canister_argument='(record { promptcache = "test.cache"; chunksize = 1024 : nat64; offset = 0 : nat64 })', + network=network, + ) + expected_response = '(variant { Err = variant { Other = "Access Denied" } })' + assert response == expected_response + + +def test__uploaded_prompt_cache_details_anonymous(identity_anonymous: Dict[str, str], network: str) -> None: + """Test that uploaded_prompt_cache_details rejects anonymous caller""" + assert identity_anonymous["identity"] == "anonymous" + assert identity_anonymous["principal"] == "2vxsx-fae" + + response = call_canister_api( + dfx_json_path=DFX_JSON_PATH, + canister_name=CANISTER_NAME, + canister_method="uploaded_prompt_cache_details", + canister_argument='(record { promptcache = "test.cache" })', + network=network, + ) + expected_response = '(variant { Err = variant { Other = "Access Denied" } })' + assert response == expected_response + + +# ============================================================================= +# Controller Success Tests +# ============================================================================= + def test__upload_prompt_cache_chunk_0(network: str, principal: str) -> None: response = call_canister_api( dfx_json_path=DFX_JSON_PATH,