Some small improvements

1 files changed, 12 insertions, 13 deletions

diff --git a/oae2.go b/oae2.go
index 928eca3..b0868a0 100644
--- a/oae2.go
+++ b/oae2.go
@@ -102,6 +102,7 @@ func NewWriter(w io.Writer, key []byte, segmentSize int) *Writer {
 		w:           w,
 		segmentSize: segmentSize,
 		encrypter:   segmentEncrypter{key: key},
+		buf:         make([]byte, 0, segmentSize+aeadOverhead),
 	}
 }
 
@@ -112,11 +113,8 @@ func (w *Writer) initialize() error {
 	if w.err = w.encrypter.init(salt); w.err != nil {
 		return w.err
 	}
-	if _, w.err = w.w.Write(salt); w.err != nil {
-		return w.err
-	}
-	w.buf = make([]byte, 0, w.segmentSize+aeadOverhead)
-	return nil
+	_, w.err = w.w.Write(salt)
+	return w.err
 }
 
 func (w *Writer) init() error {
@@ -186,7 +184,7 @@ func (w *Writer) ReadFrom(r io.Reader) (int64, error) {
 	}
 }
 
-var errClosed = errors.New("oae2.Writer.Close: already closed")
+var errClosed = errors.New("oae2.Writer: writer closed")
 
 // Close flushes the final segment. Failure to call Close will result in a
 // truncated stream.
@@ -247,6 +245,7 @@ func NewReader(r io.Reader, key []byte, segmentSize int) *Reader {
 		r:           bufReader{r: r},
 		segmentSize: segmentSize,
 		decrypter:   segmentEncrypter{key: key},
+		buf:         make([]byte, 0, segmentSize+aeadOverhead+1),
 	}
 }
 
@@ -259,11 +258,8 @@ func (r *Reader) initialize() error {
 		}
 		return r.err
 	}
-	if r.err = r.decrypter.init(salt); r.err != nil {
-		return r.err
-	}
-	r.buf = make([]byte, 0, r.segmentSize+aeadOverhead+1)
-	return nil
+	r.err = r.decrypter.init(salt)
+	return r.err
 }
 
 func (r *Reader) init() error {
@@ -274,10 +270,13 @@ func (r *Reader) init() error {
 }
 
 func (r *Reader) fillBuf() error {
+	if r.readLastChunk {
+		return io.EOF
+	}
 	n, err := io.ReadFull(&r.r, r.buf[:r.segmentSize+aeadOverhead+1])
 	if err != nil && err != io.ErrUnexpectedEOF {
-		if err == io.EOF && !r.readLastChunk {
-			return io.ErrUnexpectedEOF
+		if err == io.EOF {
+			err = io.ErrUnexpectedEOF
 		}
 		r.err = err
 		return err