From f38ef7212f7b5c04e67862aff326bb49ec5e5022 Mon Sep 17 00:00:00 2001 From: iksnerd Date: Sat, 13 Jun 2026 13:50:17 +0300 Subject: [PATCH] Fix shutdown data loss, bound untrusted metadata_size, harden Docker build MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three fixes from the 2026-06-13 code review: - cmd/tracker: graceful shutdown only called FlushToDB(), dropping the FlushUsers() + DrainBacklog() steps the periodic flusher runs every tick, so final usage deltas and queued backlog were lost on SIGTERM (the normal Cloud Run stop signal). Now mirrors the tick, and flushes *after* srv.Shutdown() so announces served during the drain window are captured. - internal/client: FetchMetadata allocated make([]byte, MetadataSize) from a peer-supplied (untrusted) value with only a == 0 check — a memory-exhaustion DoS. Cap at 16 MiB and reject <= 0. Adds a regression test. - Dockerfile: ARG TARGETARCH had no default, so a non-BuildKit `docker build` expanded an empty arch into the litestream URL and 404'd. Default to amd64. (Includes the in-progress multi-arch curl + run.sh wait-loop changes the review validated.) --- Dockerfile | 12 ++++++++---- cmd/tracker/main.go | 6 +++++- internal/client/metadata.go | 11 ++++++++++- internal/client/metadata_test.go | 23 +++++++++++++++++++++++ scripts/run.sh | 7 ++++--- 5 files changed, 50 insertions(+), 9 deletions(-) diff --git a/Dockerfile b/Dockerfile index 177624d..163f675 100644 --- a/Dockerfile +++ b/Dockerfile @@ -14,10 +14,14 @@ RUN CGO_ENABLED=0 go build \ # Use a lightweight alpine image for the final container FROM alpine:latest -# Install Litestream -RUN apk add --no-cache ca-certificates -ADD https://github.com/benbjohnson/litestream/releases/download/v0.3.13/litestream-v0.3.13-linux-amd64.tar.gz /tmp/litestream.tar.gz -RUN tar -C /usr/local/bin -xzf /tmp/litestream.tar.gz +# Install Litestream (multi-arch: amd64 or arm64). +# TARGETARCH is auto-set by BuildKit; default to amd64 so a plain `docker build` still works. +ARG TARGETARCH=amd64 +RUN apk add --no-cache ca-certificates curl && \ + curl -fsSL "https://github.com/benbjohnson/litestream/releases/download/v0.3.13/litestream-v0.3.13-linux-${TARGETARCH}.tar.gz" \ + -o /tmp/litestream.tar.gz && \ + tar -C /usr/local/bin -xzf /tmp/litestream.tar.gz && \ + rm /tmp/litestream.tar.gz # Create data directory RUN mkdir -p /data diff --git a/cmd/tracker/main.go b/cmd/tracker/main.go index 533df96..ee3dd4d 100644 --- a/cmd/tracker/main.go +++ b/cmd/tracker/main.go @@ -79,12 +79,16 @@ func main() { <-sigChan log.Println("Shutting down... stopping tickers and flushing final state to DB") close(done) - tracker.State.FlushToDB() ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) defer cancel() + // Stop accepting new announces first so the final flush captures all state, + // then mirror the periodic flusher (state + users + backlog). if err := srv.Shutdown(ctx); err != nil { log.Printf("Shutdown error: %v", err) } + tracker.State.FlushToDB() + tracker.State.FlushUsers() + tracker.State.DrainBacklog() }() fmt.Printf("Weightless Tracker %s (%s) live on :%s\n", version, commit, port) diff --git a/internal/client/metadata.go b/internal/client/metadata.go index 466b890..5524967 100644 --- a/internal/client/metadata.go +++ b/internal/client/metadata.go @@ -10,11 +10,20 @@ import ( wbencode "weightless/internal/bencode" ) +// maxMetadataSize bounds the info dictionary a peer may advertise via BEP 9. +// metadata_size is attacker-controlled (it comes straight from the peer's +// extended handshake), so cap it before allocating to avoid a memory-exhaustion +// DoS. 16 MiB is far larger than any real torrent info dict. +const maxMetadataSize = 16 << 20 + // FetchMetadata fetches the info dictionary from a peer using BEP 9 metadata exchange. func (p *PeerConn) FetchMetadata(ctx context.Context, infoHash []byte) ([]byte, error) { - if p.MetadataSize == 0 { + if p.MetadataSize <= 0 { return nil, fmt.Errorf("peer did not provide metadata_size") } + if p.MetadataSize > maxMetadataSize { + return nil, fmt.Errorf("metadata_size %d exceeds max %d", p.MetadataSize, maxMetadataSize) + } numPieces := (p.MetadataSize + 16383) / 16384 metadata := make([]byte, p.MetadataSize) diff --git a/internal/client/metadata_test.go b/internal/client/metadata_test.go index 95c225d..279ea42 100644 --- a/internal/client/metadata_test.go +++ b/internal/client/metadata_test.go @@ -2,12 +2,35 @@ package client import ( "bytes" + "context" "crypto/sha1" "testing" "github.com/zeebo/bencode" ) +func TestFetchMetadataRejectsBadSize(t *testing.T) { + t.Parallel() + // The size guards run before any network I/O, so a bare PeerConn is enough. + tests := []struct { + name string + size int + }{ + {"zero", 0}, + {"negative", -1}, + {"too large", maxMetadataSize + 1}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + p := &PeerConn{MetadataSize: tt.size} + if _, err := p.FetchMetadata(context.Background(), make([]byte, 20)); err == nil { + t.Fatalf("FetchMetadata(size=%d) = nil error, want rejection", tt.size) + } + }) + } +} + func TestFindBencodeEnd(t *testing.T) { t.Parallel() tests := []struct { diff --git a/scripts/run.sh b/scripts/run.sh index 29086a8..f78ee41 100755 --- a/scripts/run.sh +++ b/scripts/run.sh @@ -1,5 +1,4 @@ #!/bin/sh -set -e # Restore the database from Litestream if it doesn't exist if [ ! -f /data/weightless.db ]; then @@ -33,5 +32,7 @@ _term() { # Trap signals trap _term SIGTERM SIGINT -# Wait for weightless to exit (or for trap to catch signal) -wait "$TRACKER_PID" +# Wait for weightless to exit, re-entering wait if interrupted by other child exits +while kill -0 "$TRACKER_PID" 2>/dev/null; do + wait "$TRACKER_PID" 2>/dev/null || true +done