From 6160747cba6801973843bbd37c64e46353c48cf8 Mon Sep 17 00:00:00 2001 From: Adrien CABARBAYE Date: Tue, 26 Aug 2025 10:48:36 +0100 Subject: [PATCH] :sparkles: `[safeio]` Add the ability to perform safe concatenations like `cat` --- changes/20250826104813.feature | 1 + utils/safeio/copy.go | 14 ++++- utils/safeio/copy_test.go | 96 +++++++++++++++++++++++++++++++--- utils/safeio/read.go | 8 +++ 4 files changed, 111 insertions(+), 8 deletions(-) create mode 100644 changes/20250826104813.feature diff --git a/changes/20250826104813.feature b/changes/20250826104813.feature new file mode 100644 index 0000000000..69b47dcd31 --- /dev/null +++ b/changes/20250826104813.feature @@ -0,0 +1 @@ +:sparkles: `[safeio]` Add the ability to perform safe concatenations like `cat` diff --git a/utils/safeio/copy.go b/utils/safeio/copy.go index 41e7e207ef..0f2f4038e7 100644 --- a/utils/safeio/copy.go +++ b/utils/safeio/copy.go @@ -8,16 +8,26 @@ import ( "github.com/ARM-software/golang-utils/utils/parallelisation" ) -// CopyDataWithContext copies from src to dst similarly to io.Copy but with context control to stop when asked to. +// CopyDataWithContext copies from src to dst similarly to io.Copy but with context control to stop when asked. func CopyDataWithContext(ctx context.Context, src io.Reader, dst io.Writer) (copied int64, err error) { return copyDataWithContext(ctx, src, dst, io.Copy) } -// CopyNWithContext copies n bytes from src to dst similarly to io.CopyN but with context control to stop when asked to. +// CopyNWithContext copies n bytes from src to dst similarly to io.CopyN but with context control to stop when asked. func CopyNWithContext(ctx context.Context, src io.Reader, dst io.Writer, n int64) (copied int64, err error) { return copyDataWithContext(ctx, src, dst, func(dst io.Writer, src io.Reader) (int64, error) { return io.CopyN(dst, src, n) }) } +// CatN concatenates n bytes from multiple sources to dst. It is intended to provide functionality quite similar to `cat` posix command but with context control. +func CatN(ctx context.Context, dst io.Writer, n int64, src ...io.Reader) (copied int64, err error) { + return CopyNWithContext(ctx, NewContextualMultipleReader(ctx, src...), dst, n) +} + +// Cat concatenates bytes from multiple sources to dst. It is intended to provide functionality quite similar to `cat` posix command but with context control. +func Cat(ctx context.Context, dst io.Writer, src ...io.Reader) (copied int64, err error) { + return CopyDataWithContext(ctx, NewContextualMultipleReader(ctx, src...), dst) +} + func copyDataWithContext(ctx context.Context, src io.Reader, dst io.Writer, copyFunc func(io.Writer, io.Reader) (int64, error)) (copied int64, err error) { err = parallelisation.DetermineContextError(ctx) if err != nil { diff --git a/utils/safeio/copy_test.go b/utils/safeio/copy_test.go index ec71cd88d4..ccae39860d 100644 --- a/utils/safeio/copy_test.go +++ b/utils/safeio/copy_test.go @@ -11,6 +11,7 @@ import ( "github.com/ARM-software/golang-utils/utils/commonerrors" "github.com/ARM-software/golang-utils/utils/commonerrors/errortest" + "github.com/ARM-software/golang-utils/utils/safecast" ) func TestCopyDataWithContext(t *testing.T) { @@ -23,7 +24,7 @@ func TestCopyDataWithContext(t *testing.T) { n2, err := CopyDataWithContext(context.Background(), &buf1, &buf2) require.NoError(t, err) require.NotZero(t, n2) - assert.Equal(t, int64(len(text)), n2) + assert.Equal(t, safecast.ToInt64(len(text)), n2) assert.Equal(t, text, buf2.String()) ctx, cancel := context.WithCancel(context.Background()) @@ -49,10 +50,10 @@ func TestCopyNWithContext(t *testing.T) { require.NoError(t, err) require.NotZero(t, n) assert.Equal(t, len(text), n) - n2, err := CopyNWithContext(context.Background(), &buf1, &buf2, int64(len(text))) + n2, err := CopyNWithContext(context.Background(), &buf1, &buf2, safecast.ToInt64(len(text))) require.NoError(t, err) require.NotZero(t, n2) - assert.Equal(t, int64(len(text)), n2) + assert.Equal(t, safecast.ToInt64(len(text)), n2) assert.Equal(t, text, buf2.String()) ctx, cancel := context.WithCancel(context.Background()) @@ -64,7 +65,7 @@ func TestCopyNWithContext(t *testing.T) { assert.Equal(t, len(text), n) cancel() - n2, err = CopyNWithContext(ctx, &buf1, &buf2, int64(len(text))) + n2, err = CopyNWithContext(ctx, &buf1, &buf2, safecast.ToInt64(len(text))) require.Error(t, err) errortest.AssertError(t, err, commonerrors.ErrCancelled) assert.Zero(t, n2) @@ -74,8 +75,91 @@ func TestCopyNWithContext(t *testing.T) { require.NoError(t, err) require.NotZero(t, n) assert.Equal(t, len(text), n) - n2, err = CopyNWithContext(context.Background(), &buf1, &buf2, int64(len(text)-1)) + n2, err = CopyNWithContext(context.Background(), &buf1, &buf2, safecast.ToInt64(len(text)-1)) require.NoError(t, err) require.NotZero(t, n2) - assert.Equal(t, int64(len(text)-1), n2) + assert.Equal(t, safecast.ToInt64(len(text)-1), n2) +} + +func TestCat(t *testing.T) { + var buf1, buf2, buf3 bytes.Buffer + text1 := faker.Sentence() + text2 := faker.Paragraph() + n, err := WriteString(context.Background(), &buf1, text1) + require.NoError(t, err) + require.NotZero(t, n) + assert.Equal(t, len(text1), n) + n, err = WriteString(context.Background(), &buf2, text2) + require.NoError(t, err) + require.NotZero(t, n) + assert.Equal(t, len(text2), n) + n3, err := Cat(context.Background(), &buf3, &buf1, &buf2) + require.NoError(t, err) + require.NotZero(t, n3) + assert.Equal(t, safecast.ToInt64(len(text1)+len(text2)), n3) + assert.Equal(t, text1+text2, buf3.String()) + + ctx, cancel := context.WithCancel(context.Background()) + buf1.Reset() + buf2.Reset() + buf3.Reset() + n, err = WriteString(context.Background(), &buf1, text1) + require.NoError(t, err) + require.NotZero(t, n) + assert.Equal(t, len(text1), n) + n, err = WriteString(context.Background(), &buf2, text2) + require.NoError(t, err) + require.NotZero(t, n) + assert.Equal(t, len(text2), n) + + cancel() + n3, err = Cat(ctx, &buf3, &buf1, &buf2) + require.Error(t, err) + errortest.AssertError(t, err, commonerrors.ErrCancelled) + assert.Zero(t, n3) + assert.Empty(t, buf3.String()) +} + +func TestCatN(t *testing.T) { + var buf1, buf2, buf3 bytes.Buffer + text1 := faker.Sentence() + text2 := faker.Paragraph() + n, err := WriteString(context.Background(), &buf1, text1) + require.NoError(t, err) + require.NotZero(t, n) + assert.Equal(t, len(text1), n) + n, err = WriteString(context.Background(), &buf2, text2) + require.NoError(t, err) + require.NotZero(t, n) + assert.Equal(t, len(text2), n) + n3, err := CatN(context.Background(), &buf3, safecast.ToInt64(len(text1)+len(text2)), &buf1, &buf2) + require.NoError(t, err) + require.NotZero(t, n3) + assert.Equal(t, safecast.ToInt64(len(text1)+len(text2)), n3) + assert.Equal(t, text1+text2, buf3.String()) + + ctx, cancel := context.WithCancel(context.Background()) + buf1.Reset() + buf2.Reset() + buf3.Reset() + n, err = WriteString(context.Background(), &buf1, text1) + require.NoError(t, err) + require.NotZero(t, n) + assert.Equal(t, len(text1), n) + n, err = WriteString(context.Background(), &buf2, text2) + require.NoError(t, err) + require.NotZero(t, n) + assert.Equal(t, len(text2), n) + + cancel() + n3, err = CatN(ctx, &buf3, safecast.ToInt64(len(text1)+len(text2)), &buf1, &buf2) + require.Error(t, err) + errortest.AssertError(t, err, commonerrors.ErrCancelled) + assert.Zero(t, n3) + assert.Empty(t, buf3.String()) + + n3, err = CatN(context.Background(), &buf3, safecast.ToInt64(len(text1)+1), &buf1, &buf2) + require.NoError(t, err) + require.NotZero(t, n3) + assert.Equal(t, safecast.ToInt64(len(text1)+1), n3) } diff --git a/utils/safeio/read.go b/utils/safeio/read.go index e7ef62db2e..8044617dfa 100644 --- a/utils/safeio/read.go +++ b/utils/safeio/read.go @@ -76,6 +76,14 @@ func NewContextualReader(ctx context.Context, reader io.Reader) io.Reader { return contextio.NewReader(ctx, reader) } +func NewContextualMultipleReader(ctx context.Context, reader ...io.Reader) io.Reader { + readers := make([]io.Reader, len(reader)) + for i := range reader { + readers[i] = NewContextualReader(ctx, reader[i]) + } + return io.MultiReader(readers...) +} + // NewContextualReaderFrom returns a io.ReaderFrom which is context aware. // Context state is checked BEFORE every Read, Write, Copy. func NewContextualReaderFrom(ctx context.Context, reader io.ReaderFrom) io.ReaderFrom {