go-omni/contract/multipart_test.go

package contract

import (
	"bytes"
	"errors"
	apierrors "git.yuntuo.hk/k2safe/go-omni/errors"
	"io"
	"mime/multipart"
	"net/http"
	"net/http/httptest"
	"os"
	"strings"
	"testing"
)

type uploadInput struct {
	Query string `form:"query,optional"`
	Name  string `multipart:"name" binding:"required,max=10"`
	Meta  struct {
		Count int `json:"count"`
	} `multipartJSON:"meta"`
	Attachment File   `file:"attachment,maxBytes=12"`
	Extra      []File `file:"extra,optional,maxBytes=8,maxFiles=2"`
}
type uploadPart struct{ name, filename, value string }

func uploadRequest(t *testing.T, parts []uploadPart) (*http.Request, []byte) {
	t.Helper()
	var data bytes.Buffer
	w := multipart.NewWriter(&data)
	for _, part := range parts {
		if part.filename != "" {
			p, e := w.CreateFormFile(part.name, part.filename)
			if e != nil {
				t.Fatal(e)
			}
			_, _ = io.WriteString(p, part.value)
		} else {
			_ = w.WriteField(part.name, part.value)
		}
	}
	_ = w.Close()
	r := httptest.NewRequest("POST", "/upload?query=yes", bytes.NewReader(data.Bytes()))
	r.Header.Set("Content-Type", w.FormDataContentType())
	return r, data.Bytes()
}
func goodParts() []uploadPart {
	return []uploadPart{{"name", "", "合同"}, {"meta", "", `{"count":2}`}, {"attachment", "../contract.txt", "hello"}}
}
func assertUploadEmpty(t *testing.T, dir string) {
	t.Helper()
	files, err := os.ReadDir(dir)
	if err != nil || len(files) != 0 {
		t.Fatalf("temp files remain: %v %v", files, err)
	}
}
func TestMultipartBindingAndLifetime(t *testing.T) {
	b := MustNew[uploadInput]()
	dir := t.TempDir()
	r, _ := uploadRequest(t, append(goodParts(), uploadPart{"extra", "one", "first"}, uploadPart{"extra", "two", "second"}))
	var input uploadInput
	release, err := b.BindMultipart(r, nil, &input, MultipartOptions{TempDir: dir, Policy: RequestPolicy{UnknownQuery: "reject", UTF8: "reject"}})
	if err != nil {
		t.Fatal(err)
	}
	if input.Name != "合同" || input.Query != "yes" || input.Meta.Count != 2 || input.Attachment.Size != 5 || input.Attachment.Name != "contract.txt" || len(input.Extra) != 2 {
		t.Fatalf("%+v", input)
	}
	reader, err := input.Attachment.Open()
	if err != nil {
		t.Fatal(err)
	}
	data, err := io.ReadAll(reader)
	if err != nil || string(data) != "hello" {
		t.Fatal(string(data), err)
	}
	// An unclosed reader is owned by the upload cleanup as well.
	if err = release(); err != nil {
		t.Fatal(err)
	}
	if _, err = reader.Read(make([]byte, 1)); err == nil {
		t.Fatal("reader remains open")
	}
	if _, err = input.Attachment.Open(); err == nil {
		t.Fatal("expired upload reopened")
	}
	if err = release(); err != nil {
		t.Fatal(err)
	}
	assertUploadEmpty(t, dir)
	if err = b.Bind(r, nil, &input); err == nil {
		t.Fatal("ordinary binding accepted multipart fields")
	}
}
func TestMultipartFailuresAreTransactional(t *testing.T) {
	b := MustNew[uploadInput]()
	for _, tt := range []struct {
		name    string
		parts   []uploadPart
		options MultipartOptions
		status  int
	}{
		{"missing", goodParts()[1:], MultipartOptions{}, 400},
		{"unknown", append(goodParts(), uploadPart{"other", "", "x"}), MultipartOptions{}, 400},
		{"duplicate", append(goodParts(), uploadPart{"name", "", "x"}), MultipartOptions{}, 400},
		{"jsonDuplicate", []uploadPart{{"name", "", "n"}, {"attachment", "f", "x"}, {"meta", "", `{"count":1,"count":2}`}}, MultipartOptions{}, 400},
		{"badType", []uploadPart{{"name", "", "n"}, {"attachment", "f", "x"}, {"meta", "", `{"count":"x"}`}}, MultipartOptions{}, 400},
		{"badUTF8", append(goodParts()[:1], uploadPart{"meta", "", "{\"count\":\"\xff\"}"}, goodParts()[2]), MultipartOptions{Policy: RequestPolicy{UTF8: "reject"}}, 400},
		{"fileSize", []uploadPart{{"attachment", "x", strings.Repeat("x", 13)}}, MultipartOptions{}, 413},
		{"fieldSize", goodParts(), MultipartOptions{MaxFieldBytes: 2}, 413},
		{"partCount", goodParts(), MultipartOptions{MaxParts: 2}, 413},
		{"total", goodParts(), MultipartOptions{MaxBodyBytes: 32}, 413},
		{"fileKind", []uploadPart{{"attachment", "", "text"}}, MultipartOptions{}, 400},
		{"extraCount", append(goodParts(), uploadPart{"extra", "a", "a"}, uploadPart{"extra", "b", "b"}, uploadPart{"extra", "c", "c"}), MultipartOptions{}, 400},
	} {
		t.Run(tt.name, func(t *testing.T) {
			dir := t.TempDir()
			tt.options.TempDir = dir
			r, _ := uploadRequest(t, tt.parts)
			r.ContentLength = -1
			input := uploadInput{Name: "untouched"}
			release, err := b.BindMultipart(r, nil, &input, tt.options)
			var api *apierrors.APIError
			if !errors.As(err, &api) || api.Status != tt.status {
				t.Fatalf("error %v", err)
			}
			if release != nil || input.Name != "untouched" {
				t.Fatal("partial binding published")
			}
			assertUploadEmpty(t, dir)
		})
	}
}
func TestMultipartTruncationAndEpilogueLimit(t *testing.T) {
	b := MustNew[uploadInput]()
	for _, epilogue := range []bool{false, true} {
		dir := t.TempDir()
		r, original := uploadRequest(t, goodParts())
		data := original[:len(original)-8]
		limit := int64(len(original) + 16)
		if epilogue {
			data = append(original, bytes.Repeat([]byte("x"), 64)...)
		}
		r.Body = io.NopCloser(bytes.NewReader(data))
		r.ContentLength = -1
		var input uploadInput
		if _, err := b.BindMultipart(r, nil, &input, MultipartOptions{TempDir: dir, MaxBodyBytes: limit}); err == nil {
			t.Fatal("truncation or epilogue bypass accepted")
		}
		assertUploadEmpty(t, dir)
	}
}