Fix shutdown data loss, bound untrusted metadata_size, harden Docker build

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

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)

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)

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 {

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