fix(tdler): validate resume state against the part file

This commit is contained in:
krau
2026-08-25 15:22:59 +08:00
parent 991f454096
commit 8bdafd115c
5 changed files with 154 additions and 22 deletions
+20 -1
View File
@@ -8,6 +8,7 @@ import (
"io"
"net"
"os"
"strings"
"sync"
"github.com/gotd/td/tg"
@@ -84,7 +85,14 @@ func loadResumeBitmap(path string) (*resumeBitmap, error) {
}
var b resumeBitmap
if err := json.Unmarshal(data, &b); err != nil {
return nil, fmt.Errorf("parse resume bitmap: %w", err)
// 无法解析的位图 (外部损坏): 删除并视为不存在, 全量重下自愈。
_ = os.Remove(path)
return nil, nil
}
if b.Size <= 0 || b.PartSize <= 0 {
// 无效位图 (损坏或旧格式), 视为不存在, 全量重下。
_ = os.Remove(path)
return nil, nil
}
b.ensureBlocks()
return &b, nil
@@ -182,6 +190,17 @@ func DownloadResumable(
if err != nil {
return err
}
// 位图描述的数据文件 (bitmapPath 去掉 .bitmap 后缀) 必须存在且非空:
// 若缺失或为空, 已标记完成的块字节已丢失, 必须重置位图全量重下。
if bm != nil {
partPath := strings.TrimSuffix(bitmapPath, ".bitmap")
if stat, err := os.Stat(partPath); err != nil || stat.Size() == 0 {
if err := os.Remove(bitmapPath); err != nil && !errors.Is(err, os.ErrNotExist) {
return fmt.Errorf("reset stale resume bitmap: %w", err)
}
bm = nil
}
}
if bm == nil || bm.PartSize != tglimit.MaxPartSize || bm.Size != file.Size() {
bm = newResumeBitmap(file.Size())
if err := bm.save(bitmapPath); err != nil {
+111
View File
@@ -129,6 +129,117 @@ func TestDownloadResumableBitmapResetOnSizeChange(t *testing.T) {
}
}
// TestDownloadResumablePartMissingOrTruncated resets the bitmap: skipped
// blocks would otherwise be zero-filled (caller recreates the part file
// without its bytes), or the download would wedge forever on a stale
// complete bitmap.
func TestDownloadResumablePartMissingOrTruncated(t *testing.T) {
data := make([]byte, 5*1024*1024)
for i := range data {
data[i] = byte(i % 251)
}
dir := t.TempDir()
partPath := filepath.Join(dir, "test.bin.part")
bitmapPath := ResumeStatePath(partPath)
tests := []struct {
name string
doneBlocks []int
createPart bool
truncate bool
}{
{"part missing, partial bitmap", []int{0, 1, 2}, false, false},
{"part empty, partial bitmap", []int{0, 1, 2}, true, true},
{"part missing, complete bitmap", []int{0, 1, 2, 3, 4}, false, false},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
os.Remove(partPath)
os.Remove(bitmapPath)
bm := newResumeBitmap(int64(len(data)))
for _, block := range tt.doneBlocks {
bm.markDone(block)
}
if err := bm.save(bitmapPath); err != nil {
t.Fatal(err)
}
if tt.createPart {
// Simulate the caller re-creating the part file (truncating).
if err := os.WriteFile(partPath, nil, 0o644); err != nil {
t.Fatal(err)
}
if tt.truncate {
if err := os.WriteFile(partPath, make([]byte, 0), 0o644); err != nil {
t.Fatal(err)
}
}
}
partFile, err := os.OpenFile(partPath, os.O_CREATE|os.O_RDWR, 0o644)
if err != nil {
t.Fatal(err)
}
defer partFile.Close()
client := &serverLikeClient{data: data}
file := tfile.NewTGFile(&tg.InputDocumentFileLocation{ID: 1, AccessHash: 2}, client, int64(len(data)), "test.bin")
if err := DownloadResumable(context.Background(), file, partFile, 1, bitmapPath); err != nil {
t.Fatalf("download failed: %v", err)
}
got := make([]byte, len(data))
if _, err := partFile.ReadAt(got, 0); err != nil {
t.Fatal(err)
}
if !bytesEqual(got, data) {
t.Fatalf("downloaded data mismatch (blocks not reset)")
}
})
}
}
// TestDownloadResumableInvalidBitmap treats a corrupt bitmap as absent.
func TestDownloadResumableInvalidBitmap(t *testing.T) {
data := make([]byte, 1024*1024+7)
for i := range data {
data[i] = byte(i % 251)
}
dir := t.TempDir()
partPath := filepath.Join(dir, "test.bin.part")
bitmapPath := ResumeStatePath(partPath)
for _, content := range []string{
`{"part_size":1048576,"size":-1,"blocks":[]}`,
`{"part_size":1048576,"size":9223372036854775807,"blocks":[]}`,
`not json`,
} {
os.Remove(partPath)
if err := os.WriteFile(bitmapPath, []byte(content), 0o644); err != nil {
t.Fatal(err)
}
partFile, err := os.OpenFile(partPath, os.O_CREATE|os.O_RDWR, 0o644)
if err != nil {
t.Fatal(err)
}
client := &serverLikeClient{data: data}
file := tfile.NewTGFile(&tg.InputDocumentFileLocation{ID: 1, AccessHash: 2}, client, int64(len(data)), "test.bin")
err = DownloadResumable(context.Background(), file, partFile, 1, bitmapPath)
partFile.Close()
if err != nil {
t.Fatalf("download with corrupt bitmap %q failed: %v", content, err)
}
got := make([]byte, len(data))
f, err := os.Open(partPath)
if err != nil {
t.Fatal(err)
}
if _, err := f.ReadAt(got, 0); err != nil {
t.Fatal(err)
}
f.Close()
if !bytesEqual(got, data) {
t.Fatalf("downloaded data mismatch with corrupt bitmap %q", content)
}
}
}
func TestRemoveResumeState(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "x.bitmap")