fix: OCI Pull timeout

Signed-off-by: Lucas Fontes <[email protected]>

Author: lxfontes

Cargo.lock

@@ -43,6 +43,7 @@ dialoguer = { workspace = true, features = ["editor", "password", "zeroize", "fu
 etcetera = { workspace = true }
 figment = { workspace = true, features = ["json", "env"] }
 flate2 = { workspace = true, features = ["rust_backend"] }
+humantime = { workspace = true }
 indicatif = { workspace = true }
 notify = { workspace = true, features = ["macos_fsevent"] }
 reqwest = { workspace = true, features = ["json", "rustls-tls"] }
@@ -89,6 +90,7 @@ figment = { version = "0.10.19", default-features = false }
 futures = { version = "0.3", default-features = false }
 flate2 = { version = "1.0", default-features = false }
 hyper = { version = "1.6.0", default-features = false }
+humantime = { version = "2.3.0" }
 indicatif = { version = "0.18.0", default-features = false }
 k8s-openapi = { version = "0.25", default-features = false }
 kube = { version = "1", default-features = false }

Cargo.toml

@@ -43,6 +43,7 @@
 use std::collections::{HashMap, HashSet};
 use std::future::Future;
 use std::sync::Arc;
+use std::time::Duration;
 
 use anyhow::{Context, bail};
 use names::{Generator, Name};
@@ -597,9 +598,19 @@ impl std::fmt::Debug for Host {
 }
 
 /// Config for the [`Host`]
-#[derive(Clone, Debug, Default)]
+#[derive(Clone, Debug)]
 pub struct HostConfig {
     pub allow_oci_insecure: bool,
+    pub oci_pull_timeout: Option<Duration>,
+}
+
+impl Default for HostConfig {
+    fn default() -> Self {
+        Self {
+            allow_oci_insecure: false,
+            oci_pull_timeout: Duration::from_secs(30).into(),
+        }
+    }
 }
 
 /// Builder for the [`Host`]

crates/wash-runtime/src/host/mod.rs

@@ -34,6 +34,7 @@ pub struct ClusterHostBuilder {
     host_group: Option<String>,
     host_name: Option<String>,
     heartbeat_interval: Option<Duration>,
+    host_config: Option<HostConfig>,
 }
 
 impl ClusterHostBuilder {
@@ -47,6 +48,11 @@ impl ClusterHostBuilder {
         self
     }
 
+    pub fn with_host_config(mut self, host_config: HostConfig) -> Self {
+        self.host_config = Some(host_config);
+        self
+    }
+
     pub fn with_host_builder(mut self, host_builder: crate::host::HostBuilder) -> Self {
         self.host_builder = host_builder;
         self
@@ -85,6 +91,11 @@ impl ClusterHostBuilder {
         if let Some(host_name) = self.host_name {
             builder = builder.with_hostname(host_name)
         }
+
+        if let Some(host_config) = self.host_config {
+            builder = builder.with_config(host_config);
+        }
+
         let heartbeat_interval = self.heartbeat_interval.unwrap_or(HEARTBEAT_INTERVAL);
         let host = builder.build()?;
         Ok(ClusterHost {
@@ -242,12 +253,17 @@ async fn handle_command(
 
 /// Convert ImagePullSecret from protobuf to OciConfig
 fn image_pull_secret_to_oci_config(
+    config: &HostConfig,
     pull_secret: &Option<types::v2::ImagePullSecret>,
 ) -> oci::OciConfig {
-    match &pull_secret {
+    let mut oci_config = match &pull_secret {
         Some(creds) => oci::OciConfig::new_with_credentials(&creds.username, &creds.password),
         None => OciConfig::default(),
-    }
+    };
+    oci_config.insecure = config.allow_oci_insecure;
+    oci_config.timeout = config.oci_pull_timeout;
+
+    oci_config
 }
 
 async fn host_heartbeat(host: &impl HostApi) -> anyhow::Result<types::v2::HostHeartbeat> {
@@ -275,8 +291,8 @@ async fn workload_start(
     let (components, host_interfaces) = if let Some(wit_world) = wit_world {
         let mut pulled_components = Vec::with_capacity(wit_world.components.len());
         for component in &wit_world.components {
-            let mut oci_config = image_pull_secret_to_oci_config(&component.image_pull_secret);
-            oci_config.insecure = config.allow_oci_insecure;
+            let mut oci_config =
+                image_pull_secret_to_oci_config(config, &component.image_pull_secret);
             let bytes = match oci::pull_component(&component.image, oci_config).await {
                 Ok(bytes) => bytes,
                 Err(e) => {
@@ -316,7 +332,7 @@ async fn workload_start(
     };
 
     let service = if let Some(service) = service {
-        let oci_config = image_pull_secret_to_oci_config(&service.image_pull_secret);
+        let oci_config = image_pull_secret_to_oci_config(config, &service.image_pull_secret);
         let bytes = match oci::pull_component(&service.image, oci_config).await {
             Ok(bytes) => bytes,
             Err(e) => {
@@ -572,12 +588,18 @@ impl From<crate::types::WorkloadStatus> for types::v2::WorkloadStatus {
 
 #[cfg(test)]
 mod tests {
+    use crate::host;
+
     use super::*;
 
     #[tokio::test]
     async fn test_image_pull_secret_to_oci_config_none() {
+        let host_config = HostConfig {
+            allow_oci_insecure: false,
+            oci_pull_timeout: None,
+        };
         let secret: Option<types::v2::ImagePullSecret> = None;
-        let config = image_pull_secret_to_oci_config(&secret);
+        let config = image_pull_secret_to_oci_config(&host_config, &secret);
         assert!(config.credentials.is_none());
         assert!(!config.insecure);
     }
@@ -589,7 +611,11 @@ mod tests {
             password: "testpass".to_string(),
         });
 
-        let config = image_pull_secret_to_oci_config(&secret);
+        let host_config = HostConfig {
+            allow_oci_insecure: false,
+            oci_pull_timeout: None,
+        };
+        let config = image_pull_secret_to_oci_config(&host_config, &secret);
         assert_eq!(
             config.credentials,
             Some(("testuser".to_string(), "testpass".to_string()))

crates/wash-runtime/src/washlet/mod.rs

@@ -1,4 +1,4 @@
-use std::{net::SocketAddr, sync::Arc};
+use std::{net::SocketAddr, sync::Arc, time::Duration};
 
 use anyhow::Context as _;
 use clap::Args;
@@ -34,6 +34,14 @@ pub struct HostCommand {
     #[cfg(not(target_os = "windows"))]
     #[clap(long = "wasi-webgpu", default_value_t = false)]
     pub wasi_webgpu: bool,
+
+    /// Allow insecure OCI Registries
+    #[clap(long = "allow-insecure-registries", default_value_t = false)]
+    pub allow_insecure_registries: bool,
+
+    /// Timeout for pulling artifacts from OCI registries
+    #[clap(long = "registry-pull-timeout", value_parser = humantime::parse_duration, default_value = "30s")]
+    pub registry_pull_timeout: Duration,
 }
 
 impl CliCommand for HostCommand {
@@ -49,7 +57,13 @@ impl CliCommand for HostCommand {
                 .context("failed to connect to NATS")?;
         let data_nats_client = Arc::new(data_nats_client);
 
+        let host_config = wash_runtime::host::HostConfig {
+            allow_oci_insecure: self.allow_insecure_registries,
+            oci_pull_timeout: Some(self.registry_pull_timeout),
+        };
+
         let mut cluster_host_builder = wash_runtime::washlet::ClusterHostBuilder::default()
+            .with_host_config(host_config)
             .with_nats_client(Arc::new(scheduler_nats_client))
             .with_host_group(self.host_group.clone())
             .with_plugin(Arc::new(

crates/wash/src/cli/host.rs

@@ -0,0 +1,88 @@
+package runtime
+
+import (
+	"context"
+	"strings"
+	"time"
+
+	runtimev2 "go.wasmcloud.dev/runtime-operator/pkg/rpc/wasmcloud/runtime/v2"
+	"go.wasmcloud.dev/runtime-operator/pkg/wasmbus"
+	"google.golang.org/protobuf/encoding/protojson"
+	"google.golang.org/protobuf/proto"
+	"google.golang.org/protobuf/types/known/emptypb"
+)
+
+const HostRoundtripTimeout = 1 * time.Minute
+
+func NewWashHostClient(bus wasmbus.Bus, hostID string) *WashHostClient {
+	return &WashHostClient{
+		Bus:    bus,
+		HostID: hostID,
+	}
+}
+
+type WashHostClient struct {
+	Bus    wasmbus.Bus
+	HostID string
+}
+
+func (w *WashHostClient) subject(parts ...string) string {
+	return strings.Join(append([]string{
+		"runtime",
+		"host",
+		w.HostID,
+	}, parts...), ".")
+}
+
+func (w *WashHostClient) Heartbeat(ctx context.Context) (*runtimev2.HostHeartbeat, error) {
+	var resp runtimev2.HostHeartbeat
+	if err := RoundTrip(ctx, w.Bus, w.subject("heartbeat"), &emptypb.Empty{}, &resp); err != nil {
+		return nil, err
+	}
+	return &resp, nil
+}
+
+func (w *WashHostClient) Start(ctx context.Context, req *runtimev2.WorkloadStartRequest) (*runtimev2.WorkloadStartResponse, error) {
+	var resp runtimev2.WorkloadStartResponse
+	if err := RoundTrip(ctx, w.Bus, w.subject("workload.start"), req, &resp); err != nil {
+		return nil, err
+	}
+	return &resp, nil
+}
+
+func (w *WashHostClient) Status(ctx context.Context, req *runtimev2.WorkloadStatusRequest) (*runtimev2.WorkloadStatusResponse, error) {
+	var resp runtimev2.WorkloadStatusResponse
+	if err := RoundTrip(ctx, w.Bus, w.subject("workload.status"), req, &resp); err != nil {
+		return nil, err
+	}
+	return &resp, nil
+}
+
+func (w *WashHostClient) Stop(ctx context.Context, req *runtimev2.WorkloadStopRequest) (*runtimev2.WorkloadStopResponse, error) {
+	var resp runtimev2.WorkloadStopResponse
+	if err := RoundTrip(ctx, w.Bus, w.subject("workload.stop"), req, &resp); err != nil {
+		return nil, err
+	}
+	return &resp, nil
+}
+
+// RoundTrip sends a request and waits for a response.
+func RoundTrip[Req proto.Message, Resp proto.Message](ctx context.Context, bus wasmbus.Bus, subject string, req Req, resp Resp) error {
+	ctx, cancel := context.WithTimeout(ctx, HostRoundtripTimeout)
+	defer cancel()
+
+	json, err := protojson.Marshal(req)
+	if err != nil {
+		return err
+	}
+
+	msg := wasmbus.NewMessage(subject)
+	msg.Data = json
+
+	reply, err := bus.Request(ctx, msg)
+	if err != nil {
+		return err
+	}
+
+	return protojson.Unmarshal(reply.Data, resp)
+}

runtime-operator/internal/controller/runtime/host_client.go

@@ -41,7 +41,7 @@ func (r *HostReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.
 }
 
 func (r *HostReconciler) reconcileReporting(ctx context.Context, host *runtimev1alpha1.Host) error {
-	client := NewWorkloadClient(r.Bus, host.HostID)
+	client := NewWashHostClient(r.Bus, host.HostID)
 
 	ctx, cancel := context.WithTimeout(ctx, 5*time.Second)
 	defer cancel()