220 lines
7.8 KiB
Go
220 lines
7.8 KiB
Go
package storage
|
|
|
|
import (
|
|
"context"
|
|
"io"
|
|
"mime"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"net/url"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func TestReadOptions(t *testing.T) {
|
|
for _, name := range []string{"", "report.pdf", "报告 2026.pdf", `a";b.pdf`, "100%.pdf"} {
|
|
h, err := (ReadOptions{Disposition: DispositionInline, FileName: name, ContentType: "Application/PDF"}).ResponseHeaders()
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
disposition, params, err := mime.ParseMediaType(h.Get("Content-Disposition"))
|
|
if err != nil || disposition != "inline" || params["filename"] != name || h.Get("Content-Type") != "application/pdf" || h.Get("Cache-Control") != "private, no-store" {
|
|
t.Fatal("invalid read headers", h, err)
|
|
}
|
|
}
|
|
h, err := (ReadOptions{}).ResponseHeaders()
|
|
if err != nil || h.Get("Content-Disposition") != "attachment" || h.Get("Content-Type") != "" {
|
|
t.Fatal(h, err)
|
|
}
|
|
for _, options := range []ReadOptions{
|
|
{Disposition: "form-data"}, {Disposition: "Inline"}, {Disposition: "inline\r\nX: y"},
|
|
{FileName: "x\r\nX: y"}, {FileName: "a\x00b"}, {FileName: "a\tb"}, {FileName: "bad\xff"}, {FileName: "a/b"}, {FileName: `a\b`}, {FileName: ".."}, {FileName: strings.Repeat("a", 1025)},
|
|
{ContentType: "bad"}, {ContentType: "image/*"}, {ContentType: "application/pdf\r\nX: y"}, {ContentType: "application/pdf\x00"}, {ContentType: strings.Repeat("x", 257)},
|
|
} {
|
|
if _, err := options.ResponseHeaders(); err == nil {
|
|
t.Fatal("invalid options accepted", options)
|
|
}
|
|
}
|
|
}
|
|
|
|
type readTestStore struct {
|
|
calls, bytes int
|
|
streaming bool
|
|
}
|
|
|
|
func (*readTestStore) Put(context.Context, string, io.Reader, int64, string) (Info, error) {
|
|
panic("unused")
|
|
}
|
|
func (*readTestStore) Delete(context.Context, string) error { panic("unused") }
|
|
func (s *readTestStore) Get(context.Context, string) (Object, error) {
|
|
s.calls++
|
|
body := &readTestBody{Reader: strings.NewReader("0123456789"), store: s}
|
|
var stream io.ReadCloser = body
|
|
if s.streaming {
|
|
stream = struct {
|
|
io.Reader
|
|
io.Closer
|
|
}{body, body}
|
|
}
|
|
return Object{Info: Info{Size: 10, ContentType: "application/pdf", ModifiedAt: time.Unix(1000, 0)}, Body: stream}, nil
|
|
}
|
|
|
|
type readTestBody struct {
|
|
*strings.Reader
|
|
store *readTestStore
|
|
}
|
|
|
|
func (b *readTestBody) Read(p []byte) (int, error) {
|
|
n, err := b.Reader.Read(p)
|
|
b.store.bytes += n
|
|
return n, err
|
|
}
|
|
func (*readTestBody) Close() error { return nil }
|
|
|
|
func TestReadSignatureAndExpiry(t *testing.T) {
|
|
store := &readTestStore{}
|
|
s, _ := NewSigner(store, "https://files.example/read", []byte(strings.Repeat("s", 32)))
|
|
now := time.Unix(1800000000, 0).UTC()
|
|
s.now = func() time.Time { return now }
|
|
link, err := s.PresignGetWithOptions(context.Background(), "folder/file", 30*time.Minute, ReadOptions{Disposition: DispositionInline, FileName: "报告.pdf", ContentType: "application/pdf"})
|
|
if err != nil || !link.ExpiresAt.Equal(now.Add(30*time.Minute)) {
|
|
t.Fatal("incorrect expiry", err)
|
|
}
|
|
serve := func(method, address string, headers http.Header) *httptest.ResponseRecorder {
|
|
r := httptest.NewRequest(method, address, nil)
|
|
if headers != nil {
|
|
r.Header = headers
|
|
}
|
|
w := httptest.NewRecorder()
|
|
s.ServeHTTP(w, r)
|
|
return w
|
|
}
|
|
for _, method := range []string{"GET", "HEAD"} {
|
|
w := serve(method, link.URL, nil)
|
|
if w.Code != 200 || w.Header().Get("Content-Type") != "application/pdf" || w.Header().Get("Cache-Control") != "private, no-store" {
|
|
t.Fatal(w)
|
|
}
|
|
if method == "HEAD" && w.Body.Len() != 0 {
|
|
t.Fatal("HEAD body")
|
|
}
|
|
}
|
|
old, err := s.PresignGet(context.Background(), "folder/file", 30*time.Minute)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if w := serve("GET", old.String(), nil); w.Code != 200 || w.Header().Get("Content-Disposition") != "attachment" {
|
|
t.Fatal("legacy download changed")
|
|
}
|
|
mutations := []func(url.Values){
|
|
func(q url.Values) { q.Set("key", "other") }, func(q url.Values) { q.Set("expires", "1800001900") }, func(q url.Values) { q.Set("signature", "bad") },
|
|
func(q url.Values) { q.Set("filename", "other.pdf") }, func(q url.Values) { q.Set("disposition", "attachment") }, func(q url.Values) { q.Set("contentType", "text/html") },
|
|
func(q url.Values) { q.Set("v", "3") }, func(q url.Values) { q.Del("v") }, func(q url.Values) {
|
|
for _, k := range []string{"v", "filename", "disposition", "contentType"} {
|
|
q.Del(k)
|
|
}
|
|
},
|
|
func(q url.Values) { q.Set("extra", "x") }, func(q url.Values) { q.Set("filename", "evil\r\nX: yes") },
|
|
}
|
|
for _, key := range []string{"key", "expires", "signature", "v", "disposition", "filename", "contentType"} {
|
|
key := key
|
|
mutations = append(mutations, func(q url.Values) { q.Add(key, q.Get(key)) })
|
|
}
|
|
for _, mutate := range mutations {
|
|
u, _ := url.Parse(link.URL)
|
|
q := u.Query()
|
|
mutate(q)
|
|
u.RawQuery = q.Encode()
|
|
before := store.calls
|
|
if w := serve("GET", u.String(), http.Header{"Range": {"bytes=0-1"}}); w.Code != 403 || store.calls != before {
|
|
t.Fatal("tampered link reached storage")
|
|
}
|
|
}
|
|
u := *old
|
|
q := u.Query()
|
|
q.Set("disposition", "inline")
|
|
u.RawQuery = q.Encode()
|
|
if serve("GET", u.String(), nil).Code != 403 {
|
|
t.Fatal("legacy download upgraded by query")
|
|
}
|
|
if serve("GET", link.URL+"&bad=%ZZ", nil).Code != 403 {
|
|
t.Fatal("malformed query accepted")
|
|
}
|
|
before := store.calls
|
|
now = link.ExpiresAt.Add(-time.Nanosecond)
|
|
if serve("GET", link.URL, nil).Code != 200 {
|
|
t.Fatal("early expiration")
|
|
}
|
|
before = store.calls
|
|
now = link.ExpiresAt
|
|
for _, method := range []string{"GET", "HEAD"} {
|
|
w := serve(method, link.URL, http.Header{"Range": {"bytes=0-1"}, "If-Modified-Since": {"Fri, 01 Jan 2100 00:00:00 GMT"}})
|
|
if w.Code != 403 || store.calls != before || w.Header().Get("Cache-Control") != "private, no-store" {
|
|
t.Fatal("expired link reached storage or cache")
|
|
}
|
|
if method == "HEAD" && w.Body.Len() != 0 {
|
|
t.Fatal("HEAD error has body")
|
|
}
|
|
}
|
|
for _, ttl := range []time.Duration{0, time.Second - time.Nanosecond, time.Hour + time.Nanosecond} {
|
|
if _, err := s.PresignGetWithOptions(context.Background(), "key", ttl, ReadOptions{}); err == nil {
|
|
t.Fatal("invalid TTL")
|
|
}
|
|
}
|
|
for _, ttl := range []time.Duration{time.Second, time.Hour} {
|
|
if _, err := s.PresignGetWithOptions(context.Background(), "key", ttl, ReadOptions{}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
cancel()
|
|
if _, err := s.PresignGetWithOptions(ctx, "key", time.Minute, ReadOptions{}); err != context.Canceled {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
|
|
func TestReadRangeStreamingAndCache(t *testing.T) {
|
|
for _, streaming := range []bool{false, true} {
|
|
store := &readTestStore{streaming: streaming}
|
|
s, _ := NewSigner(store, "https://files.example/read", []byte(strings.Repeat("s", 32)))
|
|
link, _ := s.PresignGetWithOptions(context.Background(), "key", 30*time.Minute, ReadOptions{})
|
|
for _, method := range []string{"GET", "HEAD"} {
|
|
before := store.bytes
|
|
r := httptest.NewRequest(method, link.URL, nil)
|
|
r.Header.Set("Range", "bytes=2-4")
|
|
w := httptest.NewRecorder()
|
|
s.ServeHTTP(w, r)
|
|
if method == "HEAD" {
|
|
if w.Code != 200 || w.Body.Len() != 0 || store.bytes != before {
|
|
t.Fatal("HEAD read body", w)
|
|
}
|
|
} else if streaming {
|
|
if w.Code != 200 || w.Body.String() != "0123456789" || w.Header().Get("Accept-Ranges") != "" {
|
|
t.Fatal("stream fallback", w)
|
|
}
|
|
} else {
|
|
if w.Code != 206 || w.Body.String() != "234" || store.bytes-before != 3 {
|
|
t.Fatal("range buffered full body", w)
|
|
}
|
|
}
|
|
if w.Header().Get("Cache-Control") != "private, no-store" {
|
|
t.Fatal("missing cache policy")
|
|
}
|
|
}
|
|
if !streaming {
|
|
for _, h := range []http.Header{{"Range": {"bytes=99-100"}}, {"If-Modified-Since": {"Fri, 01 Jan 2100 00:00:00 GMT"}}} {
|
|
r := httptest.NewRequest("GET", link.URL, nil)
|
|
r.Header = h
|
|
w := httptest.NewRecorder()
|
|
s.ServeHTTP(w, r)
|
|
if w.Code != 416 && w.Code != 304 {
|
|
t.Fatal(w.Code)
|
|
}
|
|
if w.Header().Get("Cache-Control") != "private, no-store" {
|
|
t.Fatal("cache policy lost on error/conditional")
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|