github.com/cloudwego/hertz@v0.9.3/pkg/common/utils/ioutil.go (about)

     1  /*
     2   * Copyright 2022 CloudWeGo Authors
     3   *
     4   * Licensed under the Apache License, Version 2.0 (the "License");
     5   * you may not use this file except in compliance with the License.
     6   * You may obtain a copy of the License at
     7   *
     8   *     http://www.apache.org/licenses/LICENSE-2.0
     9   *
    10   * Unless required by applicable law or agreed to in writing, software
    11   * distributed under the License is distributed on an "AS IS" BASIS,
    12   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13   * See the License for the specific language governing permissions and
    14   * limitations under the License.
    15   */
    16  
    17  package utils
    18  
    19  import (
    20  	"io"
    21  
    22  	"github.com/cloudwego/hertz/pkg/network"
    23  )
    24  
    25  func CopyBuffer(dst network.Writer, src io.Reader, buf []byte) (written int64, err error) {
    26  	if buf != nil && len(buf) == 0 {
    27  		panic("empty buffer in io.CopyBuffer")
    28  	}
    29  	return copyBuffer(dst, src, buf)
    30  }
    31  
    32  // copyBuffer is the actual implementation of Copy and CopyBuffer.
    33  // If buf is nil, one is allocated.
    34  func copyBuffer(dst network.Writer, src io.Reader, buf []byte) (written int64, err error) {
    35  	if wt, ok := src.(io.WriterTo); ok {
    36  		if w, ok := dst.(io.Writer); ok {
    37  			return wt.WriteTo(w)
    38  		}
    39  	}
    40  
    41  	// Sendfile impl
    42  	if rf, ok := dst.(io.ReaderFrom); ok {
    43  		return rf.ReadFrom(src)
    44  	}
    45  
    46  	if buf == nil {
    47  		size := 32 * 1024
    48  		if l, ok := src.(*io.LimitedReader); ok && int64(size) > l.N {
    49  			if l.N < 1 {
    50  				size = 1
    51  			} else {
    52  				size = int(l.N)
    53  			}
    54  		}
    55  		buf = make([]byte, size)
    56  	}
    57  	for {
    58  		nr, er := src.Read(buf)
    59  		if nr > 0 {
    60  			nw, eb := dst.WriteBinary(buf[:nr])
    61  			if eb != nil {
    62  				err = eb
    63  				return
    64  			}
    65  
    66  			if nw > 0 {
    67  				written += int64(nw)
    68  			}
    69  			if nr != nw {
    70  				err = io.ErrShortWrite
    71  				return
    72  			}
    73  			if err = dst.Flush(); err != nil {
    74  				return
    75  			}
    76  		}
    77  		if er != nil {
    78  			if er != io.EOF {
    79  				err = er
    80  			}
    81  			break
    82  		}
    83  	}
    84  	return
    85  }
    86  
    87  func CopyZeroAlloc(w network.Writer, r io.Reader) (int64, error) {
    88  	vbuf := CopyBufPool.Get()
    89  	buf := vbuf.([]byte)
    90  	n, err := CopyBuffer(w, r, buf)
    91  	CopyBufPool.Put(vbuf)
    92  	return n, err
    93  }