diff --git a/ioext/sectionwriter.go b/ioext/sectionwriter.go new file mode 100644 index 0000000..d11841c --- /dev/null +++ b/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 } diff --git a/ioext/sectionwriter_test.go b/ioext/sectionwriter_test.go new file mode 100644 index 0000000..1c9a408 --- /dev/null +++ b/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)) + } +}