go-omni/storage/upload_test.go

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", "<html>hello</html>", "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")
	}
}