From b9cdf7e09d4b6dd668eb8aed25475e9802e06558 Mon Sep 17 00:00:00 2001 From: bootjp Date: Thu, 23 Jul 2026 15:04:19 +0900 Subject: [PATCH 1/4] snapshot: harden object offload validation --- cmd/elastickv-snapshot-offload/main.go | 17 +- cmd/elastickv-snapshot-offload/main_test.go | 2 +- internal/snapshotoffload/manifest.go | 106 ++++- internal/snapshotoffload/offload_test.go | 59 +++ internal/snapshotoffload/s3_store.go | 408 +++++++++++++++++++- internal/snapshotoffload/s3_store_test.go | 218 +++++++++-- internal/snapshotoffload/store.go | 4 +- 7 files changed, 761 insertions(+), 53 deletions(-) diff --git a/cmd/elastickv-snapshot-offload/main.go b/cmd/elastickv-snapshot-offload/main.go index 7f8c11462..78c4f9fbf 100644 --- a/cmd/elastickv-snapshot-offload/main.go +++ b/cmd/elastickv-snapshot-offload/main.go @@ -26,6 +26,7 @@ const ( commandRestore = "restore" storeLocal = "local" storeS3 = "s3" + s3SSEAES256 = "AES256" s3SSEAWSKMS = "aws:kms" ) @@ -177,6 +178,7 @@ func addStoreFlags(fs *flag.FlagSet, cfg *storeFlags) { cfg.storeKind = storeLocal cfg.s3Region = "us-east-1" cfg.s3PathStyle = true + cfg.s3ServerSideEncryption = string(s3SSEAES256) fs.StringVar(&cfg.storeKind, "store", storeLocal, "Object store backend: local or s3") fs.StringVar(&cfg.localRoot, "local-root", "", "Local object store root when --store=local") fs.StringVar(&cfg.s3Bucket, "s3-bucket", "", "S3 bucket when --store=s3") @@ -184,8 +186,8 @@ func addStoreFlags(fs *flag.FlagSet, cfg *storeFlags) { fs.StringVar(&cfg.s3Endpoint, "s3-endpoint", "", "S3-compatible endpoint URL") fs.StringVar(&cfg.s3Profile, "s3-profile", "", "AWS shared config profile") fs.BoolVar(&cfg.s3PathStyle, "s3-path-style", cfg.s3PathStyle, "Use path-style S3 addressing") - fs.StringVar(&cfg.s3ServerSideEncryption, "s3-sse", "", "Server-side encryption algorithm for uploaded objects, for example AES256 or aws:kms") - fs.StringVar(&cfg.s3KMSKeyID, "s3-kms-key-id", "", "KMS key ID when --s3-sse=aws:kms") + fs.StringVar(&cfg.s3ServerSideEncryption, "s3-sse", cfg.s3ServerSideEncryption, "Server-side encryption algorithm for uploaded objects: AES256 or aws:kms") + fs.StringVar(&cfg.s3KMSKeyID, "s3-kms-key-id", "", "KMS key ARN or bare key ID when --s3-sse=aws:kms; aliases are rejected") fs.BoolVar(&cfg.s3DisableChecksumHeaders, "s3-disable-checksum-headers", false, "Do not send S3 checksum headers; keep metadata and restore-time verification") } @@ -199,17 +201,12 @@ func validateStoreFlags(cfg storeFlags) error { if strings.TrimSpace(cfg.s3Bucket) == "" { return errors.New("--s3-bucket is required when --store=s3") } + if err := snapshotoffload.ValidateS3StoreEncryption(cfg.s3ServerSideEncryption, cfg.s3KMSKeyID); err != nil { + return errors.Wrap(err, "validate s3 encryption") + } default: return errors.Errorf("unknown --store %q", cfg.storeKind) } - if strings.TrimSpace(cfg.s3KMSKeyID) != "" { - if strings.TrimSpace(cfg.s3ServerSideEncryption) == "" { - return errors.New("--s3-kms-key-id requires --s3-sse") - } - if strings.TrimSpace(cfg.s3ServerSideEncryption) != s3SSEAWSKMS { - return errors.New("--s3-kms-key-id requires --s3-sse=aws:kms") - } - } return nil } diff --git a/cmd/elastickv-snapshot-offload/main_test.go b/cmd/elastickv-snapshot-offload/main_test.go index d2991b401..008e2f5f9 100644 --- a/cmd/elastickv-snapshot-offload/main_test.go +++ b/cmd/elastickv-snapshot-offload/main_test.go @@ -111,7 +111,7 @@ func TestSnapshotOffloadCLIS3KMSRequiresAWSKMS(t *testing.T) { "--data-dir", "data", "--group-id", "1", }) - require.ErrorContains(t, err, "--s3-kms-key-id requires --s3-sse=aws:kms") + require.ErrorContains(t, err, "s3 KMS key id requires aws:kms encryption") } func seedCLISnapshot(t *testing.T, root string, payload []byte, index uint64, term uint64) string { diff --git a/internal/snapshotoffload/manifest.go b/internal/snapshotoffload/manifest.go index c6a091f66..a4209bee6 100644 --- a/internal/snapshotoffload/manifest.go +++ b/internal/snapshotoffload/manifest.go @@ -20,6 +20,7 @@ const ( var ( ErrInvalidOptions = errors.New("snapshot offload: invalid options") ErrIntegrity = errors.New("snapshot offload: integrity check failed") + ErrObjectConflict = errors.New("snapshot offload: object conflict") ErrObjectNotFound = errors.New("snapshot offload: object not found") ) @@ -82,6 +83,22 @@ func DecodeManifest(data []byte) (Manifest, error) { } func validateManifest(manifest Manifest) error { + if err := validateManifestIdentity(manifest); err != nil { + return err + } + if err := validateManifestPayload(manifest.Payload); err != nil { + return err + } + if stringsTrim(manifest.ManifestKey) == "" { + return errors.Wrap(ErrInvalidOptions, "manifest key is required") + } + if len(manifest.ConfState.Voters) == 0 { + return errors.Wrap(ErrInvalidOptions, "manifest ConfState requires voters") + } + return validateManifestConfState(manifest.ConfState) +} + +func validateManifestIdentity(manifest Manifest) error { switch { case manifest.SchemaVersion != ManifestSchemaVersion: return errors.Wrapf(ErrInvalidOptions, "unsupported manifest schema version %d", manifest.SchemaVersion) @@ -91,18 +108,97 @@ func validateManifest(manifest Manifest) error { return errors.Wrap(ErrInvalidOptions, "snapshot index must be > 0") case manifest.SnapshotTerm == 0: return errors.Wrap(ErrInvalidOptions, "snapshot term must be > 0") - case manifest.Payload.Key == "": + default: + return nil + } +} + +func validateManifestPayload(payload PayloadDescriptor) error { + switch { + case payload.Key == "": return errors.Wrap(ErrInvalidOptions, "payload key is required") - case manifest.Payload.Bytes < 0: + case payload.Bytes < 0: return errors.Wrap(ErrInvalidOptions, "payload byte count must be >= 0") - case !isSHA256Hex(manifest.Payload.SHA256): + case !isSHA256Hex(payload.SHA256): return errors.Wrap(ErrInvalidOptions, "payload sha256 must be 64 lowercase hex characters") - case manifest.ManifestKey == "": - return errors.Wrap(ErrInvalidOptions, "manifest key is required") + default: + return nil + } +} + +func validateManifestConfState(conf ManifestConfState) error { + roles := []struct { + name string + values []uint64 + }{ + {name: "voters", values: conf.Voters}, + {name: "learners", values: conf.Learners}, + {name: "voters_outgoing", values: conf.VotersOutgoing}, + {name: "learners_next", values: conf.LearnersNext}, + } + for _, role := range roles { + for _, id := range role.values { + if id == 0 { + return errors.Wrapf(ErrInvalidOptions, "conf_state.%s contains zero", role.name) + } + } + } + return validateManifestConfStateMembership(conf) +} + +func validateManifestConfStateMembership(conf ManifestConfState) error { + learners := uint64Set(conf.Learners) + outgoing := uint64Set(conf.VotersOutgoing) + learnersNext := uint64Set(conf.LearnersNext) + if err := validateManifestConfStateOverlaps(conf, learners, outgoing, learnersNext); err != nil { + return err + } + for _, id := range conf.LearnersNext { + if _, ok := outgoing[id]; !ok { + return errors.Wrapf(ErrInvalidOptions, "conf_state.learners_next node %d is not an outgoing voter", id) + } + } + if len(conf.VotersOutgoing) == 0 && conf.AutoLeave { + return errors.Wrap(ErrInvalidOptions, "conf_state.auto_leave requires joint consensus") } return nil } +func validateManifestConfStateOverlaps( + conf ManifestConfState, + learners, outgoing, learnersNext map[uint64]struct{}, +) error { + for _, id := range conf.Voters { + if _, ok := learners[id]; ok { + return invalidManifestConfStateOverlap(id, "voters", "learners") + } + if _, ok := learnersNext[id]; ok { + return invalidManifestConfStateOverlap(id, "voters", "learners_next") + } + } + for _, id := range conf.Learners { + if _, ok := outgoing[id]; ok { + return invalidManifestConfStateOverlap(id, "learners", "voters_outgoing") + } + if _, ok := learnersNext[id]; ok { + return invalidManifestConfStateOverlap(id, "learners", "learners_next") + } + } + return nil +} + +func uint64Set(values []uint64) map[uint64]struct{} { + set := make(map[uint64]struct{}, len(values)) + for _, value := range values { + set[value] = struct{}{} + } + return set +} + +func invalidManifestConfStateOverlap(id uint64, first, second string) error { + return errors.Wrapf(ErrInvalidOptions, "node %d appears in conf_state.%s and conf_state.%s", id, first, second) +} + func verifyManifestSelfHash(manifest Manifest) error { if !isSHA256Hex(manifest.ManifestSHA256) { return errors.Wrap(ErrInvalidOptions, "manifest sha256 must be 64 lowercase hex characters") diff --git a/internal/snapshotoffload/offload_test.go b/internal/snapshotoffload/offload_test.go index f90bae9d3..07416e95f 100644 --- a/internal/snapshotoffload/offload_test.go +++ b/internal/snapshotoffload/offload_test.go @@ -249,6 +249,41 @@ func TestPutManifestReusesExistingManifestAfterCreateConflict(t *testing.T) { require.NotEmpty(t, candidate.ManifestSHA256) } +func TestDecodeManifestRejectsInvalidConfStateMembership(t *testing.T) { + base := testManifestForValidation(t) + testCases := []ManifestConfState{ + {Voters: nil}, + {Voters: []uint64{0}}, + {Voters: []uint64{1}, Learners: []uint64{1}}, + {Voters: []uint64{1}, VotersOutgoing: []uint64{1}, LearnersNext: []uint64{1}}, + {Voters: []uint64{1}, VotersOutgoing: []uint64{1}, LearnersNext: []uint64{2}}, + {Voters: []uint64{1}, AutoLeave: true}, + } + for _, confState := range testCases { + manifest := base + manifest.ConfState = confState + raw, _, err := manifest.MarshalCanonical() + require.NoError(t, err) + _, err = DecodeManifest(raw) + require.ErrorIs(t, err, ErrInvalidOptions) + } +} + +func TestDecodeManifestAcceptsJointConsensusConfState(t *testing.T) { + manifest := testManifestForValidation(t) + manifest.ConfState = ManifestConfState{ + Voters: []uint64{1, 2}, + VotersOutgoing: []uint64{1, 3}, + LearnersNext: []uint64{3}, + AutoLeave: true, + } + raw, _, err := manifest.MarshalCanonical() + require.NoError(t, err) + decoded, err := DecodeManifest(raw) + require.NoError(t, err) + require.Equal(t, manifest.ConfState, decoded.ConfState) +} + func TestPublishRejectsGroupZeroWithoutSourceClusterBeforePayload(t *testing.T) { ctx := context.Background() root := t.TempDir() @@ -530,6 +565,30 @@ func TestPrepareRestoreDownloadDirCreatesParentAndCleansOnlyStaleDirs(t *testing require.True(t, os.IsNotExist(err)) } +func testManifestForValidation(t *testing.T) Manifest { + t.Helper() + payloadSHA := hexSHA256Bytes([]byte("payload")) + key, err := manifestKey("cluster-a", 1, 20, 13) + require.NoError(t, err) + return Manifest{ + SchemaVersion: ManifestSchemaVersion, + CreatedAt: time.Unix(400, 0).UTC(), + SourceCluster: "cluster-a", + GroupID: 1, + SnapshotIndex: 20, + SnapshotTerm: 13, + ConfState: ManifestConfState{ + Voters: []uint64{1}, + }, + Payload: PayloadDescriptor{ + Key: "cluster-a/v1/payloads/test.fsm", + Bytes: int64(len("payload")), + SHA256: payloadSHA, + }, + ManifestKey: key, + } +} + func seedPhysicalSnapshot( t *testing.T, root string, diff --git a/internal/snapshotoffload/s3_store.go b/internal/snapshotoffload/s3_store.go index 98ee9963f..9e6a2644d 100644 --- a/internal/snapshotoffload/s3_store.go +++ b/internal/snapshotoffload/s3_store.go @@ -1,14 +1,17 @@ package snapshotoffload import ( + "bytes" "context" "crypto/sha256" "encoding/base64" "encoding/hex" "io" "strings" + "time" "github.com/aws/aws-sdk-go-v2/aws" + "github.com/aws/aws-sdk-go-v2/aws/arn" "github.com/aws/aws-sdk-go-v2/config" "github.com/aws/aws-sdk-go-v2/credentials" "github.com/aws/aws-sdk-go-v2/service/s3" @@ -20,12 +23,23 @@ import ( const ( s3MetadataSHA256 = "elastickv-sha256" s3LoadConfigOptionCapHint = 4 + s3ConditionalWriteRetries = 3 + s3MaxSinglePutBytes = int64(5 * 1024 * 1024 * 1024) + s3DefaultMultipartPart = int64(64 * 1024 * 1024) + s3MaxMultipartPart = int64(5 * 1024 * 1024 * 1024) + s3MaxMultipartParts = int64(10_000) + s3MaxObjectBytes = int64(5 * 1024 * 1024 * 1024 * 1024) + s3MultipartAbortTimeout = 30 * time.Second ) type S3ObjectClient interface { PutObject(context.Context, *s3.PutObjectInput, ...func(*s3.Options)) (*s3.PutObjectOutput, error) GetObject(context.Context, *s3.GetObjectInput, ...func(*s3.Options)) (*s3.GetObjectOutput, error) HeadObject(context.Context, *s3.HeadObjectInput, ...func(*s3.Options)) (*s3.HeadObjectOutput, error) + CreateMultipartUpload(context.Context, *s3.CreateMultipartUploadInput, ...func(*s3.Options)) (*s3.CreateMultipartUploadOutput, error) + UploadPart(context.Context, *s3.UploadPartInput, ...func(*s3.Options)) (*s3.UploadPartOutput, error) + CompleteMultipartUpload(context.Context, *s3.CompleteMultipartUploadInput, ...func(*s3.Options)) (*s3.CompleteMultipartUploadOutput, error) + AbortMultipartUpload(context.Context, *s3.AbortMultipartUploadInput, ...func(*s3.Options)) (*s3.AbortMultipartUploadOutput, error) } type S3StoreConfig struct { @@ -49,12 +63,17 @@ type S3Store struct { serverSideEncryption string sseKMSKeyID string disableChecksumHeaders bool + multipartThreshold int64 + multipartPartSize int64 } func NewS3Store(ctx context.Context, cfg S3StoreConfig) (*S3Store, error) { if stringsTrim(cfg.Bucket) == "" { return nil, errors.Wrap(ErrInvalidOptions, "s3 bucket is required") } + if err := validateS3EncryptionConfig(cfg.ServerSideEncryption, cfg.SSEKMSKeyID); err != nil { + return nil, err + } client := cfg.Client if client == nil { awsCfg, err := loadS3AWSConfig(ctx, cfg) @@ -74,6 +93,8 @@ func NewS3Store(ctx context.Context, cfg S3StoreConfig) (*S3Store, error) { serverSideEncryption: stringsTrim(cfg.ServerSideEncryption), sseKMSKeyID: stringsTrim(cfg.SSEKMSKeyID), disableChecksumHeaders: cfg.DisableChecksumHeaders, + multipartThreshold: s3MaxSinglePutBytes, + multipartPartSize: s3DefaultMultipartPart, }, nil } @@ -110,19 +131,203 @@ func (s *S3Store) PutObject(ctx context.Context, key string, body io.Reader, opt if err != nil { return ObjectInfo{}, err } - input, err := s.putObjectInput(normalized, body, opts) - if err != nil { - return ObjectInfo{}, err - } - if _, err := s.client.PutObject(ctx, input); err != nil { + if err := s.putObjectIfAbsent(ctx, normalized, body, opts); err != nil { + if errors.Is(err, ErrObjectConflict) { + return ObjectInfo{}, err + } if !isS3PreconditionFailed(err) { - return ObjectInfo{}, errors.Wrap(err, "put s3 object") + return ObjectInfo{}, err } return s.verifyS3ExistingObject(ctx, normalized, opts) } return s.verifyS3PutObject(ctx, normalized, opts) } +func (s *S3Store) putObjectIfAbsent(ctx context.Context, key string, body io.Reader, opts PutOptions) error { + return s.putObjectWithRetry(key, body, func() error { + if opts.Size > s.multipartThreshold { + return s.putMultipartIfAbsent(ctx, key, body, opts) + } + input, err := s.putObjectInput(key, body, opts) + if err != nil { + return err + } + if _, err = s.client.PutObject(ctx, input); err != nil { + return errors.Wrap(err, "put s3 object") + } + return nil + }) +} + +func (s *S3Store) putObjectWithRetry( + key string, + body io.Reader, + put func() error, +) error { + start, seeker := readerPosition(body) + var err error + for attempt := 0; attempt < s3ConditionalWriteRetries; attempt++ { + if attempt > 0 { + if seeker == nil { + return errors.Wrap(ErrObjectConflict, "retrying s3 conditional conflict requires a seekable body") + } + if _, seekErr := seeker.Seek(start, io.SeekStart); seekErr != nil { + return errors.Wrap(seekErr, "rewind s3 conditional write body") + } + } + if err = put(); err == nil || !isS3ConditionalConflict(err) { + return err + } + } + return errors.Wrapf(ErrObjectConflict, "s3 conditional write for %s conflicted after %d attempts", + key, s3ConditionalWriteRetries) +} + +func (s *S3Store) putMultipartIfAbsent( + ctx context.Context, + key string, + body io.Reader, + opts PutOptions, +) (retErr error) { + partSize, err := multipartPartSize(opts.Size, s.multipartPartSize) + if err != nil { + return err + } + uploadID, err := s.createMultipartUpload(ctx, key, opts) + if err != nil { + return err + } + completed := false + defer func() { + if !completed { + retErr = s.abortMultipartUpload(ctx, key, uploadID, retErr) + } + }() + parts, err := s.uploadParts(ctx, key, uploadID, body, opts.Size, partSize) + if err != nil { + return err + } + if err := s.completeMultipartUpload(ctx, key, uploadID, opts.Size, parts); err != nil { + return err + } + completed = true + return nil +} + +func (s *S3Store) createMultipartUpload(ctx context.Context, key string, opts PutOptions) (string, error) { + input := &s3.CreateMultipartUploadInput{ + Bucket: aws.String(s.bucket), + Key: aws.String(key), + ChecksumAlgorithm: types.ChecksumAlgorithmSha256, + Metadata: map[string]string{ + s3MetadataSHA256: opts.SHA256, + }, + ServerSideEncryption: types.ServerSideEncryption(s.serverSideEncryption), + } + if stringsTrim(opts.ContentType) != "" { + input.ContentType = aws.String(stringsTrim(opts.ContentType)) + } + if s.sseKMSKeyID != "" { + input.SSEKMSKeyId = aws.String(s.sseKMSKeyID) + } + out, err := s.client.CreateMultipartUpload(ctx, input) + if err != nil { + return "", errors.Wrap(err, "create s3 multipart upload") + } + uploadID := stringsTrim(aws.ToString(out.UploadId)) + if uploadID == "" { + return "", errors.Wrap(ErrIntegrity, "s3 multipart upload returned empty upload id") + } + return uploadID, nil +} + +func (s *S3Store) abortMultipartUpload(ctx context.Context, key, uploadID string, prior error) error { + abortCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), s3MultipartAbortTimeout) + defer cancel() + _, abortErr := s.client.AbortMultipartUpload(abortCtx, &s3.AbortMultipartUploadInput{ + Bucket: aws.String(s.bucket), + Key: aws.String(key), + UploadId: aws.String(uploadID), + }) + if abortErr == nil { + return prior + } + if prior == nil { + return errors.Wrap(abortErr, "abort s3 multipart upload") + } + return errors.WithStack(errors.CombineErrors(prior, errors.Wrap(abortErr, "abort s3 multipart upload"))) +} + +func (s *S3Store) uploadParts( + ctx context.Context, + key string, + uploadID string, + body io.Reader, + totalBytes int64, + partSize int64, +) ([]types.CompletedPart, error) { + parts := make([]types.CompletedPart, 0, (totalBytes+partSize-1)/partSize) + remaining := totalBytes + for partNumber := int32(1); remaining > 0; partNumber++ { + partBytes := min(remaining, partSize) + part := make([]byte, int(partBytes)) + if _, err := io.ReadFull(body, part); err != nil { + return nil, errors.Wrap(err, "read s3 multipart source") + } + sum := sha256.Sum256(part) + checksum := base64.StdEncoding.EncodeToString(sum[:]) + out, err := s.client.UploadPart(ctx, &s3.UploadPartInput{ + Bucket: aws.String(s.bucket), + Key: aws.String(key), + UploadId: aws.String(uploadID), + PartNumber: aws.Int32(partNumber), + Body: bytes.NewReader(part), + ContentLength: aws.Int64(partBytes), + ChecksumAlgorithm: types.ChecksumAlgorithmSha256, + ChecksumSHA256: aws.String(checksum), + }) + if err != nil { + return nil, errors.Wrap(err, "upload s3 multipart part") + } + if stringsTrim(aws.ToString(out.ETag)) == "" { + return nil, errors.Wrapf(ErrIntegrity, "s3 multipart part %d returned no etag", partNumber) + } + parts = append(parts, types.CompletedPart{ + ETag: out.ETag, + PartNumber: aws.Int32(partNumber), + ChecksumSHA256: aws.String(checksum), + }) + remaining -= partBytes + } + if err := requireNoTrailingBytes(body); err != nil { + return nil, errors.Wrap(err, "s3 multipart source length differs from declared length") + } + return parts, nil +} + +func (s *S3Store) completeMultipartUpload( + ctx context.Context, + key string, + uploadID string, + size int64, + parts []types.CompletedPart, +) error { + _, err := s.client.CompleteMultipartUpload(ctx, &s3.CompleteMultipartUploadInput{ + Bucket: aws.String(s.bucket), + Key: aws.String(key), + UploadId: aws.String(uploadID), + IfNoneMatch: aws.String("*"), + MpuObjectSize: aws.Int64(size), + MultipartUpload: &types.CompletedMultipartUpload{ + Parts: parts, + }, + }) + if err != nil { + return errors.Wrap(err, "complete s3 multipart upload") + } + return nil +} + func (s *S3Store) verifyS3PutObject(ctx context.Context, key string, opts PutOptions) (ObjectInfo, error) { info, ok, err := s.HeadObject(ctx, key) if err != nil { @@ -134,6 +339,9 @@ func (s *S3Store) verifyS3PutObject(ctx context.Context, key string, opts PutOpt if info.Size != opts.Size || (info.SHA256 != "" && info.SHA256 != opts.SHA256) { return ObjectInfo{}, errors.Wrapf(ErrIntegrity, "s3 object %s remote integrity mismatch", key) } + if err := s.validateS3ObjectEncryption(key, info); err != nil { + return ObjectInfo{}, err + } if info.SHA256 == "" { verified, err := s.verifyS3ObjectBytes(ctx, key, opts) if err != nil { @@ -193,11 +401,22 @@ func (s *S3Store) GetObject(ctx context.Context, key string) (io.ReadCloser, Obj if out.Body == nil { return nil, ObjectInfo{}, errors.Wrapf(ErrIntegrity, "s3 object %s returned empty body", normalized) } - info, err := s3ObjectInfo(normalized, out.ContentLength, out.Metadata, out.ChecksumSHA256) + info, err := s3ObjectInfo( + normalized, + out.ContentLength, + out.Metadata, + out.ChecksumSHA256, + out.ServerSideEncryption, + out.SSEKMSKeyId, + ) if err != nil { _ = out.Body.Close() return nil, ObjectInfo{}, err } + if err := s.validateS3ObjectEncryption(normalized, info); err != nil { + _ = out.Body.Close() + return nil, ObjectInfo{}, err + } return out.Body, info, nil } @@ -216,10 +435,20 @@ func (s *S3Store) HeadObject(ctx context.Context, key string) (ObjectInfo, bool, } return ObjectInfo{}, false, errors.Wrap(err, "head s3 object") } - info, err := s3ObjectInfo(normalized, out.ContentLength, out.Metadata, out.ChecksumSHA256) + info, err := s3ObjectInfo( + normalized, + out.ContentLength, + out.Metadata, + out.ChecksumSHA256, + out.ServerSideEncryption, + out.SSEKMSKeyId, + ) if err != nil { return ObjectInfo{}, false, err } + if err := s.validateS3ObjectEncryption(normalized, info); err != nil { + return ObjectInfo{}, false, err + } return info, true, nil } @@ -231,6 +460,9 @@ func (s *S3Store) verifyS3ExistingObject(ctx context.Context, key string, opts P if !ok { return ObjectInfo{}, errors.Wrapf(ErrIntegrity, "s3 object %s conflicted but is not visible", key) } + if err := s.validateS3ObjectEncryption(key, info); err != nil { + return ObjectInfo{}, err + } if info.Size != opts.Size { return ObjectInfo{}, errors.Wrapf(ErrIntegrity, "s3 object %s already exists with different content", key) } @@ -260,9 +492,68 @@ func (s *S3Store) verifyS3ObjectBytes(ctx context.Context, key string, opts PutO } info.Size = n info.SHA256 = gotSHA + if err := s.validateS3ObjectEncryption(key, info); err != nil { + return ObjectInfo{}, err + } return info, nil } +func (s *S3Store) validateS3ObjectEncryption(key string, info ObjectInfo) error { + if info.ServerSideEncryption != s.serverSideEncryption { + return errors.Wrapf(ErrIntegrity, "s3 object %s encryption %q, expected %q", + key, info.ServerSideEncryption, s.serverSideEncryption) + } + if s.serverSideEncryption != string(types.ServerSideEncryptionAwsKms) { + return nil + } + if !kmsKeyIDsMatch(info.SSEKMSKeyID, s.sseKMSKeyID) { + return errors.Wrapf(ErrIntegrity, "s3 object %s kms key %q, expected %q", + key, info.SSEKMSKeyID, s.sseKMSKeyID) + } + return nil +} + +func readerPosition(body io.Reader) (int64, io.Seeker) { + seeker, ok := body.(io.Seeker) + if !ok { + return 0, nil + } + position, err := seeker.Seek(0, io.SeekCurrent) + if err != nil { + return 0, nil + } + return position, seeker +} + +func multipartPartSize(totalBytes, configured int64) (int64, error) { + if totalBytes <= 0 || configured <= 0 { + return 0, errors.Wrap(ErrInvalidOptions, "s3 multipart length and part size must be positive") + } + if totalBytes > s3MaxObjectBytes { + return 0, errors.Wrapf(ErrInvalidOptions, "s3 object exceeds 5 TiB: bytes=%d", totalBytes) + } + partSize := max(configured, (totalBytes+s3MaxMultipartParts-1)/s3MaxMultipartParts) + if partSize > s3MaxMultipartPart { + return 0, errors.Wrapf(ErrInvalidOptions, "s3 multipart object is too large: bytes=%d", totalBytes) + } + return partSize, nil +} + +func requireNoTrailingBytes(reader io.Reader) error { + var extra [1]byte + n, err := reader.Read(extra[:]) + if n > 0 { + return errors.Wrap(ErrIntegrity, "source has trailing bytes after declared length") + } + if err == nil { + return errors.Wrap(ErrIntegrity, "source did not report EOF after declared length") + } + if errors.Is(err, io.EOF) { + return nil + } + return errors.WithStack(err) +} + func validateStoreObjectKey(key string) (string, error) { normalized := normalizeObjectKey(key) if normalized == "" || normalized == "." || normalized == ".." || strings.HasPrefix(normalized, "../") { @@ -271,7 +562,14 @@ func validateStoreObjectKey(key string) (string, error) { return normalized, nil } -func s3ObjectInfo(key string, contentLength *int64, metadata map[string]string, checksumSHA256 *string) (ObjectInfo, error) { +func s3ObjectInfo( + key string, + contentLength *int64, + metadata map[string]string, + checksumSHA256 *string, + encryption types.ServerSideEncryption, + kmsKeyID *string, +) (ObjectInfo, error) { if contentLength == nil || *contentLength < 0 { return ObjectInfo{}, errors.Wrapf(ErrIntegrity, "s3 object %s missing content length", key) } @@ -279,7 +577,13 @@ func s3ObjectInfo(key string, contentLength *int64, metadata map[string]string, if err != nil { return ObjectInfo{}, err } - return ObjectInfo{Key: key, Size: *contentLength, SHA256: sha}, nil + return ObjectInfo{ + Key: key, + Size: *contentLength, + SHA256: sha, + ServerSideEncryption: string(encryption), + SSEKMSKeyID: aws.ToString(kmsKeyID), + }, nil } func s3ObjectSHA256(key string, metadata map[string]string, checksumSHA256 *string) (string, error) { @@ -345,6 +649,80 @@ func sha256HexToBase64(sha string) (string, error) { return base64.StdEncoding.EncodeToString(decoded), nil } +func ValidateS3StoreEncryption(encryption, kmsKeyID string) error { + return validateS3EncryptionConfig(encryption, kmsKeyID) +} + +func validateS3EncryptionConfig(encryption, kmsKeyID string) error { + encryption = stringsTrim(encryption) + kmsKeyID = stringsTrim(kmsKeyID) + switch { + case encryption == "": + return errors.Wrap(ErrInvalidOptions, "s3 server-side encryption is required") + case encryption != string(types.ServerSideEncryptionAes256) && encryption != string(types.ServerSideEncryptionAwsKms): + return errors.Wrap(ErrInvalidOptions, "s3 server-side encryption must be AES256 or aws:kms") + case encryption == string(types.ServerSideEncryptionAwsKms) && kmsKeyID == "": + return errors.Wrap(ErrInvalidOptions, "s3 KMS key id is required for aws:kms") + case encryption != string(types.ServerSideEncryptionAwsKms) && kmsKeyID != "": + return errors.Wrap(ErrInvalidOptions, "s3 KMS key id requires aws:kms encryption") + } + if encryption == string(types.ServerSideEncryptionAwsKms) { + if err := ValidateKMSKeyID(kmsKeyID); err != nil { + return errors.Wrap(ErrInvalidOptions, err.Error()) + } + } + return nil +} + +// ValidateKMSKeyID accepts canonical key ARNs and bare key IDs. Aliases are +// rejected because S3 reports the resolved key identity, not the alias string. +func ValidateKMSKeyID(value string) error { + _, err := parseKMSKeyIdentity(value) + return err +} + +type kmsKeyIdentity struct { + arn string + keyID string +} + +func parseKMSKeyIdentity(value string) (kmsKeyIdentity, error) { + value = stringsTrim(value) + switch { + case value == "": + return kmsKeyIdentity{}, errors.Wrap(ErrInvalidOptions, "KMS key id is required") + case strings.HasPrefix(value, "alias/"): + return kmsKeyIdentity{}, errors.Wrap(ErrInvalidOptions, "KMS aliases are not supported; use a key ARN or bare key ID") + case arn.IsARN(value): + return parseKMSKeyARN(value) + case strings.ContainsAny(value, ":/"): + return kmsKeyIdentity{}, errors.Wrap(ErrInvalidOptions, "KMS key id must be a key ARN or bare key ID") + default: + return kmsKeyIdentity{keyID: value}, nil + } +} + +func parseKMSKeyARN(value string) (kmsKeyIdentity, error) { + parsed, err := arn.Parse(value) + if err != nil || parsed.Service != "kms" || parsed.Region == "" || parsed.AccountID == "" || + !strings.HasPrefix(parsed.Resource, "key/") || len(parsed.Resource) == len("key/") { + return kmsKeyIdentity{}, errors.Wrap(ErrInvalidOptions, "KMS key ARN is invalid or identifies an alias") + } + return kmsKeyIdentity{arn: value, keyID: strings.TrimPrefix(parsed.Resource, "key/")}, nil +} + +func kmsKeyIDsMatch(actual, expected string) bool { + actualIdentity, actualErr := parseKMSKeyIdentity(actual) + expectedIdentity, expectedErr := parseKMSKeyIdentity(expected) + if actualErr != nil || expectedErr != nil { + return false + } + if actualIdentity.arn != "" && expectedIdentity.arn != "" { + return actualIdentity.arn == expectedIdentity.arn + } + return actualIdentity.keyID == expectedIdentity.keyID +} + func isS3NotFound(err error) bool { var notFound *types.NotFound if errors.As(err, ¬Found) { @@ -366,9 +744,17 @@ func isS3PreconditionFailed(err error) bool { return false } switch apiErr.ErrorCode() { - case "PreconditionFailed", "ConditionalRequestConflict": + case "PreconditionFailed": return true default: return false } } + +func isS3ConditionalConflict(err error) bool { + var apiErr smithy.APIError + if !errors.As(err, &apiErr) { + return false + } + return apiErr.ErrorCode() == "ConditionalRequestConflict" +} diff --git a/internal/snapshotoffload/s3_store_test.go b/internal/snapshotoffload/s3_store_test.go index 5feaa4c9e..b1976b33b 100644 --- a/internal/snapshotoffload/s3_store_test.go +++ b/internal/snapshotoffload/s3_store_test.go @@ -3,6 +3,7 @@ package snapshotoffload import ( "bytes" "context" + "fmt" "io" "strings" "sync" @@ -156,6 +157,53 @@ func TestS3StoreConflictRejectsExistingObjectWithoutIntegrityMetadataMismatch(t require.Equal(t, 1, fake.getAttempts()) } +func TestS3StoreRetriesConditionalRequestConflict(t *testing.T) { + ctx := context.Background() + fake := newFakeS3Client() + fake.conditionalConflicts = 1 + store := newTestS3Store(t, fake) + body := []byte("retry-body") + sha := hexSHA256Bytes(body) + + info, err := store.PutObject(ctx, "snapshots/retry.fsm", bytes.NewReader(body), PutOptions{ + Size: int64(len(body)), + SHA256: sha, + }) + require.NoError(t, err) + require.Equal(t, sha, info.SHA256) + require.Equal(t, 2, fake.putAttempts()) +} + +func TestS3StoreUsesMultipartForLargeObject(t *testing.T) { + ctx := context.Background() + fake := newFakeS3Client() + store := newTestS3Store(t, fake) + store.multipartThreshold = 4 + store.multipartPartSize = 3 + body := []byte("multipart-body") + sha := hexSHA256Bytes(body) + + info, err := store.PutObject(ctx, "snapshots/multipart.fsm", bytes.NewReader(body), PutOptions{ + Size: int64(len(body)), + SHA256: sha, + }) + require.NoError(t, err) + require.Equal(t, sha, info.SHA256) + require.Equal(t, 1, fake.multipartCompletes) + require.Greater(t, fake.uploadedParts, 1) +} + +func TestS3StoreRejectsInvalidKMSConfig(t *testing.T) { + _, err := NewS3Store(context.Background(), S3StoreConfig{ + Client: newFakeS3Client(), + Bucket: "backup-bucket", + ServerSideEncryption: string(types.ServerSideEncryptionAwsKms), + SSEKMSKeyID: "alias/snapshot-key", + }) + require.ErrorIs(t, err, ErrInvalidOptions) + require.ErrorContains(t, err, "aliases are not supported") +} + func TestS3StoreRejectsParentDirectoryKeys(t *testing.T) { ctx := context.Background() store := newTestS3Store(t, newFakeS3Client()) @@ -172,30 +220,51 @@ func TestS3StoreRejectsParentDirectoryKeys(t *testing.T) { func newTestS3Store(t *testing.T, client *fakeS3Client) *S3Store { t.Helper() store, err := NewS3Store(context.Background(), S3StoreConfig{ - Client: client, - Bucket: "backup-bucket", - ForcePathStyle: true, + Client: client, + Bucket: "backup-bucket", + ForcePathStyle: true, + ServerSideEncryption: string(types.ServerSideEncryptionAes256), }) require.NoError(t, err) return store } type fakeS3Client struct { - mu sync.Mutex - objects map[string]fakeS3Object - lastPut types.ChecksumAlgorithm - attempts int - gets int + mu sync.Mutex + objects map[string]fakeS3Object + multipart map[string]*fakeMultipartUpload + lastPut types.ChecksumAlgorithm + attempts int + gets int + conditionalConflicts int + nextUploadID int + uploadedParts int + multipartCompletes int } type fakeS3Object struct { - body []byte - metadata map[string]string - checksum *string + body []byte + metadata map[string]string + checksum *string + serverSideEncryption types.ServerSideEncryption + kmsKeyID *string +} + +type fakeMultipartUpload struct { + bucket string + key string + metadata map[string]string + checksum *string + serverSideEncryption types.ServerSideEncryption + kmsKeyID *string + parts map[int32][]byte } func newFakeS3Client() *fakeS3Client { - return &fakeS3Client{objects: make(map[string]fakeS3Object)} + return &fakeS3Client{ + objects: make(map[string]fakeS3Object), + multipart: make(map[string]*fakeMultipartUpload), + } } func (c *fakeS3Client) PutObject(_ context.Context, input *s3.PutObjectInput, _ ...func(*s3.Options)) (*s3.PutObjectOutput, error) { @@ -204,6 +273,10 @@ func (c *fakeS3Client) PutObject(_ context.Context, input *s3.PutObjectInput, _ defer c.mu.Unlock() c.attempts++ c.lastPut = input.ChecksumAlgorithm + if c.conditionalConflicts > 0 { + c.conditionalConflicts-- + return nil, &smithy.GenericAPIError{Code: "ConditionalRequestConflict", Message: "retry"} + } if _, ok := c.objects[key]; ok && aws.ToString(input.IfNoneMatch) == "*" { return nil, &smithy.GenericAPIError{Code: "PreconditionFailed", Message: "exists"} } @@ -219,9 +292,11 @@ func (c *fakeS3Client) PutObject(_ context.Context, input *s3.PutObjectInput, _ metadata[k] = v } c.objects[key] = fakeS3Object{ - body: append([]byte(nil), body...), - metadata: metadata, - checksum: input.ChecksumSHA256, + body: append([]byte(nil), body...), + metadata: metadata, + checksum: input.ChecksumSHA256, + serverSideEncryption: input.ServerSideEncryption, + kmsKeyID: input.SSEKMSKeyId, } return &s3.PutObjectOutput{}, nil } @@ -238,9 +313,11 @@ func (c *fakeS3Client) HeadObject(_ context.Context, input *s3.HeadObjectInput, metadata[k] = v } return &s3.HeadObjectOutput{ - ContentLength: aws.Int64(int64(len(obj.body))), - Metadata: metadata, - ChecksumSHA256: obj.checksum, + ContentLength: aws.Int64(int64(len(obj.body))), + Metadata: metadata, + ChecksumSHA256: obj.checksum, + ServerSideEncryption: obj.serverSideEncryption, + SSEKMSKeyId: obj.kmsKeyID, }, nil } @@ -257,13 +334,103 @@ func (c *fakeS3Client) GetObject(_ context.Context, input *s3.GetObjectInput, _ metadata[k] = v } return &s3.GetObjectOutput{ - Body: io.NopCloser(bytes.NewReader(obj.body)), - ContentLength: aws.Int64(int64(len(obj.body))), - Metadata: metadata, - ChecksumSHA256: obj.checksum, + Body: io.NopCloser(bytes.NewReader(obj.body)), + ContentLength: aws.Int64(int64(len(obj.body))), + Metadata: metadata, + ChecksumSHA256: obj.checksum, + ServerSideEncryption: obj.serverSideEncryption, + SSEKMSKeyId: obj.kmsKeyID, }, nil } +func (c *fakeS3Client) CreateMultipartUpload( + _ context.Context, + input *s3.CreateMultipartUploadInput, + _ ...func(*s3.Options), +) (*s3.CreateMultipartUploadOutput, error) { + c.mu.Lock() + defer c.mu.Unlock() + c.nextUploadID++ + uploadID := fmt.Sprintf("upload-%d", c.nextUploadID) + metadata := make(map[string]string, len(input.Metadata)) + for k, v := range input.Metadata { + metadata[k] = v + } + c.multipart[uploadID] = &fakeMultipartUpload{ + bucket: aws.ToString(input.Bucket), + key: aws.ToString(input.Key), + metadata: metadata, + serverSideEncryption: input.ServerSideEncryption, + kmsKeyID: input.SSEKMSKeyId, + parts: make(map[int32][]byte), + } + return &s3.CreateMultipartUploadOutput{UploadId: aws.String(uploadID)}, nil +} + +func (c *fakeS3Client) UploadPart( + _ context.Context, + input *s3.UploadPartInput, + _ ...func(*s3.Options), +) (*s3.UploadPartOutput, error) { + c.mu.Lock() + defer c.mu.Unlock() + upload := c.multipart[aws.ToString(input.UploadId)] + if upload == nil { + return nil, &smithy.GenericAPIError{Code: "NoSuchUpload", Message: "missing upload"} + } + body, err := io.ReadAll(input.Body) + if err != nil { + return nil, err + } + upload.parts[aws.ToInt32(input.PartNumber)] = append([]byte(nil), body...) + c.uploadedParts++ + return &s3.UploadPartOutput{ETag: aws.String(fmt.Sprintf("etag-%d", aws.ToInt32(input.PartNumber)))}, nil +} + +func (c *fakeS3Client) CompleteMultipartUpload( + _ context.Context, + input *s3.CompleteMultipartUploadInput, + _ ...func(*s3.Options), +) (*s3.CompleteMultipartUploadOutput, error) { + c.mu.Lock() + defer c.mu.Unlock() + uploadID := aws.ToString(input.UploadId) + upload := c.multipart[uploadID] + if upload == nil { + return nil, &smithy.GenericAPIError{Code: "NoSuchUpload", Message: "missing upload"} + } + key := upload.bucket + "/" + upload.key + if _, ok := c.objects[key]; ok && aws.ToString(input.IfNoneMatch) == "*" { + return nil, &smithy.GenericAPIError{Code: "PreconditionFailed", Message: "exists"} + } + var body []byte + for _, part := range input.MultipartUpload.Parts { + partBody := upload.parts[aws.ToInt32(part.PartNumber)] + body = append(body, partBody...) + } + c.objects[key] = fakeS3Object{ + body: body, + metadata: upload.metadata, + checksum: upload.checksum, + serverSideEncryption: upload.serverSideEncryption, + kmsKeyID: upload.kmsKeyID, + } + c.multipartCompletes++ + delete(c.multipart, uploadID) + return &s3.CompleteMultipartUploadOutput{}, nil +} + +func (c *fakeS3Client) AbortMultipartUpload( + _ context.Context, + input *s3.AbortMultipartUploadInput, + _ ...func(*s3.Options), +) (*s3.AbortMultipartUploadOutput, error) { + c.mu.Lock() + defer c.mu.Unlock() + delete(c.multipart, aws.ToString(input.UploadId)) + return &s3.AbortMultipartUploadOutput{}, nil +} + func (c *fakeS3Client) lastPutChecksumAlgorithm() types.ChecksumAlgorithm { c.mu.Lock() defer c.mu.Unlock() @@ -290,9 +457,10 @@ func (c *fakeS3Client) putRawObject(bucket string, key string, body []byte, meta clonedMetadata[k] = v } c.objects[bucket+"/"+key] = fakeS3Object{ - body: append([]byte(nil), body...), - metadata: clonedMetadata, - checksum: checksum, + body: append([]byte(nil), body...), + metadata: clonedMetadata, + checksum: checksum, + serverSideEncryption: types.ServerSideEncryptionAes256, } } diff --git a/internal/snapshotoffload/store.go b/internal/snapshotoffload/store.go index d8fd4beb9..7049149b6 100644 --- a/internal/snapshotoffload/store.go +++ b/internal/snapshotoffload/store.go @@ -30,7 +30,9 @@ type ObjectInfo struct { Size int64 // SHA256 is optional for metadata-only Head/Get paths; PutObject returns it // when the writer verified the committed content. - SHA256 string + SHA256 string + ServerSideEncryption string + SSEKMSKeyID string } type LocalStore struct { From fdf55743f0892e5b922a1070d60dd87f4bfdfbc5 Mon Sep 17 00:00:00 2001 From: bootjp Date: Thu, 23 Jul 2026 17:31:31 +0900 Subject: [PATCH 2/4] snapshot: harden multipart offload validation --- internal/snapshotoffload/manifest.go | 5 + internal/snapshotoffload/offload_test.go | 4 + internal/snapshotoffload/s3_store.go | 74 ++++++++------ internal/snapshotoffload/s3_store_test.go | 119 ++++++++++++++++++++++ 4 files changed, 172 insertions(+), 30 deletions(-) diff --git a/internal/snapshotoffload/manifest.go b/internal/snapshotoffload/manifest.go index a4209bee6..f750ce0ba 100644 --- a/internal/snapshotoffload/manifest.go +++ b/internal/snapshotoffload/manifest.go @@ -137,10 +137,15 @@ func validateManifestConfState(conf ManifestConfState) error { {name: "learners_next", values: conf.LearnersNext}, } for _, role := range roles { + seen := make(map[uint64]struct{}, len(role.values)) for _, id := range role.values { if id == 0 { return errors.Wrapf(ErrInvalidOptions, "conf_state.%s contains zero", role.name) } + if _, ok := seen[id]; ok { + return errors.Wrapf(ErrInvalidOptions, "conf_state.%s contains duplicate node %d", role.name, id) + } + seen[id] = struct{}{} } } return validateManifestConfStateMembership(conf) diff --git a/internal/snapshotoffload/offload_test.go b/internal/snapshotoffload/offload_test.go index 07416e95f..24a114fce 100644 --- a/internal/snapshotoffload/offload_test.go +++ b/internal/snapshotoffload/offload_test.go @@ -254,6 +254,10 @@ func TestDecodeManifestRejectsInvalidConfStateMembership(t *testing.T) { testCases := []ManifestConfState{ {Voters: nil}, {Voters: []uint64{0}}, + {Voters: []uint64{1, 1}}, + {Voters: []uint64{1}, Learners: []uint64{2, 2}}, + {Voters: []uint64{1}, VotersOutgoing: []uint64{2, 2}, LearnersNext: []uint64{2}}, + {Voters: []uint64{1}, VotersOutgoing: []uint64{2}, LearnersNext: []uint64{2, 2}}, {Voters: []uint64{1}, Learners: []uint64{1}}, {Voters: []uint64{1}, VotersOutgoing: []uint64{1}, LearnersNext: []uint64{1}}, {Voters: []uint64{1}, VotersOutgoing: []uint64{1}, LearnersNext: []uint64{2}}, diff --git a/internal/snapshotoffload/s3_store.go b/internal/snapshotoffload/s3_store.go index 9e6a2644d..a13a460ae 100644 --- a/internal/snapshotoffload/s3_store.go +++ b/internal/snapshotoffload/s3_store.go @@ -28,7 +28,7 @@ const ( s3DefaultMultipartPart = int64(64 * 1024 * 1024) s3MaxMultipartPart = int64(5 * 1024 * 1024 * 1024) s3MaxMultipartParts = int64(10_000) - s3MaxObjectBytes = int64(5 * 1024 * 1024 * 1024 * 1024) + s3MaxObjectBytes = s3MaxMultipartPart * s3MaxMultipartParts s3MultipartAbortTimeout = 30 * time.Second ) @@ -203,7 +203,7 @@ func (s *S3Store) putMultipartIfAbsent( retErr = s.abortMultipartUpload(ctx, key, uploadID, retErr) } }() - parts, err := s.uploadParts(ctx, key, uploadID, body, opts.Size, partSize) + parts, err := s.uploadParts(ctx, key, uploadID, body, opts, partSize) if err != nil { return err } @@ -216,14 +216,16 @@ func (s *S3Store) putMultipartIfAbsent( func (s *S3Store) createMultipartUpload(ctx context.Context, key string, opts PutOptions) (string, error) { input := &s3.CreateMultipartUploadInput{ - Bucket: aws.String(s.bucket), - Key: aws.String(key), - ChecksumAlgorithm: types.ChecksumAlgorithmSha256, + Bucket: aws.String(s.bucket), + Key: aws.String(key), Metadata: map[string]string{ s3MetadataSHA256: opts.SHA256, }, ServerSideEncryption: types.ServerSideEncryption(s.serverSideEncryption), } + if !s.disableChecksumHeaders { + input.ChecksumAlgorithm = types.ChecksumAlgorithmSha256 + } if stringsTrim(opts.ContentType) != "" { input.ContentType = aws.String(stringsTrim(opts.ContentType)) } @@ -263,10 +265,12 @@ func (s *S3Store) uploadParts( key string, uploadID string, body io.Reader, - totalBytes int64, + opts PutOptions, partSize int64, ) ([]types.CompletedPart, error) { + totalBytes := opts.Size parts := make([]types.CompletedPart, 0, (totalBytes+partSize-1)/partSize) + fullSum := sha256.New() remaining := totalBytes for partNumber := int32(1); remaining > 0; partNumber++ { partBytes := min(remaining, partSize) @@ -274,34 +278,45 @@ func (s *S3Store) uploadParts( if _, err := io.ReadFull(body, part); err != nil { return nil, errors.Wrap(err, "read s3 multipart source") } - sum := sha256.Sum256(part) - checksum := base64.StdEncoding.EncodeToString(sum[:]) - out, err := s.client.UploadPart(ctx, &s3.UploadPartInput{ - Bucket: aws.String(s.bucket), - Key: aws.String(key), - UploadId: aws.String(uploadID), - PartNumber: aws.Int32(partNumber), - Body: bytes.NewReader(part), - ContentLength: aws.Int64(partBytes), - ChecksumAlgorithm: types.ChecksumAlgorithmSha256, - ChecksumSHA256: aws.String(checksum), - }) + _, _ = fullSum.Write(part) + input := &s3.UploadPartInput{ + Bucket: aws.String(s.bucket), + Key: aws.String(key), + UploadId: aws.String(uploadID), + PartNumber: aws.Int32(partNumber), + Body: bytes.NewReader(part), + ContentLength: aws.Int64(partBytes), + } + var checksum string + if !s.disableChecksumHeaders { + sum := sha256.Sum256(part) + checksum = base64.StdEncoding.EncodeToString(sum[:]) + input.ChecksumAlgorithm = types.ChecksumAlgorithmSha256 + input.ChecksumSHA256 = aws.String(checksum) + } + out, err := s.client.UploadPart(ctx, input) if err != nil { return nil, errors.Wrap(err, "upload s3 multipart part") } if stringsTrim(aws.ToString(out.ETag)) == "" { return nil, errors.Wrapf(ErrIntegrity, "s3 multipart part %d returned no etag", partNumber) } - parts = append(parts, types.CompletedPart{ - ETag: out.ETag, - PartNumber: aws.Int32(partNumber), - ChecksumSHA256: aws.String(checksum), - }) + completedPart := types.CompletedPart{ + ETag: out.ETag, + PartNumber: aws.Int32(partNumber), + } + if checksum != "" { + completedPart.ChecksumSHA256 = aws.String(checksum) + } + parts = append(parts, completedPart) remaining -= partBytes } if err := requireNoTrailingBytes(body); err != nil { return nil, errors.Wrap(err, "s3 multipart source length differs from declared length") } + if gotSHA := hex.EncodeToString(fullSum.Sum(nil)); gotSHA != opts.SHA256 { + return nil, errors.Wrapf(ErrIntegrity, "s3 multipart source sha256 %s, expected %s", gotSHA, opts.SHA256) + } return parts, nil } @@ -530,7 +545,7 @@ func multipartPartSize(totalBytes, configured int64) (int64, error) { return 0, errors.Wrap(ErrInvalidOptions, "s3 multipart length and part size must be positive") } if totalBytes > s3MaxObjectBytes { - return 0, errors.Wrapf(ErrInvalidOptions, "s3 object exceeds 5 TiB: bytes=%d", totalBytes) + return 0, errors.Wrapf(ErrInvalidOptions, "s3 object exceeds multipart limit: bytes=%d", totalBytes) } partSize := max(configured, (totalBytes+s3MaxMultipartParts-1)/s3MaxMultipartParts) if partSize > s3MaxMultipartPart { @@ -573,7 +588,7 @@ func s3ObjectInfo( if contentLength == nil || *contentLength < 0 { return ObjectInfo{}, errors.Wrapf(ErrIntegrity, "s3 object %s missing content length", key) } - sha, err := s3ObjectSHA256(key, metadata, checksumSHA256) + sha, err := s3ObjectSHA256(metadata, checksumSHA256) if err != nil { return ObjectInfo{}, err } @@ -586,20 +601,19 @@ func s3ObjectInfo( }, nil } -func s3ObjectSHA256(key string, metadata map[string]string, checksumSHA256 *string) (string, error) { +func s3ObjectSHA256(metadata map[string]string, checksumSHA256 *string) (string, error) { metadataSHA, err := s3MetadataSHA(metadata) if err != nil { return "", err } + if metadataSHA != "" { + return metadataSHA, nil + } checksumSHA, err := s3ChecksumSHA(checksumSHA256) if err != nil { return "", err } switch { - case metadataSHA != "" && checksumSHA != "" && metadataSHA != checksumSHA: - return "", errors.Wrapf(ErrIntegrity, "s3 object %s sha256 metadata/checksum mismatch", key) - case metadataSHA != "": - return metadataSHA, nil case checksumSHA != "": return checksumSHA, nil default: diff --git a/internal/snapshotoffload/s3_store_test.go b/internal/snapshotoffload/s3_store_test.go index b1976b33b..495aafe72 100644 --- a/internal/snapshotoffload/s3_store_test.go +++ b/internal/snapshotoffload/s3_store_test.go @@ -191,6 +191,77 @@ func TestS3StoreUsesMultipartForLargeObject(t *testing.T) { require.Equal(t, sha, info.SHA256) require.Equal(t, 1, fake.multipartCompletes) require.Greater(t, fake.uploadedParts, 1) + require.Equal(t, types.ChecksumAlgorithmSha256, fake.lastMultipartCreateChecksumAlgorithm()) + require.Equal(t, types.ChecksumAlgorithmSha256, fake.lastUploadPartChecksumAlgorithm()) + require.NotNil(t, fake.lastUploadPartSHA256()) + require.Equal(t, fake.uploadedParts, fake.completedPartChecksums()) +} + +func TestS3StoreMultipartHonorsDisabledChecksumHeaders(t *testing.T) { + ctx := context.Background() + fake := newFakeS3Client() + store, err := NewS3Store(ctx, S3StoreConfig{ + Client: fake, + Bucket: "backup-bucket", + ForcePathStyle: true, + ServerSideEncryption: string(types.ServerSideEncryptionAes256), + DisableChecksumHeaders: true, + }) + require.NoError(t, err) + store.multipartThreshold = 4 + store.multipartPartSize = 3 + body := []byte("multipart-body") + sha := hexSHA256Bytes(body) + + info, err := store.PutObject(ctx, "snapshots/multipart-no-checksum.fsm", bytes.NewReader(body), PutOptions{ + Size: int64(len(body)), + SHA256: sha, + }) + require.NoError(t, err) + require.Equal(t, sha, info.SHA256) + require.Equal(t, types.ChecksumAlgorithm(""), fake.lastMultipartCreateChecksumAlgorithm()) + require.Equal(t, types.ChecksumAlgorithm(""), fake.lastUploadPartChecksumAlgorithm()) + require.Nil(t, fake.lastUploadPartSHA256()) + require.Zero(t, fake.completedPartChecksums()) +} + +func TestS3StoreMultipartRejectsSourceHashMismatch(t *testing.T) { + ctx := context.Background() + fake := newFakeS3Client() + store := newTestS3Store(t, fake) + store.multipartThreshold = 4 + store.multipartPartSize = 3 + body := []byte("multipart-body") + expectedSHA := hexSHA256Bytes([]byte("different-body")) + + _, err := store.PutObject(ctx, "snapshots/multipart-mismatch.fsm", bytes.NewReader(body), PutOptions{ + Size: int64(len(body)), + SHA256: expectedSHA, + }) + require.ErrorIs(t, err, ErrIntegrity) + require.Zero(t, fake.multipartCompletes) + require.Zero(t, fake.activeMultipartUploads()) +} + +func TestS3ObjectSHA256PrefersMetadataForCompositeMultipartChecksum(t *testing.T) { + bodySHA := hexSHA256Bytes([]byte("full-object")) + compositeSHA, err := sha256HexToBase64(hexSHA256Bytes([]byte("aws-composite-checksum"))) + require.NoError(t, err) + + got, err := s3ObjectSHA256(map[string]string{ + s3MetadataSHA256: bodySHA, + }, aws.String(compositeSHA)) + require.NoError(t, err) + require.Equal(t, bodySHA, got) +} + +func TestMultipartPartSizeAllowsS3PartCapacity(t *testing.T) { + partSize, err := multipartPartSize(s3MaxObjectBytes, s3DefaultMultipartPart) + require.NoError(t, err) + require.Equal(t, s3MaxMultipartPart, partSize) + + _, err = multipartPartSize(s3MaxObjectBytes+1, s3DefaultMultipartPart) + require.ErrorIs(t, err, ErrInvalidOptions) } func TestS3StoreRejectsInvalidKMSConfig(t *testing.T) { @@ -234,6 +305,10 @@ type fakeS3Client struct { objects map[string]fakeS3Object multipart map[string]*fakeMultipartUpload lastPut types.ChecksumAlgorithm + lastMultipartCreate types.ChecksumAlgorithm + lastUploadPart types.ChecksumAlgorithm + lastUploadPartSHA *string + completedChecksums int attempts int gets int conditionalConflicts int @@ -352,6 +427,7 @@ func (c *fakeS3Client) CreateMultipartUpload( defer c.mu.Unlock() c.nextUploadID++ uploadID := fmt.Sprintf("upload-%d", c.nextUploadID) + c.lastMultipartCreate = input.ChecksumAlgorithm metadata := make(map[string]string, len(input.Metadata)) for k, v := range input.Metadata { metadata[k] = v @@ -382,6 +458,8 @@ func (c *fakeS3Client) UploadPart( if err != nil { return nil, err } + c.lastUploadPart = input.ChecksumAlgorithm + c.lastUploadPartSHA = cloneStringPtr(input.ChecksumSHA256) upload.parts[aws.ToInt32(input.PartNumber)] = append([]byte(nil), body...) c.uploadedParts++ return &s3.UploadPartOutput{ETag: aws.String(fmt.Sprintf("etag-%d", aws.ToInt32(input.PartNumber)))}, nil @@ -405,6 +483,9 @@ func (c *fakeS3Client) CompleteMultipartUpload( } var body []byte for _, part := range input.MultipartUpload.Parts { + if part.ChecksumSHA256 != nil { + c.completedChecksums++ + } partBody := upload.parts[aws.ToInt32(part.PartNumber)] body = append(body, partBody...) } @@ -437,6 +518,36 @@ func (c *fakeS3Client) lastPutChecksumAlgorithm() types.ChecksumAlgorithm { return c.lastPut } +func (c *fakeS3Client) lastMultipartCreateChecksumAlgorithm() types.ChecksumAlgorithm { + c.mu.Lock() + defer c.mu.Unlock() + return c.lastMultipartCreate +} + +func (c *fakeS3Client) lastUploadPartChecksumAlgorithm() types.ChecksumAlgorithm { + c.mu.Lock() + defer c.mu.Unlock() + return c.lastUploadPart +} + +func (c *fakeS3Client) lastUploadPartSHA256() *string { + c.mu.Lock() + defer c.mu.Unlock() + return cloneStringPtr(c.lastUploadPartSHA) +} + +func (c *fakeS3Client) completedPartChecksums() int { + c.mu.Lock() + defer c.mu.Unlock() + return c.completedChecksums +} + +func (c *fakeS3Client) activeMultipartUploads() int { + c.mu.Lock() + defer c.mu.Unlock() + return len(c.multipart) +} + func (c *fakeS3Client) putAttempts() int { c.mu.Lock() defer c.mu.Unlock() @@ -467,3 +578,11 @@ func (c *fakeS3Client) putRawObject(bucket string, key string, body []byte, meta func fakeS3ClientKey(bucket *string, key *string) string { return aws.ToString(bucket) + "/" + aws.ToString(key) } + +func cloneStringPtr(value *string) *string { + if value == nil { + return nil + } + cloned := *value + return &cloned +} From 7cea2491f01d25f9a53ec189005f02d0c0220ebf Mon Sep 17 00:00:00 2001 From: bootjp Date: Thu, 23 Jul 2026 19:33:48 +0900 Subject: [PATCH 3/4] snapshot: stream multipart upload parts --- internal/snapshotoffload/s3_store.go | 171 +++++++++++++++++----- internal/snapshotoffload/s3_store_test.go | 103 ++++++++++++- 2 files changed, 233 insertions(+), 41 deletions(-) diff --git a/internal/snapshotoffload/s3_store.go b/internal/snapshotoffload/s3_store.go index a13a460ae..202de54f1 100644 --- a/internal/snapshotoffload/s3_store.go +++ b/internal/snapshotoffload/s3_store.go @@ -1,7 +1,6 @@ package snapshotoffload import ( - "bytes" "context" "crypto/sha256" "encoding/base64" @@ -271,45 +270,30 @@ func (s *S3Store) uploadParts( totalBytes := opts.Size parts := make([]types.CompletedPart, 0, (totalBytes+partSize-1)/partSize) fullSum := sha256.New() + source, sourceStart, err := multipartSectionSource(body, !s.disableChecksumHeaders) + if err != nil { + return nil, err + } remaining := totalBytes + var uploaded int64 for partNumber := int32(1); remaining > 0; partNumber++ { partBytes := min(remaining, partSize) - part := make([]byte, int(partBytes)) - if _, err := io.ReadFull(body, part); err != nil { - return nil, errors.Wrap(err, "read s3 multipart source") - } - _, _ = fullSum.Write(part) - input := &s3.UploadPartInput{ - Bucket: aws.String(s.bucket), - Key: aws.String(key), - UploadId: aws.String(uploadID), - PartNumber: aws.Int32(partNumber), - Body: bytes.NewReader(part), - ContentLength: aws.Int64(partBytes), - } - var checksum string - if !s.disableChecksumHeaders { - sum := sha256.Sum256(part) - checksum = base64.StdEncoding.EncodeToString(sum[:]) - input.ChecksumAlgorithm = types.ChecksumAlgorithmSha256 - input.ChecksumSHA256 = aws.String(checksum) - } - out, err := s.client.UploadPart(ctx, input) + partReader, checksum, err := s.multipartPartReader(body, source, sourceStart+uploaded, partBytes, fullSum) if err != nil { - return nil, errors.Wrap(err, "upload s3 multipart part") - } - if stringsTrim(aws.ToString(out.ETag)) == "" { - return nil, errors.Wrapf(ErrIntegrity, "s3 multipart part %d returned no etag", partNumber) - } - completedPart := types.CompletedPart{ - ETag: out.ETag, - PartNumber: aws.Int32(partNumber), + return nil, err } - if checksum != "" { - completedPart.ChecksumSHA256 = aws.String(checksum) + completedPart, err := s.uploadMultipartPart(ctx, key, uploadID, partNumber, partBytes, partReader, checksum) + if err != nil { + return nil, err } parts = append(parts, completedPart) remaining -= partBytes + uploaded += partBytes + } + if source != nil { + if _, err := source.Seek(sourceStart+totalBytes, io.SeekStart); err != nil { + return nil, errors.Wrap(err, "advance s3 multipart source") + } } if err := requireNoTrailingBytes(body); err != nil { return nil, errors.Wrap(err, "s3 multipart source length differs from declared length") @@ -320,6 +304,107 @@ func (s *S3Store) uploadParts( return parts, nil } +func (s *S3Store) uploadMultipartPart( + ctx context.Context, + key string, + uploadID string, + partNumber int32, + partBytes int64, + partReader io.Reader, + checksum string, +) (types.CompletedPart, error) { + counted := &countingReader{reader: partReader} + input := &s3.UploadPartInput{ + Bucket: aws.String(s.bucket), + Key: aws.String(key), + UploadId: aws.String(uploadID), + PartNumber: aws.Int32(partNumber), + Body: counted, + ContentLength: aws.Int64(partBytes), + } + if checksum != "" { + input.ChecksumAlgorithm = types.ChecksumAlgorithmSha256 + input.ChecksumSHA256 = aws.String(checksum) + } + out, err := s.client.UploadPart(ctx, input) + if err != nil { + return types.CompletedPart{}, errors.Wrap(err, "upload s3 multipart part") + } + if counted.n != partBytes { + return types.CompletedPart{}, errors.Wrapf(ErrIntegrity, "s3 multipart part %d read %d bytes, expected %d", + partNumber, counted.n, partBytes) + } + if stringsTrim(aws.ToString(out.ETag)) == "" { + return types.CompletedPart{}, errors.Wrapf(ErrIntegrity, "s3 multipart part %d returned no etag", partNumber) + } + completedPart := types.CompletedPart{ + ETag: out.ETag, + PartNumber: aws.Int32(partNumber), + } + if checksum != "" { + completedPart.ChecksumSHA256 = aws.String(checksum) + } + return completedPart, nil +} + +type readAtSeeker interface { + io.ReaderAt + io.Seeker +} + +type countingReader struct { + reader io.Reader + n int64 +} + +func (r *countingReader) Read(p []byte) (int, error) { + n, err := r.reader.Read(p) + r.n += int64(n) + if err == nil { + return n, nil + } + if errors.Is(err, io.EOF) { + return n, io.EOF + } + return n, errors.WithStack(err) +} + +func multipartSectionSource(body io.Reader, requireSeekable bool) (readAtSeeker, int64, error) { + source, ok := body.(readAtSeeker) + if !ok { + if requireSeekable { + return nil, 0, errors.Wrap(ErrInvalidOptions, "s3 multipart checksum headers require a seekable source") + } + return nil, 0, nil + } + start, err := source.Seek(0, io.SeekCurrent) + if err != nil { + return nil, 0, errors.Wrap(err, "read s3 multipart source position") + } + return source, start, nil +} + +func (s *S3Store) multipartPartReader( + body io.Reader, + source readAtSeeker, + offset int64, + partBytes int64, + fullSum io.Writer, +) (io.Reader, string, error) { + if source == nil { + return io.TeeReader(io.LimitReader(body, partBytes), fullSum), "", nil + } + if s.disableChecksumHeaders { + return io.TeeReader(io.NewSectionReader(source, offset, partBytes), fullSum), "", nil + } + partSum := sha256.New() + if _, err := io.Copy(partSum, io.NewSectionReader(source, offset, partBytes)); err != nil { + return nil, "", errors.Wrap(err, "hash s3 multipart part") + } + checksum := base64.StdEncoding.EncodeToString(partSum.Sum(nil)) + return io.TeeReader(io.NewSectionReader(source, offset, partBytes), fullSum), checksum, nil +} + func (s *S3Store) completeMultipartUpload( ctx context.Context, key string, @@ -421,6 +506,7 @@ func (s *S3Store) GetObject(ctx context.Context, key string) (io.ReadCloser, Obj out.ContentLength, out.Metadata, out.ChecksumSHA256, + out.ChecksumType, out.ServerSideEncryption, out.SSEKMSKeyId, ) @@ -455,6 +541,7 @@ func (s *S3Store) HeadObject(ctx context.Context, key string) (ObjectInfo, bool, out.ContentLength, out.Metadata, out.ChecksumSHA256, + out.ChecksumType, out.ServerSideEncryption, out.SSEKMSKeyId, ) @@ -582,13 +669,14 @@ func s3ObjectInfo( contentLength *int64, metadata map[string]string, checksumSHA256 *string, + checksumType types.ChecksumType, encryption types.ServerSideEncryption, kmsKeyID *string, ) (ObjectInfo, error) { if contentLength == nil || *contentLength < 0 { return ObjectInfo{}, errors.Wrapf(ErrIntegrity, "s3 object %s missing content length", key) } - sha, err := s3ObjectSHA256(metadata, checksumSHA256) + sha, err := s3ObjectSHA256(metadata, checksumSHA256, checksumType) if err != nil { return ObjectInfo{}, err } @@ -601,19 +689,21 @@ func s3ObjectInfo( }, nil } -func s3ObjectSHA256(metadata map[string]string, checksumSHA256 *string) (string, error) { +func s3ObjectSHA256(metadata map[string]string, checksumSHA256 *string, checksumType types.ChecksumType) (string, error) { metadataSHA, err := s3MetadataSHA(metadata) if err != nil { return "", err } - if metadataSHA != "" { - return metadataSHA, nil - } - checksumSHA, err := s3ChecksumSHA(checksumSHA256) + checksumSHA, err := s3ChecksumSHA(checksumSHA256, checksumType) if err != nil { return "", err } + if metadataSHA != "" && checksumSHA != "" && metadataSHA != checksumSHA { + return "", errors.Wrap(ErrIntegrity, "s3 object sha256 metadata does not match full-object checksum") + } switch { + case metadataSHA != "": + return metadataSHA, nil case checksumSHA != "": return checksumSHA, nil default: @@ -640,11 +730,14 @@ func s3MetadataSHA(metadata map[string]string) (string, error) { return "", nil } -func s3ChecksumSHA(checksum *string) (string, error) { +func s3ChecksumSHA(checksum *string, checksumType types.ChecksumType) (string, error) { raw := stringsTrim(aws.ToString(checksum)) if raw == "" { return "", nil } + if checksumType == types.ChecksumTypeComposite { + return "", nil + } decoded, err := base64.StdEncoding.DecodeString(raw) if err != nil { return "", errors.Wrap(ErrIntegrity, "s3 object sha256 checksum is invalid base64") diff --git a/internal/snapshotoffload/s3_store_test.go b/internal/snapshotoffload/s3_store_test.go index 495aafe72..31fb2eeb5 100644 --- a/internal/snapshotoffload/s3_store_test.go +++ b/internal/snapshotoffload/s3_store_test.go @@ -197,6 +197,26 @@ func TestS3StoreUsesMultipartForLargeObject(t *testing.T) { require.Equal(t, fake.uploadedParts, fake.completedPartChecksums()) } +func TestS3StoreMultipartStreamsSeekableParts(t *testing.T) { + ctx := context.Background() + fake := newFakeS3Client() + store := newTestS3Store(t, fake) + store.multipartThreshold = 4 + store.multipartPartSize = 3 + body := []byte("multipart-body") + source := &spyReadAtSeeker{data: body} + sha := hexSHA256Bytes(body) + + info, err := store.PutObject(ctx, "snapshots/multipart-streamed.fsm", source, PutOptions{ + Size: int64(len(body)), + SHA256: sha, + }) + require.NoError(t, err) + require.Equal(t, sha, info.SHA256) + require.Greater(t, source.readAtCalls, 0) + require.Zero(t, source.readBytes) +} + func TestS3StoreMultipartHonorsDisabledChecksumHeaders(t *testing.T) { ctx := context.Background() fake := newFakeS3Client() @@ -250,11 +270,22 @@ func TestS3ObjectSHA256PrefersMetadataForCompositeMultipartChecksum(t *testing.T got, err := s3ObjectSHA256(map[string]string{ s3MetadataSHA256: bodySHA, - }, aws.String(compositeSHA)) + }, aws.String(compositeSHA), types.ChecksumTypeComposite) require.NoError(t, err) require.Equal(t, bodySHA, got) } +func TestS3ObjectSHA256RejectsMetadataFullObjectChecksumMismatch(t *testing.T) { + metadataSHA := hexSHA256Bytes([]byte("metadata-body")) + fullObjectChecksum, err := sha256HexToBase64(hexSHA256Bytes([]byte("full-object-body"))) + require.NoError(t, err) + + _, err = s3ObjectSHA256(map[string]string{ + s3MetadataSHA256: metadataSHA, + }, aws.String(fullObjectChecksum), types.ChecksumTypeFullObject) + require.ErrorIs(t, err, ErrIntegrity) +} + func TestMultipartPartSizeAllowsS3PartCapacity(t *testing.T) { partSize, err := multipartPartSize(s3MaxObjectBytes, s3DefaultMultipartPart) require.NoError(t, err) @@ -321,6 +352,7 @@ type fakeS3Object struct { body []byte metadata map[string]string checksum *string + checksumType types.ChecksumType serverSideEncryption types.ServerSideEncryption kmsKeyID *string } @@ -335,6 +367,57 @@ type fakeMultipartUpload struct { parts map[int32][]byte } +type spyReadAtSeeker struct { + data []byte + offset int64 + readBytes int64 + readAtCalls int +} + +func (r *spyReadAtSeeker) Read(p []byte) (int, error) { + if r.offset >= int64(len(r.data)) { + return 0, io.EOF + } + n := copy(p, r.data[r.offset:]) + r.offset += int64(n) + r.readBytes += int64(n) + if n < len(p) { + return n, io.EOF + } + return n, nil +} + +func (r *spyReadAtSeeker) ReadAt(p []byte, off int64) (int, error) { + r.readAtCalls++ + if off >= int64(len(r.data)) { + return 0, io.EOF + } + n := copy(p, r.data[off:]) + if n < len(p) { + return n, io.EOF + } + return n, nil +} + +func (r *spyReadAtSeeker) Seek(offset int64, whence int) (int64, error) { + var next int64 + switch whence { + case io.SeekStart: + next = offset + case io.SeekCurrent: + next = r.offset + offset + case io.SeekEnd: + next = int64(len(r.data)) + offset + default: + return 0, fmt.Errorf("invalid whence %d", whence) + } + if next < 0 { + return 0, fmt.Errorf("negative offset %d", next) + } + r.offset = next + return next, nil +} + func newFakeS3Client() *fakeS3Client { return &fakeS3Client{ objects: make(map[string]fakeS3Object), @@ -370,6 +453,7 @@ func (c *fakeS3Client) PutObject(_ context.Context, input *s3.PutObjectInput, _ body: append([]byte(nil), body...), metadata: metadata, checksum: input.ChecksumSHA256, + checksumType: s3ChecksumTypeForSHA(input.ChecksumSHA256), serverSideEncryption: input.ServerSideEncryption, kmsKeyID: input.SSEKMSKeyId, } @@ -391,6 +475,7 @@ func (c *fakeS3Client) HeadObject(_ context.Context, input *s3.HeadObjectInput, ContentLength: aws.Int64(int64(len(obj.body))), Metadata: metadata, ChecksumSHA256: obj.checksum, + ChecksumType: obj.checksumType, ServerSideEncryption: obj.serverSideEncryption, SSEKMSKeyId: obj.kmsKeyID, }, nil @@ -413,6 +498,7 @@ func (c *fakeS3Client) GetObject(_ context.Context, input *s3.GetObjectInput, _ ContentLength: aws.Int64(int64(len(obj.body))), Metadata: metadata, ChecksumSHA256: obj.checksum, + ChecksumType: obj.checksumType, ServerSideEncryption: obj.serverSideEncryption, SSEKMSKeyId: obj.kmsKeyID, }, nil @@ -493,6 +579,7 @@ func (c *fakeS3Client) CompleteMultipartUpload( body: body, metadata: upload.metadata, checksum: upload.checksum, + checksumType: s3ChecksumTypeForSHA(upload.checksum), serverSideEncryption: upload.serverSideEncryption, kmsKeyID: upload.kmsKeyID, } @@ -560,21 +647,33 @@ func (c *fakeS3Client) getAttempts() int { return c.gets } -func (c *fakeS3Client) putRawObject(bucket string, key string, body []byte, metadata map[string]string, checksum *string) { +func (c *fakeS3Client) putRawObject(bucket string, key string, body []byte, metadata map[string]string, checksum *string, checksumType ...types.ChecksumType) { c.mu.Lock() defer c.mu.Unlock() clonedMetadata := make(map[string]string, len(metadata)) for k, v := range metadata { clonedMetadata[k] = v } + storedChecksumType := s3ChecksumTypeForSHA(checksum) + if len(checksumType) > 0 { + storedChecksumType = checksumType[0] + } c.objects[bucket+"/"+key] = fakeS3Object{ body: append([]byte(nil), body...), metadata: clonedMetadata, checksum: checksum, + checksumType: storedChecksumType, serverSideEncryption: types.ServerSideEncryptionAes256, } } +func s3ChecksumTypeForSHA(checksum *string) types.ChecksumType { + if checksum == nil { + return "" + } + return types.ChecksumTypeFullObject +} + func fakeS3ClientKey(bucket *string, key *string) string { return aws.ToString(bucket) + "/" + aws.ToString(key) } From 7db43c8f3e7e8b5885a8bb6a944dc15818cb9cda Mon Sep 17 00:00:00 2001 From: bootjp Date: Thu, 23 Jul 2026 21:24:55 +0900 Subject: [PATCH 4/4] snapshot: harden s3 checksum requests --- internal/snapshotoffload/s3_store.go | 18 +++++-- internal/snapshotoffload/s3_store_test.go | 64 ++++++++++++++++++++--- 2 files changed, 71 insertions(+), 11 deletions(-) diff --git a/internal/snapshotoffload/s3_store.go b/internal/snapshotoffload/s3_store.go index 202de54f1..f8819d8df 100644 --- a/internal/snapshotoffload/s3_store.go +++ b/internal/snapshotoffload/s3_store.go @@ -27,7 +27,7 @@ const ( s3DefaultMultipartPart = int64(64 * 1024 * 1024) s3MaxMultipartPart = int64(5 * 1024 * 1024 * 1024) s3MaxMultipartParts = int64(10_000) - s3MaxObjectBytes = s3MaxMultipartPart * s3MaxMultipartParts + s3MaxObjectBytes = int64(5 * 1024 * 1024 * 1024 * 1024) s3MultipartAbortTimeout = 30 * time.Second ) @@ -488,10 +488,14 @@ func (s *S3Store) GetObject(ctx context.Context, key string) (io.ReadCloser, Obj if err != nil { return nil, ObjectInfo{}, err } - out, err := s.client.GetObject(ctx, &s3.GetObjectInput{ + input := &s3.GetObjectInput{ Bucket: aws.String(s.bucket), Key: aws.String(normalized), - }) + } + if !s.disableChecksumHeaders { + input.ChecksumMode = types.ChecksumModeEnabled + } + out, err := s.client.GetObject(ctx, input) if err != nil { if isS3NotFound(err) { return nil, ObjectInfo{}, errors.Wrapf(ErrObjectNotFound, "object %s", normalized) @@ -526,10 +530,14 @@ func (s *S3Store) HeadObject(ctx context.Context, key string) (ObjectInfo, bool, if err != nil { return ObjectInfo{}, false, err } - out, err := s.client.HeadObject(ctx, &s3.HeadObjectInput{ + input := &s3.HeadObjectInput{ Bucket: aws.String(s.bucket), Key: aws.String(normalized), - }) + } + if !s.disableChecksumHeaders { + input.ChecksumMode = types.ChecksumModeEnabled + } + out, err := s.client.HeadObject(ctx, input) if err != nil { if isS3NotFound(err) { return ObjectInfo{}, false, nil diff --git a/internal/snapshotoffload/s3_store_test.go b/internal/snapshotoffload/s3_store_test.go index 31fb2eeb5..010f65981 100644 --- a/internal/snapshotoffload/s3_store_test.go +++ b/internal/snapshotoffload/s3_store_test.go @@ -72,11 +72,13 @@ func TestS3StorePutHeadGetPreservesIntegrityMetadata(t *testing.T) { require.True(t, ok) require.Equal(t, int64(len(body)), head.Size) require.Equal(t, sha, head.SHA256) + require.Equal(t, types.ChecksumModeEnabled, fake.lastHeadObjectChecksumMode()) reader, gotInfo, err := store.GetObject(ctx, "snapshots/body.fsm") require.NoError(t, err) defer func() { require.NoError(t, reader.Close()) }() require.Equal(t, sha, gotInfo.SHA256) + require.Equal(t, types.ChecksumModeEnabled, fake.lastGetObjectChecksumMode()) gotBody, err := io.ReadAll(reader) require.NoError(t, err) require.Equal(t, body, gotBody) @@ -286,10 +288,35 @@ func TestS3ObjectSHA256RejectsMetadataFullObjectChecksumMismatch(t *testing.T) { require.ErrorIs(t, err, ErrIntegrity) } -func TestMultipartPartSizeAllowsS3PartCapacity(t *testing.T) { +func TestS3StoreHeadAndGetRequestChecksumModeForFullObjectValidation(t *testing.T) { + ctx := context.Background() + fake := newFakeS3Client() + store := newTestS3Store(t, fake) + key := "snapshots/checksum-mismatch.fsm" + metadataSHA := hexSHA256Bytes([]byte("metadata-body")) + fullObjectChecksum, err := sha256HexToBase64(hexSHA256Bytes([]byte("full-object-body"))) + require.NoError(t, err) + fake.putRawObject("backup-bucket", key, []byte("body"), map[string]string{ + s3MetadataSHA256: metadataSHA, + }, aws.String(fullObjectChecksum), types.ChecksumTypeFullObject) + + _, ok, err := store.HeadObject(ctx, key) + require.ErrorIs(t, err, ErrIntegrity) + require.False(t, ok) + require.Equal(t, types.ChecksumModeEnabled, fake.lastHeadObjectChecksumMode()) + + body, _, err := store.GetObject(ctx, key) + require.ErrorIs(t, err, ErrIntegrity) + require.Nil(t, body) + require.Equal(t, types.ChecksumModeEnabled, fake.lastGetObjectChecksumMode()) +} + +func TestMultipartPartSizeCapsAtS3ObjectLimit(t *testing.T) { partSize, err := multipartPartSize(s3MaxObjectBytes, s3DefaultMultipartPart) require.NoError(t, err) - require.Equal(t, s3MaxMultipartPart, partSize) + expectedPartSize := max(s3DefaultMultipartPart, (s3MaxObjectBytes+s3MaxMultipartParts-1)/s3MaxMultipartParts) + require.Equal(t, expectedPartSize, partSize) + require.LessOrEqual(t, partSize, s3MaxMultipartPart) _, err = multipartPartSize(s3MaxObjectBytes+1, s3DefaultMultipartPart) require.ErrorIs(t, err, ErrInvalidOptions) @@ -339,6 +366,8 @@ type fakeS3Client struct { lastMultipartCreate types.ChecksumAlgorithm lastUploadPart types.ChecksumAlgorithm lastUploadPartSHA *string + lastHeadChecksumMode types.ChecksumMode + lastGetChecksumMode types.ChecksumMode completedChecksums int attempts int gets int @@ -463,6 +492,7 @@ func (c *fakeS3Client) PutObject(_ context.Context, input *s3.PutObjectInput, _ func (c *fakeS3Client) HeadObject(_ context.Context, input *s3.HeadObjectInput, _ ...func(*s3.Options)) (*s3.HeadObjectOutput, error) { c.mu.Lock() defer c.mu.Unlock() + c.lastHeadChecksumMode = input.ChecksumMode obj, ok := c.objects[fakeS3ClientKey(input.Bucket, input.Key)] if !ok { return nil, &types.NotFound{} @@ -471,11 +501,12 @@ func (c *fakeS3Client) HeadObject(_ context.Context, input *s3.HeadObjectInput, for k, v := range obj.metadata { metadata[k] = v } + checksum, checksumType := fakeS3ChecksumForMode(input.ChecksumMode, obj) return &s3.HeadObjectOutput{ ContentLength: aws.Int64(int64(len(obj.body))), Metadata: metadata, - ChecksumSHA256: obj.checksum, - ChecksumType: obj.checksumType, + ChecksumSHA256: checksum, + ChecksumType: checksumType, ServerSideEncryption: obj.serverSideEncryption, SSEKMSKeyId: obj.kmsKeyID, }, nil @@ -485,6 +516,7 @@ func (c *fakeS3Client) GetObject(_ context.Context, input *s3.GetObjectInput, _ c.mu.Lock() defer c.mu.Unlock() c.gets++ + c.lastGetChecksumMode = input.ChecksumMode obj, ok := c.objects[fakeS3ClientKey(input.Bucket, input.Key)] if !ok { return nil, &types.NotFound{} @@ -493,12 +525,13 @@ func (c *fakeS3Client) GetObject(_ context.Context, input *s3.GetObjectInput, _ for k, v := range obj.metadata { metadata[k] = v } + checksum, checksumType := fakeS3ChecksumForMode(input.ChecksumMode, obj) return &s3.GetObjectOutput{ Body: io.NopCloser(bytes.NewReader(obj.body)), ContentLength: aws.Int64(int64(len(obj.body))), Metadata: metadata, - ChecksumSHA256: obj.checksum, - ChecksumType: obj.checksumType, + ChecksumSHA256: checksum, + ChecksumType: checksumType, ServerSideEncryption: obj.serverSideEncryption, SSEKMSKeyId: obj.kmsKeyID, }, nil @@ -623,6 +656,18 @@ func (c *fakeS3Client) lastUploadPartSHA256() *string { return cloneStringPtr(c.lastUploadPartSHA) } +func (c *fakeS3Client) lastHeadObjectChecksumMode() types.ChecksumMode { + c.mu.Lock() + defer c.mu.Unlock() + return c.lastHeadChecksumMode +} + +func (c *fakeS3Client) lastGetObjectChecksumMode() types.ChecksumMode { + c.mu.Lock() + defer c.mu.Unlock() + return c.lastGetChecksumMode +} + func (c *fakeS3Client) completedPartChecksums() int { c.mu.Lock() defer c.mu.Unlock() @@ -674,6 +719,13 @@ func s3ChecksumTypeForSHA(checksum *string) types.ChecksumType { return types.ChecksumTypeFullObject } +func fakeS3ChecksumForMode(mode types.ChecksumMode, obj fakeS3Object) (*string, types.ChecksumType) { + if mode != types.ChecksumModeEnabled { + return nil, "" + } + return obj.checksum, obj.checksumType +} + func fakeS3ClientKey(bucket *string, key *string) string { return aws.ToString(bucket) + "/" + aws.ToString(key) }