MeteorLight/util/http.go

package util

import (
	"errors"
	"fmt"
	"io"
	"mime"
	"net/http"
	"net/url"
	"strconv"
	"sync"
)

const rangeReaderBufferSize = 1024 * 16

type RangeReadSeekCloser struct {
	uri          *url.URL
	size         int64
	readOffset   int64
	body         io.ReadCloser
	bodyLock     sync.RWMutex
	buf          []byte
	bufferOffset int64
	closed       bool
	fileName     string
}

func NewRangeReadSeekCloser(uri string) (*RangeReadSeekCloser, error) {
	parsedUrl, err := url.Parse(uri)
	if err != nil {
		return nil, err
	}
	r := &RangeReadSeekCloser{
		uri: parsedUrl,
	}

	if err = r.getInformation(); err != nil {
		return nil, err
	}

	return r, nil
}

func (r *RangeReadSeekCloser) GetURI() string {
	return r.uri.String()
}

func (r *RangeReadSeekCloser) getClient() *http.Client {
	return &http.Client{
		CheckRedirect: func(req *http.Request, via []*http.Request) error {
			req.Header.Del("referer")
			return nil
		},
	}
}

func (r *RangeReadSeekCloser) retryConnect() error {
	r.bodyLock.Lock()
	defer r.bodyLock.Unlock()
	if r.closed {
		return errors.New("already closed")
	}

	headers := make(http.Header)
	startOffset := r.readOffset

	expectedLength := r.size - startOffset
	headers.Set("Range", fmt.Sprintf("bytes=%d-", startOffset))
	headers.Set("User-Agent", "MeteorLight/http-reader")
	response, err := r.getClient().Do(&http.Request{
		Method: "GET",
		URL:    r.uri,
		Header: headers,
	})

	if err != nil {
		return err
	}
	r.body = response.Body

	if response.StatusCode != http.StatusPartialContent {
		_ = r.body.Close()
		r.body = nil
		return fmt.Errorf("response status code %d != %d", response.StatusCode, http.StatusPartialContent)
	}

	contentLength, err := strconv.ParseInt(response.Header.Get("content-length"), 10, 0)
	if err != nil {
		_ = r.body.Close()
		r.body = nil
		return errors.New("server response does not have a valid Content-Length")
	}

	if contentLength != expectedLength {
		_ = r.body.Close()
		r.body = nil
		return fmt.Errorf("server returned %d bytes, expected %d", contentLength, expectedLength)
	}

	return nil
}

func (r *RangeReadSeekCloser) Read(p []byte) (n int, err error) {
	r.bodyLock.RLock()
	defer r.bodyLock.RUnlock()

	for retry := 0; retry < 2; retry++ {
		if r.readOffset >= r.size {
			return 0, io.EOF
		}

		if r.readOffset >= r.bufferOffset {
			bufStart := r.readOffset - r.bufferOffset

			if bufStart <= int64(len(r.buf)) {
				bufEnd := bufStart + int64(len(p))
				if bufEnd >= int64(len(r.buf)) {
					bufEnd = int64(len(r.buf))
				}

				if bufEnd-bufStart > 0 {
					copy(p, r.buf[bufStart:bufEnd])
					r.readOffset += bufEnd - bufStart
					return int(bufEnd - bufStart), nil
				}
			}
		}

		if r.body == nil {
			if err = func() error {
				r.bodyLock.RUnlock()
				defer r.bodyLock.RLock()
				return r.retryConnect()
			}(); err != nil {
				continue
			}
		}

		data := make([]byte, rangeReaderBufferSize)
		n, err = r.body.Read(data)
		if err != nil { //detect actual eof, not a disconnection
			if err == io.EOF || err == io.ErrUnexpectedEOF {
				if r.readOffset+int64(n) != r.size {
					continue
				}
			} else {
				continue
			}
		}
		r.buf = data[:n]
		r.bufferOffset = r.readOffset

		readBytes := len(p)
		if n < readBytes {
			readBytes = n
		}

		copy(p[:readBytes], data[:readBytes])
		r.readOffset += int64(readBytes)
		return readBytes, nil
	}

	return 0, errors.New("could not retry")
}

func (r *RangeReadSeekCloser) Seek(offset int64, whence int) (int64, error) {
	r.bodyLock.Lock()
	defer r.bodyLock.Unlock()

	oldOffset := r.readOffset

	switch whence {
	case io.SeekStart:
		r.readOffset = offset

	case io.SeekCurrent:
		r.readOffset += offset

	case io.SeekEnd:
		//todo: maybe without -1?
		r.readOffset = (r.size - 1) - offset
	default:
		return 0, fmt.Errorf("unknown whence %d", whence)
	}

	if oldOffset != r.readOffset {
		if r.body != nil {
			_ = r.body.Close()
		}
		r.body = nil
	}

	if r.readOffset >= r.size {
		return r.readOffset, io.EOF
	} else if r.readOffset < 0 {
		return r.readOffset, io.ErrUnexpectedEOF
	}

	return r.readOffset, nil
}

func (r *RangeReadSeekCloser) Close() error {
	r.bodyLock.Lock()
	defer r.bodyLock.Unlock()
	if r.body != nil {
		r.closed = true
		defer func() {
			r.body = nil
		}()
		return r.body.Close()
	}
	return nil
}

func (r *RangeReadSeekCloser) GetFileName() string {
	return r.fileName
}

func (r *RangeReadSeekCloser) getInformation() error {

	headers := make(http.Header)
	headers.Set("Range", "bytes=0-")
	headers.Set("User-Agent", "MeteorLight/http-reader")

	response, err := r.getClient().Do(&http.Request{
		Method: "HEAD",
		URL:    r.uri,
		Header: headers,
	})

	if err != nil {
		return err
	}
	defer response.Body.Close()
	if response.StatusCode != http.StatusPartialContent {
		return fmt.Errorf("response status code %d != %d, server might not accept Range", response.StatusCode, http.StatusPartialContent)
	}

	if r.size, err = strconv.ParseInt(response.Header.Get("content-length"), 10, 0); err != nil || r.size <= 0 {
		return errors.New("server response does not have a valid Content-Length")
	}
	if response.Header.Get("content-disposition") != "" {
		if _, params, err := mime.ParseMediaType(response.Header.Get("content-disposition")); err == nil {
			if disposition, ok := params["filename"]; ok {
				r.fileName = disposition
			}
		}
	}

	if r.fileName == "" {
		if response.Header.Get("content-type") != "" {
			if mimeType, params, err := mime.ParseMediaType(response.Header.Get("content-type")); err == nil {
				switch mimeType {
				case "audio/flac":
					r.fileName = "_.flac"
				case "audio/ogg":
					r.fileName = "_.ogg"
					if codecs, ok := params["codecs"]; ok {
						switch codecs {
						case "vorbis":
							r.fileName = "_.vorbis"
						case "opus":
							r.fileName = "_.opus"
						case "flac":
							r.fileName = "_.flac"
						}
					}
				case "audio/mpeg":
					r.fileName = "_.mp3"
					if codecs, ok := params["codecs"]; ok {
						switch codecs {
						case "mp3":
							r.fileName = "_.mp3"
						}
					}
				case "audio/mp4":
					r.fileName = "_.m4a"
					if codecs, ok := params["codecs"]; ok {
						switch codecs {
						case "alac":
							r.fileName = "_.alac"
						}
					}
				case "audio/aac":
					r.fileName = "_.aac"
				case "audio/x-tta", "audio/tta":
					r.fileName = "_.tta"
				}
			}
		}
	}

	if r.fileName == "" {
		r.fileName = r.GetURI()
	}

	return nil
}