github.com/TrenchBoot/u-root@v6.0.1-0.20200623220532-550e36da3258+incompatible/pkg/uio/alignreader.go (about)

     1  // Copyright 2019 the u-root Authors. All rights reserved
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package uio
     6  
     7  import (
     8  	"io"
     9  )
    10  
    11  // AlignReader keeps track of how many bytes were read so the reader can be
    12  // aligned at a future time.
    13  type AlignReader struct {
    14  	R io.Reader
    15  	N int
    16  }
    17  
    18  // Read reads from the underlying io.Reader.
    19  func (r *AlignReader) Read(b []byte) (int, error) {
    20  	n, err := r.R.Read(b)
    21  	r.N += n
    22  	return n, err
    23  }
    24  
    25  // ReadByte reads one byte from the underlying io.Reader.
    26  func (r *AlignReader) ReadByte() (byte, error) {
    27  	b := make([]byte, 1)
    28  	_, err := io.ReadFull(r, b)
    29  	return b[0], err
    30  }
    31  
    32  // Align aligns the reader to the given number of bytes and returns the
    33  // bytes read to pad it.
    34  func (r *AlignReader) Align(n int) ([]byte, error) {
    35  	if r.N%n == 0 {
    36  		return []byte{}, nil
    37  	}
    38  	pad := make([]byte, n-r.N%n)
    39  	m, err := io.ReadFull(r, pad)
    40  	return pad[:m], err
    41  }