Merge pull request #482 from projectdiscovery/introduce_safewriter

introduce safewriter

io/io.go

@@ -0,0 +1,38 @@
+package ioutil
+
+import (
+	"errors"
+	"io"
+	"sync"
+)
+
+// SafeWriter is a thread-safe wrapper for io.Writer
+type SafeWriter struct {
+	writer io.Writer   // The underlying writer
+	mutex  *sync.Mutex // Mutex for ensuring thread-safety
+}
+
+// NewSafeWriter creates and returns a new SafeWriter
+func NewSafeWriter(writer io.Writer) (*SafeWriter, error) {
+	// Check if the provided writer is nil
+	if writer == nil {
+		return nil, errors.New("writer is nil")
+	}
+
+	safeWriter := &SafeWriter{
+		writer: writer,
+		mutex:  &sync.Mutex{},
+	}
+	return safeWriter, nil
+}
+
+// Write implements the io.Writer interface in a thread-safe manner
+func (sw *SafeWriter) Write(p []byte) (n int, err error) {
+	sw.mutex.Lock()
+	defer sw.mutex.Unlock()
+
+	if sw.writer == nil {
+		return 0, io.ErrClosedPipe
+	}
+	return sw.writer.Write(p)
+}

io/io_test.go

@@ -0,0 +1,25 @@
+package ioutil
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestSafeWriter(t *testing.T) {
+	t.Run("success", func(t *testing.T) {
+		var sb strings.Builder
+		sw, err := NewSafeWriter(&sb)
+		require.Nil(t, err)
+		_, err = sw.Write([]byte("test"))
+		require.Nil(t, err)
+		require.Equal(t, "test", sb.String())
+	})
+
+	t.Run("failure", func(t *testing.T) {
+		sw, err := NewSafeWriter(nil)
+		require.NotNil(t, err)
+		require.Nil(t, sw)
+	})
+}