github.com/mutagen-io/mutagen@v0.18.0-rc1/pkg/stream/cutoff_writer.go (about)

     1  package stream
     2  
     3  import (
     4  	"io"
     5  )
     6  
     7  // cutoffWriter is the io.Writer implementation that underlies NewCutoffWriter.
     8  type cutoffWriter struct {
     9  	// writer is the underlying writer.
    10  	writer io.Writer
    11  	// cutoff is the number of bytes remaining.
    12  	cutoff uint
    13  }
    14  
    15  // NewCutoffWriter creates a new io.Writer that wraps and forwards writes to the
    16  // specified writer until a specified maximum number of bytes has been written,
    17  // at which point the writing of any further bytes will succeed automatically,
    18  // but not actually write to the underlying writer.
    19  func NewCutoffWriter(writer io.Writer, cutoff uint) io.Writer {
    20  	return &cutoffWriter{
    21  		writer: writer,
    22  		cutoff: cutoff,
    23  	}
    24  }
    25  
    26  // Write implements io.Writer.Write.
    27  func (w *cutoffWriter) Write(buffer []byte) (int, error) {
    28  	// If we've already hit the cutoff, then there's no action to take, we just
    29  	// pretend like the full write succeeded.
    30  	if w.cutoff == 0 {
    31  		return len(buffer), nil
    32  	}
    33  
    34  	// If the buffer length will fit within the cutoff, then just perform the
    35  	// write, update the cutoff, and return the result.
    36  	if uint(len(buffer)) <= w.cutoff {
    37  		written, err := w.writer.Write(buffer)
    38  		w.cutoff -= uint(written)
    39  		return written, err
    40  	}
    41  
    42  	// Otherwise, perform a truncated write, update the cutoff, and return the
    43  	// count based on whether or not the truncated write failed. If the
    44  	// truncated write succeeded, then the cutoff is now zero and we pretend
    45  	// that the rest of the bytes were also written successfully.
    46  	written, err := w.writer.Write(buffer[:w.cutoff])
    47  	w.cutoff -= uint(written)
    48  	if err != nil {
    49  		return written, err
    50  	}
    51  	return len(buffer), nil
    52  }