github.com/TeaOSLab/EdgeNode@v1.3.8/internal/utils/readers/reader_closer_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 TeeReaderCloser struct {
    10  	r io.Reader
    11  	w io.Writer
    12  
    13  	onFail func(err error)
    14  	onEOF  func()
    15  
    16  	mustWrite bool
    17  }
    18  
    19  // NewTeeReaderCloser
    20  // mustWrite - ensure writing MUST be successfully
    21  func NewTeeReaderCloser(reader io.Reader, writer io.Writer, mustWrite bool) *TeeReaderCloser {
    22  	return &TeeReaderCloser{
    23  		r:         reader,
    24  		w:         writer,
    25  		mustWrite: mustWrite,
    26  	}
    27  }
    28  
    29  func (this *TeeReaderCloser) Read(p []byte) (n int, err error) {
    30  	n, err = this.r.Read(p)
    31  	if n > 0 {
    32  		_, wErr := this.w.Write(p[:n])
    33  		if (err == nil || err == io.EOF) && wErr != nil {
    34  			if this.mustWrite {
    35  				err = wErr
    36  			} else {
    37  				if this.onFail != nil {
    38  					this.onFail(wErr)
    39  				}
    40  			}
    41  		}
    42  	}
    43  	if err != nil {
    44  		if err == io.EOF {
    45  			if this.onEOF != nil {
    46  				this.onEOF()
    47  			}
    48  		} else {
    49  			if this.onFail != nil {
    50  				this.onFail(err)
    51  			}
    52  		}
    53  	}
    54  	return
    55  }
    56  
    57  func (this *TeeReaderCloser) Close() error {
    58  	r, ok := this.r.(io.Closer)
    59  	if ok {
    60  		return r.Close()
    61  	}
    62  	return nil
    63  }
    64  
    65  func (this *TeeReaderCloser) OnFail(onFail func(err error)) {
    66  	this.onFail = onFail
    67  }
    68  
    69  func (this *TeeReaderCloser) OnEOF(onEOF func()) {
    70  	this.onEOF = onEOF
    71  }