forked from TrueCloudLab/restic
9aa2eff384
This enables the backends to request the calculation of a backend-specific hash. For the currently supported backends this will always be MD5. The hash calculation happens as early as possible, for pack files this is during assembly of the pack file. That way the hash would even capture corruptions of the temporary pack file on disk.
109 lines
2.8 KiB
Go
109 lines
2.8 KiB
Go
package limiter
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"crypto/rand"
|
|
"fmt"
|
|
"io"
|
|
"testing"
|
|
|
|
"github.com/restic/restic/internal/mock"
|
|
"github.com/restic/restic/internal/restic"
|
|
rtest "github.com/restic/restic/internal/test"
|
|
)
|
|
|
|
func randomBytes(t *testing.T, size int) []byte {
|
|
data := make([]byte, size)
|
|
_, err := io.ReadFull(rand.Reader, data)
|
|
rtest.OK(t, err)
|
|
return data
|
|
}
|
|
|
|
func TestLimitBackendSave(t *testing.T) {
|
|
testHandle := restic.Handle{Type: restic.PackFile, Name: "test"}
|
|
data := randomBytes(t, 1234)
|
|
|
|
be := mock.NewBackend()
|
|
be.SaveFn = func(ctx context.Context, h restic.Handle, rd restic.RewindReader) error {
|
|
buf := new(bytes.Buffer)
|
|
_, err := io.Copy(buf, rd)
|
|
if err != nil {
|
|
return nil
|
|
}
|
|
if !bytes.Equal(data, buf.Bytes()) {
|
|
return fmt.Errorf("data mismatch")
|
|
}
|
|
return nil
|
|
}
|
|
limiter := NewStaticLimiter(42*1024, 42*1024)
|
|
limbe := LimitBackend(be, limiter)
|
|
|
|
rd := restic.NewByteReader(data, nil)
|
|
err := limbe.Save(context.TODO(), testHandle, rd)
|
|
rtest.OK(t, err)
|
|
}
|
|
|
|
type tracedReadWriteToCloser struct {
|
|
io.Reader
|
|
io.WriterTo
|
|
Traced bool
|
|
}
|
|
|
|
func newTracedReadWriteToCloser(rd *bytes.Reader) *tracedReadWriteToCloser {
|
|
return &tracedReadWriteToCloser{Reader: rd, WriterTo: rd}
|
|
}
|
|
|
|
func (r *tracedReadWriteToCloser) WriteTo(w io.Writer) (n int64, err error) {
|
|
r.Traced = true
|
|
return r.WriterTo.WriteTo(w)
|
|
}
|
|
|
|
func (r *tracedReadWriteToCloser) Close() error {
|
|
return nil
|
|
}
|
|
|
|
func TestLimitBackendLoad(t *testing.T) {
|
|
testHandle := restic.Handle{Type: restic.PackFile, Name: "test"}
|
|
data := randomBytes(t, 1234)
|
|
|
|
for _, test := range []struct {
|
|
innerWriteTo, outerWriteTo bool
|
|
}{{false, false}, {false, true}, {true, false}, {true, true}} {
|
|
be := mock.NewBackend()
|
|
src := newTracedReadWriteToCloser(bytes.NewReader(data))
|
|
be.OpenReaderFn = func(ctx context.Context, h restic.Handle, length int, offset int64) (io.ReadCloser, error) {
|
|
if length != 0 || offset != 0 {
|
|
return nil, fmt.Errorf("Not supported")
|
|
}
|
|
// test both code paths in WriteTo of limitedReadCloser
|
|
if test.innerWriteTo {
|
|
return src, nil
|
|
}
|
|
return newTracedReadCloser(src), nil
|
|
}
|
|
limiter := NewStaticLimiter(42*1024, 42*1024)
|
|
limbe := LimitBackend(be, limiter)
|
|
|
|
err := limbe.Load(context.TODO(), testHandle, 0, 0, func(rd io.Reader) error {
|
|
dataRead := new(bytes.Buffer)
|
|
// test both Read and WriteTo
|
|
if !test.outerWriteTo {
|
|
rd = newTracedReadCloser(rd)
|
|
}
|
|
_, err := io.Copy(dataRead, rd)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if !bytes.Equal(data, dataRead.Bytes()) {
|
|
return fmt.Errorf("read broken data")
|
|
}
|
|
|
|
return nil
|
|
})
|
|
rtest.OK(t, err)
|
|
rtest.Assert(t, src.Traced == (test.innerWriteTo && test.outerWriteTo),
|
|
"unexpected/missing writeTo call innerWriteTo %v outerWriteTo %v",
|
|
test.innerWriteTo, test.outerWriteTo)
|
|
}
|
|
}
|