github.com/demonoid81/moby@v0.0.0-20200517203328-62dd8e17c460/pkg/ioutils/readers.go (about)

     1  package ioutils // import "github.com/demonoid81/moby/pkg/ioutils"
     2  
     3  import (
     4  	"context"
     5  	"crypto/sha256"
     6  	"encoding/hex"
     7  	"io"
     8  )
     9  
    10  // ReadCloserWrapper wraps an io.Reader, and implements an io.ReadCloser
    11  // It calls the given callback function when closed. It should be constructed
    12  // with NewReadCloserWrapper
    13  type ReadCloserWrapper struct {
    14  	io.Reader
    15  	closer func() error
    16  }
    17  
    18  // Close calls back the passed closer function
    19  func (r *ReadCloserWrapper) Close() error {
    20  	return r.closer()
    21  }
    22  
    23  // NewReadCloserWrapper returns a new io.ReadCloser.
    24  func NewReadCloserWrapper(r io.Reader, closer func() error) io.ReadCloser {
    25  	return &ReadCloserWrapper{
    26  		Reader: r,
    27  		closer: closer,
    28  	}
    29  }
    30  
    31  type readerErrWrapper struct {
    32  	reader io.Reader
    33  	closer func()
    34  }
    35  
    36  func (r *readerErrWrapper) Read(p []byte) (int, error) {
    37  	n, err := r.reader.Read(p)
    38  	if err != nil {
    39  		r.closer()
    40  	}
    41  	return n, err
    42  }
    43  
    44  // NewReaderErrWrapper returns a new io.Reader.
    45  func NewReaderErrWrapper(r io.Reader, closer func()) io.Reader {
    46  	return &readerErrWrapper{
    47  		reader: r,
    48  		closer: closer,
    49  	}
    50  }
    51  
    52  // HashData returns the sha256 sum of src.
    53  func HashData(src io.Reader) (string, error) {
    54  	h := sha256.New()
    55  	if _, err := io.Copy(h, src); err != nil {
    56  		return "", err
    57  	}
    58  	return "sha256:" + hex.EncodeToString(h.Sum(nil)), nil
    59  }
    60  
    61  // OnEOFReader wraps an io.ReadCloser and a function
    62  // the function will run at the end of file or close the file.
    63  type OnEOFReader struct {
    64  	Rc io.ReadCloser
    65  	Fn func()
    66  }
    67  
    68  func (r *OnEOFReader) Read(p []byte) (n int, err error) {
    69  	n, err = r.Rc.Read(p)
    70  	if err == io.EOF {
    71  		r.runFunc()
    72  	}
    73  	return
    74  }
    75  
    76  // Close closes the file and run the function.
    77  func (r *OnEOFReader) Close() error {
    78  	err := r.Rc.Close()
    79  	r.runFunc()
    80  	return err
    81  }
    82  
    83  func (r *OnEOFReader) runFunc() {
    84  	if fn := r.Fn; fn != nil {
    85  		fn()
    86  		r.Fn = nil
    87  	}
    88  }
    89  
    90  // cancelReadCloser wraps an io.ReadCloser with a context for cancelling read
    91  // operations.
    92  type cancelReadCloser struct {
    93  	cancel func()
    94  	pR     *io.PipeReader // Stream to read from
    95  	pW     *io.PipeWriter
    96  }
    97  
    98  // NewCancelReadCloser creates a wrapper that closes the ReadCloser when the
    99  // context is cancelled. The returned io.ReadCloser must be closed when it is
   100  // no longer needed.
   101  func NewCancelReadCloser(ctx context.Context, in io.ReadCloser) io.ReadCloser {
   102  	pR, pW := io.Pipe()
   103  
   104  	// Create a context used to signal when the pipe is closed
   105  	doneCtx, cancel := context.WithCancel(context.Background())
   106  
   107  	p := &cancelReadCloser{
   108  		cancel: cancel,
   109  		pR:     pR,
   110  		pW:     pW,
   111  	}
   112  
   113  	go func() {
   114  		_, err := io.Copy(pW, in)
   115  		select {
   116  		case <-ctx.Done():
   117  			// If the context was closed, p.closeWithError
   118  			// was already called. Calling it again would
   119  			// change the error that Read returns.
   120  		default:
   121  			p.closeWithError(err)
   122  		}
   123  		in.Close()
   124  	}()
   125  	go func() {
   126  		for {
   127  			select {
   128  			case <-ctx.Done():
   129  				p.closeWithError(ctx.Err())
   130  			case <-doneCtx.Done():
   131  				return
   132  			}
   133  		}
   134  	}()
   135  
   136  	return p
   137  }
   138  
   139  // Read wraps the Read method of the pipe that provides data from the wrapped
   140  // ReadCloser.
   141  func (p *cancelReadCloser) Read(buf []byte) (n int, err error) {
   142  	return p.pR.Read(buf)
   143  }
   144  
   145  // closeWithError closes the wrapper and its underlying reader. It will
   146  // cause future calls to Read to return err.
   147  func (p *cancelReadCloser) closeWithError(err error) {
   148  	p.pW.CloseWithError(err)
   149  	p.cancel()
   150  }
   151  
   152  // Close closes the wrapper its underlying reader. It will cause
   153  // future calls to Read to return io.EOF.
   154  func (p *cancelReadCloser) Close() error {
   155  	p.closeWithError(io.EOF)
   156  	return nil
   157  }