github.com/laof/lite-speed-test@v0.0.0-20230930011949-1f39b7037845/common/io.go (about)

     1  package common
     2  
     3  import (
     4  	"bufio"
     5  	"io"
     6  	"net"
     7  
     8  	"github.com/laof/lite-speed-test/log"
     9  )
    10  
    11  type RewindReader struct {
    12  	rawReader  io.Reader
    13  	buf        []byte
    14  	bufReadIdx int
    15  	rewound    bool
    16  	buffering  bool
    17  	bufferSize int
    18  }
    19  
    20  func (r *RewindReader) Read(p []byte) (int, error) {
    21  	if r.rewound {
    22  		if len(r.buf) > r.bufReadIdx {
    23  			n := copy(p, r.buf[r.bufReadIdx:])
    24  			r.bufReadIdx += n
    25  			return n, nil
    26  		}
    27  		r.rewound = false //all buffering content has been read
    28  	}
    29  	n, err := r.rawReader.Read(p)
    30  	if r.buffering {
    31  		r.buf = append(r.buf, p[:n]...)
    32  		if len(r.buf) > r.bufferSize*2 {
    33  			log.D("read too many bytes!")
    34  		}
    35  	}
    36  	return n, err
    37  }
    38  
    39  func (r *RewindReader) ReadByte() (byte, error) {
    40  	buf := [1]byte{}
    41  	_, err := r.Read(buf[:])
    42  	return buf[0], err
    43  }
    44  
    45  func (r *RewindReader) Discard(n int) (int, error) {
    46  	buf := [128]byte{}
    47  	if n < 128 {
    48  		return r.Read(buf[:n])
    49  	}
    50  	for discarded := 0; discarded+128 < n; discarded += 128 {
    51  		_, err := r.Read(buf[:])
    52  		if err != nil {
    53  			return discarded, err
    54  		}
    55  	}
    56  	if rest := n % 128; rest != 0 {
    57  		return r.Read(buf[:rest])
    58  	}
    59  	return n, nil
    60  }
    61  
    62  func (r *RewindReader) Rewind() {
    63  	if r.bufferSize == 0 {
    64  		panic("no buffer")
    65  	}
    66  	r.rewound = true
    67  	r.bufReadIdx = 0
    68  }
    69  
    70  func (r *RewindReader) StopBuffering() {
    71  	r.buffering = false
    72  }
    73  
    74  func (r *RewindReader) SetBufferSize(size int) {
    75  	if size == 0 { //disable buffering
    76  		if !r.buffering {
    77  			panic("reader is disabled")
    78  		}
    79  		r.buffering = false
    80  		r.buf = nil
    81  		r.bufReadIdx = 0
    82  		r.bufferSize = 0
    83  	} else {
    84  		if r.buffering {
    85  			panic("reader is buffering")
    86  		}
    87  		r.buffering = true
    88  		r.bufReadIdx = 0
    89  		r.bufferSize = size
    90  		r.buf = make([]byte, 0, size)
    91  	}
    92  }
    93  
    94  type RewindConn struct {
    95  	net.Conn
    96  	*RewindReader
    97  }
    98  
    99  func (c *RewindConn) Read(p []byte) (int, error) {
   100  	return c.RewindReader.Read(p)
   101  }
   102  
   103  func NewRewindConn(conn net.Conn) *RewindConn {
   104  	return &RewindConn{
   105  		Conn: conn,
   106  		RewindReader: &RewindReader{
   107  			rawReader: conn,
   108  		},
   109  	}
   110  }
   111  
   112  type BufferdConn struct {
   113  	net.Conn
   114  	Br *bufio.Reader
   115  }
   116  
   117  func (c *BufferdConn) Read(b []byte) (int, error) {
   118  	return c.Br.Read(b)
   119  }
   120  
   121  func CopyBuffer(dst io.Writer, src io.Reader, buf []byte) (written int64, err error) {
   122  
   123  	for {
   124  		nr, er := src.Read(buf)
   125  		if nr > 0 {
   126  			nw, ew := dst.Write(buf[0:nr])
   127  			if nw > 0 {
   128  				written += int64(nw)
   129  			}
   130  			if ew != nil {
   131  				err = ew
   132  				break
   133  			}
   134  			if nr != nw {
   135  				err = io.ErrShortWrite
   136  				break
   137  			}
   138  		}
   139  		if er != nil {
   140  			if er != io.EOF {
   141  				err = er
   142  			}
   143  			break
   144  		}
   145  	}
   146  	return written, err
   147  }