fix(civitai): validate downloads by content-length and content-type

An absolute 1024-byte floor rejected valid complete downloads, since
some safetensors are legitimately tiny (a lone projector-diff tensor is
~160 bytes). Verify completeness against the declared Content-Length
instead, accepting any size, and keep the floor only as a fallback when
it is missing. Reject text/* responses up front so an HTML error or
login page served with HTTP 200 is never saved as a model file.
This commit is contained in:
CalamitousFelicitousness
2026-07-03 20:16:42 +01:00
parent 98a7d17207
commit 3c0ca8f483
+23 -9
View File
@@ -178,6 +178,16 @@ class DownloadManager:
item.completed_at = datetime.now()
return
# A text/* response is an error or login page served with HTTP 200,
# not a model. Reject before writing.
content_type = r.headers.get('content-type', '').lower()
if content_type.startswith('text/'):
item.status = "failed"
item.error = f'invalid content-type: {content_type}'
item.completed_at = datetime.now()
log.warning(f'CivitAI download invalid content-type: id={item.id} content-type="{content_type}"')
return
total_size = int(r.headers.get('content-length', 0))
item.bytes_total = starting_pos + total_size
@@ -216,8 +226,18 @@ class DownloadManager:
item.progress = written / item.bytes_total
pbar.update(task, completed=written)
# Validate minimum size
if written < 1024:
# Complete = received bytes match Content-Length, at any size. Some
# safetensors are legitimately tiny (~160 bytes), so no size floor;
# the floor is only a fallback for when Content-Length is missing.
expected = starting_pos + total_size
if total_size > 0:
if written != expected:
item.status = "failed"
item.error = f'incomplete: expected={expected} got={written}'
item.completed_at = datetime.now()
log.warning(f'CivitAI download incomplete: id={item.id} expected={expected} got={written}')
return
elif written < 1024:
try:
os.remove(temp_file)
except OSError:
@@ -225,13 +245,7 @@ class DownloadManager:
item.status = "failed"
item.error = f'download too small: {written} bytes'
item.completed_at = datetime.now()
return
# Check for incomplete download
if starting_pos + total_size != written:
item.status = "failed"
item.error = f'incomplete: expected={starting_pos + total_size} got={written}'
item.completed_at = datetime.now()
log.warning(f'CivitAI download too small: id={item.id} file="{item.filename}" bytes={written}')
return
except Exception as e: