photos-api/cmd/photos-api-cli/upload.go

package main

import (
	"crypto/sha1"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"sync"

	"github.com/go-resty/resty/v2"
	"github.com/schollz/progressbar/v3"
	photosapi "gitlab.celogeek.com/photos/api/internal/photos/api"
)

type UploadCommand struct {
	Url     string `short:"u" long:"url" description:"Url of the instance" required:"true"`
	Token   string `short:"t" long:"token" description:"Token of the instance" required:"true"`
	File    string `short:"f" long:"file" description:"File to upload" required:"true"`
	Workers uint32 `short:"w" long:"workers" description:"Number of workers for uploading chunks" default:"4"`
}

type UploadError struct {
	Error string
}

type UploadFileRequest struct {
	Name     string
	Checksum string
	Chunks   []string
}

type UploadFileResponse struct {
	Sum      string
	NbChunks uint32
	Size     uint64
}

func (c *UploadCommand) Cli() *resty.Client {
	return resty.New().SetBaseURL(c.Url).SetAuthScheme("Private").SetAuthToken(c.Token)
}

func (c *UploadCommand) FileExists() (string, error) {
	f, err := os.Open(c.File)
	if err != nil {
		return "", err
	}
	defer f.Close()

	st, err := f.Stat()
	if err != nil {
		return "", err
	}

	progress := progressbar.DefaultBytes(st.Size(), fmt.Sprintf("Checking %s", filepath.Base(c.File)))
	defer progress.Close()
	tee := io.TeeReader(f, progress)
	checksum := sha1.New()
	io.Copy(checksum, tee)
	sum := hex.EncodeToString(checksum.Sum(nil))

	resp, err := c.Cli().R().Head(fmt.Sprintf("/file/%s", sum))
	if err != nil {
		return "", err
	}
	if resp.IsSuccess() {
		return "", errors.New("file already exists")
	}
	return sum, nil
}

func (c *UploadCommand) FileUpload(sum string) error {
	f, err := os.Open(c.File)
	if err != nil {
		return err
	}
	defer f.Close()

	st, err := f.Stat()
	if err != nil {
		return err
	}
	nbChunks := st.Size() / photosapi.CHUNK_SIZE
	if st.Size()%photosapi.CHUNK_SIZE > 0 {
		nbChunks++
	}

	uploadFile := &UploadFileRequest{
		Name:     filepath.Base(c.File),
		Chunks:   make([]string, nbChunks),
		Checksum: sum,
	}

	cli := c.Cli()
	progress := progressbar.DefaultBytes(st.Size(), fmt.Sprintf("Uploading %s", uploadFile.Name))
	defer progress.Close()

	wg := sync.WaitGroup{}
	mu := sync.Mutex{}
	wg.Add(4)
	wgErrors := make([]error, c.Workers)
	i := int64(0)
	for w := uint32(0); w < c.Workers; w++ {
		go func(w uint32) {
			defer wg.Done()
			b := make([]byte, photosapi.CHUNK_SIZE)
			for {
				mu.Lock()
				part := i
				i++
				n, err := f.Read(b)
				mu.Unlock()
				if n == 0 {
					if err == io.EOF {
						break
					}
					wgErrors[w] = err
					return
				}

				checksum := sha1.New()
				checksum.Write(b[0:n])
				sum := hex.EncodeToString(checksum.Sum(nil))

				resp, err := cli.R().Head(fmt.Sprintf("/file/chunk/%s", sum))
				if err != nil {
					wgErrors[w] = err
					return
				}

				if resp.IsSuccess() {
					uploadFile.Chunks[part] = sum
					progress.Add(n)
					continue
				}

				resp, err = cli.R().SetError(&UploadError{}).SetBody(b[0:n]).Post("/file/chunk")
				if err != nil {
					wgErrors[w] = err
					return
				}

				if err, ok := resp.Error().(*UploadError); ok {
					wgErrors[w] = errors.New(err.Error)
					return
				}

				uploadFile.Chunks[part] = sum
				progress.Add(n)
			}
		}(w)
	}
	wg.Wait()
	for _, err := range wgErrors {
		if err != nil {
			return err
		}
	}

	resp, err := cli.R().SetBody(uploadFile).SetError(&UploadError{}).SetResult(&UploadFileResponse{}).Post("/file")
	if err != nil {
		return err
	}

	if err, ok := resp.Error().(*UploadError); ok {
		logger.Println("Upload failed")
		return errors.New(err.Error)
	}

	if result, ok := resp.Result().(*UploadFileResponse); ok {
		fmt.Printf("Upload succeed\nSum: %s\nNbChunks: %d\nSize: %d\n", result.Sum, result.NbChunks, result.Size)
	}

	return nil
}

func (c *UploadCommand) Execute(args []string) error {
	sum, err := c.FileExists()
	if err != nil {
		return err
	}
	return c.FileUpload(sum)
}

func init() {
	parser.AddCommand("upload", "Upload a file", "", &UploadCommand{})
}