tsk/docker/
template_manager.rs

1//! Docker template management system
2//!
3//! This module provides a manager for Docker layer templates that follows
4//! the same layered approach as the AssetManager, checking for templates
5//! in project, user, and embedded locations in priority order.
6
7use anyhow::{Context, Result};
8use std::path::Path;
9use std::sync::Arc;
10
11use crate::assets::AssetManager;
12use crate::docker::layers::{DockerImageConfig, DockerLayer, DockerLayerContent, DockerLayerType};
13use crate::storage::xdg::XdgDirectories;
14
15/// Manages Docker templates and layer composition
16pub struct DockerTemplateManager {
17    asset_manager: Arc<dyn AssetManager>,
18    xdg_dirs: Arc<XdgDirectories>,
19}
20
21impl DockerTemplateManager {
22    /// Creates a new DockerTemplateManager
23    pub fn new(asset_manager: Arc<dyn AssetManager>, xdg_dirs: Arc<XdgDirectories>) -> Self {
24        Self {
25            asset_manager,
26            xdg_dirs,
27        }
28    }
29
30    /// Get the content of a specific Docker layer
31    pub fn get_layer_content(
32        &self,
33        layer: &DockerLayer,
34        project_root: Option<&Path>,
35    ) -> Result<DockerLayerContent> {
36        let layer_path = format!("dockerfiles/{}", layer.asset_path());
37
38        // For project layers, check filesystem first if project_root is provided
39        let dockerfile_content = if layer.layer_type == DockerLayerType::Project {
40            if let Some(root) = project_root {
41                let project_dockerfile = root
42                    .join(".tsk")
43                    .join("dockerfiles")
44                    .join("project")
45                    .join(&layer.name)
46                    .join("Dockerfile");
47
48                if project_dockerfile.exists() {
49                    // Read directly from filesystem for project dockerfiles
50                    std::fs::read(&project_dockerfile).with_context(|| {
51                        format!(
52                            "Failed to read project Dockerfile: {}",
53                            project_dockerfile.display()
54                        )
55                    })?
56                } else {
57                    // Fall back to asset manager
58                    let dockerfile_path = format!("{layer_path}/Dockerfile");
59                    self.get_docker_file_content(&dockerfile_path)
60                        .with_context(|| format!("Failed to get Dockerfile for layer {layer}"))?
61                }
62            } else {
63                // No project root, use asset manager
64                let dockerfile_path = format!("{layer_path}/Dockerfile");
65                self.get_docker_file_content(&dockerfile_path)
66                    .with_context(|| format!("Failed to get Dockerfile for layer {layer}"))?
67            }
68        } else {
69            // Non-project layers always use asset manager
70            let dockerfile_path = format!("{layer_path}/Dockerfile");
71            self.get_docker_file_content(&dockerfile_path)
72                .with_context(|| format!("Failed to get Dockerfile for layer {layer}"))?
73        };
74
75        // Try to get additional files if they exist
76        let mut additional_files = Vec::new();
77
78        // Check for common additional files (this could be extended or made configurable)
79        let potential_files = vec![
80            "requirements.txt",
81            "package.json",
82            "Cargo.toml",
83            "config.json",
84        ];
85
86        // For project layers with filesystem access, check filesystem first
87        if layer.layer_type == DockerLayerType::Project {
88            if let Some(root) = project_root {
89                let project_layer_dir = root
90                    .join(".tsk")
91                    .join("dockerfiles")
92                    .join("project")
93                    .join(&layer.name);
94
95                for file_name in &potential_files {
96                    let file_path = project_layer_dir.join(file_name);
97                    if file_path.exists() {
98                        if let Ok(content) = std::fs::read(&file_path) {
99                            additional_files.push((file_name.to_string(), content));
100                        }
101                    }
102                }
103            }
104        }
105
106        // If no files found on filesystem (or for non-project layers), check asset manager
107        if additional_files.is_empty() {
108            for file_name in potential_files {
109                let file_path = format!("{layer_path}/{file_name}");
110                if let Ok(content) = self.get_docker_file_content(&file_path) {
111                    additional_files.push((file_name.to_string(), content));
112                }
113            }
114        }
115
116        Ok(DockerLayerContent::with_files(
117            String::from_utf8(dockerfile_content)?,
118            additional_files,
119        ))
120    }
121
122    /// Compose a complete Dockerfile from multiple layers
123    pub fn compose_dockerfile(
124        &self,
125        config: &DockerImageConfig,
126        project_root: Option<&Path>,
127    ) -> Result<String> {
128        let layers = config.get_layers();
129        let mut composed_dockerfile = String::new();
130        let mut has_from_instruction = false;
131        let mut cmd_instruction: Option<String> = None;
132        let mut entrypoint_instruction: Option<String> = None;
133
134        // Add header comment
135        composed_dockerfile.push_str(&format!(
136            "# TSK Composed Dockerfile\n# Tech Stack: {}\n# Agent: {}\n# Project: {}\n\n",
137            config.tech_stack, config.agent, config.project
138        ));
139
140        for (index, layer) in layers.iter().enumerate() {
141            // Skip layers that don't exist (except base which is required)
142            let layer_content = match self.get_layer_content(layer, project_root) {
143                Ok(content) => content,
144                Err(_) if layer.layer_type != DockerLayerType::Base => {
145                    // Non-base layers are optional
146                    continue;
147                }
148                Err(e) => return Err(e),
149            };
150
151            // Add layer boundary comment
152            composed_dockerfile.push_str(&format!(
153                "\n###### BEGIN {} LAYER: {} ######\n",
154                layer.layer_type.to_string().to_uppercase(),
155                layer.name
156            ));
157
158            // Process the Dockerfile content
159            let (processed_content, cmd, entrypoint) = self.process_layer_content(
160                &layer_content.dockerfile_content,
161                index == 0,
162                &mut has_from_instruction,
163            )?;
164
165            // Update CMD and ENTRYPOINT if found in this layer
166            if let Some(cmd) = cmd {
167                cmd_instruction = Some(cmd);
168            }
169            if let Some(entrypoint) = entrypoint {
170                entrypoint_instruction = Some(entrypoint);
171            }
172
173            composed_dockerfile.push_str(&processed_content);
174            composed_dockerfile.push('\n');
175
176            // Add layer boundary comment
177            composed_dockerfile.push_str(&format!(
178                "###### END {} LAYER: {} ######\n",
179                layer.layer_type.to_string().to_uppercase(),
180                layer.name
181            ));
182        }
183
184        // Ensure we have at least one FROM instruction
185        if !has_from_instruction {
186            return Err(anyhow::anyhow!(
187                "No FROM instruction found in any layer. At least the base layer must contain a FROM instruction."
188            ));
189        }
190
191        // Add ENTRYPOINT and CMD at the end if they exist
192        if let Some(entrypoint) = entrypoint_instruction {
193            composed_dockerfile.push_str("\n# Final ENTRYPOINT\n");
194            composed_dockerfile.push_str(&entrypoint);
195            composed_dockerfile.push('\n');
196        }
197        if let Some(cmd) = cmd_instruction {
198            composed_dockerfile.push_str("\n# Final CMD\n");
199            composed_dockerfile.push_str(&cmd);
200            composed_dockerfile.push('\n');
201        }
202
203        Ok(composed_dockerfile)
204    }
205
206    /// List available layers of a specific type
207    pub fn list_available_layers(
208        &self,
209        layer_type: DockerLayerType,
210        project_root: Option<&Path>,
211    ) -> Vec<String> {
212        let mut layers = std::collections::HashSet::new();
213
214        let layer_dir = match layer_type {
215            DockerLayerType::Base => "dockerfiles/base".to_string(),
216            _ => format!("dockerfiles/{layer_type}"),
217        };
218
219        // Check embedded assets
220        let dockerfiles = self.asset_manager.list_dockerfiles();
221        for dockerfile in dockerfiles {
222            if dockerfile.starts_with(&layer_dir) {
223                if let Some(name) = self.extract_layer_name(&dockerfile, &layer_type) {
224                    layers.insert(name);
225                }
226            }
227        }
228
229        // Check user directory
230        let user_docker_dir = self.xdg_dirs.config_dir().join("dockerfiles");
231        if user_docker_dir.exists() {
232            self.scan_directory_for_layers(&user_docker_dir, &layer_type, &mut layers);
233        }
234
235        // Check project directory
236        if let Some(project_root) = project_root {
237            let project_docker_dir = project_root.join(".tsk").join("dockerfiles");
238            if project_docker_dir.exists() {
239                eprintln!(
240                    "Scanning project dockerfiles directory: {}",
241                    project_docker_dir.display()
242                );
243                self.scan_directory_for_layers(&project_docker_dir, &layer_type, &mut layers);
244            } else {
245                eprintln!(
246                    "No project dockerfiles directory found at: {}",
247                    project_docker_dir.display()
248                );
249            }
250        }
251
252        let mut result: Vec<String> = layers.into_iter().collect();
253        result.sort();
254        result
255    }
256
257    /// Process layer content to handle special cases
258    fn process_layer_content(
259        &self,
260        content: &str,
261        is_first_layer: bool,
262        has_from_instruction: &mut bool,
263    ) -> Result<(String, Option<String>, Option<String>)> {
264        let mut processed = String::new();
265        let mut cmd_instruction: Option<String> = None;
266        let mut entrypoint_instruction: Option<String> = None;
267        let mut seen_user_root = false;
268
269        for line in content.lines() {
270            let trimmed = line.trim();
271
272            // Handle FROM instructions
273            if trimmed.starts_with("FROM ") {
274                if *has_from_instruction && !is_first_layer {
275                    // Skip subsequent FROM instructions in non-base layers
276                    continue;
277                } else {
278                    *has_from_instruction = true;
279                }
280            }
281
282            // Track if we've seen USER root in this layer
283            if trimmed == "USER root" {
284                seen_user_root = true;
285            }
286
287            // Extract CMD and ENTRYPOINT instructions
288            if trimmed.starts_with("CMD ") {
289                cmd_instruction = Some(line.to_string());
290                continue;
291            }
292            if trimmed.starts_with("ENTRYPOINT ") {
293                entrypoint_instruction = Some(line.to_string());
294                continue;
295            }
296
297            // Skip USER agent only if we haven't seen USER root in this layer
298            // This allows switching back to agent after root operations
299            if !is_first_layer && trimmed == "USER agent" && !seen_user_root {
300                continue;
301            }
302
303            processed.push_str(line);
304            processed.push('\n');
305        }
306
307        Ok((processed, cmd_instruction, entrypoint_instruction))
308    }
309
310    /// Get Docker file content from the asset manager
311    fn get_docker_file_content(&self, path: &str) -> Result<Vec<u8>> {
312        // For now, we'll use the existing dockerfile methods
313        // In the future, this could be extended to support the new layer structure
314        if let Some(dockerfile_name) = path.strip_prefix("dockerfiles/") {
315            if let Some((name, file_path)) = dockerfile_name.split_once('/') {
316                if file_path == "Dockerfile" {
317                    self.asset_manager.get_dockerfile(name)
318                } else {
319                    self.asset_manager.get_dockerfile_file(name, file_path)
320                }
321            } else {
322                self.asset_manager.get_dockerfile(dockerfile_name)
323            }
324        } else {
325            Err(anyhow::anyhow!("Invalid dockerfile path: {}", path))
326        }
327    }
328
329    /// Extract layer name from a dockerfile path
330    fn extract_layer_name(&self, path: &str, layer_type: &DockerLayerType) -> Option<String> {
331        let parts: Vec<&str> = path.split('/').collect();
332        match layer_type {
333            DockerLayerType::Base => {
334                if parts.len() >= 2 && parts[1] == "base" {
335                    Some("base".to_string())
336                } else {
337                    None
338                }
339            }
340            _ => {
341                if parts.len() >= 3 && parts[1] == layer_type.to_string() {
342                    Some(parts[2].to_string())
343                } else {
344                    None
345                }
346            }
347        }
348    }
349
350    /// Scan a directory for layers of a specific type
351    fn scan_directory_for_layers(
352        &self,
353        dir: &Path,
354        layer_type: &DockerLayerType,
355        layers: &mut std::collections::HashSet<String>,
356    ) {
357        let layer_dir = match layer_type {
358            DockerLayerType::Base => dir.join("base"),
359            _ => dir.join(layer_type.to_string()),
360        };
361
362        if layer_dir.exists() {
363            if layer_type == &DockerLayerType::Base {
364                // Base layer is a special case - just check if Dockerfile exists
365                if layer_dir.join("Dockerfile").exists() {
366                    layers.insert("base".to_string());
367                }
368            } else {
369                // For other layer types, each subdirectory is a layer
370                if let Ok(entries) = std::fs::read_dir(&layer_dir) {
371                    for entry in entries.flatten() {
372                        if entry.path().is_dir() {
373                            if let Some(name) = entry.file_name().to_str() {
374                                if entry.path().join("Dockerfile").exists() {
375                                    layers.insert(name.to_string());
376                                }
377                            }
378                        }
379                    }
380                }
381            }
382        }
383    }
384}
385
386#[cfg(test)]
387mod tests {
388    use super::*;
389    use crate::assets::embedded::EmbeddedAssetManager;
390    use std::fs;
391    use tempfile::TempDir;
392
393    fn create_test_manager() -> DockerTemplateManager {
394        let temp_dir = TempDir::new().unwrap();
395        let xdg_dirs = XdgDirectories::new_with_paths(
396            temp_dir.path().to_path_buf(),
397            temp_dir.path().to_path_buf(),
398            temp_dir.path().to_path_buf(),
399            temp_dir.path().to_path_buf(),
400        );
401
402        DockerTemplateManager::new(Arc::new(EmbeddedAssetManager), Arc::new(xdg_dirs))
403    }
404
405    #[test]
406    fn test_docker_image_config_layers() {
407        let config = DockerImageConfig::new(
408            "rust".to_string(),
409            "claude-code".to_string(),
410            "web-api".to_string(),
411        );
412
413        let layers = config.get_layers();
414        assert_eq!(layers.len(), 4);
415        assert_eq!(layers[0].name, "base");
416        assert_eq!(layers[1].name, "rust");
417        assert_eq!(layers[2].name, "claude-code");
418        assert_eq!(layers[3].name, "web-api");
419    }
420
421    #[test]
422    fn test_process_layer_content() {
423        let manager = create_test_manager();
424        let mut has_from = false;
425
426        // Test first layer processing
427        let content = "FROM ubuntu:22.04\nRUN apt-get update\nWORKDIR /workspace\nUSER agent\nCMD [\"/bin/bash\"]";
428        let (processed, cmd, entrypoint) = manager
429            .process_layer_content(content, true, &mut has_from)
430            .unwrap();
431        assert!(processed.contains("FROM ubuntu:22.04"));
432        assert!(processed.contains("WORKDIR /workspace"));
433        assert!(processed.contains("USER agent"));
434        assert!(!processed.contains("CMD")); // CMD should be extracted
435        assert_eq!(cmd, Some("CMD [\"/bin/bash\"]".to_string()));
436        assert_eq!(entrypoint, None);
437        assert!(has_from);
438
439        // Test non-first layer processing
440        let content2 = "FROM alpine\nRUN apk add git\nWORKDIR /workspace\nUSER agent";
441        let (processed2, cmd2, entrypoint2) = manager
442            .process_layer_content(content2, false, &mut has_from)
443            .unwrap();
444        assert!(!processed2.contains("FROM alpine")); // Should skip additional FROM
445        assert!(processed2.contains("WORKDIR /workspace")); // Should keep WORKDIR now
446        assert!(!processed2.contains("USER agent")); // Should skip USER agent when not after root
447        assert!(processed2.contains("RUN apk add git")); // Should keep RUN
448        assert_eq!(cmd2, None);
449        assert_eq!(entrypoint2, None);
450    }
451
452    #[test]
453    fn test_process_layer_content_user_switching() {
454        let manager = create_test_manager();
455        let mut has_from = false;
456
457        // Test layer with USER root switching back to USER agent
458        let content = "# Switch to root\nUSER root\nRUN apt-get update\n# Switch back\nUSER agent\nRUN echo test";
459        let (processed, _, _) = manager
460            .process_layer_content(content, false, &mut has_from)
461            .unwrap();
462
463        // Should keep both USER instructions when switching from root to agent
464        assert!(processed.contains("USER root"));
465        assert!(processed.contains("USER agent"));
466        assert!(processed.contains("RUN apt-get update"));
467        assert!(processed.contains("RUN echo test"));
468    }
469
470    #[test]
471    fn test_cmd_and_entrypoint_extraction() {
472        let manager = create_test_manager();
473        let mut has_from = false;
474
475        // Test CMD extraction
476        let content1 = "RUN echo test\nCMD [\"default\"]";
477        let (processed1, cmd1, entrypoint1) = manager
478            .process_layer_content(content1, false, &mut has_from)
479            .unwrap();
480        assert!(processed1.contains("RUN echo test"));
481        assert!(!processed1.contains("CMD"));
482        assert_eq!(cmd1, Some("CMD [\"default\"]".to_string()));
483        assert_eq!(entrypoint1, None);
484
485        // Test ENTRYPOINT extraction
486        let content2 = "RUN echo test2\nENTRYPOINT [\"/entrypoint.sh\"]\nCMD [\"arg\"]";
487        let (processed2, cmd2, entrypoint2) = manager
488            .process_layer_content(content2, false, &mut has_from)
489            .unwrap();
490        assert!(processed2.contains("RUN echo test2"));
491        assert!(!processed2.contains("ENTRYPOINT"));
492        assert!(!processed2.contains("CMD"));
493        assert_eq!(cmd2, Some("CMD [\"arg\"]".to_string()));
494        assert_eq!(
495            entrypoint2,
496            Some("ENTRYPOINT [\"/entrypoint.sh\"]".to_string())
497        );
498    }
499
500    #[test]
501    fn test_compose_dockerfile_cmd_placement() {
502        let _manager = create_test_manager();
503
504        // Create a config for testing
505        let _config = DockerImageConfig {
506            tech_stack: "rust".to_string(),
507            agent: "claude".to_string(),
508            project: "default".to_string(),
509        };
510
511        // This test will use the embedded dockerfiles
512        // We can't easily test the full compose without real files,
513        // but we've tested the components thoroughly
514    }
515
516    #[test]
517    fn test_extract_layer_name() {
518        let manager = create_test_manager();
519
520        assert_eq!(
521            manager.extract_layer_name("dockerfiles/base/Dockerfile", &DockerLayerType::Base),
522            Some("base".to_string())
523        );
524
525        assert_eq!(
526            manager.extract_layer_name(
527                "dockerfiles/tech-stack/rust/Dockerfile",
528                &DockerLayerType::TechStack
529            ),
530            Some("rust".to_string())
531        );
532
533        assert_eq!(
534            manager.extract_layer_name(
535                "dockerfiles/agent/claude-code/Dockerfile",
536                &DockerLayerType::Agent
537            ),
538            Some("claude-code".to_string())
539        );
540
541        assert_eq!(
542            manager.extract_layer_name(
543                "dockerfiles/project/web-api/Dockerfile",
544                &DockerLayerType::Project
545            ),
546            Some("web-api".to_string())
547        );
548    }
549
550    #[test]
551    fn test_scan_directory_for_layers() {
552        let temp_dir = TempDir::new().unwrap();
553        let docker_dir = temp_dir.path().join("dockerfiles");
554
555        // Create test layer directories
556        fs::create_dir_all(docker_dir.join("base")).unwrap();
557        fs::write(docker_dir.join("base/Dockerfile"), "FROM ubuntu").unwrap();
558
559        fs::create_dir_all(docker_dir.join("tech-stack/rust")).unwrap();
560        fs::write(
561            docker_dir.join("tech-stack/rust/Dockerfile"),
562            "RUN install rust",
563        )
564        .unwrap();
565
566        fs::create_dir_all(docker_dir.join("tech-stack/python")).unwrap();
567        fs::write(
568            docker_dir.join("tech-stack/python/Dockerfile"),
569            "RUN install python",
570        )
571        .unwrap();
572
573        let manager = create_test_manager();
574        let mut layers = std::collections::HashSet::new();
575
576        // Test base layer scanning
577        manager.scan_directory_for_layers(&docker_dir, &DockerLayerType::Base, &mut layers);
578        assert!(layers.contains("base"));
579        layers.clear();
580
581        // Test tech-stack layer scanning
582        manager.scan_directory_for_layers(&docker_dir, &DockerLayerType::TechStack, &mut layers);
583        assert!(layers.contains("rust"));
584        assert!(layers.contains("python"));
585    }
586}
tsk/docker/template_manager.rs

tsk/docker/
template_manager.rs
