refactor: update logic refactor to make it more testable

Introduced logic to test for zip slip (path traversal)

v2/internal/runner/update.go

@@ -316,60 +316,42 @@ func (r *Runner) compareAndWriteTemplates(zipReader *zip.Reader) (*templateUpdat
 	// If the path isn't found in new update after being read from the previous checksum,
 	// it is removed. This allows us fine-grained control over the download process
 	// as well as solves a long problem with nuclei-template updates.
-	checksumFile := filepath.Join(r.templatesConfig.TemplatesDirectory, ".checksum")
+	configuredTemplateDirectory := r.templatesConfig.TemplatesDirectory
+	checksumFile := filepath.Join(configuredTemplateDirectory, ".checksum")
 	templateChecksumsMap, _ := createTemplateChecksumsMap(checksumFile)
 	for _, zipTemplateFile := range zipReader.File {
-		directory, name := filepath.Split(zipTemplateFile.Name)
+		templateAbsolutePath, skipFile, err := calculateTemplateAbsolutePath(zipTemplateFile.Name, configuredTemplateDirectory)
-		if name == "" {
+		if err != nil {
+			return nil, err
+		}
+		if skipFile {
 			continue
 		}
-		paths := strings.Split(directory, string(os.PathSeparator))
-		finalPath := filepath.Join(paths[1:]...)
-
-		if strings.HasPrefix(name, ".") || strings.HasPrefix(finalPath, ".") || strings.EqualFold(name, "README.md") {
-			continue
-		}
-		results.totalCount++
-		templateDirectory := filepath.Join(r.templatesConfig.TemplatesDirectory, finalPath)
-		if err := os.MkdirAll(templateDirectory, 0755); err != nil {
-			return nil, fmt.Errorf("failed to create template folder %s : %s", templateDirectory, err)
-		}
-
-		templatePath := filepath.Join(templateDirectory, name)
 
 		isAddition := false
-		if _, statErr := os.Stat(templatePath); os.IsNotExist(statErr) {
+		if _, statErr := os.Stat(templateAbsolutePath); os.IsNotExist(statErr) {
 			isAddition = true
 		}
-		templateFile, err := os.OpenFile(templatePath, os.O_TRUNC|os.O_CREATE|os.O_WRONLY, 0644)
+
+		newTemplateChecksum, err := writeUnZippedTemplateFile(err, templateAbsolutePath, zipTemplateFile)
 		if err != nil {
-			templateFile.Close()
+			return nil, err
-			return nil, fmt.Errorf("could not create uncompressed file: %s", err)
 		}
 
-		zipTemplateFileReader, err := zipTemplateFile.Open()
+		oldTemplateChecksum, checksumOk := templateChecksumsMap[templateAbsolutePath]
+
+		relativeTemplatePath, err := filepath.Rel(configuredTemplateDirectory, templateAbsolutePath)
 		if err != nil {
-			templateFile.Close()
+			return nil, fmt.Errorf("could not calculate relative path for template: %s. %s", templateAbsolutePath, err)
-			return nil, fmt.Errorf("could not open archive to extract file: %s", err)
 		}
-		hasher := md5.New()
 
-		// Save file and also read into hasher for md5
-		if _, err := io.Copy(templateFile, io.TeeReader(zipTemplateFileReader, hasher)); err != nil {
-			templateFile.Close()
-			return nil, fmt.Errorf("could not write template file: %s", err)
-		}
-		templateFile.Close()
-
-		oldChecksum, checksumOK := templateChecksumsMap[templatePath]
-
-		checksum := hex.EncodeToString(hasher.Sum(nil))
 		if isAddition {
-			results.additions = append(results.additions, filepath.Join(finalPath, name))
+			results.additions = append(results.additions, relativeTemplatePath)
-		} else if checksumOK && oldChecksum[0] != checksum {
+		} else if checksumOk && oldTemplateChecksum[0] != newTemplateChecksum {
-			results.modifications = append(results.modifications, filepath.Join(finalPath, name))
+			results.modifications = append(results.modifications, relativeTemplatePath)
 		}
-		results.checksums[templatePath] = checksum
+		results.checksums[templateAbsolutePath] = newTemplateChecksum
+		results.totalCount++
 	}
 
 	// If we don't find the previous file in the newly downloaded list,
@@ -378,12 +360,63 @@ func (r *Runner) compareAndWriteTemplates(zipReader *zip.Reader) (*templateUpdat
 		_, ok := results.checksums[templatePath]
 		if !ok && templateChecksums[0] == templateChecksums[1] {
 			_ = os.Remove(templatePath)
-			results.deletions = append(results.deletions, strings.TrimPrefix(strings.TrimPrefix(templatePath, r.templatesConfig.TemplatesDirectory), string(os.PathSeparator)))
+			results.deletions = append(results.deletions, strings.TrimPrefix(strings.TrimPrefix(templatePath, configuredTemplateDirectory), string(os.PathSeparator)))
 		}
 	}
 	return results, nil
 }
 
+func writeUnZippedTemplateFile(err error, templateAbsolutePath string, zipTemplateFile *zip.File) (string, error) {
+	templateFile, err := os.OpenFile(templateAbsolutePath, os.O_TRUNC|os.O_CREATE|os.O_WRONLY, 0644)
+	if err != nil {
+		return "", fmt.Errorf("could not create template file: %s", err)
+	}
+
+	zipTemplateFileReader, err := zipTemplateFile.Open()
+	if err != nil {
+		_ = templateFile.Close()
+		return "", fmt.Errorf("could not open archive to extract file: %s", err)
+	}
+
+	md5Hash := md5.New()
+
+	// Save file and also read into hash.Hash for md5
+	if _, err := io.Copy(templateFile, io.TeeReader(zipTemplateFileReader, md5Hash)); err != nil {
+		_ = templateFile.Close()
+		return "", fmt.Errorf("could not write template file: %s", err)
+	}
+
+	if err := templateFile.Close(); err != nil {
+		return "", fmt.Errorf("could not close file newly created template file: %s", err)
+	}
+
+	checksum := hex.EncodeToString(md5Hash.Sum(nil))
+	return checksum, nil
+}
+
+func calculateTemplateAbsolutePath(zipFilePath, configuredTemplateDirectory string) (string, bool, error) {
+	directory, fileName := filepath.Split(zipFilePath)
+
+	if strings.TrimSpace(fileName) == "" || strings.HasPrefix(fileName, ".") || strings.EqualFold(fileName, "README.md") {
+		return "", true, nil
+	}
+
+	directoryPathChunks := strings.Split(directory, string(os.PathSeparator))
+	relativeDirectoryPathWithoutZipRoot := filepath.Join(directoryPathChunks[1:]...)
+
+	if strings.HasPrefix(relativeDirectoryPathWithoutZipRoot, ".") {
+		return "", true, nil
+	}
+
+	templateDirectory := filepath.Join(configuredTemplateDirectory, relativeDirectoryPathWithoutZipRoot)
+
+	if err := os.MkdirAll(templateDirectory, 0755); err != nil {
+		return "", false, fmt.Errorf("failed to create template folder: %s. %s", templateDirectory, err)
+	}
+
+	return filepath.Join(templateDirectory, fileName), false, nil
+}
+
 // createTemplateChecksumsMap reads the previous checksum file from the disk.
 // Creates a map of template paths and their previous and currently calculated checksums as values.
 func createTemplateChecksumsMap(checksumsFilePath string) (map[string][2]string, error) {

v2/internal/runner/update_test.go

@@ -119,6 +119,43 @@ func TestDownloadReleaseAndUnzipDeletion(t *testing.T) {
 	require.Equal(t, "base.yaml", results.deletions[0], "could not get correct new deletions")
 }
 
+func TestCalculateTemplateAbsolutePath(t *testing.T) {
+	configuredTemplateDirectory := filepath.Join(os.TempDir(), "templates")
+	defer os.RemoveAll(configuredTemplateDirectory)
+
+	t.Run("positive scenarios", func(t *testing.T) {
+		zipFilePathsExpectedPathsMap := map[string]string{
+			"nuclei-templates/cve/test.yaml":      filepath.Join(configuredTemplateDirectory, "cve/test.yaml"),
+			"nuclei-templates/cve/test/test.yaml": filepath.Join(configuredTemplateDirectory, "cve/test/test.yaml"),
+		}
+
+		for filePathFromZip, expectedTemplateAbsPath := range zipFilePathsExpectedPathsMap {
+			calculatedTemplateAbsPath, skipFile, err := calculateTemplateAbsolutePath(filePathFromZip, configuredTemplateDirectory)
+			require.Nil(t, err)
+			require.Equal(t, expectedTemplateAbsPath, calculatedTemplateAbsPath)
+			require.False(t, skipFile)
+		}
+	})
+
+	t.Run("negative scenarios", func(t *testing.T) {
+		filePathsFromZip := []string{
+			"./../nuclei-templates/../cve/test.yaml",
+			"nuclei-templates/../cve/test.yaml",
+			"nuclei-templates/cve/../test.yaml",
+			"nuclei-templates/././../cve/test.yaml",
+			"nuclei-templates/.././../cve/test.yaml",
+			"nuclei-templates/.././../cve/../test.yaml",
+		}
+
+		for _, filePathFromZip := range filePathsFromZip {
+			calculatedTemplateAbsPath, skipFile, err := calculateTemplateAbsolutePath(filePathFromZip, configuredTemplateDirectory)
+			require.Nil(t, err)
+			require.True(t, skipFile)
+			require.Equal(t, "", calculatedTemplateAbsPath)
+		}
+	})
+}
+
 func zipFromDirectory(zipPath, directory string) error {
 	file, err := os.Create(zipPath)
 	if err != nil {