diff --git a/openapi/Swarm.yaml b/openapi/Swarm.yaml index 4168e63f80e..6bf0e89fc3c 100644 --- a/openapi/Swarm.yaml +++ b/openapi/Swarm.yaml @@ -319,7 +319,18 @@ paths: - Chunk parameters: - $ref: "SwarmCommon.yaml#/components/parameters/SwarmTagParameter" - - $ref: "SwarmCommon.yaml#/components/parameters/SwarmPostageBatchId" + - in: query + name: swarm-tag + schema: + $ref: "SwarmCommon.yaml#/components/schemas/Uid" + required: false + description: "Associate upload with an existing Tag UID (use when WebSocket client cannot set custom headers)" + - in: header + name: swarm-postage-batch-id + description: "ID of Postage Batch that is used to upload data with. Optional when chunks include pre-signed postage stamps." + required: false + schema: + $ref: "SwarmCommon.yaml#/components/schemas/SwarmAddress" responses: "200": description: "Connection established" diff --git a/pkg/api/api.go b/pkg/api/api.go index 2b713bc76d4..93595168cc3 100644 --- a/pkg/api/api.go +++ b/pkg/api/api.go @@ -845,7 +845,18 @@ func (s *Service) newStampedPutter(ctx context.Context, opts putterOptions, stam return nil, errInvalidPostageBatch } + return s.newStampedPutterWithBatch(ctx, opts, stamp, storedBatch) +} + +// newStampedPutterWithBatch creates a stamped putter using a pre-fetched batch. +// This avoids the database lookup when batch info is already cached. +func (s *Service) newStampedPutterWithBatch(ctx context.Context, opts putterOptions, stamp *postage.Stamp, storedBatch *postage.Batch) (storer.PutterSession, error) { + if !opts.Deferred && s.beeMode == DevMode { + return nil, errUnsupportedDevNodeOperation + } + var session storer.PutterSession + var err error if opts.Deferred || opts.Pin { session, err = s.storer.Upload(ctx, opts.Pin, opts.TagID) if err != nil { diff --git a/pkg/api/chunk_stream.go b/pkg/api/chunk_stream.go index 2f91939f21a..be648b964be 100644 --- a/pkg/api/chunk_stream.go +++ b/pkg/api/chunk_stream.go @@ -8,6 +8,7 @@ import ( "context" "errors" "net/http" + "strconv" "time" "github.com/ethersphere/bee/v2/pkg/cac" @@ -28,7 +29,7 @@ func (s *Service) chunkUploadStreamHandler(w http.ResponseWriter, r *http.Reques logger := s.logger.WithName("chunks_stream").Build() headers := struct { - BatchID []byte `map:"Swarm-Postage-Batch-Id" validate:"required"` + BatchID []byte `map:"Swarm-Postage-Batch-Id"` // Optional: omit if caller provides pre-signed stamps per chunk SwarmTag uint64 `map:"Swarm-Tag"` }{} if response := s.mapStructure(r.Header, &headers); response != nil { @@ -36,6 +37,19 @@ func (s *Service) chunkUploadStreamHandler(w http.ResponseWriter, r *http.Reques return } + // Fallback: read tag from query parameter (browser WebSocket can't set headers) + if headers.SwarmTag == 0 { + if qTag := r.URL.Query().Get("swarm-tag"); qTag != "" { + parsed, err := strconv.ParseUint(qTag, 10, 64) + if err != nil { + logger.Debug("invalid swarm-tag query parameter", "value", qTag, "error", err) + jsonhttp.BadRequest(w, "invalid swarm-tag query parameter") + return + } + headers.SwarmTag = parsed + } + } + var ( tag uint64 err error @@ -55,29 +69,36 @@ func (s *Service) chunkUploadStreamHandler(w http.ResponseWriter, r *http.Reques } } - // if tag not specified use direct upload - // Using context.Background here because the putter's lifetime extends beyond that of the HTTP request. - putter, err := s.newStamperPutter(context.Background(), putterOptions{ - BatchID: headers.BatchID, - TagID: tag, - Deferred: tag != 0, - }) - if err != nil { - logger.Debug("get putter failed", "error", err) - logger.Error(nil, "get putter failed") - switch { - case errors.Is(err, errBatchUnusable) || errors.Is(err, postage.ErrNotUsable): - jsonhttp.UnprocessableEntity(w, "batch not usable yet or does not exist") - case errors.Is(err, postage.ErrNotFound): - jsonhttp.NotFound(w, "batch with id not found") - case errors.Is(err, errInvalidPostageBatch): - jsonhttp.BadRequest(w, "invalid batch id") - case errors.Is(err, errUnsupportedDevNodeOperation): - jsonhttp.BadRequest(w, errUnsupportedDevNodeOperation) - default: - jsonhttp.BadRequest(w, nil) + // Create connection-level putter only if BatchID is provided. + // If BatchID is not provided, the API caller is expected to provide + // pre-signed stamps with each chunk (and is also expected to keep + // track of stamp state over time). + var putter storer.PutterSession + if len(headers.BatchID) > 0 { + // if tag not specified use direct upload + // Using context.Background here because the putter's lifetime extends beyond that of the HTTP request. + putter, err = s.newStamperPutter(context.Background(), putterOptions{ + BatchID: headers.BatchID, + TagID: tag, + Deferred: tag != 0, + }) + if err != nil { + logger.Debug("get putter failed", "error", err) + logger.Error(nil, "get putter failed") + switch { + case errors.Is(err, errBatchUnusable) || errors.Is(err, postage.ErrNotUsable): + jsonhttp.UnprocessableEntity(w, "batch not usable yet or does not exist") + case errors.Is(err, postage.ErrNotFound): + jsonhttp.NotFound(w, "batch with id not found") + case errors.Is(err, errInvalidPostageBatch): + jsonhttp.BadRequest(w, "invalid batch id") + case errors.Is(err, errUnsupportedDevNodeOperation): + jsonhttp.BadRequest(w, errUnsupportedDevNodeOperation) + default: + jsonhttp.BadRequest(w, nil) + } + return } - return } upgrader := websocket.Upgrader{ @@ -95,13 +116,46 @@ func (s *Service) chunkUploadStreamHandler(w http.ResponseWriter, r *http.Reques } s.wsWg.Add(1) - go s.handleUploadStream(logger, wsConn, putter) + var decode chunkDecoder + if len(headers.BatchID) > 0 { + decode = decodeChunkWithoutStamp + } else { + decode = decodeChunkWithStamp + } + go s.handleUploadStream(logger, wsConn, putter, tag, decode) +} + +// chunkDecoder extracts chunk data and optionally a stamp from a websocket message. +// When BatchID is provided in headers, decodeChunkWithoutStamp is used (no stamp in message). +// When BatchID is not provided, decodeChunkWithStamp is used (stamp prepended to chunk data). +type chunkDecoder func(msg []byte) (chunkData []byte, stamp *postage.Stamp, err error) + +// decodeChunkWithoutStamp returns the message as-is (used when BatchID provided in headers). +func decodeChunkWithoutStamp(msg []byte) ([]byte, *postage.Stamp, error) { + return msg, nil, nil +} + +// decodeChunkWithStamp extracts a stamp from the first 113 bytes of the message. +// Returns an error if the message is too small or the stamp is invalid. +func decodeChunkWithStamp(msg []byte) ([]byte, *postage.Stamp, error) { + if len(msg) < postage.StampSize+swarm.SpanSize { + return nil, nil, errors.New("message too small for stamp + chunk") + } + + stamp := &postage.Stamp{} + if err := stamp.UnmarshalBinary(msg[:postage.StampSize]); err != nil { + return nil, nil, errors.New("invalid stamp") + } + + return msg[postage.StampSize:], stamp, nil } func (s *Service) handleUploadStream( logger log.Logger, conn *websocket.Conn, putter storer.PutterSession, + tag uint64, + decode chunkDecoder, ) { defer s.wsWg.Done() @@ -111,11 +165,23 @@ func (s *Service) handleUploadStream( gone = make(chan struct{}) err error ) + + // Cache for batch validation to avoid database lookups for every chunk + // Key: batch ID, Value: stored batch info + // This avoids the expensive batchStore.Get() call for each chunk + batchCache := make(map[string]*postage.Batch) + defer func() { cancel() _ = conn.Close() - if err = putter.Done(swarm.ZeroAddress); err != nil { - logger.Error(err, "chunk upload stream: syncing chunks failed") + + // No cleanup needed for batch cache - it's just metadata + + // Only call Done on connection-level putter if it exists + if putter != nil { + if err = putter.Done(swarm.ZeroAddress); err != nil { + logger.Error(err, "chunk upload stream: syncing chunks failed") + } } }() @@ -190,17 +256,80 @@ func (s *Service) handleUploadStream( return } - chunk, err := cac.NewWithDataSpan(msg) + // Decode the message using the appropriate decoder + chunkData, stamp, err := decode(msg) + if err != nil { + logger.Debug("chunk upload stream: decode failed", "error", err) + logger.Error(nil, "chunk upload stream: "+err.Error()) + sendErrorClose(websocket.CloseInternalServerErr, err.Error()) + return + } + + // Determine the putter to use + var ( + chunk swarm.Chunk + chunkPutter = putter + ) + + // If stamp was extracted, create a per-chunk putter + if stamp != nil { + batchID := stamp.BatchID() + batchIDKey := string(batchID) + + storedBatch, exists := batchCache[batchIDKey] + if !exists { + storedBatch, err = s.batchStore.Get(batchID) + if err != nil { + logger.Debug("chunk upload stream: batch validation failed", "error", err) + logger.Error(nil, "chunk upload stream: batch validation failed") + if errors.Is(err, storage.ErrNotFound) { + sendErrorClose(websocket.CloseInternalServerErr, "batch not found") + } else { + sendErrorClose(websocket.CloseInternalServerErr, "batch validation failed") + } + return + } + batchCache[batchIDKey] = storedBatch + } + + chunkPutter, err = s.newStampedPutterWithBatch(ctx, putterOptions{ + BatchID: batchID, + TagID: tag, + Deferred: tag != 0, + }, stamp, storedBatch) + if err != nil { + logger.Debug("chunk upload stream: failed to create stamped putter", "error", err) + logger.Error(nil, "chunk upload stream: failed to create stamped putter") + switch { + case errors.Is(err, errBatchUnusable) || errors.Is(err, postage.ErrNotUsable): + sendErrorClose(websocket.CloseInternalServerErr, "batch not usable") + case errors.Is(err, postage.ErrNotFound): + sendErrorClose(websocket.CloseInternalServerErr, "batch not found") + default: + sendErrorClose(websocket.CloseInternalServerErr, "stamped putter creation failed") + } + return + } + } + + chunk, err = cac.NewWithDataSpan(chunkData) if err != nil { - logger.Debug("chunk upload stream: create chunk failed", "error", err) + logger.Debug("chunk upload stream: create chunk failed", "error", err, "chunk_size", len(chunkData)) logger.Error(nil, "chunk upload stream: create chunk failed") + if chunkPutter != putter { + _ = chunkPutter.Cleanup() + } + sendErrorClose(websocket.CloseInternalServerErr, "invalid chunk data") return } - err = putter.Put(ctx, chunk) + err = chunkPutter.Put(ctx, chunk) if err != nil { logger.Debug("chunk upload stream: write chunk failed", "address", chunk.Address(), "error", err) logger.Error(nil, "chunk upload stream: write chunk failed") + if chunkPutter != putter { + _ = chunkPutter.Cleanup() + } switch { case errors.Is(err, postage.ErrBucketFull): sendErrorClose(websocket.CloseInternalServerErr, "batch is overissued") @@ -210,6 +339,13 @@ func (s *Service) handleUploadStream( return } + // Clean up per-chunk putter + if chunkPutter != putter { + if err := chunkPutter.Done(swarm.ZeroAddress); err != nil { + logger.Error(err, "chunk upload stream: failed to finalize per-chunk putter") + } + } + err = sendMsg(websocket.BinaryMessage, successWsMsg) if err != nil { s.logger.Debug("chunk upload stream: sending success message failed", "error", err) diff --git a/pkg/api/chunk_stream_test.go b/pkg/api/chunk_stream_test.go index 5ba62cdd76c..add7165fa89 100644 --- a/pkg/api/chunk_stream_test.go +++ b/pkg/api/chunk_stream_test.go @@ -11,7 +11,11 @@ import ( "time" "github.com/ethersphere/bee/v2/pkg/api" + "github.com/ethersphere/bee/v2/pkg/crypto" + "github.com/ethersphere/bee/v2/pkg/postage" + mockbatchstore "github.com/ethersphere/bee/v2/pkg/postage/batchstore/mock" mockpost "github.com/ethersphere/bee/v2/pkg/postage/mock" + testingpostage "github.com/ethersphere/bee/v2/pkg/postage/testing" "github.com/ethersphere/bee/v2/pkg/spinlock" testingc "github.com/ethersphere/bee/v2/pkg/storage/testing" mockstorer "github.com/ethersphere/bee/v2/pkg/storer/mock" @@ -37,7 +41,7 @@ func TestChunkUploadStream(t *testing.T) { ) t.Run("upload and verify", func(t *testing.T) { - chsToGet := []swarm.Chunk{} + chsToGet := make([]swarm.Chunk, 0, 5) for range 5 { ch := testingc.GenerateTestRandomChunk() @@ -104,3 +108,142 @@ func TestChunkUploadStream(t *testing.T) { } }) } + +// nolint:paralleltest +func TestChunkUploadStreamWithStamp(t *testing.T) { + // Generate signer and batch for pre-signed stamps + key, err := crypto.GenerateSecp256k1Key() + if err != nil { + t.Fatal(err) + } + signer := crypto.NewDefaultSigner(key) + owner, err := signer.EthereumAddress() + if err != nil { + t.Fatal(err) + } + + // Generate chunks and their pre-signed stamps + chunks := make([]swarm.Chunk, 5) + stampBytes := make([][]byte, 5) + + for i := range 5 { + chunks[i] = testingc.GenerateTestRandomChunk() + stamp := testingpostage.MustNewValidStamp(signer, chunks[i].Address()) + sb, err := stamp.MarshalBinary() + if err != nil { + t.Fatal(err) + } + stampBytes[i] = sb + } + + // Mock batch store: accept all batch IDs and return a batch with the correct owner + batchStore := mockbatchstore.New( + mockbatchstore.WithAcceptAllExistsFunc(), + mockbatchstore.WithBatch(&postage.Batch{ + Owner: owner.Bytes(), + }), + ) + + // No Swarm-Postage-Batch-Id header — triggers per-chunk stamp mode + wsHeaders := http.Header{} + wsHeaders.Set(api.ContentTypeHeader, "application/octet-stream") + + var ( + storerMock = mockstorer.New() + _, wsConn, _, chanStorer = newTestServer(t, testServerOptions{ + Storer: storerMock, + Post: mockpost.New(mockpost.WithAcceptAll()), + BatchStore: batchStore, + WsPath: "/chunks/stream", + WsHeaders: wsHeaders, + DirectUpload: true, + }) + ) + + t.Run("upload with pre-signed stamps", func(t *testing.T) { + for i := range 5 { + // Prepend stamp bytes to chunk data + msg := append(stampBytes[i], chunks[i].Data()...) + + err := wsConn.SetWriteDeadline(time.Now().Add(time.Second)) + if err != nil { + t.Fatal(err) + } + + err = wsConn.WriteMessage(websocket.BinaryMessage, msg) + if err != nil { + t.Fatal(err) + } + + err = wsConn.SetReadDeadline(time.Now().Add(time.Second)) + if err != nil { + t.Fatal(err) + } + + mt, msg, err := wsConn.ReadMessage() + if err != nil { + t.Fatal(err) + } + + if mt != websocket.BinaryMessage || !bytes.Equal(msg, api.SuccessWsMsg) { + t.Fatal("invalid response", mt, string(msg)) + } + } + + for _, c := range chunks { + err := spinlock.Wait(100*time.Millisecond, func() bool { return chanStorer.Has(c.Address()) }) + if err != nil { + t.Fatal(err) + } + } + }) +} + +// nolint:paralleltest +func TestChunkUploadStreamInvalidStamp(t *testing.T) { + // No Swarm-Postage-Batch-Id header — triggers per-chunk stamp mode + wsHeaders := http.Header{} + wsHeaders.Set(api.ContentTypeHeader, "application/octet-stream") + + var ( + storerMock = mockstorer.New() + _, wsConn, _, _ = newTestServer(t, testServerOptions{ + Storer: storerMock, + Post: mockpost.New(mockpost.WithAcceptAll()), + WsPath: "/chunks/stream", + WsHeaders: wsHeaders, + DirectUpload: true, + }) + ) + + t.Run("message too small for stamp", func(t *testing.T) { + // Send a message smaller than StampSize + SpanSize + tooSmall := make([]byte, postage.StampSize) + + err := wsConn.SetWriteDeadline(time.Now().Add(time.Second)) + if err != nil { + t.Fatal(err) + } + + err = wsConn.WriteMessage(websocket.BinaryMessage, tooSmall) + if err != nil { + t.Fatal(err) + } + + err = wsConn.SetReadDeadline(time.Now().Add(time.Second)) + if err != nil { + t.Fatal(err) + } + + _, _, err = wsConn.ReadMessage() + if err == nil { + t.Fatal("expected failure on read") + } + // nolint:errorlint + if cerr, ok := err.(*websocket.CloseError); !ok { + t.Fatal("invalid error on read") + } else if cerr.Text != "message too small for stamp + chunk" { + t.Fatalf("incorrect response on error, exp: (message too small for stamp + chunk) got (%s)", cerr.Text) + } + }) +}