package ioutils import ( "crypto/sha256" "encoding/hex" "io" "golang.org/x/net/context" ) type readCloserWrapper struct { io.Reader closer func() error } func (r *readCloserWrapper) Close() error { return r.closer() } // NewReadCloserWrapper returns a new io.ReadCloser. func NewReadCloserWrapper(r io.Reader, closer func() error) io.ReadCloser { return &readCloserWrapper{ Reader: r, closer: closer, } } type readerErrWrapper struct { reader io.Reader closer func() } func (r *readerErrWrapper) Read(p []byte) (int, error) { n, err := r.reader.Read(p) if err != nil { r.closer() } return n, err } // NewReaderErrWrapper returns a new io.Reader. func NewReaderErrWrapper(r io.Reader, closer func()) io.Reader { return &readerErrWrapper{ reader: r, closer: closer, } } // HashData returns the sha256 sum of src. func HashData(src io.Reader) (string, error) { h := sha256.New() if _, err := io.Copy(h, src); err != nil { return "", err } return "sha256:" + hex.EncodeToString(h.Sum(nil)), nil } // OnEOFReader wraps a io.ReadCloser and a function // the function will run at the end of file or close the file. type OnEOFReader struct { Rc io.ReadCloser Fn func() } func (r *OnEOFReader) Read(p []byte) (n int, err error) { n, err = r.Rc.Read(p) if err == io.EOF { r.runFunc() } return } // Close closes the file and run the function. func (r *OnEOFReader) Close() error { err := r.Rc.Close() r.runFunc() return err } func (r *OnEOFReader) runFunc() { if fn := r.Fn; fn != nil { fn() r.Fn = nil } } // cancelReadCloser wraps an io.ReadCloser with a context for cancelling read // operations. type cancelReadCloser struct { cancel func() pR *io.PipeReader // Stream to read from pW *io.PipeWriter } // NewCancelReadCloser creates a wrapper that closes the ReadCloser when the // context is cancelled. The returned io.ReadCloser must be closed when it is // no longer needed. func NewCancelReadCloser(ctx context.Context, in io.ReadCloser) io.ReadCloser { pR, pW := io.Pipe() // Create a context used to signal when the pipe is closed doneCtx, cancel := context.WithCancel(context.Background()) p := &cancelReadCloser{ cancel: cancel, pR: pR, pW: pW, } go func() { _, err := io.Copy(pW, in) select { case <-ctx.Done(): // If the context was closed, p.closeWithError // was already called. Calling it again would // change the error that Read returns. default: p.closeWithError(err) } in.Close() }() go func() { for { select { case <-ctx.Done(): p.closeWithError(ctx.Err()) case <-doneCtx.Done(): return } } }() return p } // Read wraps the Read method of the pipe that provides data from the wrapped // ReadCloser. func (p *cancelReadCloser) Read(buf []byte) (n int, err error) { return p.pR.Read(buf) } // closeWithError closes the wrapper and its underlying reader. It will // cause future calls to Read to return err. func (p *cancelReadCloser) closeWithError(err error) { p.pW.CloseWithError(err) p.cancel() } // Close closes the wrapper its underlying reader. It will cause // future calls to Read to return io.EOF. func (p *cancelReadCloser) Close() error { p.closeWithError(io.EOF) return nil }