1package io
 2
 3import (
 4	"io"
 5	"sync"
 6)
 7
 8// NewSafeReadCloser returns a new safeReadCloser that wraps readCloser.
 9func 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.
24type safeWriteToReadCloser struct {
25	*safeReadCloser
26}
27
28// WriteTo implements the io.WriteTo interface.
29func (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.
45type 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.
52func (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.
64func (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}