github.com/karrick/gorill@v1.10.3/short.go (about)

     1  package gorill
     2  
     3  import "io"
     4  
     5  // ShortReadWriteCloser wraps a io.ReadWriteCloser, but the Read and Write operations cannot exceed
     6  // the MaxRead and MaxWrite sizes.
     7  type ShortReadWriteCloser struct {
     8  	io.Reader
     9  	io.WriteCloser
    10  	MaxRead  int
    11  	MaxWrite int
    12  }
    13  
    14  // Read reads from the wrapped io.Reader, but returns EOF if attempts to read beyond the MaxRead.
    15  func (s ShortReadWriteCloser) Read(buf []byte) (int, error) {
    16  	var short bool
    17  	index := len(buf)
    18  	if index > s.MaxRead {
    19  		index = s.MaxRead
    20  		short = true
    21  	}
    22  	n, err := s.Reader.Read(buf[:index])
    23  	if short {
    24  		return n, io.ErrUnexpectedEOF
    25  	}
    26  	return n, err
    27  }
    28  
    29  // Write writes from the wrapped io.Writer, but returns EOF if attempts to write beyond the MaxWrite.
    30  func (s ShortReadWriteCloser) Write(data []byte) (int, error) {
    31  	var short bool
    32  	index := len(data)
    33  	if index > s.MaxWrite {
    34  		index = s.MaxWrite
    35  		short = true
    36  	}
    37  	n, err := s.WriteCloser.Write(data[:index])
    38  	if short {
    39  		return n, io.ErrShortWrite
    40  	}
    41  	return n, err
    42  }
    43  
    44  // ShortReadWriter wraps a io.Reader and io.Writer, but the Read and Write operations cannot exceed
    45  // the MaxRead and MaxWrite sizes.
    46  type ShortReadWriter struct {
    47  	io.Reader
    48  	io.Writer
    49  	MaxRead  int
    50  	MaxWrite int
    51  }
    52  
    53  // Read reads from the wrapped io.Reader, but returns EOF if attempts to read beyond the MaxRead.
    54  func (s ShortReadWriter) Read(buf []byte) (int, error) {
    55  	var short bool
    56  	index := len(buf)
    57  	if index > s.MaxRead {
    58  		index = s.MaxRead
    59  		short = true
    60  	}
    61  	n, err := s.Reader.Read(buf[:index])
    62  	if short {
    63  		return n, io.ErrUnexpectedEOF
    64  	}
    65  	return n, err
    66  }
    67  
    68  // Write writes from the wrapped io.Writer, but returns EOF if attempts to write beyond the MaxWrite.
    69  func (s ShortReadWriter) Write(data []byte) (int, error) {
    70  	var short bool
    71  	index := len(data)
    72  	if index > s.MaxWrite {
    73  		index = s.MaxWrite
    74  		short = true
    75  	}
    76  	n, err := s.Writer.Write(data[:index])
    77  	if short {
    78  		return n, io.ErrShortWrite
    79  	}
    80  	return n, err
    81  }
    82  
    83  // ShortWriter returns a structure that wraps an io.Writer, but returns io.ErrShortWrite when the
    84  // number of bytes to write exceeds a preset limit.
    85  //
    86  //   bb := gorill.NopCloseBuffer()
    87  //   sw := gorill.ShortWriter(bb, 16)
    88  //
    89  //   n, err := sw.Write([]byte("short write"))
    90  //   // n == 11, err == nil
    91  //
    92  //   n, err := sw.Write([]byte("a somewhat longer write"))
    93  //   // n == 16, err == io.ErrShortWrite
    94  func ShortWriter(w io.Writer, max int) io.Writer {
    95  	return shortWriter{Writer: w, max: max}
    96  }
    97  
    98  func (s shortWriter) Write(data []byte) (int, error) {
    99  	var short bool
   100  	index := len(data)
   101  	if index > s.max {
   102  		index = s.max
   103  		short = true
   104  	}
   105  	n, err := s.Writer.Write(data[:index])
   106  	if short {
   107  		return n, io.ErrShortWrite
   108  	}
   109  	return n, err
   110  }
   111  
   112  type shortWriter struct {
   113  	io.Writer
   114  	max int
   115  }
   116  
   117  // ShortWriteCloser returns a structure that wraps an io.WriteCloser, but returns io.ErrShortWrite
   118  // when the number of bytes to write exceeds a preset limit.
   119  //
   120  //   bb := gorill.NopCloseBuffer()
   121  //   sw := gorill.ShortWriteCloser(bb, 16)
   122  //
   123  //   n, err := sw.Write([]byte("short write"))
   124  //   // n == 11, err == nil
   125  //
   126  //   n, err := sw.Write([]byte("a somewhat longer write"))
   127  //   // n == 16, err == io.ErrShortWrite
   128  func ShortWriteCloser(iowc io.WriteCloser, max int) io.WriteCloser {
   129  	return shortWriteCloser{WriteCloser: iowc, max: max}
   130  }
   131  
   132  func (s shortWriteCloser) Write(data []byte) (int, error) {
   133  	var short bool
   134  	index := len(data)
   135  	if index > s.max {
   136  		index = s.max
   137  		short = true
   138  	}
   139  	n, err := s.WriteCloser.Write(data[:index])
   140  	if short {
   141  		return n, io.ErrShortWrite
   142  	}
   143  	return n, err
   144  }
   145  
   146  type shortWriteCloser struct {
   147  	io.WriteCloser
   148  	max int
   149  }