introduce safewriter

projectdiscovery · Jul 30, 2024 · c6d7a0b · c6d7a0b
1 parent a336453
commit c6d7a0b
Show file tree

Hide file tree

Showing 2 changed files with 54 additions and 0 deletions.
diff --git a/io/io.go b/io/io.go
@@ -0,0 +1,26 @@
+package ioutil
+
+import (
+	"io"
+	"sync"
+)
+
+type SafeWriter struct {
+	writer io.Writer
+	mutex  sync.Mutex
+}
+
+func NewSafeWriter(writer io.Writer) *SafeWriter {
+	return &SafeWriter{
+		writer: writer,
+	}
+}
+
+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)
+}
diff --git a/io/io_test.go b/io/io_test.go
@@ -0,0 +1,28 @@
+package ioutil
+
+import (
+	"strings"
+	"testing"
+)
+
+func TestSafeWriter(t *testing.T) {
+	t.Run("success", func(t *testing.T) {
+		var sb strings.Builder
+		sw := NewSafeWriter(&sb)
+		_, err := sw.Write([]byte("test"))
+		if err != nil {
+			t.Fatalf("expected no error, got %v", err)
+		}
+		if sb.String() != "test" {
+			t.Fatalf("expected 'test', got '%s'", sb.String())
+		}
+	})
+
+	t.Run("failure", func(t *testing.T) {
+		sw := NewSafeWriter(nil)
+		_, err := sw.Write([]byte("test"))
+		if err == nil {
+			t.Fatalf("expected error, got nil")
+		}
+	})
+}