package storage_test import ( "bytes" "context" "errors" "io" "net/http" "net/http/httptest" "net/url" "strings" "testing" "time" "git.yuntuo.hk/k2safe/go-omni/storage" "git.yuntuo.hk/k2safe/go-omni/storage/local" ) func TestUploadPolicyActualBytes(t *testing.T) { ctx := context.Background() files, _ := local.New(t.TempDir(), 1024) uploader, err := storage.NewUploader(files, storage.UploadPolicy{MaxBytes: 64, ContentTypes: []string{"text/plain"}}) if err != nil { t.Fatal(err) } for _, tc := range []struct { name, body, contentType string size int64 want error }{ {"valid", "hello", "text/plain; charset=utf-8", 5, nil}, {"tooLong", "hello!", "text/plain", 5, storage.ErrSizeMismatch}, {"tooShort", "hello", "text/plain", 6, storage.ErrSizeMismatch}, {"emptyMismatch", "x", "text/plain", 0, storage.ErrSizeMismatch}, {"overLimit", "hello", "text/plain", 65, storage.ErrTooLarge}, {"declaredType", "hello", "image/png", 5, storage.ErrContentType}, {"forgedType", "hello", "text/plain", 18, storage.ErrContentType}, } { t.Run(tc.name, func(t *testing.T) { _, err := uploader.Put(ctx, tc.name, strings.NewReader(tc.body), tc.size, tc.contentType) if !errors.Is(err, tc.want) { t.Fatalf("%v want %v", err, tc.want) } object, getErr := files.Get(ctx, tc.name) if getErr == nil { object.Body.Close() } if tc.want != nil && !errors.Is(getErr, storage.ErrNotFound) { t.Fatal("rejected object committed", getErr) } }) } canceled, cancel := context.WithCancel(ctx) cancel() if _, err := uploader.Put(canceled, "cancel", strings.NewReader("x"), 1, "text/plain"); !errors.Is(err, context.Canceled) { t.Fatal(err) } if _, err := uploader.CreateMultipart(ctx, "local", "text/plain"); !errors.Is(err, storage.ErrUnsupported) { t.Fatal(err) } if _, err := uploader.PresignPut(ctx, "local", 1, "text/plain", time.Minute); !errors.Is(err, storage.ErrUnsupported) { t.Fatal(err) } } func TestSignedUploadConstraintsAndAtomicity(t *testing.T) { files, _ := local.New(t.TempDir(), 1024) signer, err := storage.NewUploadSigner(files, "https://files.example/upload", bytes.Repeat([]byte{1}, 32), storage.UploadPolicy{MaxBytes: 64, ContentTypes: []string{"text/plain"}}) if err != nil { t.Fatal(err) } link, err := signer.PresignPut(context.Background(), "file", 5, "text/plain", time.Minute) if err != nil { t.Fatal(err) } for _, tc := range []struct { name, body string mutate func(*http.Request) status int }{ {"success", "hello", nil, 201}, {"key", "hello", func(r *http.Request) { q := r.URL.Query(); q.Set("key", "other"); r.URL.RawQuery = q.Encode() }, 403}, {"expiry", "hello", func(r *http.Request) { q := r.URL.Query(); q.Set("expires", "1"); r.URL.RawQuery = q.Encode() }, 403}, {"type", "hello", func(r *http.Request) { r.Header.Set("Content-Type", "image/png") }, 400}, {"length", "hello", func(r *http.Request) { r.ContentLength = 6 }, 400}, {"extraBytes", "longer", func(r *http.Request) { r.ContentLength = 5 }, 413}, {"shortBody", "no", func(r *http.Request) { r.ContentLength = 5 }, 400}, {"duplicateKey", "hello", func(r *http.Request) { q := r.URL.Query(); q.Add("key", "file"); r.URL.RawQuery = q.Encode() }, 403}, } { t.Run(tc.name, func(t *testing.T) { req := httptest.NewRequest("PUT", link.URL, strings.NewReader(tc.body)) req.Header = link.Headers.Clone() if tc.mutate != nil { tc.mutate(req) } w := httptest.NewRecorder() signer.ServeHTTP(w, req) if w.Code != tc.status { t.Fatalf("status=%d want %d, %s", w.Code, tc.status, w.Body.String()) } }) } object, err := files.Get(context.Background(), "file") if err != nil { t.Fatal(err) } data, _ := io.ReadAll(object.Body) object.Body.Close() if string(data) != "hello" { t.Fatal("failed upload overwrote good object") } download, _ := signer.PresignGet(context.Background(), "file", time.Minute) req := httptest.NewRequest("PUT", download.String(), strings.NewReader("hello")) req.Header = link.Headers.Clone() w := httptest.NewRecorder() signer.ServeHTTP(w, req) if w.Code != 403 { t.Fatal("GET signature reused for PUT") } // Original read-only constructor must never accept signed uploads. readOnly, _ := storage.NewSigner(files, "https://files.example/upload", bytes.Repeat([]byte{1}, 32)) if _, err = readOnly.PresignPut(context.Background(), "file", 1, "text/plain", time.Minute); !errors.Is(err, storage.ErrUnsupported) { t.Fatal(err) } u, _ := url.Parse(link.URL) if u.Query().Get("signature") == "" { t.Fatal("unsigned URL") } }