frostfs-http-gw/internal/handler/download.go

package handler

import (
	"archive/tar"
	"archive/zip"
	"bufio"
	"compress/gzip"
	"context"
	"fmt"
	"io"
	"net/url"
	"time"

	"git.frostfs.info/TrueCloudLab/frostfs-http-gw/internal/data"
	"git.frostfs.info/TrueCloudLab/frostfs-http-gw/internal/logs"
	"git.frostfs.info/TrueCloudLab/frostfs-http-gw/response"
	"git.frostfs.info/TrueCloudLab/frostfs-http-gw/utils"
	cid "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/container/id"
	"git.frostfs.info/TrueCloudLab/frostfs-sdk-go/object"
	oid "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/object/id"
	"github.com/valyala/fasthttp"
	"go.uber.org/zap"
)

// DownloadByAddressOrBucketName handles download requests using simple cid/oid or bucketname/key format.
func (h *Handler) DownloadByAddressOrBucketName(c *fasthttp.RequestCtx) {
	oidURLParam := c.UserValue("oid").(string)
	downloadQueryParam := c.QueryArgs().GetBool("download")

	switch {
	case isObjectID(oidURLParam):
		h.byNativeAddress(c, h.receiveFile)
	case !isContainerRoot(oidURLParam) && (downloadQueryParam || !isDir(oidURLParam)):
		h.byS3Path(c, h.receiveFile)
	default:
		h.browseIndex(c)
	}
}

// DownloadByAttribute handles attribute-based download requests.
func (h *Handler) DownloadByAttribute(c *fasthttp.RequestCtx) {
	h.byAttribute(c, h.receiveFile)
}

func (h *Handler) search(ctx context.Context, cnrID cid.ID, key, val string, op object.SearchMatchType) (ResObjectSearch, error) {
	filters := object.NewSearchFilters()
	filters.AddRootFilter()
	filters.AddFilter(key, val, op)

	prm := PrmObjectSearch{
		PrmAuth: PrmAuth{
			BearerToken: bearerToken(ctx),
		},
		Container: cnrID,
		Filters:   filters,
	}

	return h.frostfs.SearchObjects(ctx, prm)
}

// DownloadZipped handles zip by prefix requests.
func (h *Handler) DownloadZipped(c *fasthttp.RequestCtx) {
	scid, _ := c.UserValue("cid").(string)

	ctx := utils.GetContextFromRequest(c)
	log := utils.GetReqLogOrDefault(ctx, h.log)
	bktInfo, err := h.getBucketInfo(ctx, scid, log)
	if err != nil {
		logAndSendBucketError(c, log, err)
		return
	}
	resSearch, err := h.searchObjectsByPrefix(c, log, bktInfo.CID)
	if err != nil {
		return
	}

	c.Response.Header.Set(fasthttp.HeaderContentType, "application/zip")
	c.Response.Header.Set(fasthttp.HeaderContentDisposition, "attachment; filename=\"archive.zip\"")

	c.SetBodyStreamWriter(h.getZipResponseWriter(ctx, log, resSearch, bktInfo))
}

func (h *Handler) getZipResponseWriter(ctx context.Context, log *zap.Logger, resSearch ResObjectSearch, bktInfo *data.BucketInfo) func(w *bufio.Writer) {
	return func(w *bufio.Writer) {
		defer resSearch.Close()

		zipWriter := zip.NewWriter(w)
		var bufZip []byte

		errIter := resSearch.Iterate(h.putObjectToArchive(ctx, log, bktInfo.CID, &bufZip,
			func(obj *object.Object) (io.Writer, error) {
				return h.createZipFile(zipWriter, obj)
			}),
		)
		if errIter != nil {
			log.Error(logs.IteratingOverSelectedObjectsFailed, zap.Error(errIter))
		} else if bufZip == nil {
			log.Error(logs.ObjectsNotFound)
		}
		if err := zipWriter.Close(); err != nil {
			log.Error(logs.CloseZipWriter, zap.Error(err))
		}
	}
}

func (h *Handler) createZipFile(zw *zip.Writer, obj *object.Object) (io.Writer, error) {
	method := zip.Store
	if h.config.ZipCompression() {
		method = zip.Deflate
	}

	filePath := getFilePath(obj)
	if len(filePath) == 0 || filePath[len(filePath)-1] == '/' {
		return nil, fmt.Errorf("invalid filepath '%s'", filePath)
	}

	return zw.CreateHeader(&zip.FileHeader{
		Name:     filePath,
		Method:   method,
		Modified: time.Now(),
	})
}

// DownloadTar forms tar.gz from objects by prefix.
func (h *Handler) DownloadTar(c *fasthttp.RequestCtx) {
	scid, _ := c.UserValue("cid").(string)

	ctx := utils.GetContextFromRequest(c)
	log := utils.GetReqLogOrDefault(ctx, h.log)
	bktInfo, err := h.getBucketInfo(ctx, scid, log)
	if err != nil {
		logAndSendBucketError(c, log, err)
		return
	}
	resSearch, err := h.searchObjectsByPrefix(c, log, bktInfo.CID)
	if err != nil {
		return
	}

	c.Response.Header.Set(fasthttp.HeaderContentType, "application/x-gzip")
	c.Response.Header.Set(fasthttp.HeaderContentDisposition, "attachment; filename=\"archive.tar.gz\"")

	c.SetBodyStreamWriter(h.getTarResponseWriter(ctx, log, resSearch, bktInfo))
}

func (h *Handler) getTarResponseWriter(ctx context.Context, log *zap.Logger, resSearch ResObjectSearch, bktInfo *data.BucketInfo) func(w *bufio.Writer) {
	return func(w *bufio.Writer) {
		defer resSearch.Close()

		var gzipWriter *gzip.Writer
		if h.config.ZipCompression() {
			gzipWriter, _ = gzip.NewWriterLevel(w, gzip.DefaultCompression)
		} else {
			gzipWriter, _ = gzip.NewWriterLevel(w, gzip.NoCompression)
		}
		defer func() {
			if err := gzipWriter.Close(); err != nil {
				log.Error(logs.CloseGzipWriter, zap.Error(err))
			}
		}()

		tarWriter := tar.NewWriter(gzipWriter)
		defer func() {
			if err := tarWriter.Close(); err != nil {
				log.Error(logs.CloseTarWriter, zap.Error(err))
			}
		}()

		var bufZip []byte

		errIter := resSearch.Iterate(h.putObjectToArchive(ctx, log, bktInfo.CID, &bufZip,
			func(obj *object.Object) (io.Writer, error) {
				return h.createTarFile(tarWriter, obj)
			}),
		)
		if errIter != nil {
			log.Error(logs.IteratingOverSelectedObjectsFailed, zap.Error(errIter))
		} else if bufZip == nil {
			log.Error(logs.ObjectsNotFound)
		}
	}
}

func (h *Handler) createTarFile(tw *tar.Writer, obj *object.Object) (io.Writer, error) {
	filePath := getFilePath(obj)
	if len(filePath) == 0 || filePath[len(filePath)-1] == '/' {
		return nil, fmt.Errorf("invalid filepath '%s'", filePath)
	}

	return tw, tw.WriteHeader(&tar.Header{
		Name: filePath,
		Mode: 0655,
		Size: int64(obj.PayloadSize()),
	})
}

func (h *Handler) putObjectToArchive(ctx context.Context, log *zap.Logger, cnrID cid.ID, bufZip *[]byte, createArchiveHeader func(obj *object.Object) (io.Writer, error)) func(id oid.ID) bool {
	return func(id oid.ID) bool {
		log := log.With(zap.String("oid", id.EncodeToString()))

		if *bufZip == nil {
			*bufZip = make([]byte, 3<<20) // the same as for upload
		}

		prm := PrmObjectGet{
			PrmAuth: PrmAuth{
				BearerToken: bearerToken(ctx),
			},
			Address: newAddress(cnrID, id),
		}

		resGet, err := h.frostfs.GetObject(ctx, prm)
		if err != nil {
			log.Error(logs.FailedToGetObject, zap.Error(err))
			return false
		}

		fileWriter, err := createArchiveHeader(&resGet.Header)
		if err != nil {
			log.Error(logs.FailedToAddObjectToArchive, zap.Error(err))
			return false
		}

		if err = h.writeToArchive(resGet, fileWriter, *bufZip); err != nil {
			log.Error(logs.FailedToAddObjectToArchive, zap.Error(err))
			return false
		}

		return false
	}
}

func (h *Handler) searchObjectsByPrefix(c *fasthttp.RequestCtx, log *zap.Logger, cnrID cid.ID) (ResObjectSearch, error) {
	scid := cnrID.EncodeToString()
	prefix, _ := c.UserValue("prefix").(string)

	ctx := utils.GetContextFromRequest(c)

	prefix, err := url.QueryUnescape(prefix)
	if err != nil {
		log.Error(logs.FailedToUnescapeQuery, zap.String("cid", scid), zap.String("prefix", prefix), zap.Error(err))
		response.Error(c, "could not unescape prefix: "+err.Error(), fasthttp.StatusBadRequest)
		return nil, err
	}

	log = log.With(zap.String("cid", scid), zap.String("prefix", prefix))

	resSearch, err := h.search(ctx, cnrID, object.AttributeFilePath, prefix, object.MatchCommonPrefix)
	if err != nil {
		log.Error(logs.CouldNotSearchForObjects, zap.Error(err))
		response.Error(c, "could not search for objects: "+err.Error(), fasthttp.StatusBadRequest)
		return nil, err
	}
	return resSearch, nil
}

func (h *Handler) writeToArchive(resGet *Object, objWriter io.Writer, bufZip []byte) error {
	var err error
	if _, err = io.CopyBuffer(objWriter, resGet.Payload, bufZip); err != nil {
		return fmt.Errorf("copy object payload to zip file: %v", err)
	}

	if err = resGet.Payload.Close(); err != nil {
		return fmt.Errorf("object body close error: %w", err)
	}

	return nil
}

func getFilePath(obj *object.Object) string {
	for _, attr := range obj.Attributes() {
		if attr.Key() == object.AttributeFilePath {
			return attr.Value()
		}
	}

	return ""
}