github.com/Prakhar-Agarwal-byte/moby@v0.0.0-20231027092010-a14e3e8ab87e/pkg/ioutils/readers.go (about)

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