diff --git a/server.py b/server.py index 174d38af1..7b179838f 100644 --- a/server.py +++ b/server.py @@ -59,7 +59,7 @@ class PromptServer(): def __init__(self, loop): PromptServer.instance = self - mimetypes.init(); + mimetypes.init() mimetypes.types_map['.js'] = 'application/javascript; charset=utf-8' self.prompt_queue = None self.loop = loop @@ -181,18 +181,43 @@ class PromptServer(): post = await request.post() return image_upload(post) + @routes.post("/upload/mask") async def upload_mask(request): post = await request.post() def image_save_function(image, post, filepath): - original_pil = Image.open(post.get("original_image").file).convert('RGBA') - mask_pil = Image.open(image.file).convert('RGBA') + original_ref = json.loads(post.get("original_ref")) + filename, output_dir = folder_paths.annotated_filepath(original_ref['filename']) - # alpha copy - new_alpha = mask_pil.getchannel('A') - original_pil.putalpha(new_alpha) - original_pil.save(filepath, compress_level=4) + # validation for security: prevent accessing arbitrary path + if filename[0] == '/' or '..' in filename: + return web.Response(status=400) + + if output_dir is None: + type = original_ref.get("type", "output") + output_dir = folder_paths.get_directory_by_type(type) + + if output_dir is None: + return web.Response(status=400) + + if original_ref.get("subfolder", "") != "": + full_output_dir = os.path.join(output_dir, original_ref["subfolder"]) + if os.path.commonpath((os.path.abspath(full_output_dir), output_dir)) != output_dir: + return web.Response(status=403) + output_dir = full_output_dir + + file = os.path.join(output_dir, filename) + + if os.path.isfile(file): + with Image.open(file) as original_pil: + original_pil = original_pil.convert('RGBA') + mask_pil = Image.open(image.file).convert('RGBA') + + # alpha copy + new_alpha = mask_pil.getchannel('A') + original_pil.putalpha(new_alpha) + original_pil.save(filepath, compress_level=4) return image_upload(post, image_save_function) @@ -226,7 +251,6 @@ class PromptServer(): if 'preview' in request.rel_url.query: with Image.open(file) as img: preview_info = request.rel_url.query['preview'].split(';') - image_format = preview_info[0] if image_format not in ['webp', 'jpeg']: image_format = 'webp' diff --git a/web/extensions/core/maskeditor.js b/web/extensions/core/maskeditor.js index 764164d5e..503c45f0e 100644 --- a/web/extensions/core/maskeditor.js +++ b/web/extensions/core/maskeditor.js @@ -346,7 +346,6 @@ class MaskEditorDialog extends ComfyDialog { const rgb_url = new URL(ComfyApp.clipspace.imgs[ComfyApp.clipspace['selectedIndex']].src); rgb_url.searchParams.delete('channel'); - rgb_url.searchParams.delete('preview'); rgb_url.searchParams.set('channel', 'rgb'); orig_image.src = rgb_url; this.image = orig_image; @@ -618,10 +617,20 @@ class MaskEditorDialog extends ComfyDialog { const dataURL = this.backupCanvas.toDataURL(); const blob = dataURLToBlob(dataURL); - const original_blob = loadedImageToBlob(this.image); + let original_url = new URL(this.image.src); + + const original_ref = { filename: original_url.searchParams.get('filename') }; + + let original_subfolder = original_url.searchParams.get("subfolder"); + if(original_subfolder) + original_ref.subfolder = original_subfolder; + + let original_type = original_url.searchParams.get("type"); + if(original_type) + original_ref.type = original_type; formData.append('image', blob, filename); - formData.append('original_image', original_blob); + formData.append('original_ref', JSON.stringify(original_ref)); formData.append('type', "input"); formData.append('subfolder', "clipspace"); diff --git a/web/scripts/app.js b/web/scripts/app.js index 385a54579..5ab9fed08 100644 --- a/web/scripts/app.js +++ b/web/scripts/app.js @@ -165,7 +165,7 @@ export class ComfyApp { if(ComfyApp.clipspace.widgets) { ComfyApp.clipspace.widgets.forEach(({ type, name, value }) => { const prop = Object.values(node.widgets).find(obj => obj.type === type && obj.name === name); - if (prop && prop.type != 'image') { + if (prop && prop.type != 'button') { if(typeof prop.value == "string" && value.filename) { prop.value = (value.subfolder?value.subfolder+'/':'') + value.filename + (value.type?` [${value.type}]`:''); } @@ -174,10 +174,6 @@ export class ComfyApp { prop.callback(value); } } - else if (prop && prop.type != 'button') { - prop.value = value; - prop.callback(value); - } }); } }