github.com/u-root/u-root@v7.0.1-0.20200915234505-ad7babab0a8e+incompatible/pkg/uio/alignwriter.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  	"bytes"
     9  	"io"
    10  )
    11  
    12  // AlignWriter keeps track of how many bytes were written so the writer can be
    13  // aligned at a future time.
    14  type AlignWriter struct {
    15  	W io.Writer
    16  	N int
    17  }
    18  
    19  // Write writes to the underlying io.Writew.
    20  func (w *AlignWriter) Write(b []byte) (int, error) {
    21  	n, err := w.W.Write(b)
    22  	w.N += n
    23  	return n, err
    24  }
    25  
    26  // Align aligns the writer to the given number of bytes using the given pad
    27  // value.
    28  func (w *AlignWriter) Align(n int, pad byte) error {
    29  	if w.N%n == 0 {
    30  		return nil
    31  	}
    32  	_, err := w.Write(bytes.Repeat([]byte{pad}, n-w.N%n))
    33  	return err
    34  }