From 340efa0847764ed31515d9358e91ae842bbe216d Mon Sep 17 00:00:00 2001
From: icpp <icpp@icpp.world>
Date: Fri, 17 Jan 2025 14:21:59 -0500
Subject: [PATCH] check_access

---
 README.md         |  5 ++++-
 src/auth.cpp      | 32 +++++++++++++++++++++-----------
 src/auth.h        |  1 +
 src/llama_cpp.did |  1 +
 4 files changed, 27 insertions(+), 12 deletions(-)

diff --git a/README.md b/README.md
index 68538dc..6da23e7 100644
--- 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
 ```
\ No newline at end of file
diff --git a/src/auth.cpp b/src/auth.cpp
index d47c40e..434cf8a 100644
--- 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});
 }
\ No newline at end of file
diff --git a/src/auth.h b/src/auth.h
index b8bb6ee..27d8181 100644
--- 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);
\ No newline at end of file
diff --git a/src/llama_cpp.did b/src/llama_cpp.did
index 01d8007..d6a7a49 100644
--- 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;