Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 16 additions & 14 deletions src/main/zimage.ts
Original file line number Diff line number Diff line change
Expand Up @@ -68,7 +68,7 @@ export async function getZImageModels(_event: IpcMainInvokeEvent, customModelDir
async function runSingleZImage(_event: IpcMainEvent, args: string[], executablePath: string): Promise<number> {
return new Promise((resolve) => {
zImageChild = spawn(executablePath, args, {
shell: true,
shell: false, // to allow prompt string with special characters
})

zImageChild.stdout.on('data', (data) => {
Expand Down Expand Up @@ -114,14 +114,6 @@ export async function runZImageCommand(_event: IpcMainEvent, options: ZImageOpti

// Construct arguments
const args: string[] = []
args.push('-p', options.prompt)

if (options.negativePrompt) {
args.push('-n', options.negativePrompt)
}

// output
args.push('-o', outputPath)

if (options.width && options.height) {
args.push('-s', `${options.width},${options.height}`)
Expand All @@ -135,15 +127,25 @@ export async function runZImageCommand(_event: IpcMainEvent, options: ZImageOpti
args.push('-r', `${options.seed}`)
}

if (options.model) {
const modelPath = join(options.modelDir, options.model)
args.push('-m', normalize(modelPath))
}

if (options.gpuId && options.gpuId !== 'auto') {
args.push('-g', `${options.gpuId}`)
}

// model path
const modelPath = join(options.modelDir, options.model)
args.push('-m', normalize(modelPath))
Comment on lines +135 to +136

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

security-high high

A high-severity path traversal vulnerability exists here: the modelPath is constructed using options.modelDir and options.model from the renderer process without sufficient validation. This allows for path traversal (e.g., using ../../ sequences) leading to arbitrary file access, even with normalize. It's crucial to validate that the resolved path remains within its intended base directory. Additionally, while model is required in ZImageOptions, it could still be an empty string. Backend validation is needed to ensure options.model is a non-empty string to prevent an invalid -m argument.

Suggested change
const modelPath = join(options.modelDir, options.model)
args.push('-m', normalize(modelPath))
const modelPath = join(options.modelDir, options.model)
const normalizedModelPath = normalize(modelPath)
if (!normalizedModelPath.startsWith(normalize(options.modelDir))) {
throw new Error('Invalid model path')
}
args.push('-m', normalizedModelPath)


// output
args.push('-o', outputPath)

// negative prompt
if (options.negativePrompt) {
args.push('-n', options.negativePrompt)
}

// prompt
args.push('-p', options.prompt)

const current_start_args = `[Batch ${i + 1}/${count}] Executing ZImage: ${getCorePath()} ${args.join(' ')} \n\n`
console.log(current_start_args)
_event.sender.send(IpcChannelOn.COMMAND_STDOUT, current_start_args)
Expand Down
2 changes: 1 addition & 1 deletion src/shared/type/zimage.ts
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ export interface ZImageOptions {
height?: number
steps?: number | 'auto'
seed?: number | 'rand'
model?: string
model: string
gpuId?: number | 'auto'
count?: number
modelDir: string
Expand Down