add sectionwriter

ioext/sectionwriter.go

@@ -0,0 +1,69 @@
+package ioext
+
+import (
+	"errors"
+	"io"
+)
+
+var ErrExceedLimit = errors.New("write exceed limit")
+
+func NewSectionWriter(w io.WriterAt, off int64, n int64) *SectionWriter {
+	return &SectionWriter{w, off, off, off + n}
+}
+
+type SectionWriter struct {
+	w     io.WriterAt
+	base  int64
+	off   int64
+	limit int64
+}
+
+func (s *SectionWriter) Write(p []byte) (n int, err error) {
+	if s.off >= s.limit {
+		return 0, ErrExceedLimit
+	}
+
+	if max := s.limit - s.off; int64(len(p)) > max {
+		return 0, ErrExceedLimit
+	}
+
+	n, err = s.w.WriteAt(p, s.off)
+	s.off += int64(n)
+	return
+}
+
+var errWhence = errors.New("Seek: invalid whence")
+var errOffset = errors.New("Seek: invalid offset")
+
+func (s *SectionWriter) Seek(offset int64, whence int) (int64, error) {
+	switch whence {
+	default:
+		return 0, errWhence
+	case 0:
+		offset += s.base
+	case 1:
+		offset += s.off
+	case 2:
+		offset += s.limit
+	}
+	if offset < s.base {
+		return 0, errOffset
+	}
+	s.off = offset
+	return offset - s.base, nil
+}
+
+func (s *SectionWriter) WriteAt(p []byte, off int64) (n int, err error) {
+	if off < 0 || off >= s.limit-s.base {
+		return 0, errOffset
+	}
+	off += s.base
+	if max := s.limit - off; int64(len(p)) > max {
+		return 0, ErrExceedLimit
+	}
+
+	return s.w.WriteAt(p, off)
+}
+
+// Size returns the size of the section in bytes.
+func (s *SectionWriter) Size() int64 { return s.limit - s.base }

ioext/sectionwriter_test.go

@@ -0,0 +1,56 @@
+package ioext
+
+import (
+	"io/ioutil"
+	"os"
+	"testing"
+)
+
+func TestSectionWriter(t *testing.T) {
+	f, err := ioutil.TempFile(".", "test_")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	defer func() {
+		n := f.Name()
+		f.Close()
+		os.Remove(n)
+	}()
+
+	f.Truncate(3)
+
+	rw := NewSectionWriter(f, 0, 1)
+
+	_, err = rw.Write([]byte{'1'})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	_, err = rw.Write([]byte{'1'})
+	if err == nil {
+		t.Fatal("must err")
+	}
+
+	rw = NewSectionWriter(f, 1, 2)
+
+	_, err = rw.Write([]byte{'2', '3', '4'})
+	if err == nil {
+		t.Fatal("must err")
+	}
+
+	_, err = rw.Write([]byte{'2', '3'})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	buf := make([]byte, 3)
+	_, err = f.ReadAt(buf, 0)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if string(buf) != "123" {
+		t.Fatal(string(buf))
+	}
+}