vendor/github.com/git-lfs/gitobj/object_writer.go

   1 package gitobj
   2
   3 import (
   4         "compress/zlib"
   5         "crypto/sha1"
   6         "fmt"
   7         "hash"
   8         "io"
   9         "sync/atomic"
  10 )
  11
  12 // ObjectWriter provides an implementation of io.Writer that compresses and
  13 // writes data given to it, and keeps track of the SHA1 hash of the data as it
  14 // is written.
  15 type ObjectWriter struct {
  16         // members managed via sync/atomic must be aligned at the top of this
  17         // structure (see: https://github.com/git-lfs/git-lfs/pull/2880).
  18
  19         // wroteHeader is a uint32 managed by the sync/atomic package. It is 1
  20         // if the header was written, and 0 otherwise.
  21         wroteHeader uint32
  22
  23         // w is the underling writer that this ObjectWriter is writing to.
  24         w io.Writer
  25         // sum is the in-progress hash calculation.
  26         sum hash.Hash
  27
  28         // closeFn supplies an optional function that, when called, frees an
  29         // resources (open files, memory, etc) held by this instance of the
  30         // *ObjectWriter.
  31         //
  32         // closeFn returns any error encountered when closing/freeing resources
  33         // held.
  34         //
  35         // It is allowed to be nil.
  36         closeFn func() error
  37 }
  38
  39 // nopCloser provides a no-op implementation of the io.WriteCloser interface by
  40 // taking an io.Writer and wrapping it with a Close() method that returns nil.
  41 type nopCloser struct {
  42         // Writer is an embedded io.Writer that receives the Write() method
  43         // call.
  44         io.Writer
  45 }
  46
  47 // Close implements the io.Closer interface by returning nil.
  48 func (n *nopCloser) Close() error {
  49         return nil
  50 }
  51
  52 // NewObjectWriter returns a new *ObjectWriter instance that drains incoming
  53 // writes into the io.Writer given, "w".
  54 func NewObjectWriter(w io.Writer) *ObjectWriter {
  55         return NewObjectWriteCloser(&nopCloser{w})
  56 }
  57
  58 // NewObjectWriter returns a new *ObjectWriter instance that drains incoming
  59 // writes into the io.Writer given, "w".
  60 //
  61 // Upon closing, it calls the given Close() function of the io.WriteCloser.
  62 func NewObjectWriteCloser(w io.WriteCloser) *ObjectWriter {
  63         zw := zlib.NewWriter(w)
  64         sum := sha1.New()
  65
  66         return &ObjectWriter{
  67                 w:   io.MultiWriter(zw, sum),
  68                 sum: sum,
  69
  70                 closeFn: func() error {
  71                         if err := zw.Close(); err != nil {
  72                                 return err
  73                         }
  74                         if err := w.Close(); err != nil {
  75                                 return err
  76                         }
  77                         return nil
  78                 },
  79         }
  80 }
  81
  82 // WriteHeader writes object header information and returns the number of
  83 // uncompressed bytes written, or any error that was encountered along the way.
  84 //
  85 // WriteHeader MUST be called only once, or a panic() will occur.
  86 func (w *ObjectWriter) WriteHeader(typ ObjectType, len int64) (n int, err error) {
  87         if !atomic.CompareAndSwapUint32(&w.wroteHeader, 0, 1) {
  88                 panic("gitobj: cannot write headers more than once")
  89         }
  90         return fmt.Fprintf(w, "%s %d\x00", typ, len)
  91 }
  92
  93 // Write writes the given buffer "p" of uncompressed bytes into the underlying
  94 // data-stream, returning the number of uncompressed bytes written, along with
  95 // any error encountered along the way.
  96 //
  97 // A call to WriteHeaders MUST occur before calling Write, or a panic() will
  98 // occur.
  99 func (w *ObjectWriter) Write(p []byte) (n int, err error) {
 100         if atomic.LoadUint32(&w.wroteHeader) != 1 {
 101                 panic("gitobj: cannot write data without header")
 102         }
 103         return w.w.Write(p)
 104 }
 105
 106 // Sha returns the in-progress SHA1 of the compressed object contents.
 107 func (w *ObjectWriter) Sha() []byte {
 108         return w.sum.Sum(nil)
 109 }
 110
 111 // Close closes the ObjectWriter and frees any resources held by it, including
 112 // flushing the zlib-compressed content to the underling writer. It must be
 113 // called before discarding of the Writer instance.
 114 //
 115 // If any error occurred while calling close, it will be returned immediately,
 116 // otherwise nil.
 117 func (w *ObjectWriter) Close() error {
 118         if w.closeFn == nil {
 119                 return nil
 120         }
 121         return w.closeFn()
 122 }