github.com/insolar/vanilla@v0.0.0-20201023172447-248fdf805322/iokit/tee.go (about)

     1  // Copyright 2020 Insolar Network Ltd.
     2  // All rights reserved.
     3  // This material is licensed under the Insolar License version 1.0,
     4  // available at https://github.com/insolar/assured-ledger/blob/master/LICENSE.md.
     5  
     6  package iokit
     7  
     8  import (
     9  	"io"
    10  )
    11  
    12  func NewTeeWriter(main, copy io.Writer) io.Writer {
    13  	return &TeeWriter{main, teeTemplate{copy, 0}}
    14  }
    15  
    16  func NewTeeWriterWithSkip(main, copy io.Writer, zeroLeadingBytes int) io.Writer {
    17  	return &TeeWriter{main, teeTemplate{copy, zeroLeadingBytes}}
    18  }
    19  
    20  func NewTeeReader(main io.Reader, copy io.Writer) io.Reader {
    21  	return &TeeReader{main, teeTemplate{copy, 0}}
    22  }
    23  
    24  func NewTeeReaderWithSkip(main io.Reader, copy io.Writer, skipLeadingBytes int) io.Reader {
    25  	return &TeeReader{main, teeTemplate{copy, skipLeadingBytes}}
    26  }
    27  
    28  type teeTemplate struct {
    29  	CopyTo   io.Writer
    30  	CopySkip int
    31  }
    32  
    33  func (w *teeTemplate) teeWrite(n int, b []byte) {
    34  	switch {
    35  	case w.CopyTo == nil:
    36  		return
    37  	case w.CopySkip <= 0:
    38  		_, _ = w.CopyTo.Write(b[:n])
    39  	case w.CopySkip >= n:
    40  		w.CopySkip -= n
    41  	default:
    42  		_, _ = w.CopyTo.Write(b[w.CopySkip:n])
    43  		w.CopySkip = 0
    44  	}
    45  }
    46  
    47  /***********************************************/
    48  
    49  var _ io.Writer = &TeeWriter{}
    50  
    51  type TeeWriter struct {
    52  	W io.Writer
    53  	teeTemplate
    54  }
    55  
    56  func (w *TeeWriter) Write(b []byte) (n int, err error) {
    57  	n, err = w.W.Write(b)
    58  	w.teeWrite(n, b)
    59  	return n, err
    60  }
    61  
    62  /***********************************************/
    63  
    64  var _ io.Reader = &TeeReader{}
    65  
    66  type TeeReader struct {
    67  	R io.Reader
    68  	teeTemplate
    69  }
    70  
    71  func (w *TeeReader) Read(b []byte) (n int, err error) {
    72  	n, err = w.R.Read(b)
    73  	w.teeWrite(n, b)
    74  	return n, err
    75  }