dragonfly/client/dfget/dfget.go

/*
 *     Copyright 2020 The Dragonfly Authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package dfget

import (
	"context"
	"fmt"
	"io"
	"net/url"
	"os"
	"path"
	"path/filepath"
	"regexp"
	"strings"
	"time"

	"github.com/go-http-utils/headers"
	"github.com/pkg/errors"
	"github.com/schollz/progressbar/v3"

	"d7y.io/dragonfly/v2/client/config"
	logger "d7y.io/dragonfly/v2/internal/dflog"
	"d7y.io/dragonfly/v2/pkg/basic"
	"d7y.io/dragonfly/v2/pkg/rpc/base"
	"d7y.io/dragonfly/v2/pkg/rpc/dfdaemon"
	daemonclient "d7y.io/dragonfly/v2/pkg/rpc/dfdaemon/client"
	"d7y.io/dragonfly/v2/pkg/source"
	"d7y.io/dragonfly/v2/pkg/util/digestutils"
	"d7y.io/dragonfly/v2/pkg/util/stringutils"
)

func Download(cfg *config.DfgetConfig, client daemonclient.DaemonClient) error {
	var (
		ctx       = context.Background()
		cancel    context.CancelFunc
		wLog      = logger.With("url", cfg.URL)
		downError error
	)

	wLog.Info("init success and start to download")
	fmt.Println("init success and start to download")

	if cfg.Timeout > 0 {
		ctx, cancel = context.WithTimeout(ctx, cfg.Timeout)
	} else {
		ctx, cancel = context.WithCancel(ctx)
	}

	go func() {
		downError = download(ctx, client, cfg, wLog)
		cancel()
	}()

	<-ctx.Done()

	if ctx.Err() == context.DeadlineExceeded {
		return errors.Errorf("download timeout(%s)", cfg.Timeout)
	}
	return downError
}

func download(ctx context.Context, client daemonclient.DaemonClient, cfg *config.DfgetConfig, wLog *logger.SugaredLoggerOnWith) error {
	if cfg.Recursive {
		return recursiveDownload(ctx, client, cfg)
	}
	return singleDownload(ctx, client, cfg, wLog)
}

func singleDownload(ctx context.Context, client daemonclient.DaemonClient, cfg *config.DfgetConfig, wLog *logger.SugaredLoggerOnWith) error {
	hdr := parseHeader(cfg.Header)

	if client == nil {
		return downloadFromSource(ctx, cfg, hdr)
	}

	var (
		start     = time.Now()
		stream    *daemonclient.DownResultStream
		result    *dfdaemon.DownResult
		pb        *progressbar.ProgressBar
		request   = newDownRequest(cfg, hdr)
		downError error
	)

	if stream, downError = client.Download(ctx, request); downError == nil {
		if cfg.ShowProgress {
			pb = newProgressBar(-1)
		}

		for {
			if result, downError = stream.Recv(); downError != nil {
				break
			}

			if result.CompletedLength > 0 && pb != nil {
				_ = pb.Set64(int64(result.CompletedLength))
			}

			// success
			if result.Done {
				if pb != nil {
					pb.Describe("Downloaded")
					_ = pb.Close()
				}

				wLog.Infof("download from daemon success, length: %d bytes cost: %d ms", result.CompletedLength, time.Now().Sub(start).Milliseconds())
				fmt.Printf("finish total length %d bytes\n", result.CompletedLength)

				break
			}
		}
	}

	if downError != nil && !cfg.KeepOriginalOffset {
		wLog.Warnf("daemon downloads file error: %v", downError)
		fmt.Printf("daemon downloads file error: %v\n", downError)
		downError = downloadFromSource(ctx, cfg, hdr)
	}

	return downError
}

func downloadFromSource(ctx context.Context, cfg *config.DfgetConfig, hdr map[string]string) error {
	if cfg.DisableBackSource {
		return errors.New("try to download from source but back source is disabled")
	}

	var (
		wLog     = logger.With("url", cfg.URL)
		start    = time.Now()
		target   *os.File
		response *source.Response
		err      error
		written  int64
	)

	wLog.Info("try to download from source and ignore rate limit")
	fmt.Println("try to download from source and ignore rate limit")

	if target, err = os.CreateTemp(filepath.Dir(cfg.Output), ".df_"); err != nil {
		return err
	}
	defer os.Remove(target.Name())
	defer target.Close()

	downloadRequest, err := source.NewRequestWithContext(ctx, cfg.URL, hdr)
	if err != nil {
		return err
	}
	if response, err = source.Download(downloadRequest); err != nil {
		return err
	}
	defer response.Body.Close()

	if written, err = io.Copy(target, response.Body); err != nil {
		return err
	}

	if !stringutils.IsBlank(cfg.Digest) {
		parsedHash := digestutils.Parse(cfg.Digest)
		realHash := digestutils.HashFile(target.Name(), digestutils.Algorithms[parsedHash[0]])

		if realHash != "" && realHash != parsedHash[1] {
			return errors.Errorf("%s digest is not matched: real[%s] expected[%s]", parsedHash[0], realHash, parsedHash[1])
		}
	}

	// change file owner
	if err = os.Chown(target.Name(), basic.UserID, basic.UserGroup); err != nil {
		return errors.Wrapf(err, "change file owner to uid[%d] gid[%d]", basic.UserID, basic.UserGroup)
	}

	if err = os.Rename(target.Name(), cfg.Output); err != nil {
		return err
	}

	wLog.Infof("download from source success, length: %d bytes cost: %d ms", written, time.Now().Sub(start).Milliseconds())
	fmt.Printf("finish total length %d bytes\n", written)

	return nil
}

func parseHeader(s []string) map[string]string {
	hdr := make(map[string]string)
	var key, value string
	for _, h := range s {
		idx := strings.Index(h, ":")
		if idx > 0 {
			key = strings.TrimSpace(h[:idx])
			value = strings.TrimSpace(h[idx+1:])
			hdr[key] = value
		}
	}

	return hdr
}

func newDownRequest(cfg *config.DfgetConfig, hdr map[string]string) *dfdaemon.DownRequest {
	var rg string
	if r, ok := hdr[headers.Range]; ok {
		rg = strings.TrimLeft(r, "bytes=")
	} else {
		rg = cfg.Range
	}
	return &dfdaemon.DownRequest{
		Url:               cfg.URL,
		Output:            cfg.Output,
		Timeout:           uint64(cfg.Timeout),
		Limit:             float64(cfg.RateLimit),
		DisableBackSource: cfg.DisableBackSource,
		UrlMeta: &base.UrlMeta{
			Digest: cfg.Digest,
			Tag:    cfg.Tag,
			Range:  rg,
			Filter: cfg.Filter,
			Header: hdr,
		},
		Pattern:            cfg.Pattern,
		Callsystem:         cfg.CallSystem,
		Uid:                int64(basic.UserID),
		Gid:                int64(basic.UserGroup),
		KeepOriginalOffset: cfg.KeepOriginalOffset,
	}
}

func newProgressBar(max int64) *progressbar.ProgressBar {
	return progressbar.NewOptions64(max,
		progressbar.OptionShowBytes(true),
		progressbar.OptionShowIts(),
		progressbar.OptionSetPredictTime(true),
		progressbar.OptionUseANSICodes(true),
		progressbar.OptionEnableColorCodes(true),
		progressbar.OptionFullWidth(),
		progressbar.OptionSetDescription("[cyan]Downloading...[reset]"),
		progressbar.OptionSetRenderBlankState(true),
		progressbar.OptionSetTheme(progressbar.Theme{
			Saucer:        "[green]=[reset]",
			SaucerHead:    "[green]>[reset]",
			SaucerPadding: " ",
			BarStart:      "[",
			BarEnd:        "]",
		}))
}

func accept(u string, parent, sub string, level uint, accept, reject string) bool {
	if !checkDirectoryLevel(parent, sub, level) {
		logger.Debugf("url %s not accept by level: %d", u, level)
		return false
	}
	if !acceptRegex(u, accept) {
		logger.Debugf("url %s not accept by regex: %s", u, accept)
		return false
	}
	if rejectRegex(u, reject) {
		logger.Debugf("url %s rejected by regex: %s", u, reject)
		return false
	}
	return true
}

func checkDirectoryLevel(parent, sub string, level uint) bool {
	if level == 0 {
		return true
	}
	dirs := strings.Split(strings.Trim(strings.TrimLeft(sub, parent), "/"), "/")
	return uint(len(dirs)) <= level
}

func acceptRegex(u string, accept string) bool {
	if accept == "" {
		return true
	}
	return regexp.MustCompile(accept).Match([]byte(u))
}

func rejectRegex(u string, reject string) bool {
	if reject == "" {
		return false
	}
	return regexp.MustCompile(reject).Match([]byte(u))
}

func recursiveDownload(ctx context.Context, client daemonclient.DaemonClient, cfg *config.DfgetConfig) error {
	request, err := source.NewRequestWithContext(ctx, cfg.URL, parseHeader(cfg.Header))
	if err != nil {
		return err
	}
	dirURL, err := url.Parse(cfg.URL)
	if err != nil {
		return err
	}
	logger.Debugf("dirURL: %s", cfg.URL)

	urls, err := source.List(request)
	if err != nil {
		return err
	}
	for _, u := range urls {
		// reuse dfget config
		c := *cfg
		// update some attributes
		c.Recursive, c.URL, c.Output = false, u.String(), path.Join(cfg.Output, strings.TrimPrefix(u.Path, dirURL.Path))
		if !accept(c.URL, dirURL.Path, u.Path, cfg.RecursiveLevel, cfg.RecursiveAcceptRegex, cfg.RecursiveRejectRegex) {
			logger.Debugf("url %s is not accepted, skip", c.URL)
			continue
		}
		if cfg.RecursiveList {
			fmt.Printf("%s\n", u.String())
			continue
		}
		// validate new dfget config
		if err = c.Validate(); err != nil {
			logger.Errorf("validate failed: %s", err)
			return err
		}

		logger.Debugf("download %s to %s", c.URL, c.Output)
		err = download(ctx, client, &c, logger.With("url", c.URL))
		if err != nil {
			return err
		}
	}
	return nil
}