package objfile import ( "compress/zlib" "errors" "io" "strconv" "gopkg.in/src-d/go-git.v4/plumbing" ) var ( ErrOverflow = errors.New("objfile: declared data length exceeded (overflow)") ) // Writer writes and encodes data in compressed objfile format to a provided // io.Writerº. Close should be called when finished with the Writer. Close will // not close the underlying io.Writer. type Writer struct { raw io.Writer zlib io.WriteCloser hasher plumbing.Hasher multi io.Writer closed bool pending int64 // number of unwritten bytes } // NewWriter returns a new Writer writing to w. // // The returned Writer implements io.WriteCloser. Close should be called when // finished with the Writer. Close will not close the underlying io.Writer. func NewWriter(w io.Writer) *Writer { return &Writer{ raw: w, zlib: zlib.NewWriter(w), } } // WriteHeader writes the type and the size and prepares to accept the object's // contents. If an invalid t is provided, plumbing.ErrInvalidType is returned. If a // negative size is provided, ErrNegativeSize is returned. func (w *Writer) WriteHeader(t plumbing.ObjectType, size int64) error { if !t.Valid() { return plumbing.ErrInvalidType } if size < 0 { return ErrNegativeSize } b := t.Bytes() b = append(b, ' ') b = append(b, []byte(strconv.FormatInt(size, 10))...) b = append(b, 0) defer w.prepareForWrite(t, size) _, err := w.zlib.Write(b) return err } func (w *Writer) prepareForWrite(t plumbing.ObjectType, size int64) { w.pending = size w.hasher = plumbing.NewHasher(t, size) w.multi = io.MultiWriter(w.zlib, w.hasher) } // Write writes the object's contents. Write returns the error ErrOverflow if // more than size bytes are written after WriteHeader. func (w *Writer) Write(p []byte) (n int, err error) { if w.closed { return 0, ErrClosed } overwrite := false if int64(len(p)) > w.pending { p = p[0:w.pending] overwrite = true } n, err = w.multi.Write(p) w.pending -= int64(n) if err == nil && overwrite { err = ErrOverflow return } return } // Hash returns the hash of the object data stream that has been written so far. // It can be called before or after Close. func (w *Writer) Hash() plumbing.Hash { return w.hasher.Sum() // Not yet closed, return hash of data written so far } // Close releases any resources consumed by the Writer. // // Calling Close does not close the wrapped io.Writer originally passed to // NewWriter. func (w *Writer) Close() error { if err := w.zlib.Close(); err != nil { return err } w.closed = true return nil }