fix(stream): answer 502 when Telegram fails before the first byte

internal/api/raw.go

@@ -12,11 +12,13 @@ import (
 	"syscall"
 
 	"github.com/google/uuid"
+	"github.com/gotd/td/tgerr"
 
 	"github.com/tgdrive/teldrive/v2/internal/api/gen"
 	"github.com/tgdrive/teldrive/v2/internal/db/sqlcgen"
 	"github.com/tgdrive/teldrive/v2/internal/dbtypes"
 	"github.com/tgdrive/teldrive/v2/internal/shares"
+	"github.com/tgdrive/teldrive/v2/internal/telegramstore"
 	"github.com/tgdrive/teldrive/v2/internal/transfer"
 )
 
@@ -143,10 +145,30 @@ func (h *RawHandler) streamFile(ctx context.Context, w http.ResponseWriter, user
 		UserID: userID, FileID: fileID, Offset: rangeSpec.Offset, Length: rangeSpec.Length,
 	})
 	if err != nil {
+		if isTelegramPartGone(err) {
+			// Opening asks Telegram for part sizes older uploads didn't record.
+			slog.WarnContext(ctx, "api.stream_unreadable", "file_id", fileID, "offset", rangeSpec.Offset, "length", rangeSpec.Length, "error", err)
+			return mapStreamError(err)
+		}
 		return mapServiceError(err)
 	}
 	defer download.Reader.Close()
 
+	// Read the first bytes before committing the status line: when Telegram can't serve the
+	// part, the client gets a 502 instead of a 206 that ends early with nothing in it.
+	first := make([]byte, min(rangeSpec.Length, streamFirstRead))
+	firstN := 0
+	if len(first) > 0 {
+		firstN, err = io.ReadAtLeast(download.Reader, first, 1)
+		if firstN == 0 {
+			if isExpectedStreamEnd(ctx, err) {
+				return ctx.Err()
+			}
+			slog.WarnContext(ctx, "api.stream_unreadable", "file_id", fileID, "offset", rangeSpec.Offset, "length", rangeSpec.Length, "error", err)
+			return mapStreamError(err)
+		}
+	}
+
 	status := http.StatusOK
 	if partial {
 		status = http.StatusPartialContent
@@ -160,7 +182,10 @@ func (h *RawHandler) streamFile(ctx context.Context, w http.ResponseWriter, user
 	w.Header().Set("Last-Modified", file.ModTime.Time.UTC().Format(http.TimeFormat))
 	w.Header().Set("X-Accel-Buffering", "no")
 	w.WriteHeader(status)
-	_, err = io.CopyN(w, download.Reader, rangeSpec.Length)
+	_, err = w.Write(first[:firstN])
+	if err == nil {
+		_, err = io.CopyN(w, download.Reader, rangeSpec.Length-int64(firstN))
+	}
 	if err != nil && !isExpectedStreamEnd(ctx, err) {
 		// Headers are already committed, so a JSON error response cannot replace
 		// this stream. The Content-Length mismatch tells the client it was truncated.
@@ -169,6 +194,35 @@ func (h *RawHandler) streamFile(ctx context.Context, w http.ResponseWriter, user
 	return nil
 }
 
+// streamFirstRead is how much of a stream is read before the response status is sent.
+const streamFirstRead = 64 << 10
+
+// mapStreamError is the response for a stream that failed before its first byte: errors with
+// their own status keep it, anything else is Telegram (the storage behind the request) failing.
+// The details, which name Telegram messages, stay in the log.
+func mapStreamError(err error) error {
+	if isTelegramPartGone(err) {
+		return problem(http.StatusBadGateway, "telegram_part_unavailable", "Telegram could not serve a part of this file", err)
+	}
+	mapped := mapServiceError(err)
+	var p *Problem
+	if errors.As(mapped, &p) && p.Status == http.StatusInternalServerError {
+		return problem(http.StatusBadGateway, "telegram_read_failed", "reading this file from Telegram failed", err)
+	}
+	return mapped
+}
+
+// isTelegramPartGone: Telegram answered, but can't serve the part: its message or document is
+// gone, or -503 "Timeout" (the file DC couldn't read it), also when that was still being retried
+// as the deadline ran out.
+func isTelegramPartGone(err error) bool {
+	if errors.Is(err, telegramstore.ErrMessageNotFound) || errors.Is(err, telegramstore.ErrDocumentNotFound) {
+		return true
+	}
+	var rpc *tgerr.Error
+	return errors.As(err, &rpc) && rpc.Code == -503
+}
+
 func isExpectedStreamEnd(ctx context.Context, err error) bool {
 	return ctx.Err() != nil || isClientDisconnect(err)
 }

internal/api/raw_stream_test.go

@@ -0,0 +1,153 @@
+package api
+
+import (
+	"bytes"
+	"context"
+	"errors"
+	"fmt"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/google/uuid"
+	"github.com/gotd/td/tgerr"
+	"github.com/jackc/pgx/v5/pgtype"
+
+	"github.com/tgdrive/teldrive/v2/internal/api/gen"
+	"github.com/tgdrive/teldrive/v2/internal/db/sqlcgen"
+	"github.com/tgdrive/teldrive/v2/internal/telegramstore"
+	"github.com/tgdrive/teldrive/v2/internal/transfer"
+)
+
+func TestStreamFileAnswers502WhenTelegramFailsBeforeTheFirstByte(t *testing.T) {
+	for _, tc := range []struct {
+		name string
+		err  error
+		code string
+	}{
+		{"-503 behind the retry deadline", fmt.Errorf("download Telegram document chunk at 0: %w", errors.Join(context.DeadlineExceeded, tgerr.New(-503, "Timeout"))), "telegram_part_unavailable"},
+		{"message gone", telegramstore.ErrMessageNotFound, "telegram_part_unavailable"},
+		{"connection trouble", errors.New("engine was closed"), "telegram_read_failed"},
+	} {
+		file, parts := streamTestFile(64 * 1024)
+		handler := &RawHandler{handler: &Handler{Downloader: transfer.NewDownloader(streamTestCatalog{file, parts}, &streamTestStorage{err: tc.err}, nil)}}
+		fileID, _ := uuid.FromBytes(file.ID.Bytes[:])
+		recorder := httptest.NewRecorder()
+		err := handler.streamFile(context.Background(), recorder, file.UserID, fileID, file, gen.NewOptString("bytes=0-"), gen.OptETag{}, false)
+		var p *Problem
+		if !errors.As(err, &p) || p.Status != http.StatusBadGateway || p.Code != tc.code {
+			t.Errorf("%s: streamFile() = %v, want 502 %s", tc.name, err, tc.code)
+			continue
+		}
+		if len(recorder.Header().Get("Content-Range")) != 0 || recorder.Body.Len() != 0 {
+			t.Errorf("%s: the 206 was committed before the error", tc.name)
+		}
+	}
+}
+
+func TestStreamFileAnswers502WhenAPartWithoutSizeIsGone(t *testing.T) {
+	file, parts := streamTestFile(64 * 1024)
+	parts[0].PlainSize, parts[0].StoredSize = pgtype.Int8{}, pgtype.Int8{} // an older upload: sizes resolved on open
+	handler := &RawHandler{handler: &Handler{Downloader: transfer.NewDownloader(streamTestCatalog{file, parts}, &streamTestStorage{metadataErr: telegramstore.ErrMessageNotFound}, nil)}}
+	fileID, _ := uuid.FromBytes(file.ID.Bytes[:])
+	err := handler.streamFile(context.Background(), httptest.NewRecorder(), file.UserID, fileID, file, gen.NewOptString("bytes=0-"), gen.OptETag{}, false)
+	var p *Problem
+	if !errors.As(err, &p) || p.Status != http.StatusBadGateway || p.Code != "telegram_part_unavailable" {
+		t.Fatalf("streamFile() = %v, want 502 telegram_part_unavailable", err)
+	}
+}
+
+func TestStreamFileStreamsWholeRangeAfterTheFirstRead(t *testing.T) {
+	file, parts := streamTestFile(300 * 1024)
+	payload := bytes.Repeat([]byte("abcdef"), 50*1024)
+	handler := &RawHandler{handler: &Handler{Downloader: transfer.NewDownloader(streamTestCatalog{file, parts}, &streamTestStorage{data: payload}, nil)}}
+	fileID, _ := uuid.FromBytes(file.ID.Bytes[:])
+	recorder := httptest.NewRecorder()
+	if err := handler.streamFile(context.Background(), recorder, file.UserID, fileID, file, gen.NewOptString("bytes=10-"), gen.OptETag{}, false); err != nil {
+		t.Fatalf("streamFile() = %v", err)
+	}
+	if recorder.Code != http.StatusPartialContent || !bytes.Equal(recorder.Body.Bytes(), payload[10:]) {
+		t.Fatalf("code %d, %d bytes, want 206 with %d bytes", recorder.Code, recorder.Body.Len(), len(payload)-10)
+	}
+}
+
+func TestMapStreamErrorKeepsOwnStatuses(t *testing.T) {
+	var p *Problem
+	if err := mapStreamError(context.DeadlineExceeded); !errors.As(err, &p) || p.Status != http.StatusGatewayTimeout {
+		t.Fatalf("deadline: %v, want 504", err)
+	}
+	if err := mapStreamError(io.ErrUnexpectedEOF); !errors.As(err, &p) || p.Status != http.StatusBadGateway {
+		t.Fatalf("short read: %v, want 502", err)
+	}
+}
+
+func streamTestFile(size int64) (*sqlcgen.File, []*sqlcgen.FilePart) {
+	file := &sqlcgen.File{
+		ID: pgtype.UUID{Bytes: uuid.New(), Valid: true}, UserID: 7, Name: "film.mkv",
+		Kind: sqlcgen.FileKindFile, Size: pgtype.Int8{Int64: size, Valid: true}, Status: sqlcgen.FileStatusActive,
+	}
+	return file, []*sqlcgen.FilePart{{
+		PartNo: 1, ChannelID: 11, MessageID: 101,
+		PlainSize: pgtype.Int8{Int64: size, Valid: true}, StoredSize: pgtype.Int8{Int64: size, Valid: true},
+	}}
+}
+
+type streamTestCatalog struct {
+	file  *sqlcgen.File
+	parts []*sqlcgen.FilePart
+}
+
+func (c streamTestCatalog) Get(context.Context, int64, uuid.UUID) (*sqlcgen.File, error) {
+	return c.file, nil
+}
+
+func (c streamTestCatalog) Parts(context.Context, int64, uuid.UUID) ([]*sqlcgen.FilePart, error) {
+	return c.parts, nil
+}
+
+type streamTestStorage struct {
+	data          []byte
+	err           error
+	metadataErr   error
+	metadataCalls int
+}
+
+func (s *streamTestStorage) OpenRange(_ context.Context, request telegramstore.RangeRequest) (io.ReadCloser, error) {
+	if s.err != nil {
+		return io.NopCloser(failingReader{s.err}), nil
+	}
+	end := int64(len(s.data))
+	if request.Length >= 0 && request.Offset+request.Length < end {
+		end = request.Offset + request.Length
+	}
+	return io.NopCloser(bytes.NewReader(s.data[request.Offset:end])), nil
+}
+
+func (s *streamTestStorage) Upload(context.Context, telegramstore.UploadRequest) (telegramstore.StoredPart, error) {
+	return telegramstore.StoredPart{}, nil
+}
+
+func (s *streamTestStorage) DeleteMessages(context.Context, int64, int64, []int64) error { return nil }
+
+func (s *streamTestStorage) CopyPart(context.Context, int64, int64, int64, int64) (telegramstore.StoredPart, error) {
+	return telegramstore.StoredPart{}, nil
+}
+
+func (s *streamTestStorage) CreateChannel(context.Context, int64, string) (telegramstore.Channel, error) {
+	return telegramstore.Channel{}, nil
+}
+
+func (s *streamTestStorage) DeleteChannel(context.Context, int64, int64) error { return nil }
+
+type failingReader struct{ err error }
+
+func (r failingReader) Read([]byte) (int, error) { return 0, r.err }
+
+func (s *streamTestStorage) Metadata(_ context.Context, request telegramstore.MetadataRequest) (telegramstore.StoredPart, error) {
+	s.metadataCalls++
+	if s.metadataErr != nil {
+		return telegramstore.StoredPart{}, s.metadataErr
+	}
+	return telegramstore.StoredPart{ChannelID: request.ChannelID, MessageID: request.MessageID, Size: int64(len(s.data))}, nil
+}

internal/telegramstore/retry.go

@@ -2,6 +2,7 @@ package telegramstore
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"log/slog"
 	"strings"
@@ -35,20 +36,22 @@ type retryMiddleware struct{ max int }
 
 func (m retryMiddleware) Handle(next tg.Invoker) telegram.InvokeFunc {
 	return func(ctx context.Context, input bin.Encoder, output bin.Decoder) error {
+		var last error // the transient Telegram error being retried
 		for attempt := 0; ; attempt++ {
 			if err := ctx.Err(); err != nil {
-				return err
+				return deadlineWithCause(err, last)
 			}
 			err := next.Invoke(ctx, input, output)
 			if err == nil {
 				return nil
 			}
 			if ctx.Err() != nil {
-				return ctx.Err()
+				return deadlineWithCause(ctx.Err(), last)
 			}
 			if !isTransientTelegramError(err) {
 				return err
 			}
+			last = err
 			clientID, _ := ClientID(ctx)
 			if attempt >= m.max {
 				slog.WarnContext(ctx, "telegram.rpc_retry_exhausted",
@@ -66,13 +69,33 @@ func (m retryMiddleware) Handle(next tg.Invoker) telegram.InvokeFunc {
 			select {
 			case <-ctx.Done():
 				timer.Stop()
-				return ctx.Err()
+				return deadlineWithCause(ctx.Err(), err)
 			case <-timer.C:
 			}
 		}
 	}
 }
 
+// deadlineWithCause keeps the Telegram error that was being retried when the caller's deadline
+// ran out, so "Timeout" (-503) and friends aren't reported as a bare context deadline. A cancel
+// stays a plain cancel: nobody is waiting for the reason.
+func deadlineWithCause(ctxErr, last error) error {
+	if errors.Is(ctxErr, context.DeadlineExceeded) && last != nil {
+		return &retryDeadlineError{last: last, deadline: ctxErr}
+	}
+	return ctxErr
+}
+
+// retryDeadlineError is a deadline that ran out while a Telegram error was being retried: it
+// reads as that error and matches both (errors.Is context.DeadlineExceeded, tgerr.As).
+type retryDeadlineError struct{ last, deadline error }
+
+func (e *retryDeadlineError) Error() string {
+	return e.last.Error() + " (retried until the deadline)"
+}
+
+func (e *retryDeadlineError) Unwrap() []error { return []error{e.deadline, e.last} }
+
 func telegramRetryDelay(attempt int) time.Duration {
 	delay := 500 * time.Millisecond
 	const maxDelay = 10 * time.Second

internal/telegramstore/retry_test.go

@@ -160,3 +160,42 @@ func TestRetryMiddlewareRetriesConnectionResetByPeer(t *testing.T) {
 		t.Fatalf("calls = %d, want 2", got)
 	}
 }
+
+func TestRetryMiddlewareKeepsTelegramErrorWhenDeadlineEnds(t *testing.T) {
+	synctest.Test(t, func(t *testing.T) {
+		refusal := tgerr.New(-503, "Timeout")
+		invoke := retryMiddleware{max: 10}.Handle(telegram.InvokeFunc(func(context.Context, bin.Encoder, bin.Decoder) error {
+			return refusal
+		}))
+		ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
+		defer cancel()
+		err := invoke(ctx, &tg.UploadGetFileRequest{}, &tg.UploadFileBox{})
+		if !errors.Is(err, context.DeadlineExceeded) {
+			t.Fatalf("error = %v, want the deadline", err)
+		}
+		if rpc, ok := tgerr.As(err); !ok || rpc.Code != -503 {
+			t.Fatalf("error = %v, want Telegram's -503 kept with it", err)
+		}
+
+		// A cancel stays a plain cancel.
+		ctx, cancel = context.WithCancel(context.Background())
+		time.AfterFunc(time.Second, cancel)
+		err = invoke(ctx, &tg.UploadGetFileRequest{}, &tg.UploadFileBox{})
+		if !errors.Is(err, context.Canceled) {
+			t.Fatalf("error = %v, want canceled", err)
+		}
+		if _, ok := tgerr.As(err); ok {
+			t.Fatalf("a cancel carried the Telegram error: %v", err)
+		}
+	})
+}
+
+func TestRetryDeadlineErrorReadsAsTheTelegramError(t *testing.T) {
+	err := deadlineWithCause(context.DeadlineExceeded, tgerr.New(-503, "Timeout"))
+	if got, want := err.Error(), "rpc error code -503: Timeout (retried until the deadline)"; got != want {
+		t.Fatalf("Error() = %q, want %q", got, want)
+	}
+	if got := deadlineWithCause(context.Canceled, tgerr.New(-503, "Timeout")); got != context.Canceled {
+		t.Fatalf("cancel = %v, want the plain cancel", got)
+	}
+}

internal/transfer/download.go

@@ -170,7 +170,7 @@ func (d *Downloader) resolveMissingPartSizes(ctx context.Context, session telegr
 			UserID: userID, ChannelID: part.ChannelID, MessageID: part.MessageID,
 		})
 		if err != nil {
-			return fmt.Errorf("resolve Telegram part %d metadata: %w", part.PartNo, err)
+			return fmt.Errorf("resolve Telegram part %d (message %d) metadata: %w", part.PartNo, part.MessageID, err)
 		}
 		plainSize := stored.Size
 		if file.Encryption {
@@ -428,7 +428,7 @@ func (r *downloadReader) openPartReader(segment downloadSegment, partOffset, spa
 		})
 	}
 	if err != nil {
-		return nil, fmt.Errorf("open Telegram part %d: %w", segment.part.PartNo, err)
+		return nil, fmt.Errorf("open Telegram part %d (message %d): %w", segment.part.PartNo, segment.part.MessageID, err)
 	}
 	return reader, nil
 }