github.com/TeaOSLab/EdgeNode@v1.3.8/internal/utils/readers/reader_tee.go (about)

     1  // Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
     2  
     3  package readers
     4  
     5  import (
     6  	"io"
     7  )
     8  
     9  type TeeReader struct {
    10  	r io.Reader
    11  	w io.Writer
    12  
    13  	onFail func(err error)
    14  	onEOF  func()
    15  }
    16  
    17  func NewTeeReader(reader io.Reader, writer io.Writer) *TeeReader {
    18  	return &TeeReader{
    19  		r: reader,
    20  		w: writer,
    21  	}
    22  }
    23  
    24  func (this *TeeReader) Read(p []byte) (n int, err error) {
    25  	n, err = this.r.Read(p)
    26  	if n > 0 {
    27  		_, wErr := this.w.Write(p[:n])
    28  		if err == nil && wErr != nil {
    29  			err = wErr
    30  		}
    31  	}
    32  	if err != nil {
    33  		if err == io.EOF {
    34  			if this.onEOF != nil {
    35  				this.onEOF()
    36  			}
    37  		} else {
    38  			if this.onFail != nil {
    39  				this.onFail(err)
    40  			}
    41  		}
    42  	}
    43  	return
    44  }
    45  
    46  func (this *TeeReader) OnFail(onFail func(err error)) {
    47  	this.onFail = onFail
    48  }
    49  
    50  func (this *TeeReader) OnEOF(onEOF func()) {
    51  	this.onEOF = onEOF
    52  }