feat: add WASM custom env vars host function (GITHUB_TOKEN support)

echobt · echobt · commit 2d7f6d1ce1df · 2026-03-15T19:58:59.000Z
- Add env_get host function to platform_sandbox WASM module
- Add host_env_get to challenge-sdk-wasm for WASM challenges
- Add custom_env_vars to InstanceConfig and RuntimeState
- Add GITHUB_TOKEN CLI arg to validator-node
- Pass GITHUB_TOKEN env to WASM challenges to avoid GitHub API rate limiting
- Add GITHUB_TOKEN to docker-compose.yml
diff --git a/bins/validator-node/src/main.rs b/bins/validator-node/src/main.rs
@@ -284,6 +284,10 @@ struct Args {
     #[arg(long, env = "CHUTES_API_KEY")]
     chutes_api_key: Option<String>,
 
+    /// GitHub token for WASM challenges to use for API requests (avoids rate limiting)
+    #[arg(long, env = "GITHUB_TOKEN")]
+    github_token: Option<String>,
+
     /// Disable RPC server
     #[arg(long)]
     no_rpc: bool,
@@ -728,6 +732,13 @@ async fn main() -> Result<()> {
             keypair.clone(),
         )),
         chutes_api_key: args.chutes_api_key.clone(),
+        custom_env_vars: {
+            let mut env = std::collections::HashMap::new();
+            if let Some(ref token) = args.github_token {
+                env.insert("GITHUB_TOKEN".to_string(), token.clone());
+            }
+            env
+        },
         distributed_storage: Some(storage_dyn),
         llm_validators_json: Arc::clone(&shared_llm_validators_json),
         registered_hotkeys_json: Arc::clone(&shared_registered_hotkeys_json),
@@ -5238,7 +5249,11 @@ async fn handle_block_event(
                                     mechanism_id, e
                                 );
                                 needs_reconnect = true;
-                                failed_submissions.push((*mechanism_id, uids.clone(), weights.clone()));
+                                failed_submissions.push((
+                                    *mechanism_id,
+                                    uids.clone(),
+                                    weights.clone(),
+                                ));
                             } else {
                                 error!(
                                     "Mechanism {} weight submission failed: {}",
@@ -5253,10 +5268,15 @@ async fn handle_block_event(
                 drop(weights_to_submit);
 
                 if needs_reconnect && !failed_submissions.is_empty() {
-                    if try_reconnect_subtensor(subtensor, subtensor_endpoint, subtensor_state_path).await {
+                    if try_reconnect_subtensor(subtensor, subtensor_endpoint, subtensor_state_path)
+                        .await
+                    {
                         if let (Some(new_st), Some(sig)) = (subtensor.as_ref(), signer.as_ref()) {
                             for (mechanism_id, uids, weights) in &failed_submissions {
-                                info!("Retrying weight submission after reconnect for mechanism {}", mechanism_id);
+                                info!(
+                                    "Retrying weight submission after reconnect for mechanism {}",
+                                    mechanism_id
+                                );
                                 match new_st
                                     .set_mechanism_weights(
                                         sig,
@@ -5277,7 +5297,10 @@ async fn handle_block_event(
                                         *last_weight_submission_epoch = epoch;
                                     }
                                     Ok(resp) => {
-                                        warn!("Mechanism {} issue after reconnect: {}", mechanism_id, resp.message);
+                                        warn!(
+                                            "Mechanism {} issue after reconnect: {}",
+                                            mechanism_id, resp.message
+                                        );
                                     }
                                     Err(e2) => {
                                         error!(
diff --git a/bins/validator-node/src/wasm_executor.rs b/bins/validator-node/src/wasm_executor.rs
@@ -26,6 +26,8 @@ pub struct WasmExecutorConfig {
     pub storage_host_config: StorageHostConfig,
     pub storage_backend: Arc<dyn StorageBackend>,
     pub chutes_api_key: Option<String>,
+    /// Custom environment variables passed to WASM challenges
+    pub custom_env_vars: std::collections::HashMap<String, String>,
     /// Optional distributed storage for loading WASM modules
     pub distributed_storage: Option<Arc<dyn platform_distributed_storage::DistributedStore>>,
     /// Shared chain state: LLM-capable validators (JSON bytes)
@@ -59,6 +61,7 @@ impl Default for WasmExecutorConfig {
             storage_host_config: StorageHostConfig::default(),
             storage_backend: Arc::new(InMemoryStorageBackend::new()),
             chutes_api_key: None,
+            custom_env_vars: std::collections::HashMap::new(),
             distributed_storage: None,
             llm_validators_json: Arc::new(parking_lot::RwLock::new(Vec::new())),
             registered_hotkeys_json: Arc::new(parking_lot::RwLock::new(Vec::new())),
@@ -929,6 +932,7 @@ impl WasmChallengeExecutor {
             },
             llm_validators_json: self.config.llm_validators_json.read().clone(),
             registered_hotkeys_json: self.config.registered_hotkeys_json.read().clone(),
+            custom_env_vars: self.config.custom_env_vars.clone(),
             ..Default::default()
         };
 
@@ -1100,7 +1104,10 @@ impl WasmChallengeExecutor {
                     );
                     return Ok(cached.clone());
                 }
-                debug!(module = module_path, "get_weights skipped: already running, no cache available");
+                debug!(
+                    module = module_path,
+                    "get_weights skipped: already running, no cache available"
+                );
                 return Ok(Vec::new());
             }
         };
@@ -1230,6 +1237,7 @@ impl WasmChallengeExecutor {
             state.consensus_state.epoch = epoch;
             state.fixed_timestamp_ms = Some(real_now_ms as i64);
             state.time_state.set_fixed_timestamp(real_now_ms);
+            state.custom_env_vars = self.config.custom_env_vars.clone();
         }
 
         // Reset fuel before each sync call
diff --git a/crates/bittensor-integration/src/block_sync.rs b/crates/bittensor-integration/src/block_sync.rs
@@ -211,8 +211,7 @@ impl BlockSync {
                         // the broadcast channel.  After a threshold of
                         // consecutive errors we proactively recreate the
                         // listener with a fresh client.
-                        let is_connection_error =
-                            matches!(&event, BlockEvent::ConnectionError(_));
+                        let is_connection_error = matches!(&event, BlockEvent::ConnectionError(_));
 
                         let should_break = BlockSync::handle_block_event(
                             event,
@@ -250,47 +249,30 @@ impl BlockSync {
                             }
 
                             let delay_secs = 5u64;
-                            tokio::time::sleep(
-                                std::time::Duration::from_secs(delay_secs),
-                            )
-                            .await;
+                            tokio::time::sleep(std::time::Duration::from_secs(delay_secs)).await;
 
-                            match BlockSync::recreate_listener(&rpc_url, &config)
-                                .await
-                            {
-                                Ok((
-                                    new_client,
-                                    new_listener,
-                                    new_rx,
-                                    epoch_info,
-                                )) => {
+                            match BlockSync::recreate_listener(&rpc_url, &config).await {
+                                Ok((new_client, new_listener, new_rx, epoch_info)) => {
                                     info!(
                                         block = epoch_info.current_block,
                                         epoch = epoch_info.epoch_number,
                                         "Bittensor client recreated after \
                                          consecutive connection errors"
                                     );
-                                    *current_block.write().await =
-                                        epoch_info.current_block;
-                                    *current_epoch.write().await =
-                                        epoch_info.epoch_number;
-                                    *current_phase.write().await =
-                                        epoch_info.phase;
+                                    *current_block.write().await = epoch_info.current_block;
+                                    *current_epoch.write().await = epoch_info.epoch_number;
+                                    *current_phase.write().await = epoch_info.phase;
                                     current_block_rx = new_rx;
                                     was_disconnected = true;
 
-                                    if let Err(e) =
-                                        new_listener.start(new_client).await
-                                    {
+                                    if let Err(e) = new_listener.start(new_client).await {
                                         warn!(
                                             "Failed to start recreated \
                                              listener: {}",
                                             e
                                         );
                                     } else {
-                                        let _ = event_tx
-                                            .send(BlockSyncEvent::Reconnected)
-                                            .await;
+                                        let _ = event_tx.send(BlockSyncEvent::Reconnected).await;
                                     }
                                 }
                                 Err(e) => {
diff --git a/crates/challenge-sdk-wasm/src/host_functions.rs b/crates/challenge-sdk-wasm/src/host_functions.rs
@@ -247,6 +247,7 @@ extern "C" {
     fn sandbox_exec(req_ptr: i32, req_len: i32, resp_ptr: i32, resp_len: i32) -> i32;
     fn get_timestamp() -> i64;
     fn log_message(level: i32, msg_ptr: i32, msg_len: i32);
+    fn env_get(key_ptr: i32, key_len: i32, val_ptr: i32, val_len: i32) -> i32;
 }
 
 pub fn host_sandbox_exec(request: &[u8]) -> Result<Vec<u8>, i32> {
@@ -274,6 +275,23 @@ pub fn host_log(level: u8, msg: &str) {
     unsafe { log_message(level as i32, msg.as_ptr() as i32, msg.len() as i32) }
 }
 
+pub fn host_env_get(key: &str) -> Option<Vec<u8>> {
+    let mut buf = vec![0u8; 4096];
+    let status = unsafe {
+        env_get(
+            key.as_ptr() as i32,
+            key.len() as i32,
+            buf.as_mut_ptr() as i32,
+            buf.len() as i32,
+        )
+    };
+    if status <= 0 {
+        return None;
+    }
+    buf.truncate(status as usize);
+    Some(buf)
+}
+
 #[link(wasm_import_module = "platform_llm")]
 extern "C" {
     fn llm_chat_completion(req_ptr: i32, req_len: i32, resp_ptr: i32, resp_len: i32) -> i32;
diff --git a/crates/wasm-runtime-interface/src/runtime.rs b/crates/wasm-runtime-interface/src/runtime.rs
@@ -141,6 +141,8 @@ pub struct InstanceConfig {
     pub llm_validators_json: Vec<u8>,
     /// JSON-encoded list of all registered hotkeys from metagraph
     pub registered_hotkeys_json: Vec<u8>,
+    /// Custom environment variables passed to WASM challenges
+    pub custom_env_vars: std::collections::HashMap<String, String>,
 }
 
 impl Default for InstanceConfig {
@@ -169,6 +171,7 @@ impl Default for InstanceConfig {
             epoch: 0,
             llm_validators_json: Vec::new(),
             registered_hotkeys_json: Vec::new(),
+            custom_env_vars: std::collections::HashMap::new(),
         }
     }
 }
@@ -208,6 +211,8 @@ pub struct RuntimeState {
     pub container_state: ContainerState,
     /// LLM state for LLM inference host operations.
     pub llm_state: LlmState,
+    /// Custom environment variables accessible from WASM
+    pub custom_env_vars: std::collections::HashMap<String, String>,
     limits: StoreLimits,
 }
 
@@ -231,6 +236,7 @@ impl RuntimeState {
         config_version: u64,
         storage_state: StorageHostState,
         fixed_timestamp_ms: Option<i64>,
+        custom_env_vars: std::collections::HashMap<String, String>,
         limits: StoreLimits,
     ) -> Self {
         Self {
@@ -251,6 +257,7 @@ impl RuntimeState {
             config_version,
             storage_state,
             fixed_timestamp_ms,
+            custom_env_vars,
             limits,
         }
     }
@@ -396,6 +403,7 @@ impl WasmRuntime {
             instance_config.config_version,
             storage_state,
             instance_config.fixed_timestamp_ms,
+            instance_config.custom_env_vars.clone(),
             limits.build(),
         );
         let mut store = Store::new(&self.engine, runtime_state);
diff --git a/crates/wasm-runtime-interface/src/sandbox.rs b/crates/wasm-runtime-interface/src/sandbox.rs
@@ -33,6 +33,7 @@ pub const HOST_SANDBOX_CONFIGURE: &str = "sandbox_configure";
 pub const HOST_SANDBOX_STATUS: &str = "sandbox_status";
 pub const HOST_SANDBOX_GET_TIMESTAMP: &str = "get_timestamp";
 pub const HOST_SANDBOX_LOG_MESSAGE: &str = "log_message";
+pub const HOST_SANDBOX_ENV_GET: &str = "env_get";
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 #[repr(i32)]
@@ -270,6 +271,26 @@ impl HostFunctionRegistrar for SandboxHostFunctions {
                 ))
             })?;
 
+        linker
+            .func_wrap(
+                HOST_SANDBOX_NAMESPACE,
+                HOST_SANDBOX_ENV_GET,
+                |mut caller: Caller<RuntimeState>,
+                 key_ptr: i32,
+                 key_len: i32,
+                 val_ptr: i32,
+                 val_len: i32|
+                 -> i32 {
+                    handle_env_get(&mut caller, key_ptr, key_len, val_ptr, val_len)
+                },
+            )
+            .map_err(|e| {
+                WasmRuntimeError::HostFunction(format!(
+                    "failed to register {}: {}",
+                    HOST_SANDBOX_ENV_GET, e
+                ))
+            })?;
+
         Ok(())
     }
 }
@@ -493,6 +514,29 @@ fn handle_log_message(caller: &mut Caller<RuntimeState>, level: i32, msg_ptr: i3
     }
 }
 
+fn handle_env_get(
+    caller: &mut Caller<RuntimeState>,
+    key_ptr: i32,
+    key_len: i32,
+    val_ptr: i32,
+    val_len: i32,
+) -> i32 {
+    let key_bytes = match read_memory(caller, key_ptr, key_len) {
+        Ok(b) => b,
+        Err(_) => return -1,
+    };
+    let key = match std::str::from_utf8(&key_bytes) {
+        Ok(s) => s,
+        Err(_) => return -1,
+    };
+
+    let value = caller.data().custom_env_vars.get(key).cloned();
+    match value {
+        Some(v) => write_bytes(caller, val_ptr, val_len, v.as_bytes()),
+        None => 0,
+    }
+}
+
 fn read_memory(caller: &mut Caller<RuntimeState>, ptr: i32, len: i32) -> Result<Vec<u8>, String> {
     if ptr < 0 || len < 0 {
         return Err("negative pointer/length".to_string());
diff --git a/docker-compose.yml b/docker-compose.yml
@@ -37,6 +37,11 @@ services:
       - WITH_BOOTNODE=${WITH_BOOTNODE:-true}
       - BOOTNODE_PORT=${P2P_PORT:-8090}
       - CHUTES_API_KEY=${CHUTES_API_KEY:-}
+      - GITHUB_TOKEN=${GITHUB_TOKEN:-}
+
+    dns:
+      - 8.8.8.8
+      - 1.1.1.1
 
     healthcheck:
       test: ["CMD-SHELL", "curl -sf http://localhost:8080/health || exit 1"]
