zkcmd.go

package main

import (
	"archive/zip"
	"bytes"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"os/signal"
	"path"
	"sort"
	"strings"
	"sync"
	"syscall"
	"time"

	log "github.com/golang/glog"
	"github.com/samuel/go-zookeeper/zk"
	"golang.org/x/crypto/ssh/terminal"
	"golang.org/x/net/context"

	"github.com/youtube/vitess/go/exit"
	"github.com/youtube/vitess/go/vt/logutil"
	"github.com/youtube/vitess/go/vt/topo/zk2topo"
	"github.com/youtube/vitess/go/vt/vtctld"
)

var doc = `
zk is a tool for wrangling the zookeeper

It tries to mimic unix file system commands wherever possible, but
there are some slight differences in flag handling.

zk -h - provide help on overriding cell selection

zk cat /zk/path
zk cat -l /zk/path1 /zk/path2 (list filename before file data)

zk chmod n-mode /zk/path
zk chmod n+mode /zk/path

zk cp /zk/path .
zk cp ./config /zk/path/config
zk cp ./config /zk/path/ (trailing slash indicates directory)

zk edit /zk/path (create a local copy, edit and write changes back to cell)

zk ls /zk
zk ls -l /zk
zk ls -ld /zk (list directory node itself)
zk ls -R /zk (recursive, expensive)

zk stat /zk/path

zk touch /zk/path
zk touch -c /zk/path (don't create, just touch timestamp)
zk touch -p /zk/path (create all parts necessary, think mkdir -p)
NOTE: there is no mkdir - just touch a node. The distinction
between file and directory is just not relevant in zookeeper.

zk rm /zk/path
zk rm -r /zk/path (recursive)
zk rm -f /zk/path (no error on nonexistent node)

zk wait /zk/path (wait for node change or creation)
zk wait /zk/path/children/ (trailing slash waits on children)

zk watch /zk/path (print changes)

zk unzip zktree.zip /
zk unzip zktree.zip /zk/prefix

zk zip /zk/root zktree.zip
NOTE: zip file can't be dumped to the file system since znodes
can have data and children.

The zk tool looks for the address of the cluster in /etc/zookeeper/zk_client.conf,
or the file specified in the ZK_CLIENT_CONFIG environment variable.

The local cell may be overridden with the ZK_CLIENT_LOCAL_CELL environment
variable.
`

const (
	timeFmt      = "2006-01-02 15:04:05"
	timeFmtMicro = "2006-01-02 15:04:05.000000"
)

type cmdFunc func(ctx context.Context, subFlags *flag.FlagSet, args []string) error

var cmdMap map[string]cmdFunc
var zconn zk2topo.Conn

func init() {
	cmdMap = map[string]cmdFunc{
		"cat":   cmdCat,
		"chmod": cmdChmod,
		"cp":    cmdCp,
		"edit":  cmdEdit,
		"ls":    cmdLs,
		"rm":    cmdRm,
		"stat":  cmdStat,
		"touch": cmdTouch,
		"unzip": cmdUnzip,
		"wait":  cmdWait,
		"watch": cmdWatch,
		"zip":   cmdZip,
	}
}

var (
	server = flag.String("server", "", "server(s) to connect to")
)

func main() {
	defer exit.Recover()
	defer logutil.Flush()
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage of %v:\n", os.Args[0])
		flag.PrintDefaults()
		fmt.Fprintf(os.Stderr, doc)
	}
	flag.Parse()
	args := flag.Args()
	if len(args) == 0 {
		flag.Usage()
		exit.Return(1)
	}

	cmdName := args[0]
	args = args[1:]
	cmd, ok := cmdMap[cmdName]
	if !ok {
		log.Fatalf("Unknown command %v", cmdName)
	}
	subFlags := flag.NewFlagSet(cmdName, flag.ExitOnError)

	// Create a context for the command, cancel it if we get a signal.
	ctx, cancel := context.WithCancel(context.Background())
	sigRecv := make(chan os.Signal, 1)
	signal.Notify(sigRecv, os.Interrupt)
	go func() {
		<-sigRecv
		cancel()
	}()

	// Connect to the server.
	zconn = zk2topo.Connect(*server)

	// Run the command.
	if err := cmd(ctx, subFlags, args); err != nil {
		log.Error(err)
		exit.Return(1)
	}
}

func fixZkPath(zkPath string) string {
	if zkPath != "/" {
		zkPath = strings.TrimRight(zkPath, "/")
	}
	return path.Clean(zkPath)
}

func isZkFile(path string) bool {
	return strings.HasPrefix(path, "/zk")
}

func cmdWait(ctx context.Context, subFlags *flag.FlagSet, args []string) error {
	var (
		exitIfExists = subFlags.Bool("e", false, "exit if the path already exists")
	)

	subFlags.Parse(args)

	if subFlags.NArg() != 1 {
		return fmt.Errorf("wait: can only wait for one path")
	}
	zkPath := subFlags.Arg(0)
	isDir := zkPath[len(zkPath)-1] == '/'
	zkPath = fixZkPath(zkPath)

	var wait <-chan zk.Event
	var err error
	if isDir {
		_, _, wait, err = zconn.ChildrenW(ctx, zkPath)
	} else {
		_, _, wait, err = zconn.GetW(ctx, zkPath)
	}
	if err != nil {
		if err == zk.ErrNoNode {
			_, _, wait, err = zconn.ExistsW(ctx, zkPath)
		} else {
			return fmt.Errorf("wait: error %v: %v", zkPath, err)
		}
	} else {
		if *exitIfExists {
			return fmt.Errorf("already exists: %v", zkPath)
		}
	}
	event := <-wait
	fmt.Printf("event: %v\n", event)
	return nil
}

// Watch for changes to the node.
func cmdWatch(ctx context.Context, subFlags *flag.FlagSet, args []string) error {
	subFlags.Parse(args)

	eventChan := make(chan zk.Event, 16)
	for _, arg := range subFlags.Args() {
		zkPath := fixZkPath(arg)
		_, _, watch, err := zconn.GetW(ctx, zkPath)
		if err != nil {
			return fmt.Errorf("watch error: %v", err)
		}
		go func() {
			eventChan <- <-watch
		}()
	}

	for {
		select {
		case <-ctx.Done():
			return nil
		case event := <-eventChan:
			log.Infof("watch: event %v: %v", event.Path, event)
			if event.Type == zk.EventNodeDataChanged {
				data, stat, watch, err := zconn.GetW(ctx, event.Path)
				if err != nil {
					return fmt.Errorf("ERROR: failed to watch %v", err)
				}
				log.Infof("watch: %v %v\n", event.Path, stat)
				println(data)
				go func() {
					eventChan <- <-watch
				}()
			} else if event.State == zk.StateDisconnected {
				return nil
			} else if event.Type == zk.EventNodeDeleted {
				log.Infof("watch: %v deleted\n", event.Path)
			} else {
				// Most likely a session event - try t
				_, _, watch, err := zconn.GetW(ctx, event.Path)
				if err != nil {
					return fmt.Errorf("ERROR: failed to watch %v", err)
				}
				go func() {
					eventChan <- <-watch
				}()
			}
		}
	}
}

func cmdLs(ctx context.Context, subFlags *flag.FlagSet, args []string) error {
	var (
		longListing      = subFlags.Bool("l", false, "long listing")
		directoryListing = subFlags.Bool("d", false, "list directory instead of contents")
		force            = subFlags.Bool("f", false, "no warning on nonexistent node")
		recursiveListing = subFlags.Bool("R", false, "recursive listing")
	)
	subFlags.Parse(args)
	if subFlags.NArg() == 0 {
		return fmt.Errorf("ls: no path specified")
	}
	resolved, err := zk2topo.ResolveWildcards(ctx, zconn, subFlags.Args())
	if err != nil {
		return fmt.Errorf("ls: invalid wildcards: %v", err)
	}
	if len(resolved) == 0 {
		// the wildcards didn't result in anything, we're
		// done.
		return nil
	}

	hasError := false
	needsHeader := len(resolved) > 1 && !*directoryListing
	for _, arg := range resolved {
		zkPath := fixZkPath(arg)
		var children []string
		var err error
		isDir := true
		if *directoryListing {
			children = []string{""}
			isDir = false
		} else if *recursiveListing {
			children, err = zk2topo.ChildrenRecursive(ctx, zconn, zkPath)
		} else {
			children, _, err = zconn.Children(ctx, zkPath)
			// Assume this is a file node if it has no children.
			if len(children) == 0 {
				children = []string{""}
				isDir = false
			}
		}
		if err != nil {
			hasError = true
			if !*force || err != zk.ErrNoNode {
				log.Warningf("ls: cannot access %v: %v", zkPath, err)
			}
		}

		// Show the full path when it helps.
		showFullPath := false
		if *recursiveListing {
			showFullPath = true
		} else if *longListing && (*directoryListing || !isDir) {
			showFullPath = true
		}
		if needsHeader {
			fmt.Printf("%v:\n", zkPath)
		}
		if len(children) > 0 {
			if *longListing && isDir {
				fmt.Printf("total: %v\n", len(children))
			}
			sort.Strings(children)
			stats := make([]*zk.Stat, len(children))
			wg := sync.WaitGroup{}
			f := func(i int) {
				localPath := path.Join(zkPath, children[i])
				_, stat, err := zconn.Exists(ctx, localPath)
				if err != nil {
					if !*force || err != zk.ErrNoNode {
						log.Warningf("ls: cannot access: %v: %v", localPath, err)
					}
				} else {
					stats[i] = stat
				}
				wg.Done()
			}
			for i := range children {
				wg.Add(1)
				go f(i)
			}
			wg.Wait()

			for i, child := range children {
				localPath := path.Join(zkPath, child)
				if stat := stats[i]; stat != nil {
					fmtPath(stat, localPath, showFullPath, *longListing)
				}
			}
		}
		if needsHeader {
			fmt.Println()
		}
	}
	if hasError {
		return fmt.Errorf("ls: some paths had errors")
	}
	return nil
}

func fmtPath(stat *zk.Stat, zkPath string, showFullPath bool, longListing bool) {
	var name, perms string

	if !showFullPath {
		name = path.Base(zkPath)
	} else {
		name = zkPath
	}

	if longListing {
		if stat.NumChildren > 0 {
			// FIXME(msolomon) do permissions check?
			perms = "drwxrwxrwx"
			if stat.DataLength > 0 {
				// give a visual indication that this node has data as well as children
				perms = "nrw-rw-rw-"
			}
		} else if stat.EphemeralOwner != 0 {
			perms = "erw-rw-rw-"
		} else {
			perms = "-rw-rw-rw-"
		}
		// always print the Local version of the time. zookeeper's
		// go / C library would return a local time anyway, but
		// might as well be sure.
		fmt.Printf("%v %v %v % 8v % 20v %v\n", perms, "zk", "zk", stat.DataLength, zk2topo.Time(stat.Mtime).Local().Format(timeFmt), name)
	} else {
		fmt.Printf("%v\n", name)
	}
}

func cmdTouch(ctx context.Context, subFlags *flag.FlagSet, args []string) error {
	var (
		createParents = subFlags.Bool("p", false, "create parents")
		touchOnly     = subFlags.Bool("c", false, "touch only - don't create")
	)

	subFlags.Parse(args)
	if subFlags.NArg() != 1 {
		return fmt.Errorf("touch: need to specify exactly one path")
	}

	zkPath := fixZkPath(subFlags.Arg(0))

	var (
		version int32 = -1
		create        = false
	)

	data, stat, err := zconn.Get(ctx, zkPath)
	switch {
	case err == nil:
		version = stat.Version
	case err == zk.ErrNoNode:
		create = true
	default:
		return fmt.Errorf("touch: cannot access %v: %v", zkPath, err)
	}

	switch {
	case !create:
		_, err = zconn.Set(ctx, zkPath, data, version)
	case *touchOnly:
		return fmt.Errorf("touch: no such path %v", zkPath)
	case *createParents:
		_, err = zk2topo.CreateRecursive(ctx, zconn, zkPath, data, 0, zk.WorldACL(zk.PermAll), 10)
	default:
		_, err = zconn.Create(ctx, zkPath, data, 0, zk.WorldACL(zk.PermAll))
	}

	if err != nil {
		return fmt.Errorf("touch: cannot modify %v: %v", zkPath, err)
	}
	return nil
}

func cmdRm(ctx context.Context, subFlags *flag.FlagSet, args []string) error {
	var (
		force             = subFlags.Bool("f", false, "no warning on nonexistent node")
		recursiveDelete   = subFlags.Bool("r", false, "recursive delete")
		forceAndRecursive = subFlags.Bool("rf", false, "shorthand for -r -f")
	)
	subFlags.Parse(args)
	*force = *force || *forceAndRecursive
	*recursiveDelete = *recursiveDelete || *forceAndRecursive

	if subFlags.NArg() == 0 {
		return fmt.Errorf("rm: no path specified")
	}

	if *recursiveDelete {
		for _, arg := range subFlags.Args() {
			zkPath := fixZkPath(arg)
			if strings.Count(zkPath, "/") < 2 {
				return fmt.Errorf("rm: overly general path: %v", zkPath)
			}
		}
	}

	resolved, err := zk2topo.ResolveWildcards(ctx, zconn, subFlags.Args())
	if err != nil {
		return fmt.Errorf("rm: invalid wildcards: %v", err)
	}
	if len(resolved) == 0 {
		// the wildcards didn't result in anything, we're done
		return nil
	}

	hasError := false
	for _, arg := range resolved {
		zkPath := fixZkPath(arg)
		var err error
		if *recursiveDelete {
			err = zk2topo.DeleteRecursive(ctx, zconn, zkPath, -1)
		} else {
			err = zconn.Delete(ctx, zkPath, -1)
		}
		if err != nil && (!*force || err != zk.ErrNoNode) {
			hasError = true
			log.Warningf("rm: cannot delete %v: %v", zkPath, err)
		}
	}
	if hasError {
		// to be consistent with the command line 'rm -f', return
		// 0 if using 'zk rm -f' and the file doesn't exist.
		return fmt.Errorf("rm: some paths had errors")
	}
	return nil
}

func cmdCat(ctx context.Context, subFlags *flag.FlagSet, args []string) error {
	var (
		longListing = subFlags.Bool("l", false, "long listing")
		force       = subFlags.Bool("f", false, "no warning on nonexistent node")
		decodeProto = subFlags.Bool("p", false, "decode proto files and display them as text")
	)
	subFlags.Parse(args)
	if subFlags.NArg() == 0 {
		return fmt.Errorf("cat: no path specified")
	}
	resolved, err := zk2topo.ResolveWildcards(ctx, zconn, subFlags.Args())
	if err != nil {
		return fmt.Errorf("cat: invalid wildcards: %v", err)
	}
	if len(resolved) == 0 {
		// the wildcards didn't result in anything, we're done
		return nil
	}

	hasError := false
	for _, arg := range resolved {
		zkPath := fixZkPath(arg)
		data, _, err := zconn.Get(ctx, zkPath)
		if err != nil {
			hasError = true
			if !*force || err != zk.ErrNoNode {
				log.Warningf("cat: cannot access %v: %v", zkPath, err)
			}
			continue
		}

		if *longListing {
			fmt.Printf("%v:\n", zkPath)
		}
		decoded := ""
		if *decodeProto {
			decoded, err = vtctld.DecodeContent(zkPath, data)
			if err != nil {
				log.Warningf("cat: cannot proto decode %v: %v", zkPath, err)
				decoded = string(data)
			}
		} else {
			decoded = string(data)
		}
		fmt.Print(decoded)
		if len(decoded) > 0 && decoded[len(decoded)-1] != '\n' && (terminal.IsTerminal(int(os.Stdout.Fd())) || *longListing) {
			fmt.Print("\n")
		}
	}
	if hasError {
		return fmt.Errorf("cat: some paths had errors")
	}
	return nil
}

func cmdEdit(ctx context.Context, subFlags *flag.FlagSet, args []string) error {
	var (
		force = subFlags.Bool("f", false, "no warning on nonexistent node")
	)
	subFlags.Parse(args)
	if subFlags.NArg() == 0 {
		return fmt.Errorf("edit: no path specified")
	}
	arg := subFlags.Arg(0)
	zkPath := fixZkPath(arg)
	data, stat, err := zconn.Get(ctx, zkPath)
	if err != nil {
		if !*force || err != zk.ErrNoNode {
			log.Warningf("edit: cannot access %v: %v", zkPath, err)
		}
		return fmt.Errorf("edit: cannot access %v: %v", zkPath, err)
	}

	name := path.Base(zkPath)
	tmpPath := fmt.Sprintf("/tmp/zk-edit-%v-%v", name, time.Now().UnixNano())
	f, err := os.Create(tmpPath)
	if err == nil {
		_, err = f.Write(data)
		f.Close()
	}
	if err != nil {
		return fmt.Errorf("edit: cannot write file %v", err)
	}

	cmd := exec.Command(os.Getenv("EDITOR"), tmpPath)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	err = cmd.Run()
	if err != nil {
		os.Remove(tmpPath)
		return fmt.Errorf("edit: cannot start $EDITOR: %v", err)
	}

	fileData, err := ioutil.ReadFile(tmpPath)
	if err != nil {
		os.Remove(tmpPath)
		return fmt.Errorf("edit: cannot read file %v", err)
	}

	if bytes.Compare(fileData, data) != 0 {
		// data changed - update if we can
		_, err = zconn.Set(ctx, zkPath, fileData, stat.Version)
		if err != nil {
			os.Remove(tmpPath)
			return fmt.Errorf("edit: cannot write zk file %v", err)
		}
	}
	os.Remove(tmpPath)
	return nil
}

func cmdStat(ctx context.Context, subFlags *flag.FlagSet, args []string) error {
	var (
		force = subFlags.Bool("f", false, "no warning on nonexistent node")
	)
	subFlags.Parse(args)

	if subFlags.NArg() == 0 {
		return fmt.Errorf("stat: no path specified")
	}

	resolved, err := zk2topo.ResolveWildcards(ctx, zconn, subFlags.Args())
	if err != nil {
		return fmt.Errorf("stat: invalid wildcards: %v", err)
	}
	if len(resolved) == 0 {
		// the wildcards didn't result in anything, we're done
		return nil
	}

	hasError := false
	for _, arg := range resolved {
		zkPath := fixZkPath(arg)
		acls, stat, err := zconn.GetACL(ctx, zkPath)
		if stat == nil {
			err = fmt.Errorf("no such node")
		}
		if err != nil {
			hasError = true
			if !*force || err != zk.ErrNoNode {
				log.Warningf("stat: cannot access %v: %v", zkPath, err)
			}
			continue
		}
		fmt.Printf("Path: %s\n", zkPath)
		fmt.Printf("Created: %s\n", zk2topo.Time(stat.Ctime).Format(timeFmtMicro))
		fmt.Printf("Modified: %s\n", zk2topo.Time(stat.Mtime).Format(timeFmtMicro))
		fmt.Printf("Size: %v\n", stat.DataLength)
		fmt.Printf("Children: %v\n", stat.NumChildren)
		fmt.Printf("Version: %v\n", stat.Version)
		fmt.Printf("Ephemeral: %v\n", stat.EphemeralOwner)
		fmt.Printf("ACL:\n")
		for _, acl := range acls {
			fmt.Printf(" %v:%v %v\n", acl.Scheme, acl.ID, fmtACL(acl))
		}
	}
	if hasError {
		return fmt.Errorf("stat: some paths had errors")
	}
	return nil
}

var charPermMap map[string]int32
var permCharMap map[int32]string

func init() {
	charPermMap = map[string]int32{
		"r": zk.PermRead,
		"w": zk.PermWrite,
		"d": zk.PermDelete,
		"c": zk.PermCreate,
		"a": zk.PermAdmin,
	}
	permCharMap = make(map[int32]string)
	for c, p := range charPermMap {
		permCharMap[p] = c
	}
}

func fmtACL(acl zk.ACL) string {
	s := ""

	for _, perm := range []int32{zk.PermRead, zk.PermWrite, zk.PermDelete, zk.PermCreate, zk.PermAdmin} {
		if acl.Perms&perm != 0 {
			s += permCharMap[perm]
		} else {
			s += "-"
		}
	}
	return s
}

func cmdChmod(ctx context.Context, subFlags *flag.FlagSet, args []string) error {
	subFlags.Parse(args)
	if subFlags.NArg() < 2 {
		return fmt.Errorf("chmod: no permission specified")
	}
	mode := subFlags.Arg(0)
	if mode[0] != 'n' {
		return fmt.Errorf("chmod: invalid mode")
	}

	addPerms := false
	if mode[1] == '+' {
		addPerms = true
	} else if mode[1] != '-' {
		return fmt.Errorf("chmod: invalid mode")
	}

	var permMask int32
	for _, c := range mode[2:] {
		permMask |= charPermMap[string(c)]
	}

	resolved, err := zk2topo.ResolveWildcards(ctx, zconn, subFlags.Args()[1:])
	if err != nil {
		return fmt.Errorf("chmod: invalid wildcards: %v", err)
	}
	if len(resolved) == 0 {
		// the wildcards didn't result in anything, we're done
		return nil
	}

	hasError := false
	for _, arg := range resolved {
		zkPath := fixZkPath(arg)
		aclv, _, err := zconn.GetACL(ctx, zkPath)
		if err != nil {
			hasError = true
			log.Warningf("chmod: cannot set access %v: %v", zkPath, err)
			continue
		}
		if addPerms {
			aclv[0].Perms |= permMask
		} else {
			aclv[0].Perms &= ^permMask
		}
		err = zconn.SetACL(ctx, zkPath, aclv, -1)
		if err != nil {
			hasError = true
			log.Warningf("chmod: cannot set access %v: %v", zkPath, err)
			continue
		}
	}
	if hasError {
		return fmt.Errorf("chmod: some paths had errors")
	}
	return nil
}

func cmdCp(ctx context.Context, subFlags *flag.FlagSet, args []string) error {
	subFlags.Parse(args)
	switch {
	case subFlags.NArg() < 2:
		return fmt.Errorf("cp: need to specify source and destination paths")
	case subFlags.NArg() == 2:
		return fileCp(ctx, args[0], args[1])
	default:
		return multiFileCp(ctx, args)
	}
}

func getPathData(ctx context.Context, filePath string) ([]byte, error) {
	if isZkFile(filePath) {
		data, _, err := zconn.Get(ctx, filePath)
		return data, err
	}
	var err error
	file, err := os.Open(filePath)
	if err == nil {
		data, err := ioutil.ReadAll(file)
		if err == nil {
			return data, err
		}
	}
	return nil, err
}

func setPathData(ctx context.Context, filePath string, data []byte) error {
	if isZkFile(filePath) {
		_, err := zconn.Set(ctx, filePath, data, -1)
		if err == zk.ErrNoNode {
			_, err = zk2topo.CreateRecursive(ctx, zconn, filePath, data, 0, zk.WorldACL(zk.PermAll), 10)
		}
		return err
	}
	return ioutil.WriteFile(filePath, []byte(data), 0666)
}

func fileCp(ctx context.Context, srcPath, dstPath string) error {
	dstIsDir := dstPath[len(dstPath)-1] == '/'
	srcPath = fixZkPath(srcPath)
	dstPath = fixZkPath(dstPath)

	if !isZkFile(srcPath) && !isZkFile(dstPath) {
		return fmt.Errorf("cp: neither src nor dst is a /zk file: exitting")
	}

	data, err := getPathData(ctx, srcPath)
	if err != nil {
		return fmt.Errorf("cp: cannot read %v: %v", srcPath, err)
	}

	// If we are copying to a local directory - say '.', make the filename
	// the same as the source.
	if !isZkFile(dstPath) {
		fileInfo, err := os.Stat(dstPath)
		if err != nil {
			if err.(*os.PathError).Err != syscall.ENOENT {
				return fmt.Errorf("cp: cannot stat %v: %v", dstPath, err)
			}
		} else if fileInfo.IsDir() {
			dstPath = path.Join(dstPath, path.Base(srcPath))
		}
	} else if dstIsDir {
		// If we are copying into zk, interpret trailing slash as treating the
		// dstPath as a directory.
		dstPath = path.Join(dstPath, path.Base(srcPath))
	}
	if err := setPathData(ctx, dstPath, data); err != nil {
		return fmt.Errorf("cp: cannot write %v: %v", dstPath, err)
	}
	return nil
}

func multiFileCp(ctx context.Context, args []string) error {
	dstPath := args[len(args)-1]
	if dstPath[len(dstPath)-1] != '/' {
		// In multifile context, dstPath must be a directory.
		dstPath += "/"
	}

	for _, srcPath := range args[:len(args)-1] {
		if err := fileCp(ctx, srcPath, dstPath); err != nil {
			return err
		}
	}
	return nil
}

type zkItem struct {
	path string
	data []byte
	stat *zk.Stat
	err  error
}

// Store a zk tree in a zip archive. This won't be immediately useful to
// zip tools since even "directories" can contain data.
func cmdZip(ctx context.Context, subFlags *flag.FlagSet, args []string) error {
	subFlags.Parse(args)
	if subFlags.NArg() < 2 {
		return fmt.Errorf("zip: need to specify source and destination paths")
	}

	dstPath := subFlags.Arg(subFlags.NArg() - 1)
	paths := subFlags.Args()[:len(args)-1]
	if !strings.HasSuffix(dstPath, ".zip") {
		return fmt.Errorf("zip: need to specify destination .zip path: %v", dstPath)
	}
	zipFile, err := os.Create(dstPath)
	if err != nil {
		return fmt.Errorf("zip: error %v", err)
	}

	wg := sync.WaitGroup{}
	items := make(chan *zkItem, 64)
	for _, arg := range paths {
		zkPath := fixZkPath(arg)
		children, err := zk2topo.ChildrenRecursive(ctx, zconn, zkPath)
		if err != nil {
			return fmt.Errorf("zip: error %v", err)
		}
		for _, child := range children {
			toAdd := path.Join(zkPath, child)
			wg.Add(1)
			go func() {
				data, stat, err := zconn.Get(ctx, toAdd)
				items <- &zkItem{toAdd, data, stat, err}
				wg.Done()
			}()
		}
	}
	go func() {
		wg.Wait()
		close(items)
	}()

	zipWriter := zip.NewWriter(zipFile)
	for item := range items {
		path, data, stat, err := item.path, item.data, item.stat, item.err
		if err != nil {
			return fmt.Errorf("zip: get failed: %v", err)
		}
		// Skip ephemerals - not sure why you would archive them.
		if stat.EphemeralOwner > 0 {
			continue
		}
		fi := &zip.FileHeader{Name: path, Method: zip.Deflate}
		fi.SetModTime(zk2topo.Time(stat.Mtime))
		f, err := zipWriter.CreateHeader(fi)
		if err != nil {
			return fmt.Errorf("zip: create failed: %v", err)
		}
		_, err = f.Write(data)
		if err != nil {
			return fmt.Errorf("zip: create failed: %v", err)
		}
	}
	err = zipWriter.Close()
	if err != nil {
		return fmt.Errorf("zip: close failed: %v", err)
	}
	zipFile.Close()
	return nil
}

func cmdUnzip(ctx context.Context, subFlags *flag.FlagSet, args []string) error {
	subFlags.Parse(args)
	if subFlags.NArg() != 2 {
		return fmt.Errorf("zip: need to specify source and destination paths")
	}

	srcPath, dstPath := subFlags.Arg(0), subFlags.Arg(1)

	if !strings.HasSuffix(srcPath, ".zip") {
		return fmt.Errorf("zip: need to specify src .zip path: %v", srcPath)
	}

	zipReader, err := zip.OpenReader(srcPath)
	if err != nil {
		return fmt.Errorf("zip: error %v", err)
	}
	defer zipReader.Close()

	for _, zf := range zipReader.File {
		rc, err := zf.Open()
		if err != nil {
			return fmt.Errorf("unzip: error %v", err)
		}
		data, err := ioutil.ReadAll(rc)
		if err != nil {
			return fmt.Errorf("unzip: failed reading archive: %v", err)
		}
		zkPath := zf.Name
		if dstPath != "/" {
			zkPath = path.Join(dstPath, zkPath)
		}
		_, err = zk2topo.CreateRecursive(ctx, zconn, zkPath, data, 0, zk.WorldACL(zk.PermAll), 10)
		if err != nil && err != zk.ErrNodeExists {
			return fmt.Errorf("unzip: zk create failed: %v", err)
		}
		_, err = zconn.Set(ctx, zkPath, data, -1)
		if err != nil {
			return fmt.Errorf("unzip: zk set failed: %v", err)
		}
		rc.Close()
	}
	return nil
}