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 }