Open-IM-Server/tools/url2im/pkg/manage.go

package pkg

import (
	"bufio"
	"context"
	"crypto/md5"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/OpenIMSDK/protocol/third"
)

type Upload struct {
	URL         string `json:"url"`
	Name        string `json:"name"`
	ContentType string `json:"contentType"`
}

type Task struct {
	Index  int
	Upload Upload
}

type PartInfo struct {
	ContentType string
	PartSize    int64
	PartNum     int
	FileMd5     string
	PartMd5     string
	PartSizes   []int64
	PartMd5s    []string
}

func Run(conf Config) error {
	m := &Manage{
		prefix: time.Now().Format("20060102150405"),
		conf:   &conf,
		ctx:    context.Background(),
	}
	return m.Run()
}

type Manage struct {
	conf      *Config
	ctx       context.Context
	api       *Api
	partLimit *third.PartLimitResp
	prefix    string
	tasks     chan Task
	id        uint64
	success   int64
	failed    int64
}

func (m *Manage) tempFilePath() string {
	return filepath.Join(m.conf.TempDir, fmt.Sprintf("%s_%d", m.prefix, atomic.AddUint64(&m.id, 1)))
}

func (m *Manage) Run() error {
	defer func(start time.Time) {
		log.Printf("run time %s\n", time.Since(start))
	}(time.Now())
	m.api = &Api{
		Api:    m.conf.Api,
		UserID: m.conf.UserID,
		Secret: m.conf.Secret,
		Client: &http.Client{Timeout: m.conf.Timeout},
	}
	var err error
	ctx := context.WithValue(m.ctx, "operationID", fmt.Sprintf("%s_init", m.prefix))
	m.api.Token, err = m.api.GetToken(ctx)
	if err != nil {
		return err
	}
	m.partLimit, err = m.api.GetPartLimit(ctx)
	if err != nil {
		return err
	}
	progress, err := ReadProgress(m.conf.ProgressPath)
	if err != nil {
		return err
	}
	progressFile, err := os.OpenFile(m.conf.ProgressPath, os.O_CREATE|os.O_RDWR|os.O_APPEND, 0666)
	if err != nil {
		return err
	}
	var mutex sync.Mutex
	writeSuccessIndex := func(index int) {
		mutex.Lock()
		defer mutex.Unlock()
		if _, err := progressFile.Write([]byte(strconv.Itoa(index) + "\n")); err != nil {
			log.Printf("write progress err: %v\n", err)
		}
	}
	file, err := os.Open(m.conf.TaskPath)
	if err != nil {
		return err
	}
	m.tasks = make(chan Task, m.conf.Concurrency*2)
	go func() {
		defer file.Close()
		defer close(m.tasks)
		scanner := bufio.NewScanner(file)
		var (
			index int
			num   int
		)
		for scanner.Scan() {
			line := strings.TrimSpace(scanner.Text())
			if line == "" {
				continue
			}
			index++
			if progress.IsUploaded(index) {
				log.Printf("index: %d already uploaded %s\n", index, line)
				continue
			}
			var upload Upload
			if err := json.Unmarshal([]byte(line), &upload); err != nil {
				log.Printf("index: %d json.Unmarshal(%s) err: %v", index, line, err)
				continue
			}
			num++
			m.tasks <- Task{
				Index:  index,
				Upload: upload,
			}
		}
		if num == 0 {
			log.Println("mark all completed")
		}
	}()
	var wg sync.WaitGroup
	wg.Add(m.conf.Concurrency)
	for i := 0; i < m.conf.Concurrency; i++ {
		go func(tid int) {
			defer wg.Done()
			for task := range m.tasks {
				var success bool
				for n := 0; n < m.conf.Retry; n++ {
					ctx := context.WithValue(m.ctx, "operationID", fmt.Sprintf("%s_%d_%d_%d", m.prefix, tid, task.Index, n+1))
					if urlRaw, err := m.RunTask(ctx, task); err == nil {
						writeSuccessIndex(task.Index)
						log.Println("index:", task.Index, "upload success", "urlRaw", urlRaw)
						success = true
						break
					} else {
						log.Printf("index: %d upload: %+v err: %v", task.Index, task.Upload, err)
					}
				}
				if success {
					atomic.AddInt64(&m.success, 1)
				} else {
					atomic.AddInt64(&m.failed, 1)
					log.Printf("index: %d upload: %+v failed", task.Index, task.Upload)
				}
			}
		}(i + 1)
	}
	wg.Wait()
	log.Printf("execution completed success %d failed %d\n", m.success, m.failed)
	return nil
}

func (m *Manage) RunTask(ctx context.Context, task Task) (string, error) {
	resp, err := m.HttpGet(ctx, task.Upload.URL)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	reader, err := NewReader(resp.Body, m.conf.CacheSize, m.tempFilePath())
	if err != nil {
		return "", err
	}
	defer reader.Close()
	part, err := m.getPartInfo(ctx, reader, reader.Size())
	if err != nil {
		return "", err
	}
	var contentType string
	if task.Upload.ContentType == "" {
		contentType = part.ContentType
	} else {
		contentType = task.Upload.ContentType
	}
	initiateMultipartUploadResp, err := m.api.InitiateMultipartUpload(ctx, &third.InitiateMultipartUploadReq{
		Hash:        part.PartMd5,
		Size:        reader.Size(),
		PartSize:    part.PartSize,
		MaxParts:    -1,
		Cause:       "batch-import",
		Name:        task.Upload.Name,
		ContentType: contentType,
	})
	if err != nil {
		return "", err
	}
	if initiateMultipartUploadResp.Upload == nil {
		return initiateMultipartUploadResp.Url, nil
	}
	if _, err := reader.Seek(0, io.SeekStart); err != nil {
		return "", err
	}
	uploadParts := make([]*third.SignPart, part.PartNum)
	for _, part := range initiateMultipartUploadResp.Upload.Sign.Parts {
		uploadParts[part.PartNumber-1] = part
	}
	for i, currentPartSize := range part.PartSizes {
		md5Reader := NewMd5Reader(io.LimitReader(reader, currentPartSize))
		if m.doPut(ctx, m.api.Client, initiateMultipartUploadResp.Upload.Sign, uploadParts[i], md5Reader, currentPartSize); err != nil {
			return "", err
		}
		if md5val := md5Reader.Md5(); md5val != part.PartMd5s[i] {
			return "", fmt.Errorf("upload part %d failed, md5 not match, expect %s, got %s", i, part.PartMd5s[i], md5val)
		}
	}
	urlRaw, err := m.api.CompleteMultipartUpload(ctx, &third.CompleteMultipartUploadReq{
		UploadID:    initiateMultipartUploadResp.Upload.UploadID,
		Parts:       part.PartMd5s,
		Name:        task.Upload.Name,
		ContentType: contentType,
		Cause:       "batch-import",
	})
	if err != nil {
		return "", err
	}
	return urlRaw, nil
}

func (m *Manage) partSize(size int64) (int64, error) {
	if size <= 0 {
		return 0, errors.New("size must be greater than 0")
	}
	if size > m.partLimit.MaxPartSize*int64(m.partLimit.MaxNumSize) {
		return 0, fmt.Errorf("size must be less than %db", m.partLimit.MaxPartSize*int64(m.partLimit.MaxNumSize))
	}
	if size <= m.partLimit.MinPartSize*int64(m.partLimit.MaxNumSize) {
		return m.partLimit.MinPartSize, nil
	}
	partSize := size / int64(m.partLimit.MaxNumSize)
	if size%int64(m.partLimit.MaxNumSize) != 0 {
		partSize++
	}
	return partSize, nil
}

func (m *Manage) partMD5(parts []string) string {
	s := strings.Join(parts, ",")
	md5Sum := md5.Sum([]byte(s))
	return hex.EncodeToString(md5Sum[:])
}

func (m *Manage) getPartInfo(ctx context.Context, r io.Reader, fileSize int64) (*PartInfo, error) {
	partSize, err := m.partSize(fileSize)
	if err != nil {
		return nil, err
	}
	partNum := int(fileSize / partSize)
	if fileSize%partSize != 0 {
		partNum++
	}
	partSizes := make([]int64, partNum)
	for i := 0; i < partNum; i++ {
		partSizes[i] = partSize
	}
	partSizes[partNum-1] = fileSize - partSize*(int64(partNum)-1)
	partMd5s := make([]string, partNum)
	buf := make([]byte, 1024*8)
	fileMd5 := md5.New()
	var contentType string
	for i := 0; i < partNum; i++ {
		h := md5.New()
		r := io.LimitReader(r, partSize)
		for {
			if n, err := r.Read(buf); err == nil {
				if contentType == "" {
					contentType = http.DetectContentType(buf[:n])
				}
				h.Write(buf[:n])
				fileMd5.Write(buf[:n])
			} else if err == io.EOF {
				break
			} else {
				return nil, err
			}
		}
		partMd5s[i] = hex.EncodeToString(h.Sum(nil))
	}
	partMd5Val := m.partMD5(partMd5s)
	fileMd5val := hex.EncodeToString(fileMd5.Sum(nil))
	return &PartInfo{
		ContentType: contentType,
		PartSize:    partSize,
		PartNum:     partNum,
		FileMd5:     fileMd5val,
		PartMd5:     partMd5Val,
		PartSizes:   partSizes,
		PartMd5s:    partMd5s,
	}, nil
}

func (m *Manage) doPut(ctx context.Context, client *http.Client, sign *third.AuthSignParts, part *third.SignPart, reader io.Reader, size int64) error {
	rawURL := part.Url
	if rawURL == "" {
		rawURL = sign.Url
	}
	if len(sign.Query)+len(part.Query) > 0 {
		u, err := url.Parse(rawURL)
		if err != nil {
			return err
		}
		query := u.Query()
		for i := range sign.Query {
			v := sign.Query[i]
			query[v.Key] = v.Values
		}
		for i := range part.Query {
			v := part.Query[i]
			query[v.Key] = v.Values
		}
		u.RawQuery = query.Encode()
		rawURL = u.String()
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodPut, rawURL, reader)
	if err != nil {
		return err
	}
	for i := range sign.Header {
		v := sign.Header[i]
		req.Header[v.Key] = v.Values
	}
	for i := range part.Header {
		v := part.Header[i]
		req.Header[v.Key] = v.Values
	}
	req.ContentLength = size
	resp, err := client.Do(req)
	if err != nil {
		return err
	}
	defer func() {
		_ = resp.Body.Close()
	}()
	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return err
	}
	if resp.StatusCode/200 != 1 {
		return fmt.Errorf("PUT %s part %d failed, status code %d, body %s", rawURL, part.PartNumber, resp.StatusCode, string(body))
	}
	return nil
}

func (m *Manage) HttpGet(ctx context.Context, url string) (*http.Response, error) {
	reqUrl := url
	for {
		request, err := http.NewRequestWithContext(ctx, http.MethodGet, reqUrl, nil)
		if err != nil {
			return nil, err
		}
		DefaultRequestHeader(request.Header)
		response, err := m.api.Client.Do(request)
		if err != nil {
			return nil, err
		}
		if response.StatusCode != http.StatusOK {
			_ = response.Body.Close()
			return nil, fmt.Errorf("http get %s status %s", url, response.Status)
		}
		return response, nil
	}
}