safe.go

Documentation: github.com/aws/smithy-go/transport/http/internal/io

     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
View as plain text