frostfs-node/cmd/frostfs-cli/modules/object/range.go

package object

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"os"
	"strconv"
	"strings"

	internalclient "git.frostfs.info/TrueCloudLab/frostfs-node/cmd/frostfs-cli/internal/client"
	"git.frostfs.info/TrueCloudLab/frostfs-node/cmd/frostfs-cli/internal/commonflags"
	"git.frostfs.info/TrueCloudLab/frostfs-node/cmd/frostfs-cli/internal/key"
	commonCmd "git.frostfs.info/TrueCloudLab/frostfs-node/cmd/internal/common"
	cid "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/container/id"
	objectSDK "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/object"
	oid "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/object/id"
	"github.com/spf13/cobra"
)

var objectRangeCmd = &cobra.Command{
	Use:   "range",
	Short: "Get payload range data of an object",
	Long:  "Get payload range data of an object",
	Run:   getObjectRange,
}

func initObjectRangeCmd() {
	commonflags.Init(objectRangeCmd)
	initFlagSession(objectRangeCmd, "RANGE")

	flags := objectRangeCmd.Flags()

	flags.String(commonflags.CIDFlag, "", commonflags.CIDFlagUsage)
	_ = objectRangeCmd.MarkFlagRequired(commonflags.CIDFlag)

	flags.String(commonflags.OIDFlag, "", commonflags.OIDFlagUsage)
	_ = objectRangeCmd.MarkFlagRequired(commonflags.OIDFlag)

	flags.String("range", "", "Range to take data from in the form offset:length")
	flags.String(fileFlag, "", "File to write object payload to. Default: stdout.")
	flags.Bool(rawFlag, false, rawFlagDesc)
}

func getObjectRange(cmd *cobra.Command, _ []string) {
	var cnr cid.ID
	var obj oid.ID

	objAddr := readObjectAddress(cmd, &cnr, &obj)

	ranges, err := getRangeList(cmd)
	commonCmd.ExitOnErr(cmd, "", err)

	if len(ranges) != 1 {
		commonCmd.ExitOnErr(cmd, "", fmt.Errorf("exactly one range must be specified, got: %d", len(ranges)))
	}

	var out io.Writer

	filename := cmd.Flag(fileFlag).Value.String()
	if filename == "" {
		out = os.Stdout
	} else {
		f, err := os.OpenFile(filename, os.O_CREATE|os.O_WRONLY, os.ModePerm)
		if err != nil {
			commonCmd.ExitOnErr(cmd, "", fmt.Errorf("can't open file '%s': %w", filename, err))
		}

		defer f.Close()

		out = f
	}

	pk := key.GetOrGenerate(cmd)

	cli := internalclient.GetSDKClientByFlag(cmd, pk, commonflags.RPC)

	var prm internalclient.PayloadRangePrm
	prm.SetClient(cli)
	Prepare(cmd, &prm)
	readSession(cmd, &prm, pk, cnr, obj)

	raw, _ := cmd.Flags().GetBool(rawFlag)
	prm.SetRawFlag(raw)
	prm.SetAddress(objAddr)
	prm.SetRange(&ranges[0])
	prm.SetPayloadWriter(out)

	_, err = internalclient.PayloadRange(cmd.Context(), prm)
	if err != nil {
		if ok := printSplitInfoErr(cmd, err); ok {
			return
		}

		commonCmd.ExitOnErr(cmd, "can't get object payload range: %w", err)
	}

	if filename != "" {
		cmd.Printf("[%s] Payload successfully saved\n", filename)
	}
}

func printSplitInfoErr(cmd *cobra.Command, err error) bool {
	var errSplitInfo *objectSDK.SplitInfoError

	ok := errors.As(err, &errSplitInfo)

	if ok {
		cmd.PrintErrln("Object is complex, split information received.")
		printSplitInfo(cmd, errSplitInfo.SplitInfo())
	}

	return ok
}

func printSplitInfo(cmd *cobra.Command, info *objectSDK.SplitInfo) {
	bs, err := marshalSplitInfo(cmd, info)
	commonCmd.ExitOnErr(cmd, "can't marshal split info: %w", err)

	cmd.Println(string(bs))
}

func marshalSplitInfo(cmd *cobra.Command, info *objectSDK.SplitInfo) ([]byte, error) {
	toJSON, _ := cmd.Flags().GetBool(commonflags.JSON)
	toProto, _ := cmd.Flags().GetBool("proto")
	switch {
	case toJSON && toProto:
		return nil, errors.New("'--json' and '--proto' flags are mutually exclusive")
	case toJSON:
		return info.MarshalJSON()
	case toProto:
		return info.Marshal()
	default:
		b := bytes.NewBuffer(nil)
		if splitID := info.SplitID(); splitID != nil {
			b.WriteString("Split ID: " + splitID.String() + "\n")
		}
		if link, ok := info.Link(); ok {
			b.WriteString("Linking object: " + link.String() + "\n")
		}
		if last, ok := info.LastPart(); ok {
			b.WriteString("Last object: " + last.String() + "\n")
		}
		return b.Bytes(), nil
	}
}

func printECInfoErr(cmd *cobra.Command, err error) bool {
	var errECInfo *objectSDK.ECInfoError

	ok := errors.As(err, &errECInfo)

	if ok {
		cmd.PrintErrln("Object is erasure-encoded, ec information received.")
		printECInfo(cmd, errECInfo.ECInfo())
	}

	return ok
}

func printECInfo(cmd *cobra.Command, info *objectSDK.ECInfo) {
	bs, err := marshalECInfo(cmd, info)
	commonCmd.ExitOnErr(cmd, "can't marshal split info: %w", err)

	cmd.Println(string(bs))
}

func marshalECInfo(cmd *cobra.Command, info *objectSDK.ECInfo) ([]byte, error) {
	toJSON, _ := cmd.Flags().GetBool(commonflags.JSON)
	toProto, _ := cmd.Flags().GetBool("proto")
	switch {
	case toJSON && toProto:
		return nil, errors.New("'--json' and '--proto' flags are mutually exclusive")
	case toJSON:
		return info.MarshalJSON()
	case toProto:
		return info.Marshal()
	default:
		b := bytes.NewBuffer(nil)
		b.WriteString("Total chunks: " + strconv.Itoa(int(info.Chunks[0].Total)))
		for _, chunk := range info.Chunks {
			var id oid.ID
			if err := id.Decode(chunk.ID.GetValue()); err != nil {
				return nil, fmt.Errorf("unable to decode chunk id: %w", err)
			}
			b.WriteString("\n  Index: " + strconv.Itoa(int(chunk.Index)) + " ID: " + id.String())
		}
		return b.Bytes(), nil
	}
}

func getRangeList(cmd *cobra.Command) ([]objectSDK.Range, error) {
	v := cmd.Flag("range").Value.String()
	if len(v) == 0 {
		return nil, nil
	}
	vs := strings.Split(v, ",")
	rs := make([]objectSDK.Range, len(vs))
	for i := range vs {
		before, after, found := strings.Cut(vs[i], rangeSep)
		if !found {
			return nil, fmt.Errorf("invalid range specifier: %s", vs[i])
		}

		offset, err := strconv.ParseUint(before, 10, 64)
		if err != nil {
			return nil, fmt.Errorf("invalid '%s' range offset specifier: %w", vs[i], err)
		}
		length, err := strconv.ParseUint(after, 10, 64)
		if err != nil {
			return nil, fmt.Errorf("invalid '%s' range length specifier: %w", vs[i], err)
		}

		if length == 0 {
			return nil, fmt.Errorf("invalid '%s' range: zero length", vs[i])
		}

		if offset+length <= offset {
			return nil, fmt.Errorf("invalid '%s' range: uint64 overflow", vs[i])
		}

		rs[i].SetOffset(offset)
		rs[i].SetLength(length)
	}
	return rs, nil
}