package container import ( "encoding/hex" "encoding/json" "errors" "fmt" "os" "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" "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/netmap" "github.com/chzyer/readline" "github.com/spf13/cobra" "github.com/spf13/viper" ) type policyPlaygroundREPL struct { cmd *cobra.Command nodes map[string]netmap.NodeInfo } func newPolicyPlaygroundREPL(cmd *cobra.Command) (*policyPlaygroundREPL, error) { return &policyPlaygroundREPL{ cmd: cmd, nodes: map[string]netmap.NodeInfo{}, }, nil } func (repl *policyPlaygroundREPL) handleLs(args []string) error { if len(args) > 0 { return fmt.Errorf("too many arguments for command 'ls': got %d, want 0", len(args)) } i := 1 for id, node := range repl.nodes { var attrs []string node.IterateAttributes(func(k, v string) { attrs = append(attrs, fmt.Sprintf("%s:%q", k, v)) }) fmt.Printf("\t%2d: id=%s attrs={%v}\n", i, id, strings.Join(attrs, " ")) i++ } return nil } func (repl *policyPlaygroundREPL) handleAdd(args []string) error { if len(args) == 0 { return fmt.Errorf("too few arguments for command 'add': got %d, want >0", len(args)) } id := args[0] key, err := hex.DecodeString(id) if err != nil { return fmt.Errorf("node id must be a hex string: got %q: %v", id, err) } node := repl.nodes[id] node.SetPublicKey(key) for _, attr := range args[1:] { kv := strings.Split(attr, ":") if len(kv) != 2 { return fmt.Errorf("node attributes must be in the format 'KEY:VALUE': got %q", attr) } node.SetAttribute(kv[0], kv[1]) } repl.nodes[id] = node return nil } func (repl *policyPlaygroundREPL) handleLoad(args []string) error { if len(args) != 1 { return fmt.Errorf("too few arguments for command 'add': got %d, want 1", len(args)) } jsonNetmap := map[string]map[string]string{} b, err := os.ReadFile(args[0]) if err != nil { return fmt.Errorf("reading netmap file %q: %v", args[0], err) } if err := json.Unmarshal(b, &jsonNetmap); err != nil { return fmt.Errorf("decoding json netmap: %v", err) } repl.nodes = make(map[string]netmap.NodeInfo) for id, attrs := range jsonNetmap { key, err := hex.DecodeString(id) if err != nil { return fmt.Errorf("node id must be a hex string: got %q: %v", id, err) } node := repl.nodes[id] node.SetPublicKey(key) for k, v := range attrs { node.SetAttribute(k, v) } repl.nodes[id] = node } return nil } func (repl *policyPlaygroundREPL) handleRemove(args []string) error { if len(args) == 0 { return fmt.Errorf("too few arguments for command 'remove': got %d, want >0", len(args)) } id := args[0] if _, exists := repl.nodes[id]; exists { delete(repl.nodes, id) return nil } return fmt.Errorf("node not found: id=%q", id) } func (repl *policyPlaygroundREPL) handleEval(args []string) error { policyStr := strings.TrimSpace(strings.Join(args, " ")) var nodes [][]netmap.NodeInfo nm := repl.netMap() if strings.HasPrefix(policyStr, "CBF") || strings.HasPrefix(policyStr, "SELECT") || strings.HasPrefix(policyStr, "FILTER") { // Assume that the input is a partial SELECT-FILTER expression. // Full inline policies always start with UNIQUE or REP keywords, // or different prefixes when it's the case of an external file. sfExpr, err := netmap.DecodeSelectFilterString(policyStr) if err != nil { return fmt.Errorf("parsing select-filter expression: %v", err) } nodes, err = nm.SelectFilterNodes(sfExpr) if err != nil { return fmt.Errorf("building select-filter nodes: %v", err) } } else { // Assume that the input is a full policy or input file otherwise. placementPolicy, err := parseContainerPolicy(repl.cmd, policyStr) if err != nil { return fmt.Errorf("parsing placement policy: %v", err) } nodes, err = nm.ContainerNodes(*placementPolicy, nil) if err != nil { return fmt.Errorf("building container nodes: %v", err) } } for i, ns := range nodes { var ids []string for _, node := range ns { ids = append(ids, hex.EncodeToString(node.PublicKey())) } fmt.Printf("\t%2d: %v\n", i+1, ids) } return nil } func (repl *policyPlaygroundREPL) netMap() netmap.NetMap { var nm netmap.NetMap var nodes []netmap.NodeInfo for _, node := range repl.nodes { nodes = append(nodes, node) } nm.SetNodes(nodes) return nm } var policyPlaygroundCompleter = readline.NewPrefixCompleter( readline.PcItem("list"), readline.PcItem("ls"), readline.PcItem("add"), readline.PcItem("load"), readline.PcItem("remove"), readline.PcItem("rm"), readline.PcItem("eval"), ) func (repl *policyPlaygroundREPL) run() error { if len(viper.GetString(commonflags.RPC)) > 0 { key := key.GetOrGenerate(repl.cmd) cli := internalclient.GetSDKClientByFlag(repl.cmd, key, commonflags.RPC) var prm internalclient.NetMapSnapshotPrm prm.SetClient(cli) resp, err := internalclient.NetMapSnapshot(repl.cmd.Context(), prm) commonCmd.ExitOnErr(repl.cmd, "unable to get netmap snapshot to populate initial netmap: %w", err) for _, node := range resp.NetMap().Nodes() { id := hex.EncodeToString(node.PublicKey()) repl.nodes[id] = node } } cmdHandlers := map[string]func([]string) error{ "list": repl.handleLs, "ls": repl.handleLs, "add": repl.handleAdd, "load": repl.handleLoad, "remove": repl.handleRemove, "rm": repl.handleRemove, "eval": repl.handleEval, } rl, err := readline.NewEx(&readline.Config{ Prompt: "> ", InterruptPrompt: "^C", AutoComplete: policyPlaygroundCompleter, }) if err != nil { return fmt.Errorf("error initializing readline: %w", err) } defer rl.Close() var exit bool for { line, err := rl.Readline() if err != nil { if errors.Is(err, readline.ErrInterrupt) { if exit { return nil } exit = true continue } return fmt.Errorf("reading line: %v", err) } exit = false parts := strings.Fields(line) if len(parts) == 0 { continue } cmd := parts[0] handler, exists := cmdHandlers[cmd] if exists { if err := handler(parts[1:]); err != nil { fmt.Printf("error: %v\n", err) } } else { fmt.Printf("error: unknown command %q\n", cmd) } } } var policyPlaygroundCmd = &cobra.Command{ Use: "policy-playground", Short: "A REPL for testing placement policies", Long: `A REPL for testing placement policies. If a wallet and endpoint is provided, the initial netmap data will be loaded from the snapshot of the node. Otherwise, an empty playground is created.`, Run: func(cmd *cobra.Command, _ []string) { repl, err := newPolicyPlaygroundREPL(cmd) commonCmd.ExitOnErr(cmd, "could not create policy playground: %w", err) commonCmd.ExitOnErr(cmd, "policy playground failed: %w", repl.run()) }, } func initContainerPolicyPlaygroundCmd() { commonflags.Init(policyPlaygroundCmd) }