check_access

onicai · Jan 17, 2025 · 340efa0 · 340efa0
1 parent 62bfca1
commit 340efa0
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 12 deletions.
diff --git a/README.md b/README.md
@@ -244,7 +244,7 @@ You can run a smoketest on the deployed LLM:
   pytest -vv test/test_qwen2.py
   ```
 
-## Securing your LLM
+## Access control
 
 By default, only a controller can call the inference endpoints:
 - new_chat
@@ -261,4 +261,7 @@ dfx canister call llama_cpp set_access '(record {level = 1 : nat16})'
 
 # Verify it worked
 dfx canister call llama_cpp get_access
+
+# A caller can check it's access rights with
+dfx canister call llama_cpp check_access
 ```
diff --git a/src/auth.cpp b/src/auth.cpp
@@ -50,17 +50,6 @@ std::string get_explanation_() {
   return access_levels[0];
 }
 
-void get_access() {
-  IC_API ic_api(CanisterUpdate{std::string(__func__)}, false);
-  if (!is_caller_a_controller(ic_api)) return;
-
-  // Return the status over the wire
-  CandidTypeRecord access_record;
-  access_record.append("level", CandidTypeNat16{access_level});
-  access_record.append("explanation", CandidTypeText{get_explanation_()});
-  ic_api.to_wire(CandidTypeVariant{"Ok", CandidTypeRecord{access_record}});
-}
-
 void set_access() {
   IC_API ic_api(CanisterUpdate{std::string(__func__)}, false);
   if (!is_caller_a_controller(ic_api)) return;
@@ -85,4 +74,25 @@ void set_access() {
   access_record.append("level", CandidTypeNat16{access_level});
   access_record.append("explanation", CandidTypeText{get_explanation_()});
   ic_api.to_wire(CandidTypeVariant{"Ok", CandidTypeRecord{access_record}});
+}
+
+void get_access() {
+  IC_API ic_api(CanisterQuery{std::string(__func__)}, false);
+  if (!is_caller_a_controller(ic_api)) return;
+
+  // Return the status over the wire
+  CandidTypeRecord access_record;
+  access_record.append("level", CandidTypeNat16{access_level});
+  access_record.append("explanation", CandidTypeText{get_explanation_()});
+  ic_api.to_wire(CandidTypeVariant{"Ok", CandidTypeRecord{access_record}});
+}
+
+void check_access() {
+  IC_API ic_api(CanisterQuery{std::string(__func__)}, false);
+  if (!is_caller_whitelisted(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});
 }
diff --git a/src/auth.h b/src/auth.h
@@ -7,6 +7,7 @@
 
 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");
 
 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);
diff --git a/src/llama_cpp.did b/src/llama_cpp.did
@@ -126,6 +126,7 @@ service : {
   // Access level
   set_access : (AccessInputRecord) -> (AccessRecordResult);
   get_access : () -> (AccessRecordResult) query;
+  check_access : () -> (StatusCodeRecordResult) query;
 
   // Other admin endpoints
   whoami : () -> (text) query;