fixedIO_test.go raw

   1  package wire
   2  
   3  import (
   4  	"bytes"
   5  	"io"
   6  )
   7  
   8  // fixedWriter implements the io.Writer interface and intentially allows testing of error paths by forcing short writes.
   9  type fixedWriter struct {
  10  	b   []byte
  11  	pos int
  12  }
  13  
  14  // Write writes the contents of p to w. When the contents of p would cause the writer to exceed the maximum allowed size
  15  // of the fixed writer, io.ErrShortWrite is returned and the writer is left unchanged. This satisfies the io.Writer
  16  // interface.
  17  func (w *fixedWriter) Write(p []byte) (n int, e error) {
  18  	lenp := len(p)
  19  	if w.pos+lenp > cap(w.b) {
  20  		return 0, io.ErrShortWrite
  21  	}
  22  	n = lenp
  23  	w.pos += copy(w.b[w.pos:], p)
  24  	return
  25  }
  26  
  27  // Hash returns the bytes already written to the fixed writer.
  28  func (w *fixedWriter) Bytes() []byte {
  29  	return w.b
  30  }
  31  
  32  // newFixedWriter returns a new io.Writer that will error once more bytes than the specified max have been written.
  33  func newFixedWriter(max int) io.Writer {
  34  	b := make([]byte, max)
  35  	fw := fixedWriter{b, 0}
  36  	return &fw
  37  }
  38  
  39  // fixedReader implements the io.Reader interface and intentially allows testing of error paths by forcing short reads.
  40  type fixedReader struct {
  41  	buf   []byte
  42  	pos   int
  43  	iobuf *bytes.Buffer
  44  }
  45  
  46  // Read reads the next len(p) bytes from the fixed reader. When the number of bytes read would exceed the maximum number
  47  // of allowed bytes to be read from the fixed writer, an error is returned. This satisfies the io.Reader interface.
  48  func (fr *fixedReader) Read(p []byte) (n int, e error) {
  49  	n, e = fr.iobuf.Read(p)
  50  	fr.pos += n
  51  	return
  52  }
  53  
  54  // newFixedReader returns a new io.Reader that will error once more bytes than the specified max have been read.
  55  func newFixedReader(max int, buf []byte) io.Reader {
  56  	b := make([]byte, max)
  57  	if buf != nil {
  58  		copy(b[:], buf)
  59  	}
  60  	iobuf := bytes.NewBuffer(b)
  61  	fr := fixedReader{b, 0, iobuf}
  62  	return &fr
  63  }
  64