refactor: implement builder pattern for OmnectDeviceServiceClient with lifecycle management

.gitignore

@@ -5,4 +5,5 @@ centrifugo
 *.tar.gz
 *.env
 *.http
-/static
+/static
+CLAUDE.md

Cargo.toml

@@ -7,7 +7,7 @@ license = "MIT OR Apache-2.0"
 name = "omnect-ui"
 readme = "README.md"
 repository = "git@github.com:omnect/omnect-ui.git"
-version = "1.0.5"
+version = "1.0.6"
 build = "src/build.rs"
 
 [dependencies]

Dockerfile

@@ -42,7 +42,7 @@ COPY src/build.rs ./omnect-ui/src/build.rs
 
 RUN --mount=type=cache,target=/usr/local/cargo/registry cd omnect-ui && cargo build ${OMNECT_UI_BUILD_ARG} --release --target-dir ./build
 
-COPY src/* ./omnect-ui/src/
+COPY src ./omnect-ui/src/
 COPY .git ./omnect-ui/.git
 RUN --mount=type=cache,target=/usr/local/cargo/registry <<EOF
   set -e

src/auth/token.rs

@@ -46,7 +46,7 @@ impl TokenManager {
         self.inner
             .key
             .authenticate(claims)
-            .map_err(|e| anyhow::anyhow!("failed to create token: {}", e))
+            .map_err(|e| anyhow::anyhow!("failed to create token: {e:#}"))
     }
 
     /// Verify a token and check if it's valid

src/certificate.rs

@@ -1,19 +1,16 @@
 #![cfg_attr(feature = "mock", allow(dead_code, unused_imports))]
 
-use crate::{
-    common::handle_http_response,
-    http_client,
-    omnect_device_service_client::{DeviceServiceClient, OmnectDeviceServiceClient},
-};
+use crate::{common::handle_http_response, http_client};
 use anyhow::{Context, Result};
 use log::info;
 use serde::{Deserialize, Serialize};
 use std::{fs::File, io::Write};
 
-#[derive(Serialize)]
-struct CreateCertPayload {
+// Public payload for passing to certificate creation
+#[derive(Debug, Serialize)]
+pub struct CreateCertPayload {
     #[serde(rename = "commonName")]
-    common_name: String,
+    pub common_name: String,
 }
 
 #[derive(Debug, Deserialize)]
@@ -42,14 +39,13 @@ pub fn key_path() -> String {
 }
 
 #[cfg(feature = "mock")]
-pub async fn create_module_certificate() -> Result<()> {
+pub async fn create_module_certificate(_payload: CreateCertPayload) -> Result<()> {
     Ok(())
 }
 
 #[cfg(not(feature = "mock"))]
-pub async fn create_module_certificate() -> Result<()> {
+pub async fn create_module_certificate(payload: CreateCertPayload) -> Result<()> {
     info!("create module certificate");
-    let ods_client = OmnectDeviceServiceClient::new(false).await?;
     let id = std::env::var("IOTEDGE_MODULEID")
         .context("failed to read IOTEDGE_MODULEID environment variable")?;
     let gen_id = std::env::var("IOTEDGE_MODULEGENERATIONID")
@@ -59,17 +55,13 @@ pub async fn create_module_certificate() -> Result<()> {
     let workload_uri = std::env::var("IOTEDGE_WORKLOADURI")
         .context("failed to read IOTEDGE_WORKLOADURI environment variable")?;
 
-    let payload = CreateCertPayload {
-        common_name: ods_client.ip_address().await?,
-    };
-
-    let path = format!("/modules/{id}/genid/{gen_id}/certificate/server?api-version={api_version}");
+    let path = format!("modules/{id}/genid/{gen_id}/certificate/server?api-version={api_version}");
 
     // Create a client for the IoT Edge workload socket
     let client = http_client::unix_socket_client(&workload_uri)?;
 
-    let url = format!("http://localhost{}", path);
-    info!("POST {url} (IoT Edge workload API)");
+    let url = format!("http://localhost/{path}");
+    info!("POST {url} (IoT Edge workload API) with payload: {payload:?}");
 
     let res = client
         .post(&url)

src/common.rs

@@ -33,6 +33,23 @@ pub fn centrifugo_config() -> Arc<CentrifugoConfig> {
         .clone()
 }
 
+pub fn centrifugo_publish_endpoint() -> crate::omnect_device_service_client::PublishEndpoint {
+    let cfg = centrifugo_config();
+    crate::omnect_device_service_client::PublishEndpoint {
+        url: format!("https://localhost:{}/api/publish", cfg.port),
+        headers: vec![
+            crate::omnect_device_service_client::HeaderKeyValue {
+                name: String::from("Content-Type"),
+                value: String::from("application/json"),
+            },
+            crate::omnect_device_service_client::HeaderKeyValue {
+                name: String::from("X-API-Key"),
+                value: cfg.api_key.clone(),
+            },
+        ],
+    }
+}
+
 macro_rules! config_path {
     () => {
         std::path::Path::new(
@@ -109,10 +126,7 @@ pub async fn handle_http_response(res: Response, context_msg: &str) -> Result<St
 
     ensure!(
         status.is_success(),
-        "{} failed with status {} and body: {}",
-        context_msg,
-        status,
-        body
+        "{context_msg} failed with status {status} and body: {body}"
     );
 
     Ok(body)

src/lib.rs

@@ -1,5 +1,6 @@
 pub mod api;
 pub mod auth;
+pub mod certificate;
 pub mod common;
 pub mod http_client;
 pub mod keycloak_client;

src/main.rs

@@ -10,10 +10,12 @@ mod omnect_device_service_client;
 use crate::{
     api::Api,
     auth::TokenManager,
-    certificate::create_module_certificate,
-    common::{centrifugo_config, config_path},
+    certificate::{CreateCertPayload, create_module_certificate},
+    common::{centrifugo_config, centrifugo_publish_endpoint, config_path},
     keycloak_client::KeycloakProvider,
-    omnect_device_service_client::{DeviceServiceClient, OmnectDeviceServiceClient},
+    omnect_device_service_client::{
+        DeviceServiceClient, OmnectDeviceServiceClient, OmnectDeviceServiceClientBuilder,
+    },
 };
 use actix_files::Files;
 use actix_multipart::form::MultipartFormConfig;
@@ -41,8 +43,7 @@ use tokio::{
 const UPLOAD_LIMIT_BYTES: usize = 250 * 1024 * 1024;
 const MEMORY_LIMIT_BYTES: usize = 10 * 1024 * 1024;
 
-#[actix_web::main]
-async fn main() {
+fn initialize() {
     log_panics::init();
 
     let mut builder = if cfg!(debug_assertions) {
@@ -69,13 +70,35 @@ async fn main() {
         env!("GIT_SHORT_REV")
     );
 
-    create_module_certificate()
-        .await
-        .expect("failed to create module certificate");
+    CryptoProvider::install_default(default_provider()).expect("failed to install crypto provider");
+
+    let Ok(true) = fs::exists("/data") else {
+        panic!("failed to find required data directory: /data is missing");
+    };
+
+    if !fs::exists(config_path!()).is_ok_and(|ok| ok) {
+        fs::create_dir_all(config_path!()).expect("failed to create config directory");
+    };
+
+    common::create_frontend_config_file().expect("failed to create frontend config file");
+}
+
+#[actix_web::main]
+async fn main() {
+    initialize();
 
     let mut sigterm = signal(SignalKind::terminate()).expect("failed to install SIGTERM handler");
     let mut centrifugo = run_centrifugo();
-    let (server_handle, server_task, service_client) = run_server().await;
+    let service_client = OmnectDeviceServiceClientBuilder::new()
+        .with_certificate_setup(|payload: CreateCertPayload| async move {
+            create_module_certificate(payload).await
+        })
+        .with_publish_endpoint(centrifugo_publish_endpoint())
+        .build()
+        .await
+        .expect("failed to create device service client");
+
+    let (server_handle, server_task) = run_server(service_client.clone()).await;
 
     tokio::select! {
         _ = tokio::signal::ctrl_c() => {
@@ -95,7 +118,7 @@ async fn main() {
     // Unified cleanup sequence - ensures consistent shutdown regardless of exit reason
     info!("shutting down...");
 
-    // 1. Shutdown service client first (unregister from omnect-device-service)
+    // 1. Shutdown service client (unregister from omnect-device-service)
     if let Err(e) = service_client.shutdown().await {
         error!("failed to shutdown service client: {e:#}");
     }
@@ -111,33 +134,7 @@ async fn main() {
     info!("centrifugo stopped");
 }
 
-async fn run_server() -> (
-    ServerHandle,
-    tokio::task::JoinHandle<Result<(), std::io::Error>>,
-    OmnectDeviceServiceClient,
-) {
-    CryptoProvider::install_default(default_provider()).expect("failed to install crypto provider");
-
-    let Ok(true) = fs::exists("/data") else {
-        panic!("failed to find required data directory: /data is missing");
-    };
-
-    if !fs::exists(config_path!()).is_ok_and(|ok| ok) {
-        fs::create_dir_all(config_path!()).expect("failed to create config directory");
-    };
-
-    common::create_frontend_config_file().expect("failed to create frontend config file");
-
-    type UiApi = Api<OmnectDeviceServiceClient, KeycloakProvider>;
-
-    let service_client = OmnectDeviceServiceClient::new(true)
-        .await
-        .expect("failed to create client to device service");
-
-    let api = UiApi::new(service_client.clone(), Default::default())
-        .await
-        .expect("failed to create api");
-
+fn load_tls_config() -> rustls::ServerConfig {
     let mut tls_certs = std::io::BufReader::new(
         std::fs::File::open(certificate::cert_path()).expect("failed to read certificate file"),
     );
@@ -149,8 +146,7 @@ async fn run_server() -> (
         .collect::<Result<Vec<_>, _>>()
         .expect("failed to parse cert pem");
 
-    // set up TLS config options
-    let tls_config = match rustls_pemfile::read_one(&mut tls_key)
+    match rustls_pemfile::read_one(&mut tls_key)
         .expect("failed to read key pem file")
         .expect("failed to parse key pem file: no valid key found")
     {
@@ -163,7 +159,22 @@ async fn run_server() -> (
             .with_single_cert(tls_certs, rustls::pki_types::PrivateKeyDer::Pkcs8(key))
             .expect("failed to create TLS config"),
         _ => panic!("failed to parse key pem file: unexpected item type found"),
-    };
+    }
+}
+
+async fn run_server(
+    service_client: OmnectDeviceServiceClient,
+) -> (
+    ServerHandle,
+    tokio::task::JoinHandle<Result<(), std::io::Error>>,
+) {
+    type UiApi = Api<OmnectDeviceServiceClient, KeycloakProvider>;
+
+    let api = UiApi::new(service_client.clone(), Default::default())
+        .await
+        .expect("failed to create api");
+
+    let tls_config = load_tls_config();
 
     let ui_port = std::env::var("UI_PORT")
         .expect("failed to read UI_PORT environment variable")
@@ -256,7 +267,7 @@ async fn run_server() -> (
     .disable_signals()
     .run();
 
-    (server.handle(), tokio::spawn(server), service_client)
+    (server.handle(), tokio::spawn(server))
 }
 
 fn run_centrifugo() -> Child {