upload.go

package tgbotapi

import (
	"fmt"
	"io"
	"mime/multipart"
	"strings"
)

type requestPayload struct {
	body        io.Reader
	closer      io.Closer
	contentType string
}

func (p requestPayload) close() {
	if p.closer != nil {
		_ = p.closer.Close()
	}
}

type requestDebug struct {
	params    Params
	fileCount int
}

type uploadPlan struct {
	files  []RequestFile
	inline Params
}

func newUploadPlan() *uploadPlan {
	return &uploadPlan{}
}

func (p *uploadPlan) AddField(name string, data RequestFileData) {
	if data == nil {
		return
	}
	if data.NeedsUpload() {
		p.addUpload(name, data)
		return
	}
	p.addInline(name, data.SendData())
}

func (p *uploadPlan) AddUploadOnly(name string, data RequestFileData) {
	if data == nil || !data.NeedsUpload() {
		return
	}
	p.addUpload(name, data)
}

func (p *uploadPlan) Apply(params Params) Params {
	if len(p.inline) == 0 {
		return params
	}
	if params == nil {
		params = make(Params)
	}
	for key, value := range p.inline {
		params[key] = value
	}
	return params
}

func (p *uploadPlan) Files() []RequestFile {
	return p.files
}

func (p *uploadPlan) NeedsUpload() bool {
	return len(p.files) > 0
}

func (p *uploadPlan) addUpload(name string, data RequestFileData) {
	p.files = append(p.files, RequestFile{
		Name: name,
		Data: data,
	})
}

func (p *uploadPlan) addInline(name, value string) {
	if p.inline == nil {
		p.inline = make(Params)
	}
	p.inline[name] = value
}

func uploadPlanFromFiles(files []RequestFile) *uploadPlan {
	plan := newUploadPlan()
	for _, file := range files {
		plan.AddField(file.Name, file.Data)
	}
	return plan
}

func requestFiles(files ...RequestFile) []RequestFile {
	out := make([]RequestFile, 0, len(files))
	for _, file := range files {
		if file.Data != nil {
			out = append(out, file)
		}
	}
	return out
}

func requestFile(name string, data RequestFileData) RequestFile {
	return RequestFile{
		Name: name,
		Data: data,
	}
}

func buildFormPayload(params Params) requestPayload {
	values := buildParams(params)
	return requestPayload{
		body:        strings.NewReader(values.Encode()),
		contentType: "application/x-www-form-urlencoded",
	}
}

func buildMultipartPayload(params Params, files []RequestFile) (requestPayload, error) {
	reader, writer := io.Pipe()
	multipartWriter := multipart.NewWriter(writer)

	go func() {
		if err := writeMultipartPayload(multipartWriter, params, files); err != nil {
			_ = writer.CloseWithError(err)
			return
		}
		if err := multipartWriter.Close(); err != nil {
			_ = writer.CloseWithError(err)
			return
		}
		_ = writer.Close()
	}()

	return requestPayload{
		body:        reader,
		closer:      reader,
		contentType: multipartWriter.FormDataContentType(),
	}, nil
}

func writeMultipartPayload(writer *multipart.Writer, params Params, files []RequestFile) error {
	for field, value := range params {
		if err := writer.WriteField(field, value); err != nil {
			return fmt.Errorf("write multipart field %q: %w", field, err)
		}
	}

	for _, file := range files {
		if file.Data == nil {
			return fmt.Errorf("multipart file %q has nil data", file.Name)
		}

		if file.Data.NeedsUpload() {
			if err := writeMultipartUpload(writer, file); err != nil {
				return err
			}
			continue
		}

		if err := writer.WriteField(file.Name, file.Data.SendData()); err != nil {
			return fmt.Errorf("write multipart file reference %q: %w", file.Name, err)
		}
	}

	return nil
}

func writeMultipartUpload(writer *multipart.Writer, file RequestFile) error {
	name, reader, err := file.Data.UploadData()
	if err != nil {
		return fmt.Errorf("open upload %q: %w", file.Name, err)
	}
	if reader == nil {
		return fmt.Errorf("open upload %q: nil reader", file.Name)
	}

	part, err := writer.CreateFormFile(file.Name, name)
	if err != nil {
		if closeErr := closeUploadReader(reader); closeErr != nil {
			return fmt.Errorf("create multipart file %q: %w; close upload: %w", file.Name, err, closeErr)
		}
		return fmt.Errorf("create multipart file %q: %w", file.Name, err)
	}

	_, copyErr := io.Copy(part, reader)
	closeErr := closeUploadReader(reader)
	if copyErr != nil {
		if closeErr != nil {
			return fmt.Errorf("copy upload %q: %w; close upload: %w", file.Name, copyErr, closeErr)
		}
		return fmt.Errorf("copy upload %q: %w", file.Name, copyErr)
	}
	if closeErr != nil {
		return fmt.Errorf("close upload %q: %w", file.Name, closeErr)
	}

	return nil
}

func closeUploadReader(reader io.Reader) error {
	if closer, ok := reader.(io.Closer); ok {
		return closer.Close()
	}
	return nil
}

func prepareInputMediaUploadPlan(inputMedia []InputMedia, prefix string) ([]InputMedia, *uploadPlan) {
	prepared := cloneMediaSlice(inputMedia)
	plan := newUploadPlan()

	for idx, media := range prepared {
		if media == nil {
			continue
		}

		name := fmt.Sprintf("%s-%d", prefix, idx)
		prepareInputMediaItem(media, name, plan)
		prepared[idx] = media
	}

	return prepared, plan
}

func prepareInputMediaItem(media InputMedia, name string, plan *uploadPlan) {
	if file := media.getMedia(); file != nil && file.NeedsUpload() {
		media.setUploadMedia("attach://" + name)
		plan.AddUploadOnly(name, file)
	}

	switch input := media.(type) {
	case *InputMediaLivePhoto:
		prepareInputMediaSecondary(input.Photo, name+"-photo", plan, func(ref string) {
			input.Photo = fileAttach(ref)
		})
	case *InputPaidMedia:
		prepareInputPaidMediaPhoto(input, name, plan)
	}

	if thumb := media.getThumb(); thumb != nil && thumb.NeedsUpload() {
		media.setUploadThumb("attach://" + name + "-thumb")
		plan.AddUploadOnly(name+"-thumb", thumb)
	}
}

func prepareInputPaidMediaPhoto(media *InputPaidMedia, name string, plan *uploadPlan) {
	if livePhoto, ok := media.Media.(*InputMediaLivePhoto); ok && livePhoto.Photo != nil {
		prepareInputMediaSecondary(livePhoto.Photo, name+"-photo", plan, func(ref string) {
			livePhoto.Photo = fileAttach(ref)
		})
		return
	}

	prepareInputMediaSecondary(media.Photo, name+"-photo", plan, func(ref string) {
		media.Photo = fileAttach(ref)
	})
}

func prepareInputMediaSecondary(data RequestFileData, name string, plan *uploadPlan, set func(string)) {
	if data == nil || !data.NeedsUpload() {
		return
	}
	set("attach://" + name)
	plan.AddUploadOnly(name, data)
}

func prepareInputProfilePhotoUploadPlan(photo InputProfilePhoto) (InputProfilePhoto, *uploadPlan) {
	prepared := cloneInputProfilePhoto(photo)
	plan := newUploadPlan()
	if prepared == nil {
		return nil, plan
	}

	if media := prepared.getMedia(); media != nil && media.NeedsUpload() {
		prepared.setUploadMedia("attach://file-0")
		plan.AddUploadOnly("file-0", media)
	}

	return prepared, plan
}

func prepareInputStoryContentUploadPlan(content InputStoryContent) (InputStoryContent, *uploadPlan) {
	prepared := cloneInputStoryContent(content)
	plan := newUploadPlan()
	if prepared == nil {
		return nil, plan
	}

	if media := prepared.getMedia(); media != nil && media.NeedsUpload() {
		prepared.setUploadMedia("attach://file-0")
		plan.AddUploadOnly("file-0", media)
	}

	return prepared, plan
}