1 package io
2 3 import (
4 "io"
5 "sync"
6 )
7 8 // NewSafeReadCloser returns a new safeReadCloser that wraps readCloser.
9 func NewSafeReadCloser(readCloser io.ReadCloser) io.ReadCloser {
10 sr := &safeReadCloser{
11 readCloser: readCloser,
12 }
13 14 if _, ok := readCloser.(io.WriterTo); ok {
15 return &safeWriteToReadCloser{safeReadCloser: sr}
16 }
17 18 return sr
19 }
20 21 // safeWriteToReadCloser wraps a safeReadCloser but exposes a WriteTo interface implementation. This will panic
22 // if the underlying io.ReadClose does not support WriteTo. Use NewSafeReadCloser to ensure the proper handling of this
23 // type.
24 type safeWriteToReadCloser struct {
25 *safeReadCloser
26 }
27 28 // WriteTo implements the io.WriteTo interface.
29 func (r *safeWriteToReadCloser) WriteTo(w io.Writer) (int64, error) {
30 r.safeReadCloser.mtx.Lock()
31 defer r.safeReadCloser.mtx.Unlock()
32 33 if r.safeReadCloser.closed {
34 return 0, io.EOF
35 }
36 37 return r.safeReadCloser.readCloser.(io.WriterTo).WriteTo(w)
38 }
39 40 // safeReadCloser wraps a io.ReadCloser and presents an io.ReadCloser interface. When Close is called on safeReadCloser
41 // the underlying Close method will be executed, and then the reference to the reader will be dropped. This type
42 // is meant to be used with the net/http library which will retain a reference to the request body for the lifetime
43 // of a goroutine connection. Wrapping in this manner will ensure that no data race conditions are falsely reported.
44 // This type is thread-safe.
45 type safeReadCloser struct {
46 readCloser io.ReadCloser
47 closed bool
48 mtx sync.Mutex
49 }
50 51 // Read reads up to len(p) bytes into p from the underlying read. If the reader is closed io.EOF will be returned.
52 func (r *safeReadCloser) Read(p []byte) (n int, err error) {
53 r.mtx.Lock()
54 defer r.mtx.Unlock()
55 if r.closed {
56 return 0, io.EOF
57 }
58 59 return r.readCloser.Read(p)
60 }
61 62 // Close calls the underlying io.ReadCloser's Close method, removes the reference to the reader, and returns any error
63 // reported from Close. Subsequent calls to Close will always return a nil error.
64 func (r *safeReadCloser) Close() error {
65 r.mtx.Lock()
66 defer r.mtx.Unlock()
67 if r.closed {
68 return nil
69 }
70 71 r.closed = true
72 rc := r.readCloser
73 r.readCloser = nil
74 return rc.Close()
75 }
76