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