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

     1  // Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
     2  
     3  package readers
     4  
     5  import (
     6  	"bytes"
     7  	"github.com/iwind/TeaGo/types"
     8  	"io"
     9  	"mime/multipart"
    10  	"net/textproto"
    11  	"regexp"
    12  	"strings"
    13  )
    14  
    15  type OnPartReadHandler func(start int64, end int64, total int64, data []byte, header textproto.MIMEHeader)
    16  
    17  var contentRangeRegexp = regexp.MustCompile(`^(\d+)-(\d+)/(\d+|\*)`)
    18  
    19  type ByteRangesReaderCloser struct {
    20  	BaseReader
    21  
    22  	rawReader io.ReadCloser
    23  	boundary  string
    24  
    25  	mReader *multipart.Reader
    26  	part    *multipart.Part
    27  
    28  	buf   *bytes.Buffer
    29  	isEOF bool
    30  
    31  	onPartReadHandler OnPartReadHandler
    32  	rangeStart        int64
    33  	rangeEnd          int64
    34  	total             int64
    35  
    36  	isStarted bool
    37  	nl        string
    38  }
    39  
    40  func NewByteRangesReaderCloser(reader io.ReadCloser, boundary string) *ByteRangesReaderCloser {
    41  	return &ByteRangesReaderCloser{
    42  		rawReader: reader,
    43  		mReader:   multipart.NewReader(reader, boundary),
    44  		boundary:  boundary,
    45  		buf:       &bytes.Buffer{},
    46  		nl:        "\r\n",
    47  	}
    48  }
    49  
    50  func (this *ByteRangesReaderCloser) Read(p []byte) (n int, err error) {
    51  	n, err = this.read(p)
    52  	return
    53  }
    54  
    55  func (this *ByteRangesReaderCloser) Close() error {
    56  	return this.rawReader.Close()
    57  }
    58  
    59  func (this *ByteRangesReaderCloser) OnPartRead(handler OnPartReadHandler) {
    60  	this.onPartReadHandler = handler
    61  }
    62  
    63  func (this *ByteRangesReaderCloser) read(p []byte) (n int, err error) {
    64  	// read from buffer
    65  	n, err = this.buf.Read(p)
    66  	if !this.isEOF {
    67  		err = nil
    68  	}
    69  	if n > 0 {
    70  		return
    71  	}
    72  	if this.isEOF {
    73  		return
    74  	}
    75  
    76  	if this.part == nil {
    77  		part, partErr := this.mReader.NextPart()
    78  		if partErr != nil {
    79  			if partErr == io.EOF {
    80  				this.buf.WriteString(this.nl + "--" + this.boundary + "--" + this.nl)
    81  				this.isEOF = true
    82  				n, _ = this.buf.Read(p)
    83  				return
    84  			}
    85  
    86  			return 0, partErr
    87  		}
    88  
    89  		if !this.isStarted {
    90  			this.isStarted = true
    91  			this.buf.WriteString("--" + this.boundary + this.nl)
    92  		} else {
    93  			this.buf.WriteString(this.nl + "--" + this.boundary + this.nl)
    94  		}
    95  
    96  		// Headers
    97  		var hasRange = false
    98  		for k, v := range part.Header {
    99  			for _, v1 := range v {
   100  				this.buf.WriteString(k + ": " + v1 + this.nl)
   101  
   102  				// parse range
   103  				if k == "Content-Range" {
   104  					var bytesPrefix = "bytes "
   105  					if strings.HasPrefix(v1, bytesPrefix) {
   106  						var r = v1[len(bytesPrefix):]
   107  						var matches = contentRangeRegexp.FindStringSubmatch(r)
   108  						if len(matches) > 2 {
   109  							var start = types.Int64(matches[1])
   110  							var end = types.Int64(matches[2])
   111  							var total int64 = 0
   112  							if matches[3] != "*" {
   113  								total = types.Int64(matches[3])
   114  							}
   115  							if start <= end {
   116  								hasRange = true
   117  								this.rangeStart = start
   118  								this.rangeEnd = end
   119  								this.total = total
   120  							}
   121  						}
   122  					}
   123  				}
   124  			}
   125  		}
   126  
   127  		if !hasRange {
   128  			this.rangeStart = -1
   129  			this.rangeEnd = -1
   130  		}
   131  
   132  		this.buf.WriteString(this.nl)
   133  		this.part = part
   134  
   135  		n, _ = this.buf.Read(p)
   136  		return
   137  	}
   138  
   139  	n, err = this.part.Read(p)
   140  
   141  	if this.onPartReadHandler != nil && n > 0 && this.rangeStart >= 0 && this.rangeEnd >= 0 {
   142  		this.onPartReadHandler(this.rangeStart, this.rangeEnd, this.total, p[:n], this.part.Header)
   143  		this.rangeStart += int64(n)
   144  	}
   145  
   146  	if err == io.EOF {
   147  		this.part = nil
   148  		err = nil
   149  
   150  		// 如果没有读取到内容，则直接跳到下一个Part
   151  		if n == 0 {
   152  			return this.read(p)
   153  		}
   154  	}
   155  
   156  	return
   157  }