Compare commits

..

3 Commits

Author SHA1 Message Date
krau
784888f44c fix(tdler): treat offset past EOF as end of file
gotd's downloader is size-unaware: when a file's size is an exact
multiple of the 1MiB part size, it issues one final upload.getFile
at offset == size, which Telegram rejects with 400 OFFSET_INVALID,
failing the whole download. Wrap the client to answer such requests
with an empty chunk, matching the EOF semantics the downloader
expects. Regression test uses a fake client with real server
behavior (OFFSET_INVALID past EOF).
2026-08-25 14:00:47 +08:00
krau
c86867ff8b test(bot): added wrapper invocation regression test 2026-08-25 13:26:44 +08:00
krau
7c4c7ef3c7 fix(bot): stopped swallowing callbacks on permission pass 2026-08-25 13:26:44 +08:00
4 changed files with 174 additions and 3 deletions

View File

@@ -1,6 +1,8 @@
package handlers
import (
"errors"
"github.com/celestix/gotgproto/dispatcher"
"github.com/celestix/gotgproto/ext"
"github.com/duke-git/lancet/v2/slice"
@@ -37,10 +39,12 @@ func checkPermission(ctx *ext.Context, update *ext.Update) error {
}
// withPermission wraps a callback handler with the same whitelist check used
// for message handlers (checkPermission).
// for message handlers (checkPermission). ContinueGroups is the dispatcher's
// success sentinel, not an error: only real failures and EndGroups stop the
// chain before the wrapped handler runs.
func withPermission(handler func(*ext.Context, *ext.Update) error) func(*ext.Context, *ext.Update) error {
return func(ctx *ext.Context, update *ext.Update) error {
if err := checkPermission(ctx, update); err != nil {
if err := checkPermission(ctx, update); err != nil && !errors.Is(err, dispatcher.ContinueGroups) {
return err
}
return handler(ctx, update)

View File

@@ -1,13 +1,20 @@
package handlers
import (
"os"
"path/filepath"
"testing"
"github.com/celestix/gotgproto/ext"
"github.com/celestix/gotgproto/types"
"github.com/gotd/td/tg"
"github.com/krau/SaveAny-Bot/config"
)
// Regression: callback queries usually arrive as updateShort without entity
// maps, so resolving the sender through the entity map yields ID 0 and every
// click was denied by the whitelist check. Callback updates must use the
// native UserID field.
func TestResponsibleUserID(t *testing.T) {
tests := []struct {
name string
@@ -49,3 +56,30 @@ func TestResponsibleUserID(t *testing.T) {
})
}
}
// Regression: withPermission must treat ContinueGroups (the dispatcher's
// success sentinel) as a pass and invoke the wrapped handler. v0.60.1 treated
// it as an error, so every permitted callback was swallowed before the real
// handler ran.
func TestWithPermissionInvokesHandler(t *testing.T) {
path := filepath.Join(t.TempDir(), "config.toml")
if err := os.WriteFile(path, []byte("workers = 2\n\n[[users]]\nid = 42\n"), 0o644); err != nil {
t.Fatal(err)
}
if err := config.Init(t.Context(), path); err != nil {
t.Fatal(err)
}
update := &ext.Update{CallbackQuery: &tg.UpdateBotCallbackQuery{UserID: 42}}
called := false
handler := withPermission(func(ctx *ext.Context, u *ext.Update) error {
called = true
return nil
})
if err := handler(&ext.Context{}, update); err != nil {
t.Fatalf("withPermission returned error: %v", err)
}
if !called {
t.Fatal("withPermission did not invoke the wrapped handler")
}
}

View File

@@ -1,7 +1,10 @@
package tdler
import (
"context"
"github.com/gotd/td/telegram/downloader"
"github.com/gotd/td/tg"
"github.com/krau/SaveAny-Bot/common/utils/dlutil"
"github.com/krau/SaveAny-Bot/config"
"github.com/krau/SaveAny-Bot/pkg/consts/tglimit"
@@ -10,5 +13,23 @@ import (
func NewDownloader(file tfile.TGFile) *downloader.Builder {
return downloader.NewDownloader().WithPartSize(tglimit.MaxPartSize).
Download(file.Dler(), file.Location()).WithThreads(dlutil.BestThreads(file.Size(), config.C().Threads))
Download(eofAwareClient{Client: file.Dler(), size: file.Size()}, file.Location()).
WithThreads(dlutil.BestThreads(file.Size(), config.C().Threads))
}
// eofAwareClient answers upload.getFile requests at or past the end of the
// file with an empty chunk. gotd's downloader is size-unaware: for files
// whose size is an exact multiple of the part size it issues one final
// request at offset == size and expects an empty chunk, but Telegram rejects
// it with 400 OFFSET_INVALID and the whole download fails.
type eofAwareClient struct {
downloader.Client
size int64
}
func (c eofAwareClient) UploadGetFile(ctx context.Context, req *tg.UploadGetFileRequest) (tg.UploadFileClass, error) {
if req.Offset >= c.size {
return &tg.UploadFile{}, nil
}
return c.Client.UploadGetFile(ctx, req)
}

112
common/tdler/dler_test.go Normal file
View File

@@ -0,0 +1,112 @@
package tdler
import (
"bytes"
"context"
"sync"
"testing"
"github.com/gotd/td/tg"
"github.com/gotd/td/tgerr"
"github.com/krau/SaveAny-Bot/pkg/tfile"
)
// serverLikeClient mimics real Telegram upload.getFile behavior: it returns
// up to limit bytes per chunk, and answers any offset at or past the end of
// the file with 400 OFFSET_INVALID.
type serverLikeClient struct {
data []byte
mu sync.Mutex
maxOffset int64
}
func (c *serverLikeClient) UploadGetFile(_ context.Context, req *tg.UploadGetFileRequest) (tg.UploadFileClass, error) {
c.mu.Lock()
if req.Offset > c.maxOffset {
c.maxOffset = req.Offset
}
c.mu.Unlock()
if req.Offset >= int64(len(c.data)) {
return nil, tgerr.New(400, "OFFSET_INVALID")
}
end := min(len(c.data), int(req.Offset)+req.Limit)
return &tg.UploadFile{Bytes: c.data[req.Offset:end]}, nil
}
func (c *serverLikeClient) UploadGetFileHashes(context.Context, *tg.UploadGetFileHashesRequest) ([]tg.FileHash, error) {
return nil, nil
}
func (c *serverLikeClient) UploadReuploadCDNFile(context.Context, *tg.UploadReuploadCDNFileRequest) ([]tg.FileHash, error) {
return nil, nil
}
func (c *serverLikeClient) UploadGetCDNFileHashes(context.Context, *tg.UploadGetCDNFileHashesRequest) ([]tg.FileHash, error) {
return nil, nil
}
func (c *serverLikeClient) UploadGetWebFile(context.Context, *tg.UploadGetWebFileRequest) (*tg.UploadWebFile, error) {
return nil, nil
}
type memWriterAt struct {
b []byte
}
func (w *memWriterAt) WriteAt(p []byte, off int64) (int, error) {
copy(w.b[off:], p)
return len(p), nil
}
func TestDownloadServerLikeEOF(t *testing.T) {
const partSize = 1024 * 1024
tests := []struct {
name string
size int
parallel bool
}{
{"stream exact multiple of part size", 2 * partSize, false},
{"stream non-multiple", 2*partSize + 12345, false},
{"stream smaller than part size", 1234, false},
{"parallel exact multiple of part size", 2 * partSize, true},
{"parallel non-multiple", 2*partSize + 12345, true},
{"parallel smaller than part size", 1234, true},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
data := make([]byte, tt.size)
for i := range data {
data[i] = byte(i % 251)
}
client := &serverLikeClient{data: data}
file := tfile.NewTGFile(
&tg.InputDocumentFileLocation{ID: 1, AccessHash: 2},
client, int64(tt.size), "test.bin",
)
dl := NewDownloader(file)
var got []byte
var err error
if tt.parallel {
buf := make([]byte, tt.size)
_, err = dl.WithThreads(4).Parallel(context.Background(), &memWriterAt{b: buf})
got = buf
} else {
var buf bytes.Buffer
_, err = dl.Stream(context.Background(), &buf)
got = buf.Bytes()
}
if err != nil {
t.Fatalf("download failed: %v", err)
}
if !bytes.Equal(got, data) {
t.Fatalf("downloaded %d bytes, want %d matching bytes", len(got), len(data))
}
if client.maxOffset >= int64(tt.size) {
t.Fatalf("requested offset %d at or past EOF (size %d)", client.maxOffset, tt.size)
}
})
}
}