github.com/searKing/golang/go@v1.2.117/io/seeker.go (about)

     1  // Copyright 2021 The searKing Author. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package io
     6  
     7  import (
     8  	"errors"
     9  	"io"
    10  )
    11  
    12  var errWhence = errors.New("Seek: invalid whence")
    13  var errOffset = errors.New("Seek: invalid offset")
    14  var errSeeker = errors.New("Seek: can't seek")
    15  
    16  // SeekerLen returns the length of the file and an error, if any.
    17  func SeekerLen(s io.Seeker) (int64, error) {
    18  	curOffset, err := s.Seek(0, io.SeekCurrent)
    19  	if err != nil {
    20  		return 0, err
    21  	}
    22  
    23  	endOffset, err := s.Seek(0, io.SeekEnd)
    24  	if err != nil {
    25  		return 0, err
    26  	}
    27  
    28  	_, err = s.Seek(curOffset, io.SeekStart)
    29  	if err != nil {
    30  		return 0, err
    31  	}
    32  
    33  	return endOffset - curOffset, nil
    34  }
    35  
    36  // SniffCopy copies the seekable reader to an io.Writer
    37  func SniffCopy(dst io.Writer, src io.ReadSeeker) (int64, error) {
    38  	curPos, err := src.Seek(0, io.SeekCurrent)
    39  	if err != nil {
    40  		return 0, err
    41  	}
    42  
    43  	// copy errors may be assumed to be from the body.
    44  	n, err := io.Copy(dst, src)
    45  	if err != nil {
    46  		return n, err
    47  	}
    48  
    49  	// seek back to the first position after reading to reset
    50  	// the body for transmission.
    51  	_, err = src.Seek(curPos, io.SeekStart)
    52  	if err != nil {
    53  		return n, err
    54  	}
    55  
    56  	return n, nil
    57  }
    58  
    59  // SniffRead reads up to len(p) bytes into p.
    60  func SniffRead(p []byte, src io.ReadSeeker) (int, error) {
    61  	curPos, err := src.Seek(0, io.SeekCurrent)
    62  	if err != nil {
    63  		return 0, err
    64  	}
    65  
    66  	// copy errors may be assumed to be from the body.
    67  	n, err := src.Read(p)
    68  	if err != nil {
    69  		return n, err
    70  	}
    71  
    72  	// seek back to the first position after reading to reset
    73  	// the body for transmission.
    74  	_, err = src.Seek(curPos, io.SeekStart)
    75  	if err != nil {
    76  		return n, err
    77  	}
    78  
    79  	return n, nil
    80  }
    81  
    82  // LimitReadSeeker returns a Reader that reads from r
    83  // but stops with EOF after n bytes.
    84  // The underlying implementation is a *LimitedReader.
    85  func LimitReadSeeker(r io.ReadSeeker, n int64) io.ReadSeeker { return &LimitedReadSeeker{r, n} }
    86  
    87  // LimitedReadSeeker A LimitSeekable reads from R but limits the size of the file N bytes.
    88  // Read returns EOF when N <= 0 or when the underlying R returns EOF.
    89  type LimitedReadSeeker struct {
    90  	rs    io.ReadSeeker // underlying readSeeker
    91  	limit int64         // max bytes remaining
    92  }
    93  
    94  func (l *LimitedReadSeeker) Read(p []byte) (n int, err error) {
    95  	// speedup
    96  	if l.limit <= 0 {
    97  		return 0, io.EOF
    98  	}
    99  
   100  	offset, err := l.rs.Seek(0, io.SeekCurrent)
   101  	if err != nil {
   102  		return 0, errOffset
   103  	}
   104  
   105  	readLimit := l.limit - offset
   106  
   107  	if readLimit <= 0 {
   108  		return 0, io.EOF
   109  	}
   110  
   111  	if int64(len(p)) > readLimit {
   112  		p = p[0:readLimit]
   113  	}
   114  	n, err = l.rs.Read(p)
   115  	return
   116  }
   117  
   118  func (l *LimitedReadSeeker) Seek(offset int64, whence int) (int64, error) {
   119  	lastPos, err := l.rs.Seek(0, io.SeekCurrent)
   120  	if err != nil {
   121  		return 0, errOffset
   122  	}
   123  
   124  	size, err := l.rs.Seek(offset, whence)
   125  	if err != nil {
   126  		return 0, errSeeker
   127  	}
   128  	if size > l.limit {
   129  		// recover if overflow
   130  		_, _ = l.rs.Seek(lastPos, io.SeekStart)
   131  		return 0, errOffset
   132  	}
   133  	return size, nil
   134  }