437 lines
14 KiB
Plaintext
437 lines
14 KiB
Plaintext
import t, c
|
||
from stdint import *
|
||
import string
|
||
import stdio
|
||
import stdlib
|
||
import requests
|
||
import requests.https as rhttps
|
||
import vthreading as vt
|
||
import w32.win32file as wf
|
||
import w32.win32base as wb
|
||
import w32.win32sync as wsync
|
||
|
||
|
||
# ============================================================
|
||
# requests/download.vp — 大文件切片下载（Range + 串行/多线程）
|
||
#
|
||
# 用法:
|
||
# import requests.download as rd
|
||
# rc: INT = rd.download("https://host/big.bin", "big.bin", threads=4,
|
||
# chunk_size=262144, verify=True, ca_file="ca.crt")
|
||
# if rc != 0:
|
||
# print(rd.last_error())
|
||
#
|
||
# 策略:
|
||
# threads == 1 串行 Range 切片（顺序拉取并写入文件）
|
||
# threads > 1 多线程：文件按线程数均分为若干带(band)，每线程一个
|
||
# 连接顺序下载自己带的切片，写入同一文件的对应偏移。
|
||
#
|
||
# 说明:
|
||
# - 通过 Content-Range 探测总大小与服务器 Range 支持；
|
||
# - http:// 走 requests（纯 socket），https:// 走 requests.https（openssl）；
|
||
# - 每片请求不带 Accept-Encoding，响应为原始字节，直接落盘。
|
||
# ============================================================
|
||
|
||
DL_MAX_THREADS: t.CDefine = 8
|
||
DL_DEFAULT_CHUNK: t.CDefine = 262144
|
||
DL_REDIRECT_MAX: t.CDefine = 5
|
||
|
||
_dl_err_store: t.CArray[t.CChar, 512]
|
||
|
||
|
||
def last_error() -> str:
|
||
"""返回最近一次下载失败原因文本（无则空串）。"""
|
||
return _dl_err_store
|
||
|
||
|
||
def _dl_set_err(text: str):
|
||
if text is None:
|
||
return
|
||
_dl_err_store[0] = 0
|
||
stdio.snprintf(_dl_err_store, 512, "%s", text)
|
||
|
||
|
||
def _dl_strdup(s: str) -> str:
|
||
"""stdlib 堆上复制字符串（调用方 free），None 失败。"""
|
||
if s is None:
|
||
return None
|
||
n: t.CSizeT = string.strlen(s)
|
||
out: str = stdlib.malloc(n + 1)
|
||
if out is None:
|
||
return None
|
||
if n > 0:
|
||
string.memcpy(out, s, n)
|
||
out[n] = 0
|
||
return out
|
||
|
||
|
||
# ============================================================
|
||
# 单切片拉取（含重定向跟随，≤5 跳；http/https 按 scheme 分发）
|
||
# ============================================================
|
||
|
||
def _dl_slice(url: str, pos: t.CSizeT, want: INT, verify: bool, ca_file: str,
|
||
dst: str, dst_cap: INT, out_total: t.CSizeT | t.CPtr) -> INT:
|
||
"""拉取 [pos, pos+want) 一片，返回体字节数写入 dst（dst 内正文起始需
|
||
调用方用 _dl_body_off 定位）；-2 重定向超限；-3 网络/服务器失败。"""
|
||
if url is None or dst is None or out_total is None:
|
||
return -1
|
||
cur: str = _dl_strdup(url)
|
||
if cur is None:
|
||
return -1
|
||
hops: INT = 0
|
||
rc: INT = 0
|
||
while True:
|
||
is_https: INT = 1 if (string.strncmp(cur, "https://", 8) == 0) else 0
|
||
st: INT = 0
|
||
tot: t.CSizeT = 0
|
||
loc: t.CArray[t.CChar, 512]
|
||
if is_https != 0:
|
||
rc = rhttps._DlOnce(cur, pos, want, verify, ca_file, dst, dst_cap,
|
||
c.Addr(st), loc, 512, c.Addr(tot))
|
||
else:
|
||
rc = requests._DlOnce(cur, pos, want, dst, dst_cap,
|
||
c.Addr(st), loc, 512, c.Addr(tot))
|
||
if rc == -2 and loc[0] != 0:
|
||
hops += 1
|
||
if hops > DL_REDIRECT_MAX:
|
||
_dl_set_err("too many redirects")
|
||
stdlib.free(cur)
|
||
return -2
|
||
new_url: str = stdlib.malloc(2048)
|
||
if new_url is None:
|
||
stdlib.free(cur)
|
||
return -1
|
||
if requests._ResolveRedirectUrl(cur, loc, new_url, 2048) != 0:
|
||
stdlib.free(new_url)
|
||
stdlib.free(cur)
|
||
_dl_set_err("bad redirect location")
|
||
return -3
|
||
stdlib.free(cur)
|
||
cur = new_url
|
||
continue
|
||
break
|
||
stdlib.free(cur)
|
||
if rc < 0:
|
||
return rc
|
||
c.Set(c.Deref(out_total), tot)
|
||
return rc
|
||
|
||
|
||
def _dl_body_off(dst: str, dst_cap: INT) -> INT:
|
||
"""定位响应体（头结束处）在 dst 中的偏移，找不到返回 -1。"""
|
||
return requests._FindBodyStart(dst, dst_cap)
|
||
|
||
|
||
# ============================================================
|
||
# 文件写入（按偏移落盘，支持多线程并发写不同区段）
|
||
# ============================================================
|
||
|
||
def _dl_write_at(path: str, src: str, length: INT, offset: t.CSizeT) -> INT:
|
||
"""打开文件（不存在则创建），Seek 到 offset 写 length 字节。0 成功；-1 失败。"""
|
||
if path is None or src is None or length <= 0:
|
||
return -1
|
||
h: wb.HANDLE = wf.CreateFileA(
|
||
path, wf.GENERIC_WRITE,
|
||
wf.FILE_SHARE_READ | wf.FILE_SHARE_WRITE | wf.FILE_SHARE_DELETE,
|
||
None, wf.OPEN_ALWAYS, wf.FILE_ATTRIBUTE_NORMAL, None)
|
||
if h == wb.INVALID_HANDLE_VALUE:
|
||
return -1
|
||
li: wb.LARGE_INTEGER = wb.LARGE_INTEGER()
|
||
li.QuadPart = t.CLongLong(offset)
|
||
ok_seek: BOOL = wf.SetFilePointerEx(h, c.Addr(li), None, wf.FILE_BEGIN)
|
||
if ok_seek == 0:
|
||
wb.CloseHandle(h)
|
||
return -1
|
||
written: ULONG = 0
|
||
ok_w: BOOL = wf.WriteFile(h, src, t.CUnsignedLong(length), c.Addr(written), None)
|
||
wb.CloseHandle(h)
|
||
if ok_w == 0:
|
||
return -1
|
||
if t.CInt(written) != length:
|
||
return -1
|
||
return 0
|
||
|
||
|
||
def _dl_truncate(path: str) -> INT:
|
||
"""CREATE_ALWAYS 建（截断）空文件。0 成功；-1 失败。"""
|
||
h: wb.HANDLE = wf.CreateFileA(
|
||
path, wf.GENERIC_WRITE,
|
||
wf.FILE_SHARE_READ | wf.FILE_SHARE_WRITE | wf.FILE_SHARE_DELETE,
|
||
None, wf.CREATE_ALWAYS, wf.FILE_ATTRIBUTE_NORMAL, None)
|
||
if h == wb.INVALID_HANDLE_VALUE:
|
||
return -1
|
||
wb.CloseHandle(h)
|
||
return 0
|
||
|
||
|
||
# ============================================================
|
||
# 探测：总大小 / Range 支持（首片 bytes=0-0）
|
||
# ============================================================
|
||
|
||
def _dl_probe(url: str, verify: bool, ca_file: str,
|
||
out_total: t.CSizeT | t.CPtr) -> INT:
|
||
"""探测服务器总大小与 Range 支持。
|
||
返回 0=支持 Range 且 out_total 有效；1=服务器返回 200（不支持 Range，
|
||
退化为整段下载，out_total=Content-Length）；-3 失败。"""
|
||
c.Set(c.Deref(out_total), 0)
|
||
buf: str = stdlib.malloc(8192)
|
||
if buf is None:
|
||
_dl_set_err("out of memory")
|
||
return -3
|
||
tot: t.CSizeT = 0
|
||
rc: INT = _dl_slice(url, 0, 0, verify, ca_file, buf, 8192, c.Addr(tot))
|
||
stdlib.free(buf)
|
||
if rc < 0:
|
||
_dl_set_err("probe failed")
|
||
return -3
|
||
if tot > 0:
|
||
c.Set(c.Deref(out_total), tot)
|
||
return 0
|
||
_dl_set_err("cannot determine size")
|
||
return -3
|
||
|
||
|
||
# ============================================================
|
||
# 下载一个区间带（[b_start, b_end)），逐片拉取并落盘
|
||
# ============================================================
|
||
|
||
def _dl_download_band(path: str, url: str, is_https: INT, verify: bool, ca_file: str,
|
||
b_start: t.CSizeT, b_end: t.CSizeT, chunk: INT,
|
||
buf: str, buf_cap: INT,
|
||
p_done: t.CSizeT | t.CPtr) -> INT:
|
||
"""0 成功；-1 网络/服务器；-2 写入失败。p_done 累计已写字节。"""
|
||
pos: t.CSizeT = b_start
|
||
while pos < b_end:
|
||
want: INT = chunk
|
||
if b_end - pos < t.CSizeT(want):
|
||
want = INT(b_end - pos)
|
||
tot: t.CSizeT = 0
|
||
rc: INT = _dl_slice(url, pos, want, verify, ca_file, buf, buf_cap, c.Addr(tot))
|
||
if rc <= 0:
|
||
_dl_set_err("slice fetch failed")
|
||
return -1
|
||
bo: INT = _dl_body_off(buf, buf_cap)
|
||
if bo < 0:
|
||
_dl_set_err("bad slice response")
|
||
return -1
|
||
got: INT = rc
|
||
if bo + got > buf_cap:
|
||
got = buf_cap - bo
|
||
if got <= 0:
|
||
_dl_set_err("slice empty")
|
||
return -1
|
||
if _dl_write_at(path, buf + bo, got, pos) != 0:
|
||
_dl_set_err("write file failed")
|
||
return -2
|
||
c.Set(c.Deref(p_done), c.Deref(p_done) + t.CSizeT(got))
|
||
pos += t.CSizeT(got)
|
||
if got < want:
|
||
# 服务端提前结束（正常收尾时仅末片 < want）
|
||
if pos >= b_end:
|
||
break
|
||
_dl_set_err("short slice")
|
||
return -1
|
||
return 0
|
||
|
||
|
||
# ============================================================
|
||
# 多线程 worker
|
||
# ============================================================
|
||
|
||
@t.Object
|
||
class DlJob:
|
||
path: str
|
||
url: str
|
||
is_https: INT
|
||
verify: INT
|
||
ca_file: str
|
||
start: t.CSizeT
|
||
end: t.CSizeT
|
||
chunk: INT
|
||
buf: str
|
||
buf_cap: INT
|
||
done: t.CSizeT
|
||
status: INT
|
||
|
||
|
||
def _dl_worker(arg: VOIDPTR) -> t.CPtr:
|
||
if arg is None:
|
||
return None
|
||
job: DlJob | t.CPtr = arg
|
||
rc: INT = _dl_download_band(
|
||
job.path, job.url, job.is_https,
|
||
True if job.verify != 0 else False,
|
||
job.ca_file,
|
||
job.start, job.end, job.chunk, job.buf, job.buf_cap,
|
||
c.Addr(job.done))
|
||
if rc != 0:
|
||
job.status = rc
|
||
return None
|
||
|
||
|
||
# ============================================================
|
||
# 公共 API
|
||
# ============================================================
|
||
|
||
def download(url: str, path: str, threads: INT = 1,
|
||
chunk_size: INT = DL_DEFAULT_CHUNK,
|
||
verify: bool = True, ca_file: str = None) -> INT:
|
||
"""大文件切片下载。
|
||
|
||
Args:
|
||
url: 支持 http:// 与 https://（Range 分段拉取）
|
||
path: 保存路径（不存在自动创建；已存在则截断重建）
|
||
threads: 1=串行切片；2~8=多线程并行分带下载
|
||
chunk_size: 单片大小（字节，默认 256KB）
|
||
verify / ca_file: https 证书校验参数（对齐 requests.https）
|
||
Returns:
|
||
0 成功；-1 参数错误；-2 文件创建失败；-3 服务器不支持/下载失败；
|
||
-4 写入失败。last_error() 可查原因文本。
|
||
"""
|
||
_dl_err_store[0] = 0
|
||
if url is None or path is None or url[0] == 0 or path[0] == 0:
|
||
_dl_set_err("url/path empty")
|
||
return -1
|
||
if chunk_size < 1024:
|
||
chunk_size = DL_DEFAULT_CHUNK
|
||
if threads < 1:
|
||
threads = 1
|
||
if threads > DL_MAX_THREADS:
|
||
threads = DL_MAX_THREADS
|
||
if string.strncmp(url, "http://", 7) != 0 \
|
||
and string.strncmp(url, "https://", 8) != 0:
|
||
_dl_set_err("unsupported scheme")
|
||
return -1
|
||
|
||
# 探测总大小与 Range 支持
|
||
total: t.CSizeT = 0
|
||
prc: INT = _dl_probe(url, verify, ca_file, c.Addr(total))
|
||
if prc < 0:
|
||
return -3
|
||
if total == 0:
|
||
_dl_set_err("cannot determine size")
|
||
return -3
|
||
|
||
# 建立（截断）目标文件
|
||
if _dl_truncate(path) != 0:
|
||
_dl_set_err("create file failed")
|
||
return -2
|
||
|
||
stdio.printf("[download] %s -> %s total=%zu bytes threads=%d chunk=%d\n",
|
||
url, path, total, threads, chunk_size)
|
||
|
||
slab: INT = chunk_size + 8192
|
||
|
||
if threads == 1:
|
||
# ---- 串行切片 ----
|
||
buf: str = stdlib.malloc(t.CSizeT(slab))
|
||
if buf is None:
|
||
_dl_set_err("out of memory")
|
||
return -3
|
||
done: t.CSizeT = 0
|
||
rc: INT = _dl_download_band(path, url, 0, verify, ca_file,
|
||
0, total, chunk_size, buf, slab, c.Addr(done))
|
||
stdlib.free(buf)
|
||
if rc != 0:
|
||
return -3 if rc == -1 else -4
|
||
stdio.printf("[download] done: %zu/%zu bytes\n", done, total)
|
||
return 0
|
||
|
||
# ---- 多线程并行分带 ----
|
||
band: t.CSizeT = (total + t.CSizeT(threads) - 1) / t.CSizeT(threads)
|
||
ctxs: t.CArray[VOIDPTR, DL_MAX_THREADS]
|
||
handles: t.CArray[VOIDPTR, DL_MAX_THREADS]
|
||
threads_ok: INT = 0
|
||
|
||
i: INT = 0
|
||
while i < threads:
|
||
b_start: t.CSizeT = t.CSizeT(i) * band
|
||
if b_start >= total:
|
||
break
|
||
b_end: t.CSizeT = b_start + band
|
||
if b_end > total:
|
||
b_end = total
|
||
job: DlJob | t.CPtr = stdlib.calloc(1, DlJob.__sizeof__())
|
||
if job is None:
|
||
break
|
||
job.path = _dl_strdup(path)
|
||
job.url = _dl_strdup(url)
|
||
if job.path is None or job.url is None:
|
||
stdlib.free(job.path)
|
||
stdlib.free(job.url)
|
||
stdlib.free(t.CPtr(t.CUInt64T(job)))
|
||
break
|
||
job.is_https = 1 if (string.strncmp(url, "https://", 8) == 0) else 0
|
||
job.verify = 1 if verify else 0
|
||
job.ca_file = _dl_strdup(ca_file) if ca_file is not None else None
|
||
job.start = b_start
|
||
job.end = b_end
|
||
job.chunk = chunk_size
|
||
job.buf_cap = slab
|
||
job.buf = stdlib.malloc(t.CSizeT(slab))
|
||
if job.buf is None:
|
||
stdlib.free(job.path)
|
||
stdlib.free(job.url)
|
||
if job.ca_file is not None:
|
||
stdlib.free(job.ca_file)
|
||
stdlib.free(t.CPtr(t.CUInt64T(job)))
|
||
break
|
||
job.done = 0
|
||
job.status = 0
|
||
ctxs[i] = t.CPtr(t.CUInt64T(job))
|
||
|
||
th: rawptr[vt.Thread] = vt.Thread(_dl_worker, ctxs[i])
|
||
if th is None:
|
||
stdlib.free(job.buf)
|
||
stdlib.free(job.path)
|
||
stdlib.free(job.url)
|
||
if job.ca_file is not None:
|
||
stdlib.free(job.ca_file)
|
||
stdlib.free(t.CPtr(t.CUInt64T(job)))
|
||
break
|
||
th.start()
|
||
handles[i] = th.handle
|
||
threads_ok += 1
|
||
i += 1
|
||
|
||
if threads_ok == 0:
|
||
_dl_set_err("thread create failed")
|
||
return -3
|
||
|
||
# 等待全部线程结束
|
||
j: INT = 0
|
||
while j < threads_ok:
|
||
wsync.WaitForSingleObject(handles[j], 0xFFFFFFFF)
|
||
wb.CloseHandle(handles[j])
|
||
j += 1
|
||
|
||
# 汇总结果
|
||
total_done: t.CSizeT = 0
|
||
failed_job: INT = 0
|
||
k: INT = 0
|
||
while k < threads_ok:
|
||
jobk: DlJob | t.CPtr = ctxs[k]
|
||
total_done += jobk.done
|
||
if jobk.status != 0:
|
||
failed_job = 1
|
||
k += 1
|
||
|
||
# 释放线程上下文
|
||
m: INT = 0
|
||
while m < threads_ok:
|
||
jobm: DlJob | t.CPtr = ctxs[m]
|
||
stdlib.free(jobm.path)
|
||
stdlib.free(jobm.url)
|
||
if jobm.ca_file is not None:
|
||
stdlib.free(jobm.ca_file)
|
||
stdlib.free(jobm.buf)
|
||
stdlib.free(t.CPtr(t.CUInt64T(jobm)))
|
||
m += 1
|
||
|
||
if failed_job != 0 or total_done != total:
|
||
stdio.printf("[download] incomplete: %zu/%zu bytes\n", total_done, total)
|
||
_dl_set_err("download incomplete")
|
||
return -3
|
||
stdio.printf("[download] done: %zu/%zu bytes\n", total_done, total)
|
||
return 0
|