<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"os"
	"strings"

	log "github.com/Sirupsen/logrus"
	cobra "github.com/spf13/cobra"
)

type handler func(*cobra.Command, []string) error

func genCommand(use, short string, fn handler) *cobra.Command {
	run := func(cmd *cobra.Command, args []string) {
		if err := fn(cmd, args); err != nil {
			log.Fatal(err)
		} else {
			log.Debug("success")
		}
	}
	return &cobra.Command{
		Use:   use,
		Short: short,
		Run:   run,
	}

}

func main() {
	dummyPersistentPreRun := func(cmd *cobra.Command, args []string) {
		// Use PersistentPreRun here to replace the root one
	}
	root := &cobra.Command{
		Use:  "gozfs",
		Long: "gozfs is the cli for testing the go zfs ioctl api",
		PersistentPreRun: func(cmd *cobra.Command, args []string) {
			if cmd.Use == "gozfs" {
				return
			}
			name, err := cmd.Flags().GetString("name")
			if err != nil {
				log.Fatal(err)
			}
			if name == "" {
				log.Fatal("missing name")
			}
		},
		Run: help,
	}
	root.PersistentFlags().StringP("name", "n", "", "dataset name")

	cmdExists := genCommand("exists", "Test for dataset existence.",
		func(cmd *cobra.Command, args []string) error {
			name, _ := cmd.Flags().GetString("name")
			return exists(name)
		})

	cmdDestroy := genCommand("destroy", "Destroys a dataset or volume.",
		func(cmd *cobra.Command, args []string) error {
			name, _ := cmd.Flags().GetString("name")
			recursive, _ := cmd.Flags().GetBool("recursive")
			recursiveClones, _ := cmd.Flags().GetBool("recursiveclones")
			forceUnmount, _ := cmd.Flags().GetBool("forceunmount")
			deferDestroy, _ := cmd.Flags().GetBool("defer")

			d, err := GetDataset(name)
			if err != nil {
				return err
			}

			opts := &DestroyOptions{
				Recursive:       recursive,
				RecursiveClones: recursiveClones,
				ForceUnmount:    forceUnmount,
				Defer:           deferDestroy,
			}
			return d.Destroy(opts)
		})
	cmdDestroy.Flags().BoolP("defer", "d", false, "defer destroy")
	cmdDestroy.Flags().BoolP("recursive", "r", false, "recursively destroy datasets")
	cmdDestroy.Flags().BoolP("recursiveclones", "c", false, "recursively destroy clones")
	cmdDestroy.Flags().BoolP("forceunmount", "f", false, "force unmount")

	cmdHolds := genCommand("holds", "Retrieve list of user holds on the specified snapshot.",
		func(cmd *cobra.Command, args []string) error {
			name, _ := cmd.Flags().GetString("name")
			holds, err := holds(name)
			if err == nil {
				fmt.Println(holds)
			}
			return err
		})

	cmdSnapshot := genCommand("snapshot", "Creates a snapshot of a dataset or volume",
		func(cmd *cobra.Command, snaps []string) error {
			recursive, _ := cmd.Flags().GetBool("recursive")
			name, _ := cmd.Flags().GetString("name")
			nameParts := strings.Split(name, "@")
			if len(nameParts) != 2 {
				log.Fatal("invalid snapshot name")
			}

			var props map[string]string
			propJSON, _ := cmd.Flags().GetString("props")
			if err := json.Unmarshal([]byte(propJSON), &props); err != nil {
				log.Fatal("bad prop json")
			}

			d, err := GetDataset(nameParts[0])
			if err != nil {
				return err
			}

			return d.Snapshot(nameParts[1], recursive)
		},
	)
	cmdSnapshot.Flags().StringP("props", "p", "{}", "snapshot properties")
	cmdSnapshot.Flags().BoolP("recursive", "r", false, "recurisvely create snapshots")

	cmdRollback := genCommand("rollback", "Rollback this filesystem or volume to its most recent snapshot.",
		func(cmd *cobra.Command, args []string) error {
			destroyMoreRecent, _ := cmd.Flags().GetBool("destroyrecent")
			name, _ := cmd.Flags().GetString("name")

			d, err := GetDataset(name)
			if err != nil {
				return err
			}

			return d.Rollback(destroyMoreRecent)
		},
	)
	cmdRollback.Flags().BoolP("destroyrecent", "r", false, "destroy more recent snapshots and their clones")

	cmdCreate := genCommand("create", "Create a ZFS dataset or volume",
		func(cmd *cobra.Command, args []string) error {
			name, _ := cmd.Flags().GetString("name")
			createTypeName, _ := cmd.Flags().GetString("type")
			createType, err := getDMUType(createTypeName)
			if err != nil {
				log.Fatal("invalid type")
			}

			propJSON, _ := cmd.Flags().GetString("props")
			var props map[string]interface{}
			if err := json.Unmarshal([]byte(propJSON), &props); err != nil {
				log.Fatal("bad prop json")
			}

			if volsize, ok := props["volsize"]; ok {
				if volsizeFloat, ok := volsize.(float64); ok {
					props["volsize"] = uint64(volsizeFloat)
				}
			}

			if volblocksize, ok := props["volblocksize"]; ok {
				if volblocksizeFloat, ok := volblocksize.(float64); ok {
					props["volblocksize"] = uint64(volblocksizeFloat)
				}
			}

			return create(name, createType, props)
		},
	)
	cmdCreate.Flags().StringP("type", "t", "zfs", "zfs or zvol")
	cmdCreate.Flags().StringP("props", "p", "{}", "create properties")

	cmdSend := genCommand("send", "Generate a send stream from a snapshot",
		func(cmd *cobra.Command, args []string) error {
			name, _ := cmd.Flags().GetString("name")
			largeBlockOK, _ := cmd.Flags().GetBool("largeblock")
			embedOK, _ := cmd.Flags().GetBool("embed")
			fromSnap, _ := cmd.Flags().GetString("fromsnap")

			outputFD := os.Stdout.Fd()
			output, _ := cmd.Flags().GetString("output")
			if output != "/dev/stdout" {
				outputFile, err := os.Create(output)
				if err != nil {
					return err
				}
				defer outputFile.Close()

				outputFD = outputFile.Fd()
			} else {
				// If sending on stdout, don't log anything else unless there's
				// an error
				log.SetLevel(log.ErrorLevel)
			}

			return send(name, outputFD, fromSnap, largeBlockOK, embedOK)
		},
	)
	cmdSend.Flags().StringP("output", "o", "/dev/stdout", "output file")
	cmdSend.Flags().StringP("fromsnap", "f", "", "full snap name to send incremental from")
	cmdSend.Flags().BoolP("embed", "e", false, "embed data")
	cmdSend.Flags().BoolP("largeblock", "l", false, "large block")

	cmdClone := genCommand("clone", "Creates a clone from a snapshot",
		func(cmd *cobra.Command, args []string) error {
			name, _ := cmd.Flags().GetString("name")
			origin, _ := cmd.Flags().GetString("origin")
			if origin == "" {
				log.Fatal("missing origin snapshot")
			}

			var props map[string]interface{}
			propJSON, _ := cmd.Flags().GetString("props")
			if err := json.Unmarshal([]byte(propJSON), &props); err != nil {
				log.Fatal("bad prop json")
			}

			return clone(name, origin, props)
		},
	)
	cmdClone.Flags().StringP("origin", "o", "", "name of origin snapshot")
	cmdClone.Flags().StringP("props", "p", "{}", "snapshot properties")

	cmdRename := genCommand("rename", "Rename a dataset",
		func(cmd *cobra.Command, args []string) error {
			name, _ := cmd.Flags().GetString("name")
			newName, _ := cmd.Flags().GetString("newname")
			recursive, _ := cmd.Flags().GetBool("recursive")

			failedName, err := rename(name, newName, recursive)
			if failedName != "" {
				log.Error(failedName)
			}
			return err
		},
	)
	cmdRename.Flags().StringP("newname", "o", "", "new name of dataset")
	cmdRename.Flags().BoolP("recursive", "r", false, "recursively rename snapshots")

	cmdList := genCommand("list", "List filesystems, volumes, snapshots and bookmarks.",
		func(cmd *cobra.Command, args []string) error {
			name, _ := cmd.Flags().GetString("name")
			recurse, _ := cmd.Flags().GetBool("recurse")
			depth, _ := cmd.Flags().GetUint64("depth")
			typesList, _ := cmd.Flags().GetString("types")

			types := map[string]bool{}
			if typesList != "" {
				for _, t := range strings.Split(typesList, ",") {
					types[t] = true
				}
			}

			m, err := list(name, types, recurse, depth)
			if err != nil {
				return err
			}
			out, err := json.MarshalIndent(m, "", "  ")
			if err != nil {
				return err
			}
			fmt.Printf("%s\n", out)
			return nil
		})
	cmdList.PersistentPreRun = dummyPersistentPreRun
	cmdList.Flags().StringP("types", "t", "", "Comma seperated list of dataset types to list.")
	cmdList.Flags().BoolP("recurse", "r", false, "Recurse to all levels.")
	cmdList.Flags().Uint64P("depth", "d", 0, "Recursion depth limit, 0 for unlimited")

	cmdGet := genCommand("get", "Get dataset properties",
		func(cmd *cobra.Command, args []string) error {
			name, _ := cmd.Flags().GetString("name")
			recurse, _ := cmd.Flags().GetBool("recurse")
			depth, _ := cmd.Flags().GetUint64("depth")
			typesList, _ := cmd.Flags().GetString("types")

			types := map[string]bool{}
			if typesList != "" {
				for _, t := range strings.Split(typesList, ",") {
					types[t] = true
				}
			}

			m, err := properties(name, types, recurse, depth)
			if err != nil {
				return err
			}
			out, err := json.MarshalIndent(m, "", "  ")
			if err != nil {
				return err
			}
			fmt.Printf("%s\n", out)
			return nil
		})
	cmdGet.Flags().StringP("types", "t", "", "Comma seperated list of dataset types to list.")
	cmdGet.Flags().BoolP("recurse", "r", false, "Recurse to all levels.")
	cmdGet.Flags().Uint64P("depth", "d", 0, "Recursion depth limit, 0 for unlimited")

	root.AddCommand(
		cmdExists,
		cmdDestroy,
		cmdHolds,
		cmdSnapshot,
		cmdRollback,
		cmdCreate,
		cmdSend,
		cmdClone,
		cmdRename,
		cmdList,
		cmdGet,
	)
	if err := root.Execute(); err != nil {
		log.Fatal("root execute failed:", err)
	}
}

func help(cmd *cobra.Command, _ []string) {
	if err := cmd.Help(); err != nil {
		log.Fatal("help failed:", err)
	}
}
<commit_msg>Use public API Dataset.Create[Volume|Filesystem] in command<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"os"
	"strings"

	log "github.com/Sirupsen/logrus"
	cobra "github.com/spf13/cobra"
)

type handler func(*cobra.Command, []string) error

func genCommand(use, short string, fn handler) *cobra.Command {
	run := func(cmd *cobra.Command, args []string) {
		if err := fn(cmd, args); err != nil {
			log.Fatal(err)
		} else {
			log.Debug("success")
		}
	}
	return &cobra.Command{
		Use:   use,
		Short: short,
		Run:   run,
	}

}

func main() {
	dummyPersistentPreRun := func(cmd *cobra.Command, args []string) {
		// Use PersistentPreRun here to replace the root one
	}
	root := &cobra.Command{
		Use:  "gozfs",
		Long: "gozfs is the cli for testing the go zfs ioctl api",
		PersistentPreRun: func(cmd *cobra.Command, args []string) {
			if cmd.Use == "gozfs" {
				return
			}
			name, err := cmd.Flags().GetString("name")
			if err != nil {
				log.Fatal(err)
			}
			if name == "" {
				log.Fatal("missing name")
			}
		},
		Run: help,
	}
	root.PersistentFlags().StringP("name", "n", "", "dataset name")

	cmdExists := genCommand("exists", "Test for dataset existence.",
		func(cmd *cobra.Command, args []string) error {
			name, _ := cmd.Flags().GetString("name")
			return exists(name)
		})

	cmdDestroy := genCommand("destroy", "Destroys a dataset or volume.",
		func(cmd *cobra.Command, args []string) error {
			name, _ := cmd.Flags().GetString("name")
			recursive, _ := cmd.Flags().GetBool("recursive")
			recursiveClones, _ := cmd.Flags().GetBool("recursiveclones")
			forceUnmount, _ := cmd.Flags().GetBool("forceunmount")
			deferDestroy, _ := cmd.Flags().GetBool("defer")

			d, err := GetDataset(name)
			if err != nil {
				return err
			}

			opts := &DestroyOptions{
				Recursive:       recursive,
				RecursiveClones: recursiveClones,
				ForceUnmount:    forceUnmount,
				Defer:           deferDestroy,
			}
			return d.Destroy(opts)
		})
	cmdDestroy.Flags().BoolP("defer", "d", false, "defer destroy")
	cmdDestroy.Flags().BoolP("recursive", "r", false, "recursively destroy datasets")
	cmdDestroy.Flags().BoolP("recursiveclones", "c", false, "recursively destroy clones")
	cmdDestroy.Flags().BoolP("forceunmount", "f", false, "force unmount")

	cmdHolds := genCommand("holds", "Retrieve list of user holds on the specified snapshot.",
		func(cmd *cobra.Command, args []string) error {
			name, _ := cmd.Flags().GetString("name")
			holds, err := holds(name)
			if err == nil {
				fmt.Println(holds)
			}
			return err
		})

	cmdSnapshot := genCommand("snapshot", "Creates a snapshot of a dataset or volume",
		func(cmd *cobra.Command, snaps []string) error {
			recursive, _ := cmd.Flags().GetBool("recursive")
			name, _ := cmd.Flags().GetString("name")
			nameParts := strings.Split(name, "@")
			if len(nameParts) != 2 {
				log.Fatal("invalid snapshot name")
			}

			var props map[string]string
			propJSON, _ := cmd.Flags().GetString("props")
			if err := json.Unmarshal([]byte(propJSON), &props); err != nil {
				log.Fatal("bad prop json")
			}

			d, err := GetDataset(nameParts[0])
			if err != nil {
				return err
			}

			return d.Snapshot(nameParts[1], recursive)
		},
	)
	cmdSnapshot.Flags().StringP("props", "p", "{}", "snapshot properties")
	cmdSnapshot.Flags().BoolP("recursive", "r", false, "recurisvely create snapshots")

	cmdRollback := genCommand("rollback", "Rollback this filesystem or volume to its most recent snapshot.",
		func(cmd *cobra.Command, args []string) error {
			destroyMoreRecent, _ := cmd.Flags().GetBool("destroyrecent")
			name, _ := cmd.Flags().GetString("name")

			d, err := GetDataset(name)
			if err != nil {
				return err
			}

			return d.Rollback(destroyMoreRecent)
		},
	)
	cmdRollback.Flags().BoolP("destroyrecent", "r", false, "destroy more recent snapshots and their clones")

	cmdCreate := genCommand("create", "Create a ZFS dataset or volume",
		func(cmd *cobra.Command, args []string) error {
			name, _ := cmd.Flags().GetString("name")
			volsize, _ := cmd.Flags().GetUint64("volsize")
			createTypeName, _ := cmd.Flags().GetString("type")

			propJSON, _ := cmd.Flags().GetString("props")
			var props map[string]interface{}
			if err := json.Unmarshal([]byte(propJSON), &props); err != nil {
				log.Fatal("bad prop json")
			}

			if createTypeName == "zvol" {
				if volblocksize, ok := props["volblocksize"]; ok {
					if volblocksizeFloat, ok := volblocksize.(float64); ok {
						props["volblocksize"] = uint64(volblocksizeFloat)
					}
				}

				_, err := CreateVolume(name, volsize, props)
				return err
			} else {
				_, err := CreateFilesystem(name, props)
				return err
			}
		},
	)
	cmdCreate.Flags().StringP("type", "t", "zfs", "zfs or zvol")
	cmdCreate.Flags().Uint64P("volsize", "s", 0, "volume size")
	cmdCreate.Flags().StringP("props", "p", "{}", "create properties")

	cmdSend := genCommand("send", "Generate a send stream from a snapshot",
		func(cmd *cobra.Command, args []string) error {
			name, _ := cmd.Flags().GetString("name")
			largeBlockOK, _ := cmd.Flags().GetBool("largeblock")
			embedOK, _ := cmd.Flags().GetBool("embed")
			fromSnap, _ := cmd.Flags().GetString("fromsnap")

			outputFD := os.Stdout.Fd()
			output, _ := cmd.Flags().GetString("output")
			if output != "/dev/stdout" {
				outputFile, err := os.Create(output)
				if err != nil {
					return err
				}
				defer outputFile.Close()

				outputFD = outputFile.Fd()
			} else {
				// If sending on stdout, don't log anything else unless there's
				// an error
				log.SetLevel(log.ErrorLevel)
			}

			return send(name, outputFD, fromSnap, largeBlockOK, embedOK)
		},
	)
	cmdSend.Flags().StringP("output", "o", "/dev/stdout", "output file")
	cmdSend.Flags().StringP("fromsnap", "f", "", "full snap name to send incremental from")
	cmdSend.Flags().BoolP("embed", "e", false, "embed data")
	cmdSend.Flags().BoolP("largeblock", "l", false, "large block")

	cmdClone := genCommand("clone", "Creates a clone from a snapshot",
		func(cmd *cobra.Command, args []string) error {
			name, _ := cmd.Flags().GetString("name")
			origin, _ := cmd.Flags().GetString("origin")
			if origin == "" {
				log.Fatal("missing origin snapshot")
			}

			var props map[string]interface{}
			propJSON, _ := cmd.Flags().GetString("props")
			if err := json.Unmarshal([]byte(propJSON), &props); err != nil {
				log.Fatal("bad prop json")
			}

			return clone(name, origin, props)
		},
	)
	cmdClone.Flags().StringP("origin", "o", "", "name of origin snapshot")
	cmdClone.Flags().StringP("props", "p", "{}", "snapshot properties")

	cmdRename := genCommand("rename", "Rename a dataset",
		func(cmd *cobra.Command, args []string) error {
			name, _ := cmd.Flags().GetString("name")
			newName, _ := cmd.Flags().GetString("newname")
			recursive, _ := cmd.Flags().GetBool("recursive")

			failedName, err := rename(name, newName, recursive)
			if failedName != "" {
				log.Error(failedName)
			}
			return err
		},
	)
	cmdRename.Flags().StringP("newname", "o", "", "new name of dataset")
	cmdRename.Flags().BoolP("recursive", "r", false, "recursively rename snapshots")

	cmdList := genCommand("list", "List filesystems, volumes, snapshots and bookmarks.",
		func(cmd *cobra.Command, args []string) error {
			name, _ := cmd.Flags().GetString("name")
			recurse, _ := cmd.Flags().GetBool("recurse")
			depth, _ := cmd.Flags().GetUint64("depth")
			typesList, _ := cmd.Flags().GetString("types")

			types := map[string]bool{}
			if typesList != "" {
				for _, t := range strings.Split(typesList, ",") {
					types[t] = true
				}
			}

			m, err := list(name, types, recurse, depth)
			if err != nil {
				return err
			}
			out, err := json.MarshalIndent(m, "", "  ")
			if err != nil {
				return err
			}
			fmt.Printf("%s\n", out)
			return nil
		})
	cmdList.PersistentPreRun = dummyPersistentPreRun
	cmdList.Flags().StringP("types", "t", "", "Comma seperated list of dataset types to list.")
	cmdList.Flags().BoolP("recurse", "r", false, "Recurse to all levels.")
	cmdList.Flags().Uint64P("depth", "d", 0, "Recursion depth limit, 0 for unlimited")

	cmdGet := genCommand("get", "Get dataset properties",
		func(cmd *cobra.Command, args []string) error {
			name, _ := cmd.Flags().GetString("name")
			recurse, _ := cmd.Flags().GetBool("recurse")
			depth, _ := cmd.Flags().GetUint64("depth")
			typesList, _ := cmd.Flags().GetString("types")

			types := map[string]bool{}
			if typesList != "" {
				for _, t := range strings.Split(typesList, ",") {
					types[t] = true
				}
			}

			m, err := properties(name, types, recurse, depth)
			if err != nil {
				return err
			}
			out, err := json.MarshalIndent(m, "", "  ")
			if err != nil {
				return err
			}
			fmt.Printf("%s\n", out)
			return nil
		})
	cmdGet.Flags().StringP("types", "t", "", "Comma seperated list of dataset types to list.")
	cmdGet.Flags().BoolP("recurse", "r", false, "Recurse to all levels.")
	cmdGet.Flags().Uint64P("depth", "d", 0, "Recursion depth limit, 0 for unlimited")

	root.AddCommand(
		cmdExists,
		cmdDestroy,
		cmdHolds,
		cmdSnapshot,
		cmdRollback,
		cmdCreate,
		cmdSend,
		cmdClone,
		cmdRename,
		cmdList,
		cmdGet,
	)
	if err := root.Execute(); err != nil {
		log.Fatal("root execute failed:", err)
	}
}

func help(cmd *cobra.Command, _ []string) {
	if err := cmd.Help(); err != nil {
		log.Fatal("help failed:", err)
	}
}
<|endoftext|>


<commit_before>// Package c2go contains the main function for running the executable.
package main

import (
	"bytes"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"regexp"
	"strings"

	"github.com/elliotchance/c2go/ast"
	"github.com/elliotchance/c2go/program"
	"github.com/elliotchance/c2go/transpiler"
)

// Version can be requested through the command line with:
//
//     c2go -v
//
// See https://github.com/elliotchance/c2go/wiki/Release-Process
const Version = "0.13.3"

// ProgramArgs - arguments of program
type ProgramArgs struct {
	verbose     bool
	ast         bool
	inputFile   string
	outputFile  string
	packageName string
}

func readAST(data []byte) []string {
	uncolored := regexp.MustCompile(`\x1b\[[\d;]+m`).ReplaceAll(data, []byte{})
	return strings.Split(string(uncolored), "\n")
}

type treeNode struct {
	indent int
	node   ast.Node
}

func convertLinesToNodes(lines []string) []treeNode {
	nodes := []treeNode{}
	for _, line := range lines {
		if strings.TrimSpace(line) == "" {
			continue
		}

		// It is tempting to discard null AST nodes, but these may
		// have semantic importance: for example, they represent omitted
		// for-loop conditions, as in for(;;).
		line = strings.Replace(line, "<<<NULL>>>", "NullStmt", 1)

		indentAndType := regexp.MustCompile("^([|\\- `]*)(\\w+)").FindStringSubmatch(line)
		if len(indentAndType) == 0 {
			panic(fmt.Sprintf("Cannot understand line '%s'", line))
		}

		offset := len(indentAndType[1])
		node := ast.Parse(line[offset:])

		indentLevel := len(indentAndType[1]) / 2
		nodes = append(nodes, treeNode{indentLevel, node})
	}

	return nodes
}

// buildTree convert an array of nodes, each prefixed with a depth into a tree.
func buildTree(nodes []treeNode, depth int) []ast.Node {
	if len(nodes) == 0 {
		return []ast.Node{}
	}

	// Split the list into sections, treat each section as a a tree with its own root.
	sections := [][]treeNode{}
	for _, node := range nodes {
		if node.indent == depth {
			sections = append(sections, []treeNode{node})
		} else {
			sections[len(sections)-1] = append(sections[len(sections)-1], node)
		}
	}

	results := []ast.Node{}
	for _, section := range sections {
		slice := []treeNode{}
		for _, n := range section {
			if n.indent > depth {
				slice = append(slice, n)
			}
		}

		children := buildTree(slice, depth+1)
		for _, child := range children {
			section[0].node.AddChild(child)
		}
		results = append(results, section[0].node)
	}

	return results
}

/* Dead code
// ToJSON - tree convert to JSON
func ToJSON(tree []interface{}) []map[string]interface{} {
	r := make([]map[string]interface{}, len(tree))

	for j, n := range tree {
		rn := reflect.ValueOf(n).Elem()
		r[j] = make(map[string]interface{})
		r[j]["node"] = rn.Type().Name()

		for i := 0; i < rn.NumField(); i++ {
			name := strings.ToLower(rn.Type().Field(i).Name)
			value := rn.Field(i).Interface()

			if name == "children" {
				v := value.([]interface{})

				if len(v) == 0 {
					continue
				}

				value = ToJSON(v)
			}

			r[j][name] = value
		}
	}

	return r
}
*/

// Start - base function
func Start(args ProgramArgs) error {
	if os.Getenv("GOPATH") == "" {
		return fmt.Errorf("The $GOPATH must be set")
	}

	// 1. Compile it first (checking for errors)
	_, err := os.Stat(args.inputFile)
	if err != nil {
		return fmt.Errorf("Input file is not found")
	}

	// 2. Preprocess
	var pp []byte
	{
		// See : https://clang.llvm.org/docs/CommandGuide/clang.html
		// clang -E <file>    Run the preprocessor stage.
		cmd := exec.Command("clang", "-E", args.inputFile)
		var out bytes.Buffer
		var stderr bytes.Buffer
		cmd.Stdout = &out
		cmd.Stderr = &stderr
		err = cmd.Run()
		if err != nil {
			return fmt.Errorf(fmt.Sprintf("preprocess failed: %v\nStdErr = %v\n", err, stderr.String())
		}
		pp = []byte(out.String())
	}

	ppFilePath := path.Join(os.TempDir(), "pp.c")
	err = ioutil.WriteFile(ppFilePath, pp, 0644)
	if err != nil {
		return fmt.Errorf("writing to /tmp/pp.c failed: %v", err)
	}

	// 3. Generate JSON from AST
	astPP, err := exec.Command("clang", "-Xclang", "-ast-dump", "-fsyntax-only", ppFilePath).Output()
	if err != nil {
		// If clang fails it still prints out the AST, so we have to run it
		// again to get the real error.
		errBody, _ := exec.Command("clang", ppFilePath).CombinedOutput()

		panic("clang failed: " + err.Error() + ":\n\n" + string(errBody))
	}

	lines := readAST(astPP)
	if args.ast {
		for _, l := range lines {
			fmt.Println(l)
		}
		fmt.Println()
	}

	nodes := convertLinesToNodes(lines)
	tree := buildTree(nodes, 0)

	p := program.NewProgram()
	p.Verbose = args.verbose

	err = transpiler.TranspileAST(args.inputFile, args.packageName, p, tree[0].(ast.Node))
	if err != nil {
		panic(err)
	}

	outputFilePath := args.outputFile

	if outputFilePath == "" {
		cleanFileName := filepath.Clean(filepath.Base(args.inputFile))
		extension := filepath.Ext(args.inputFile)
		outputFilePath = cleanFileName[0:len(cleanFileName)-len(extension)] + ".go"
	}

	err = ioutil.WriteFile(outputFilePath, []byte(p.String()), 0755)
	if err != nil {
		return fmt.Errorf("writing C output file failed: %v", err)
	}

	return nil
}

// newTempFile - returns temp file
func newTempFile(dir, prefix, suffix string) (*os.File, error) {
	for index := 1; index < 10000; index++ {
		path := filepath.Join(dir, fmt.Sprintf("%s%03d%s", prefix, index, suffix))
		if _, err := os.Stat(path); err != nil {
			return os.Create(path)
		}
	}
	return nil, fmt.Errorf("could not create file: %s%03d%s", prefix, 1, suffix)
}

func main() {
	var (
		versionFlag       = flag.Bool("v", false, "print the version and exit")
		transpileCommand  = flag.NewFlagSet("transpile", flag.ContinueOnError)
		verboseFlag       = transpileCommand.Bool("V", false, "print progress as comments")
		outputFlag        = transpileCommand.String("o", "", "output Go generated code to the specified file")
		packageFlag       = transpileCommand.String("p", "main", "set the name of the generated package")
		transpileHelpFlag = transpileCommand.Bool("h", false, "print help information")
		astCommand        = flag.NewFlagSet("ast", flag.ContinueOnError)
		astHelpFlag       = astCommand.Bool("h", false, "print help information")
	)

	flag.Usage = func() {
		usage := "Usage: %s [-v] [<command>] [<flags>] file.c\n\n"
		usage += "Commands:\n"
		usage += "  transpile\ttranspile an input C source file to Go\n"
		usage += "  ast\t\tprint AST before translated Go code\n\n"

		usage += "Flags:\n"
		fmt.Fprintf(os.Stderr, usage, os.Args[0])
		flag.PrintDefaults()

		// print flags of transpile command
		fmt.Println("\nFlags of transpile command:")
		transpileCommand.PrintDefaults()

		// print flags of ast command
		fmt.Println("\nFlags of ast command:")
		astCommand.PrintDefaults()

		// examples
		fmt.Println("\nExamples of flag using:")
		fmt.Println("\nc2go -h\n\treturn the help", "")
		fmt.Println("\nc2go transpile -o source.go source.c\n\ttranspiling file source.c to Go file with name source.go")
	}

	flag.Parse()

	if *versionFlag {
		// Simply print out the version and exit.
		fmt.Println(Version)
		return
	}

	if flag.NArg() < 1 {
		flag.Usage()
		os.Exit(1)
	}

	args := ProgramArgs{verbose: *verboseFlag, ast: false}

	switch os.Args[1] {
	case "ast":
		err := astCommand.Parse(os.Args[2:])
		if err != nil {
			fmt.Printf("ast command cannot parse: %v", err)
			os.Exit(1)
		}

		if *astHelpFlag || astCommand.NArg() == 0 {
			fmt.Fprintf(os.Stderr, "Usage: %s ast file.c\n", os.Args[0])
			astCommand.PrintDefaults()
			os.Exit(1)
		}

		args.ast = true
		args.inputFile = astCommand.Arg(0)
	case "transpile":
		err := transpileCommand.Parse(os.Args[2:])
		if err != nil {
			fmt.Printf("transpile command cannot parse: %v", err)
			os.Exit(1)
		}

		if *transpileHelpFlag || transpileCommand.NArg() == 0 {
			fmt.Fprintf(os.Stderr, "Usage: %s transpile [-V] [-o file.go] [-p package] file.c\n", os.Args[0])
			transpileCommand.PrintDefaults()
			os.Exit(1)
		}

		args.inputFile = transpileCommand.Arg(0)
		args.outputFile = *outputFlag
		args.packageName = *packageFlag
	default:
		flag.Usage()
		os.Exit(1)
	}

	if err := Start(args); err != nil {
		fmt.Printf("Error: %v", err)
		os.Exit(1)
	}
}
<commit_msg>Prepare for pull request<commit_after>// Package c2go contains the main function for running the executable.
package main

import (
	"bytes"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"regexp"
	"strings"

	"github.com/elliotchance/c2go/ast"
	"github.com/elliotchance/c2go/program"
	"github.com/elliotchance/c2go/transpiler"
)

// Version can be requested through the command line with:
//
//     c2go -v
//
// See https://github.com/elliotchance/c2go/wiki/Release-Process
const Version = "0.13.3"

// ProgramArgs - arguments of program
type ProgramArgs struct {
	verbose     bool
	ast         bool
	inputFile   string
	outputFile  string
	packageName string
}

func readAST(data []byte) []string {
	uncolored := regexp.MustCompile(`\x1b\[[\d;]+m`).ReplaceAll(data, []byte{})
	return strings.Split(string(uncolored), "\n")
}

type treeNode struct {
	indent int
	node   ast.Node
}

func convertLinesToNodes(lines []string) []treeNode {
	nodes := []treeNode{}
	for _, line := range lines {
		if strings.TrimSpace(line) == "" {
			continue
		}

		// It is tempting to discard null AST nodes, but these may
		// have semantic importance: for example, they represent omitted
		// for-loop conditions, as in for(;;).
		line = strings.Replace(line, "<<<NULL>>>", "NullStmt", 1)

		indentAndType := regexp.MustCompile("^([|\\- `]*)(\\w+)").FindStringSubmatch(line)
		if len(indentAndType) == 0 {
			panic(fmt.Sprintf("Cannot understand line '%s'", line))
		}

		offset := len(indentAndType[1])
		node := ast.Parse(line[offset:])

		indentLevel := len(indentAndType[1]) / 2
		nodes = append(nodes, treeNode{indentLevel, node})
	}

	return nodes
}

// buildTree convert an array of nodes, each prefixed with a depth into a tree.
func buildTree(nodes []treeNode, depth int) []ast.Node {
	if len(nodes) == 0 {
		return []ast.Node{}
	}

	// Split the list into sections, treat each section as a a tree with its own root.
	sections := [][]treeNode{}
	for _, node := range nodes {
		if node.indent == depth {
			sections = append(sections, []treeNode{node})
		} else {
			sections[len(sections)-1] = append(sections[len(sections)-1], node)
		}
	}

	results := []ast.Node{}
	for _, section := range sections {
		slice := []treeNode{}
		for _, n := range section {
			if n.indent > depth {
				slice = append(slice, n)
			}
		}

		children := buildTree(slice, depth+1)
		for _, child := range children {
			section[0].node.AddChild(child)
		}
		results = append(results, section[0].node)
	}

	return results
}

/* Dead code
// ToJSON - tree convert to JSON
func ToJSON(tree []interface{}) []map[string]interface{} {
	r := make([]map[string]interface{}, len(tree))

	for j, n := range tree {
		rn := reflect.ValueOf(n).Elem()
		r[j] = make(map[string]interface{})
		r[j]["node"] = rn.Type().Name()

		for i := 0; i < rn.NumField(); i++ {
			name := strings.ToLower(rn.Type().Field(i).Name)
			value := rn.Field(i).Interface()

			if name == "children" {
				v := value.([]interface{})

				if len(v) == 0 {
					continue
				}

				value = ToJSON(v)
			}

			r[j][name] = value
		}
	}

	return r
}
*/

// Start - base function
func Start(args ProgramArgs) error {
	if os.Getenv("GOPATH") == "" {
		return fmt.Errorf("The $GOPATH must be set")
	}

	// 1. Compile it first (checking for errors)
	_, err := os.Stat(args.inputFile)
	if err != nil {
		return fmt.Errorf("Input file is not found")
	}

	// 2. Preprocess
	var pp []byte
	{
		// See : https://clang.llvm.org/docs/CommandGuide/clang.html
		// clang -E <file>    Run the preprocessor stage.
		cmd := exec.Command("clang", "-E", args.inputFile)
		var out bytes.Buffer
		var stderr bytes.Buffer
		cmd.Stdout = &out
		cmd.Stderr = &stderr
		err = cmd.Run()
		if err != nil {
			return fmt.Errorf("preprocess failed: %v\nStdErr = %v\n", err, stderr.String())
		}
		pp = []byte(out.String())
	}

	ppFilePath := path.Join(os.TempDir(), "pp.c")
	err = ioutil.WriteFile(ppFilePath, pp, 0644)
	if err != nil {
		return fmt.Errorf("writing to /tmp/pp.c failed: %v", err)
	}

	// 3. Generate JSON from AST
	astPP, err := exec.Command("clang", "-Xclang", "-ast-dump", "-fsyntax-only", ppFilePath).Output()
	if err != nil {
		// If clang fails it still prints out the AST, so we have to run it
		// again to get the real error.
		errBody, _ := exec.Command("clang", ppFilePath).CombinedOutput()

		panic("clang failed: " + err.Error() + ":\n\n" + string(errBody))
	}

	lines := readAST(astPP)
	if args.ast {
		for _, l := range lines {
			fmt.Println(l)
		}
		fmt.Println()
	}

	nodes := convertLinesToNodes(lines)
	tree := buildTree(nodes, 0)

	p := program.NewProgram()
	p.Verbose = args.verbose

	err = transpiler.TranspileAST(args.inputFile, args.packageName, p, tree[0].(ast.Node))
	if err != nil {
		panic(err)
	}

	outputFilePath := args.outputFile

	if outputFilePath == "" {
		cleanFileName := filepath.Clean(filepath.Base(args.inputFile))
		extension := filepath.Ext(args.inputFile)
		outputFilePath = cleanFileName[0:len(cleanFileName)-len(extension)] + ".go"
	}

	err = ioutil.WriteFile(outputFilePath, []byte(p.String()), 0755)
	if err != nil {
		return fmt.Errorf("writing C output file failed: %v", err)
	}

	return nil
}

// newTempFile - returns temp file
func newTempFile(dir, prefix, suffix string) (*os.File, error) {
	for index := 1; index < 10000; index++ {
		path := filepath.Join(dir, fmt.Sprintf("%s%03d%s", prefix, index, suffix))
		if _, err := os.Stat(path); err != nil {
			return os.Create(path)
		}
	}
	return nil, fmt.Errorf("could not create file: %s%03d%s", prefix, 1, suffix)
}

func main() {
	var (
		versionFlag       = flag.Bool("v", false, "print the version and exit")
		transpileCommand  = flag.NewFlagSet("transpile", flag.ContinueOnError)
		verboseFlag       = transpileCommand.Bool("V", false, "print progress as comments")
		outputFlag        = transpileCommand.String("o", "", "output Go generated code to the specified file")
		packageFlag       = transpileCommand.String("p", "main", "set the name of the generated package")
		transpileHelpFlag = transpileCommand.Bool("h", false, "print help information")
		astCommand        = flag.NewFlagSet("ast", flag.ContinueOnError)
		astHelpFlag       = astCommand.Bool("h", false, "print help information")
	)

	flag.Usage = func() {
		usage := "Usage: %s [-v] [<command>] [<flags>] file.c\n\n"
		usage += "Commands:\n"
		usage += "  transpile\ttranspile an input C source file to Go\n"
		usage += "  ast\t\tprint AST before translated Go code\n\n"

		usage += "Flags:\n"
		fmt.Fprintf(os.Stderr, usage, os.Args[0])
		flag.PrintDefaults()

		// print flags of transpile command
		fmt.Println("\nFlags of transpile command:")
		transpileCommand.PrintDefaults()

		// print flags of ast command
		fmt.Println("\nFlags of ast command:")
		astCommand.PrintDefaults()

		// examples
		fmt.Println("\nExamples of flag using:")
		fmt.Println("\nc2go -h\n\treturn the help", "")
		fmt.Println("\nc2go transpile -o source.go source.c\n\ttranspiling file source.c to Go file with name source.go")
	}

	flag.Parse()

	if *versionFlag {
		// Simply print out the version and exit.
		fmt.Println(Version)
		return
	}

	if flag.NArg() < 1 {
		flag.Usage()
		os.Exit(1)
	}

	args := ProgramArgs{verbose: *verboseFlag, ast: false}

	switch os.Args[1] {
	case "ast":
		err := astCommand.Parse(os.Args[2:])
		if err != nil {
			fmt.Printf("ast command cannot parse: %v", err)
			os.Exit(1)
		}

		if *astHelpFlag || astCommand.NArg() == 0 {
			fmt.Fprintf(os.Stderr, "Usage: %s ast file.c\n", os.Args[0])
			astCommand.PrintDefaults()
			os.Exit(1)
		}

		args.ast = true
		args.inputFile = astCommand.Arg(0)
	case "transpile":
		err := transpileCommand.Parse(os.Args[2:])
		if err != nil {
			fmt.Printf("transpile command cannot parse: %v", err)
			os.Exit(1)
		}

		if *transpileHelpFlag || transpileCommand.NArg() == 0 {
			fmt.Fprintf(os.Stderr, "Usage: %s transpile [-V] [-o file.go] [-p package] file.c\n", os.Args[0])
			transpileCommand.PrintDefaults()
			os.Exit(1)
		}

		args.inputFile = transpileCommand.Arg(0)
		args.outputFile = *outputFlag
		args.packageName = *packageFlag
	default:
		flag.Usage()
		os.Exit(1)
	}

	if err := Start(args); err != nil {
		fmt.Printf("Error: %v", err)
		os.Exit(1)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"os/signal"
	"regexp"
	"strings"
	"sync"

	"github.com/atotto/clipboard"
)

// re matches FILE:LINE.
var re = regexp.MustCompile(`\S+\.go:\d+`)

var matched bool
var mu sync.Mutex

func main() {
	log.SetFlags(0)

	// Execute "go" command with the same arguments.
	cmd := exec.Command("go", os.Args[1:]...)

	// Pass through standard input.
	stdin, err := cmd.StdinPipe()
	if err != nil {
		log.Fatal(err)
	}
	go io.Copy(stdin, os.Stdin)

	// Create a wait group for stdout/stderr.
	var wg sync.WaitGroup
	wg.Add(2)

	// Pass through standard out.
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		log.Fatal(err)
	}
	go func() {
		processPipe(os.Stdout, stdout)
		wg.Done()
	}()

	// Read through stderr and decorate.
	stderr, err := cmd.StderrPipe()
	if err != nil {
		log.Fatal(err)
	}
	go func() {
		processPipe(os.Stderr, stderr)
		wg.Done()
	}()

	// Execute command.
	if err := cmd.Start(); err != nil {
		log.Fatal(err)
	}

	// Forward signals to command.
	go func() {
		c := make(chan os.Signal, 1)
		signal.Notify(c)
		for sig := range c {
			cmd.Process.Signal(sig)
		}
	}()

	// Wait for pipes to finish reading and then wait for command to exit.
	wg.Wait()
	if err = cmd.Wait(); err != nil {
		log.Fatal(err)
	}
}

// processPipe scans the src by line and attempts to match the first FILE:LINE.
func processPipe(dst io.Writer, src io.Reader) {
	scanner := bufio.NewScanner(src)
	for scanner.Scan() {
		line := scanner.Text()
		func() {
			mu.Lock()
			defer mu.Unlock()

			if !matched {
				if m := re.FindString(line); m != "" && !strings.Contains(m, "testing.go") {
					// Remove "./" prefix.
					m = strings.TrimPrefix(m, "./")

					// Remove present working directory prefix.
					if pwd, _ := os.Getwd(); pwd != "" {
						m = strings.TrimPrefix(m, pwd+"/")
					}

					// Copy match.
					clipboard.WriteAll(m)

					// Bold line.
					line = "\033[1m" + line + "\033[0m"
					matched = true
				}
			}
			fmt.Fprintln(dst, line)
		}()
	}
}
<commit_msg>add visual beep<commit_after>package main

import (
	"bufio"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"os/signal"
	"regexp"
	"strings"
	"sync"

	"github.com/atotto/clipboard"
)

// re matches FILE:LINE.
var re = regexp.MustCompile(`\S+\.go:\d+`)

var matched bool
var mu sync.Mutex

func main() {
	log.SetFlags(0)

	// Execute "go" command with the same arguments.
	cmd := exec.Command("go", os.Args[1:]...)

	// Pass through standard input.
	stdin, err := cmd.StdinPipe()
	if err != nil {
		log.Fatal(err)
	}
	go io.Copy(stdin, os.Stdin)

	// Create a wait group for stdout/stderr.
	var wg sync.WaitGroup
	wg.Add(2)

	// Pass through standard out.
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		log.Fatal(err)
	}
	go func() {
		processPipe(os.Stdout, stdout)
		wg.Done()
	}()

	// Read through stderr and decorate.
	stderr, err := cmd.StderrPipe()
	if err != nil {
		log.Fatal(err)
	}
	go func() {
		processPipe(os.Stderr, stderr)
		wg.Done()
	}()

	// Execute command.
	if err := cmd.Start(); err != nil {
		log.Fatal(err)
	}

	// Forward signals to command.
	go func() {
		c := make(chan os.Signal, 1)
		signal.Notify(c)
		for sig := range c {
			cmd.Process.Signal(sig)
		}
	}()

	// Wait for pipes to finish reading and then wait for command to exit.
	wg.Wait()

	// Print a visual beep to cause the dock icon to bounce.
	fmt.Print("\x07")

	if err := cmd.Wait(); err != nil {
		os.Exit(1)
	}
}

// processPipe scans the src by line and attempts to match the first FILE:LINE.
func processPipe(dst io.Writer, src io.Reader) {
	scanner := bufio.NewScanner(src)
	for scanner.Scan() {
		line := scanner.Text()
		func() {
			mu.Lock()
			defer mu.Unlock()

			if !matched {
				if m := re.FindString(line); m != "" && !strings.Contains(m, "testing.go") {
					// Remove "./" prefix.
					m = strings.TrimPrefix(m, "./")

					// Remove present working directory prefix.
					if pwd, _ := os.Getwd(); pwd != "" {
						m = strings.TrimPrefix(m, pwd+"/")
					}

					// Copy match.
					clipboard.WriteAll(m)

					// Bold line.
					line = "\033[1m" + line + "\033[0m"
					matched = true
				}
			}
			fmt.Fprintln(dst, line)
		}()
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"errors"
	"io/ioutil"
	"log"
	"os"
	"os/user"
	"path"
	"path/filepath"
	"sort"
	"strings"
	"syscall"

	"github.com/codegangsta/cli"
	"github.com/hpcloud/tail"
)

type EdwardConfiguration struct {
	Dir       string
	LogDir    string
	PidDir    string
	ScriptDir string
}

var EdwardConfig EdwardConfiguration = EdwardConfiguration{}

func createDirIfNeeded(path string) {
	if _, err := os.Stat(path); os.IsNotExist(err) {
		os.MkdirAll(path, 0777)
	}
}

func (e *EdwardConfiguration) initialize() error {
	user, err := user.Current()
	if err != nil {
		return err
	}
	e.Dir = path.Join(user.HomeDir, ".edward")
	e.LogDir = path.Join(e.Dir, "logs")
	e.PidDir = path.Join(e.Dir, "pidFiles")
	e.ScriptDir = path.Join(e.Dir, "scriptFiles")
	createDirIfNeeded(e.Dir)
	createDirIfNeeded(e.LogDir)
	createDirIfNeeded(e.PidDir)
	createDirIfNeeded(e.ScriptDir)
	return nil
}

var groups map[string]*ServiceGroupConfig
var services map[string]*ServiceConfig

func thirdPartyService(name string, startCommand string, stopCommand string, started string) *ServiceConfig {
	pathStr := "$ALPHA"
	return &ServiceConfig{
		Name: name,
		Path: &pathStr,
		Env:  []string{"YEXT_RABBITMQ=localhost"},
		Commands: ServiceConfigCommands{
			Launch: startCommand,
			Stop:   stopCommand,
		},
		Properties: ServiceConfigProperties{
			Started: started,
		},
	}
}

func getAlpha() string {
	for _, env := range os.Environ() {
		pair := strings.Split(env, "=")
		if pair[0] == "ALPHA" {
			return pair[1]
		}
	}
	return ""
}

func addFoundServices() {
	foundServices, _, err := generateServices(getAlpha())
	if err != nil {
		log.Fatal(err)
	}

	for _, s := range foundServices {
		if _, found := services[s.Name]; !found {
			services[s.Name] = s
		}
	}
}

func getConfigPath() string {
	return filepath.Join(EdwardConfig.Dir, "edward.json")
}

func loadConfig() {
	groups = make(map[string]*ServiceGroupConfig)
	services = make(map[string]*ServiceConfig)

	configPath := getConfigPath()

	if _, err := os.Stat(configPath); err == nil {
		println("Loading configuration from", configPath)
		r, err := os.Open(configPath)
		if err != nil {
			log.Fatal(err)
		}
		config, err := LoadConfig(r)
		if err != nil {
			log.Fatal(err)
		}

		services = config.ServiceMap
		groups = config.GroupMap
		return
	} else {
		addFoundServices()
		applyHardCodedServicesAndGroups()
	}

}

func getServicesOrGroups(names []string) ([]ServiceOrGroup, error) {
	var outSG []ServiceOrGroup
	for _, name := range names {
		sg, err := getServiceOrGroup(name)
		if err != nil {
			return nil, err
		}
		outSG = append(outSG, sg)
	}
	return outSG, nil
}

func getServiceOrGroup(name string) (ServiceOrGroup, error) {
	if group, ok := groups[name]; ok {
		return group, nil
	}
	if service, ok := services[name]; ok {
		return service, nil
	}
	return nil, errors.New("Service or group not found")
}

func list(c *cli.Context) error {

	var groupNames []string
	var serviceNames []string
	for name, _ := range groups {
		groupNames = append(groupNames, name)
	}
	for name, _ := range services {
		serviceNames = append(serviceNames, name)
	}

	sort.Strings(groupNames)
	sort.Strings(serviceNames)

	println("Services and groups")
	println("Groups:")
	for _, name := range groupNames {
		println("\t", name)
	}
	println("Services:")
	for _, name := range serviceNames {
		println("\t", name)
	}

	return nil
}

func generate(c *cli.Context) error {

	// Add any new services to the config as appropriate
	addFoundServices()

	configPath := getConfigPath()

	if err := generateConfigFile(configPath); err != nil {
		return err
	}
	println("Wrote to", configPath)

	return nil
}

func allStatus() {
	var statuses []ServiceStatus
	for _, service := range services {
		statuses = append(statuses, service.GetStatus()...)
	}
	for _, status := range statuses {
		if status.Status != "STOPPED" {
			println(status.Service.Name, ":", status.Status)
		}
	}
}

func status(c *cli.Context) error {

	if len(c.Args()) == 0 {
		allStatus()
		return nil
	}

	sgs, err := getServicesOrGroups(c.Args())
	if err != nil {
		return err
	}
	for _, s := range sgs {
		statuses := s.GetStatus()
		for _, status := range statuses {
			println(status.Service.Name, ":", status.Status)
		}
	}
	return nil
}

func messages(c *cli.Context) error {
	return errors.New("Unimplemented")
}

func start(c *cli.Context) error {
	sgs, err := getServicesOrGroups(c.Args())
	if err != nil {
		return err
	}
	for _, s := range sgs {
		println("==== Build Phase ====")
		err = s.Build()
		if err != nil {
			return errors.New("Error building " + s.GetName() + ": " + err.Error())
		}
		println("==== Launch Phase ====")
		err = s.Start()
		if err != nil {
			return errors.New("Error launching " + s.GetName() + ": " + err.Error())
		}
	}
	return nil
}

func allServices() []ServiceOrGroup {
	var as []ServiceOrGroup
	for _, service := range services {
		as = append(as, service)
	}
	return as
}

func stop(c *cli.Context) error {
	var sgs []ServiceOrGroup
	var err error
	if len(c.Args()) == 0 {
		sgs = allServices()
	} else {
		sgs, err = getServicesOrGroups(c.Args())
		if err != nil {
			return err
		}
	}
	for _, s := range sgs {
		_ = s.Stop()
	}
	return nil
}

func restart(c *cli.Context) error {
	sgs, err := getServicesOrGroups(c.Args())
	if err != nil {
		return err
	}
	for _, s := range sgs {
		_ = s.Stop()
		err = s.Build()
		if err != nil {
			return err
		}
		err = s.Start()
		if err != nil {
			return err
		}
	}
	return nil
}

func doLog(c *cli.Context) error {
	if len(c.Args()) > 1 {
		return errors.New("Cannot output multiple service logs")
	}
	name := c.Args()[0]
	if _, ok := groups[name]; ok {
		return errors.New("Cannot output group logs")
	}
	if service, ok := services[name]; ok {
		command := service.GetCommand()
		runLog := command.Logs.Run
		t, err := tail.TailFile(runLog, tail.Config{Follow: true})
		if err != nil {
			return nil
		}
		for line := range t.Lines {
			println(line.Text)
		}
		return nil
	}
	return errors.New("Service not found: " + name)
}

func checkNotSudo() {
	user, err := user.Current()
	if err != nil {
		log.Fatal(err)
	}
	if user.Uid == "0" {
		log.Fatal("edward should not be run with sudo")
	}
}

func createScriptFile(suffix string, content string) (*os.File, error) {
	file, err := ioutil.TempFile(os.TempDir(), suffix)
	if err != nil {
		return nil, err
	}
	file.WriteString(content)
	file.Close()

	err = os.Chmod(file.Name(), 0777)
	if err != nil {
		return nil, err
	}

	return file, nil
}

func ensureSudoAble() {
	var buffer bytes.Buffer

	buffer.WriteString("#!/bin/bash\n")
	buffer.WriteString("sudo echo Test > /dev/null\n")
	buffer.WriteString("ISCHILD=YES ")
	buffer.WriteString(strings.Join(os.Args, " "))
	buffer.WriteString("\n")

	file, err := createScriptFile("sudoAbility", buffer.String())
	if err != nil {
		log.Fatal(err)
	}

	err = syscall.Exec(file.Name(), []string{file.Name()}, os.Environ())
	if err != nil {
		log.Fatal(err)
	}
}

func prepareForSudo() {
	checkNotSudo()

	isChild := os.Getenv("ISCHILD")
	if isChild == "" {
		ensureSudoAble()
		return
	}
}

func main() {

	app := cli.NewApp()
	app.Name = "Edward"
	app.Usage = "Manage local microservices"
	app.Before = func(c *cli.Context) error {
		command := c.Args().First()
		if command == "start" || command == "stop" || command == "restart" {
			prepareForSudo()
		}

		err := EdwardConfig.initialize()
		if err != nil {
			return err
		}
		loadConfig()
		return nil
	}
	app.Commands = []cli.Command{
		{
			Name:   "list",
			Usage:  "List available services",
			Action: list,
		},
		{
			Name:   "generate",
			Usage:  "Generate Edward config for a source tree",
			Action: generate,
		},
		{
			Name:   "status",
			Usage:  "Display service status",
			Action: status,
		},
		{
			Name:   "messages",
			Usage:  "Show messages from services",
			Action: messages,
		},
		{
			Name:   "start",
			Usage:  "Build and launch a service",
			Action: start,
		},
		{
			Name:   "stop",
			Usage:  "Stop a service",
			Action: stop,
		},
		{
			Name:   "restart",
			Usage:  "Rebuild and relaunch a service",
			Action: restart,
		},
		{
			Name:   "log",
			Usage:  "Tail the log for a service",
			Action: doLog,
		},
	}

	err := app.Run(os.Args)
	if err != nil {
		log.Fatal(err)
	}
}
<commit_msg>Make 'tail' an alias for 'log'<commit_after>package main

import (
	"bytes"
	"errors"
	"io/ioutil"
	"log"
	"os"
	"os/user"
	"path"
	"path/filepath"
	"sort"
	"strings"
	"syscall"

	"github.com/codegangsta/cli"
	"github.com/hpcloud/tail"
)

type EdwardConfiguration struct {
	Dir       string
	LogDir    string
	PidDir    string
	ScriptDir string
}

var EdwardConfig EdwardConfiguration = EdwardConfiguration{}

func createDirIfNeeded(path string) {
	if _, err := os.Stat(path); os.IsNotExist(err) {
		os.MkdirAll(path, 0777)
	}
}

func (e *EdwardConfiguration) initialize() error {
	user, err := user.Current()
	if err != nil {
		return err
	}
	e.Dir = path.Join(user.HomeDir, ".edward")
	e.LogDir = path.Join(e.Dir, "logs")
	e.PidDir = path.Join(e.Dir, "pidFiles")
	e.ScriptDir = path.Join(e.Dir, "scriptFiles")
	createDirIfNeeded(e.Dir)
	createDirIfNeeded(e.LogDir)
	createDirIfNeeded(e.PidDir)
	createDirIfNeeded(e.ScriptDir)
	return nil
}

var groups map[string]*ServiceGroupConfig
var services map[string]*ServiceConfig

func thirdPartyService(name string, startCommand string, stopCommand string, started string) *ServiceConfig {
	pathStr := "$ALPHA"
	return &ServiceConfig{
		Name: name,
		Path: &pathStr,
		Env:  []string{"YEXT_RABBITMQ=localhost"},
		Commands: ServiceConfigCommands{
			Launch: startCommand,
			Stop:   stopCommand,
		},
		Properties: ServiceConfigProperties{
			Started: started,
		},
	}
}

func getAlpha() string {
	for _, env := range os.Environ() {
		pair := strings.Split(env, "=")
		if pair[0] == "ALPHA" {
			return pair[1]
		}
	}
	return ""
}

func addFoundServices() {
	foundServices, _, err := generateServices(getAlpha())
	if err != nil {
		log.Fatal(err)
	}

	for _, s := range foundServices {
		if _, found := services[s.Name]; !found {
			services[s.Name] = s
		}
	}
}

func getConfigPath() string {
	return filepath.Join(EdwardConfig.Dir, "edward.json")
}

func loadConfig() {
	groups = make(map[string]*ServiceGroupConfig)
	services = make(map[string]*ServiceConfig)

	configPath := getConfigPath()

	if _, err := os.Stat(configPath); err == nil {
		println("Loading configuration from", configPath)
		r, err := os.Open(configPath)
		if err != nil {
			log.Fatal(err)
		}
		config, err := LoadConfig(r)
		if err != nil {
			log.Fatal(err)
		}

		services = config.ServiceMap
		groups = config.GroupMap
		return
	} else {
		addFoundServices()
		applyHardCodedServicesAndGroups()
	}

}

func getServicesOrGroups(names []string) ([]ServiceOrGroup, error) {
	var outSG []ServiceOrGroup
	for _, name := range names {
		sg, err := getServiceOrGroup(name)
		if err != nil {
			return nil, err
		}
		outSG = append(outSG, sg)
	}
	return outSG, nil
}

func getServiceOrGroup(name string) (ServiceOrGroup, error) {
	if group, ok := groups[name]; ok {
		return group, nil
	}
	if service, ok := services[name]; ok {
		return service, nil
	}
	return nil, errors.New("Service or group not found")
}

func list(c *cli.Context) error {

	var groupNames []string
	var serviceNames []string
	for name, _ := range groups {
		groupNames = append(groupNames, name)
	}
	for name, _ := range services {
		serviceNames = append(serviceNames, name)
	}

	sort.Strings(groupNames)
	sort.Strings(serviceNames)

	println("Services and groups")
	println("Groups:")
	for _, name := range groupNames {
		println("\t", name)
	}
	println("Services:")
	for _, name := range serviceNames {
		println("\t", name)
	}

	return nil
}

func generate(c *cli.Context) error {

	// Add any new services to the config as appropriate
	addFoundServices()

	configPath := getConfigPath()

	if err := generateConfigFile(configPath); err != nil {
		return err
	}
	println("Wrote to", configPath)

	return nil
}

func allStatus() {
	var statuses []ServiceStatus
	for _, service := range services {
		statuses = append(statuses, service.GetStatus()...)
	}
	for _, status := range statuses {
		if status.Status != "STOPPED" {
			println(status.Service.Name, ":", status.Status)
		}
	}
}

func status(c *cli.Context) error {

	if len(c.Args()) == 0 {
		allStatus()
		return nil
	}

	sgs, err := getServicesOrGroups(c.Args())
	if err != nil {
		return err
	}
	for _, s := range sgs {
		statuses := s.GetStatus()
		for _, status := range statuses {
			println(status.Service.Name, ":", status.Status)
		}
	}
	return nil
}

func messages(c *cli.Context) error {
	return errors.New("Unimplemented")
}

func start(c *cli.Context) error {
	sgs, err := getServicesOrGroups(c.Args())
	if err != nil {
		return err
	}
	for _, s := range sgs {
		println("==== Build Phase ====")
		err = s.Build()
		if err != nil {
			return errors.New("Error building " + s.GetName() + ": " + err.Error())
		}
		println("==== Launch Phase ====")
		err = s.Start()
		if err != nil {
			return errors.New("Error launching " + s.GetName() + ": " + err.Error())
		}
	}
	return nil
}

func allServices() []ServiceOrGroup {
	var as []ServiceOrGroup
	for _, service := range services {
		as = append(as, service)
	}
	return as
}

func stop(c *cli.Context) error {
	var sgs []ServiceOrGroup
	var err error
	if len(c.Args()) == 0 {
		sgs = allServices()
	} else {
		sgs, err = getServicesOrGroups(c.Args())
		if err != nil {
			return err
		}
	}
	for _, s := range sgs {
		_ = s.Stop()
	}
	return nil
}

func restart(c *cli.Context) error {
	sgs, err := getServicesOrGroups(c.Args())
	if err != nil {
		return err
	}
	for _, s := range sgs {
		_ = s.Stop()
		err = s.Build()
		if err != nil {
			return err
		}
		err = s.Start()
		if err != nil {
			return err
		}
	}
	return nil
}

func doLog(c *cli.Context) error {
	if len(c.Args()) > 1 {
		return errors.New("Cannot output multiple service logs")
	}
	name := c.Args()[0]
	if _, ok := groups[name]; ok {
		return errors.New("Cannot output group logs")
	}
	if service, ok := services[name]; ok {
		command := service.GetCommand()
		runLog := command.Logs.Run
		t, err := tail.TailFile(runLog, tail.Config{Follow: true})
		if err != nil {
			return nil
		}
		for line := range t.Lines {
			println(line.Text)
		}
		return nil
	}
	return errors.New("Service not found: " + name)
}

func checkNotSudo() {
	user, err := user.Current()
	if err != nil {
		log.Fatal(err)
	}
	if user.Uid == "0" {
		log.Fatal("edward should not be run with sudo")
	}
}

func createScriptFile(suffix string, content string) (*os.File, error) {
	file, err := ioutil.TempFile(os.TempDir(), suffix)
	if err != nil {
		return nil, err
	}
	file.WriteString(content)
	file.Close()

	err = os.Chmod(file.Name(), 0777)
	if err != nil {
		return nil, err
	}

	return file, nil
}

func ensureSudoAble() {
	var buffer bytes.Buffer

	buffer.WriteString("#!/bin/bash\n")
	buffer.WriteString("sudo echo Test > /dev/null\n")
	buffer.WriteString("ISCHILD=YES ")
	buffer.WriteString(strings.Join(os.Args, " "))
	buffer.WriteString("\n")

	file, err := createScriptFile("sudoAbility", buffer.String())
	if err != nil {
		log.Fatal(err)
	}

	err = syscall.Exec(file.Name(), []string{file.Name()}, os.Environ())
	if err != nil {
		log.Fatal(err)
	}
}

func prepareForSudo() {
	checkNotSudo()

	isChild := os.Getenv("ISCHILD")
	if isChild == "" {
		ensureSudoAble()
		return
	}
}

func main() {

	app := cli.NewApp()
	app.Name = "Edward"
	app.Usage = "Manage local microservices"
	app.Before = func(c *cli.Context) error {
		command := c.Args().First()
		if command == "start" || command == "stop" || command == "restart" {
			prepareForSudo()
		}

		err := EdwardConfig.initialize()
		if err != nil {
			return err
		}
		loadConfig()
		return nil
	}
	app.Commands = []cli.Command{
		{
			Name:   "list",
			Usage:  "List available services",
			Action: list,
		},
		{
			Name:   "generate",
			Usage:  "Generate Edward config for a source tree",
			Action: generate,
		},
		{
			Name:   "status",
			Usage:  "Display service status",
			Action: status,
		},
		{
			Name:   "messages",
			Usage:  "Show messages from services",
			Action: messages,
		},
		{
			Name:   "start",
			Usage:  "Build and launch a service",
			Action: start,
		},
		{
			Name:   "stop",
			Usage:  "Stop a service",
			Action: stop,
		},
		{
			Name:   "restart",
			Usage:  "Rebuild and relaunch a service",
			Action: restart,
		},
		{
			Name:    "log",
			Aliases: []string{"tail"},
			Usage:   "Tail the log for a service",
			Action:  doLog,
		},
	}

	err := app.Run(os.Args)
	if err != nil {
		log.Fatal(err)
	}
}
<|endoftext|>


<commit_before>// Copyright 2016-2017 Authors of Cilium
//
// 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 labels

import (
	"encoding/json"
	"fmt"
	"os"
	"regexp"
	"strings"

	k8sConst "github.com/cilium/cilium/pkg/k8s/apis/cilium.io"
	"github.com/cilium/cilium/pkg/lock"
	"github.com/cilium/cilium/pkg/logging"
	"github.com/cilium/cilium/pkg/logging/logfields"
)

var (
	log                  = logging.DefaultLogger.WithField(logfields.LogSubsys, "labels-filter")
	validLabelPrefixesMU lock.RWMutex
	validLabelPrefixes   *labelPrefixCfg // Label prefixes used to filter from all labels
)

const (
	// LPCfgFileVersion represents the version of a Label Prefix Configuration File
	LPCfgFileVersion = 1
)

// LabelPrefix is the cilium's representation of a container label.
// +k8s:deepcopy-gen=false
// +k8s:openapi-gen=false
type LabelPrefix struct {
	// Ignore if true will cause this prefix to be ignored insted of being accepted
	Ignore bool   `json:"invert"`
	Prefix string `json:"prefix"`
	Source string `json:"source"`
	expr   *regexp.Regexp
}

// String returns a human readable representation of the LabelPrefix
func (p LabelPrefix) String() string {
	s := fmt.Sprintf("%s:%s", p.Source, p.Prefix)
	if p.Ignore {
		s = "!" + s
	}

	return s
}

// matches returns true and the length of the matched section if the label is
// matched by the LabelPrefix. The Ignore flag has no effect at this point.
func (p LabelPrefix) matches(l Label) (bool, int) {
	if p.Source != "" && p.Source != l.Source {
		return false, 0
	}

	// If no regular expression is available, fall back to prefix matching
	if p.expr == nil {
		return strings.HasPrefix(l.Key, p.Prefix), len(p.Prefix)
	}

	res := p.expr.FindStringIndex(l.Key)

	// No match if regexp was not found
	if res == nil {
		return false, 0
	}

	// Otherwise match if match was found at start of key
	return res[0] == 0, res[1]
}

// parseLabelPrefix returns a LabelPrefix created from the string label parameter.
func parseLabelPrefix(label string) (*LabelPrefix, error) {
	labelPrefix := LabelPrefix{}
	i := strings.IndexByte(label, ':')
	if i >= 0 {
		labelPrefix.Source = label[:i]
		labelPrefix.Prefix = label[i+1:]
	} else {
		labelPrefix.Prefix = label
	}

	if labelPrefix.Prefix[0] == '!' {
		labelPrefix.Ignore = true
		labelPrefix.Prefix = labelPrefix.Prefix[1:]
	}

	r, err := regexp.Compile(labelPrefix.Prefix)
	if err != nil {
		return nil, fmt.Errorf("unable to compile regexp: %s", err)
	}
	labelPrefix.expr = r

	return &labelPrefix, nil
}

// ParseLabelPrefixCfg parses valid label prefixes from a file and from a slice
// of valid prefixes. Both are optional. If both are provided, both list are
// appended together.
func ParseLabelPrefixCfg(prefixes []string, file string) error {
	cfg, err := readLabelPrefixCfgFrom(file)
	if err != nil {
		return fmt.Errorf("unable to read label prefix file: %s", err)
	}

	for _, label := range prefixes {
		p, err := parseLabelPrefix(label)
		if err != nil {
			return err
		}

		if !p.Ignore {
			cfg.whitelist = true
		}

		cfg.LabelPrefixes = append(cfg.LabelPrefixes, p)
	}

	validLabelPrefixes = cfg

	log.Info("Valid label prefix configuration:")
	for _, l := range validLabelPrefixes.LabelPrefixes {
		log.Infof(" - %s", l)
	}

	return nil
}

// labelPrefixCfg is the label prefix configuration to filter labels of started
// containers.
// +k8s:openapi-gen=false
type labelPrefixCfg struct {
	Version       int            `json:"version"`
	LabelPrefixes []*LabelPrefix `json:"valid-prefixes"`
	// whitelist if true, indicates that an inclusive rule has to match
	// in order for the label to be considered
	whitelist bool
}

// defaultLabelPrefixCfg returns a default LabelPrefixCfg using the latest
// LPCfgFileVersion
func defaultLabelPrefixCfg() *labelPrefixCfg {
	cfg := &labelPrefixCfg{
		Version:       LPCfgFileVersion,
		LabelPrefixes: []*LabelPrefix{},
	}

	expressions := []string{
		k8sConst.PodNamespaceLabel,      // include io.kubernetes.pod.namespace
		k8sConst.PodNamespaceMetaLabels, // include all namespace labels
		k8sConst.AppKubernetes,          // include app.kubernetes.io
		"!io.kubernetes",                // ignore all other io.kubernetes labels
		"!kubernetes.io",                // ignore all other kubernetes.io labels
		"!.*beta.kubernetes.io",         // ignore all beta.kubernetes.io labels
		"!k8s.io",                       // ignore all k8s.io labels
		"!pod-template-generation",      // ignore pod-template-generation
		"!pod-template-hash",            // ignore pod-template-hash
		"!controller-revision-hash",     // ignore controller-revision-hash
		"!annotation.*",                 // ignore all annotation labels
		"!etcd_node",                    // ignore etcd_node label
	}

	for _, e := range expressions {
		p, err := parseLabelPrefix(e)
		if err != nil {
			msg := fmt.Sprintf("BUG: Unable to parse default label prefix '%s': %s", e, err)
			panic(msg)
		}
		cfg.LabelPrefixes = append(cfg.LabelPrefixes, p)
	}

	return cfg
}

// readLabelPrefixCfgFrom reads a label prefix configuration file from fileName. If the
// version is not supported by us it returns an error.
func readLabelPrefixCfgFrom(fileName string) (*labelPrefixCfg, error) {
	// if not file is specified, the default is empty
	if fileName == "" {
		return defaultLabelPrefixCfg(), nil
	}

	f, err := os.Open(fileName)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	lpc := labelPrefixCfg{}
	err = json.NewDecoder(f).Decode(&lpc)
	if err != nil {
		return nil, err
	}
	if lpc.Version != LPCfgFileVersion {
		return nil, fmt.Errorf("unsupported version %d", lpc.Version)
	}
	for _, lp := range lpc.LabelPrefixes {
		if lp.Prefix == "" {
			return nil, fmt.Errorf("invalid label prefix file: prefix was empty")
		}
		if lp.Source == "" {
			return nil, fmt.Errorf("invalid label prefix file: source was empty")
		}
		if !lp.Ignore {
			lpc.whitelist = true
		}
	}
	return &lpc, nil
}

func (cfg *labelPrefixCfg) filterLabels(lbls Labels) (identityLabels, informationLabels Labels) {
	if lbls == nil {
		return nil, nil
	}

	validLabelPrefixesMU.RLock()
	defer validLabelPrefixesMU.RUnlock()

	identityLabels = Labels{}
	informationLabels = Labels{}
	for k, v := range lbls {
		included, ignored := 0, 0

		for _, p := range cfg.LabelPrefixes {
			if m, len := p.matches(v); m {
				if p.Ignore {
					// save length of shortest matching ignore
					if ignored == 0 || len < ignored {
						ignored = len
					}
				} else {
					// save length of longest matching include
					if len > included {
						included = len
					}
				}
			}
		}

		// A label is accepted if :
		// - No inclusive LabelPrefix (Ignore flag not set) is
		//   configured and label is not ignored.
		// - An inclusive LabelPrefix matches the label
		// - If both an inclusive and ignore LabelPrefix match, the
		//   label is accepted if the matching section in the label
		//   is greater than the ignored matching section in label,
		//   e.g. when evaluating the label foo.bar, the prefix rules
		//   {!foo, foo.bar} will cause the label to be accepted
		//   because the inclusive prefix matches over a longer section.
		if (!cfg.whitelist && ignored == 0) || included > ignored {
			// Just want to make sure we don't have labels deleted in
			// on side and disappearing in the other side...
			identityLabels[k] = v
		} else {
			informationLabels[k] = v
		}
	}
	return identityLabels, informationLabels
}

// FilterLabels returns Labels from the given labels that have the same source and the
// same prefix as one of lpc valid prefixes, as well as labels that do not match
// the aforementioned filtering criteria.
func FilterLabels(lbls Labels) (identityLabels, informationLabels Labels) {
	return validLabelPrefixes.filterLabels(lbls)
}
<commit_msg>pkg/labels check if the map is zero instead of nil<commit_after>// Copyright 2016-2019 Authors of Cilium
//
// 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 labels

import (
	"encoding/json"
	"fmt"
	"os"
	"regexp"
	"strings"

	k8sConst "github.com/cilium/cilium/pkg/k8s/apis/cilium.io"
	"github.com/cilium/cilium/pkg/lock"
	"github.com/cilium/cilium/pkg/logging"
	"github.com/cilium/cilium/pkg/logging/logfields"
)

var (
	log                  = logging.DefaultLogger.WithField(logfields.LogSubsys, "labels-filter")
	validLabelPrefixesMU lock.RWMutex
	validLabelPrefixes   *labelPrefixCfg // Label prefixes used to filter from all labels
)

const (
	// LPCfgFileVersion represents the version of a Label Prefix Configuration File
	LPCfgFileVersion = 1
)

// LabelPrefix is the cilium's representation of a container label.
// +k8s:deepcopy-gen=false
// +k8s:openapi-gen=false
type LabelPrefix struct {
	// Ignore if true will cause this prefix to be ignored insted of being accepted
	Ignore bool   `json:"invert"`
	Prefix string `json:"prefix"`
	Source string `json:"source"`
	expr   *regexp.Regexp
}

// String returns a human readable representation of the LabelPrefix
func (p LabelPrefix) String() string {
	s := fmt.Sprintf("%s:%s", p.Source, p.Prefix)
	if p.Ignore {
		s = "!" + s
	}

	return s
}

// matches returns true and the length of the matched section if the label is
// matched by the LabelPrefix. The Ignore flag has no effect at this point.
func (p LabelPrefix) matches(l Label) (bool, int) {
	if p.Source != "" && p.Source != l.Source {
		return false, 0
	}

	// If no regular expression is available, fall back to prefix matching
	if p.expr == nil {
		return strings.HasPrefix(l.Key, p.Prefix), len(p.Prefix)
	}

	res := p.expr.FindStringIndex(l.Key)

	// No match if regexp was not found
	if res == nil {
		return false, 0
	}

	// Otherwise match if match was found at start of key
	return res[0] == 0, res[1]
}

// parseLabelPrefix returns a LabelPrefix created from the string label parameter.
func parseLabelPrefix(label string) (*LabelPrefix, error) {
	labelPrefix := LabelPrefix{}
	i := strings.IndexByte(label, ':')
	if i >= 0 {
		labelPrefix.Source = label[:i]
		labelPrefix.Prefix = label[i+1:]
	} else {
		labelPrefix.Prefix = label
	}

	if labelPrefix.Prefix[0] == '!' {
		labelPrefix.Ignore = true
		labelPrefix.Prefix = labelPrefix.Prefix[1:]
	}

	r, err := regexp.Compile(labelPrefix.Prefix)
	if err != nil {
		return nil, fmt.Errorf("unable to compile regexp: %s", err)
	}
	labelPrefix.expr = r

	return &labelPrefix, nil
}

// ParseLabelPrefixCfg parses valid label prefixes from a file and from a slice
// of valid prefixes. Both are optional. If both are provided, both list are
// appended together.
func ParseLabelPrefixCfg(prefixes []string, file string) error {
	cfg, err := readLabelPrefixCfgFrom(file)
	if err != nil {
		return fmt.Errorf("unable to read label prefix file: %s", err)
	}

	for _, label := range prefixes {
		p, err := parseLabelPrefix(label)
		if err != nil {
			return err
		}

		if !p.Ignore {
			cfg.whitelist = true
		}

		cfg.LabelPrefixes = append(cfg.LabelPrefixes, p)
	}

	validLabelPrefixes = cfg

	log.Info("Valid label prefix configuration:")
	for _, l := range validLabelPrefixes.LabelPrefixes {
		log.Infof(" - %s", l)
	}

	return nil
}

// labelPrefixCfg is the label prefix configuration to filter labels of started
// containers.
// +k8s:openapi-gen=false
type labelPrefixCfg struct {
	Version       int            `json:"version"`
	LabelPrefixes []*LabelPrefix `json:"valid-prefixes"`
	// whitelist if true, indicates that an inclusive rule has to match
	// in order for the label to be considered
	whitelist bool
}

// defaultLabelPrefixCfg returns a default LabelPrefixCfg using the latest
// LPCfgFileVersion
func defaultLabelPrefixCfg() *labelPrefixCfg {
	cfg := &labelPrefixCfg{
		Version:       LPCfgFileVersion,
		LabelPrefixes: []*LabelPrefix{},
	}

	expressions := []string{
		k8sConst.PodNamespaceLabel,      // include io.kubernetes.pod.namespace
		k8sConst.PodNamespaceMetaLabels, // include all namespace labels
		k8sConst.AppKubernetes,          // include app.kubernetes.io
		"!io.kubernetes",                // ignore all other io.kubernetes labels
		"!kubernetes.io",                // ignore all other kubernetes.io labels
		"!.*beta.kubernetes.io",         // ignore all beta.kubernetes.io labels
		"!k8s.io",                       // ignore all k8s.io labels
		"!pod-template-generation",      // ignore pod-template-generation
		"!pod-template-hash",            // ignore pod-template-hash
		"!controller-revision-hash",     // ignore controller-revision-hash
		"!annotation.*",                 // ignore all annotation labels
		"!etcd_node",                    // ignore etcd_node label
	}

	for _, e := range expressions {
		p, err := parseLabelPrefix(e)
		if err != nil {
			msg := fmt.Sprintf("BUG: Unable to parse default label prefix '%s': %s", e, err)
			panic(msg)
		}
		cfg.LabelPrefixes = append(cfg.LabelPrefixes, p)
	}

	return cfg
}

// readLabelPrefixCfgFrom reads a label prefix configuration file from fileName. If the
// version is not supported by us it returns an error.
func readLabelPrefixCfgFrom(fileName string) (*labelPrefixCfg, error) {
	// if not file is specified, the default is empty
	if fileName == "" {
		return defaultLabelPrefixCfg(), nil
	}

	f, err := os.Open(fileName)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	lpc := labelPrefixCfg{}
	err = json.NewDecoder(f).Decode(&lpc)
	if err != nil {
		return nil, err
	}
	if lpc.Version != LPCfgFileVersion {
		return nil, fmt.Errorf("unsupported version %d", lpc.Version)
	}
	for _, lp := range lpc.LabelPrefixes {
		if lp.Prefix == "" {
			return nil, fmt.Errorf("invalid label prefix file: prefix was empty")
		}
		if lp.Source == "" {
			return nil, fmt.Errorf("invalid label prefix file: source was empty")
		}
		if !lp.Ignore {
			lpc.whitelist = true
		}
	}
	return &lpc, nil
}

func (cfg *labelPrefixCfg) filterLabels(lbls Labels) (identityLabels, informationLabels Labels) {
	if len(lbls) == 0 {
		return nil, nil
	}

	validLabelPrefixesMU.RLock()
	defer validLabelPrefixesMU.RUnlock()

	identityLabels = Labels{}
	informationLabels = Labels{}
	for k, v := range lbls {
		included, ignored := 0, 0

		for _, p := range cfg.LabelPrefixes {
			if m, len := p.matches(v); m {
				if p.Ignore {
					// save length of shortest matching ignore
					if ignored == 0 || len < ignored {
						ignored = len
					}
				} else {
					// save length of longest matching include
					if len > included {
						included = len
					}
				}
			}
		}

		// A label is accepted if :
		// - No inclusive LabelPrefix (Ignore flag not set) is
		//   configured and label is not ignored.
		// - An inclusive LabelPrefix matches the label
		// - If both an inclusive and ignore LabelPrefix match, the
		//   label is accepted if the matching section in the label
		//   is greater than the ignored matching section in label,
		//   e.g. when evaluating the label foo.bar, the prefix rules
		//   {!foo, foo.bar} will cause the label to be accepted
		//   because the inclusive prefix matches over a longer section.
		if (!cfg.whitelist && ignored == 0) || included > ignored {
			// Just want to make sure we don't have labels deleted in
			// on side and disappearing in the other side...
			identityLabels[k] = v
		} else {
			informationLabels[k] = v
		}
	}
	return identityLabels, informationLabels
}

// FilterLabels returns Labels from the given labels that have the same source and the
// same prefix as one of lpc valid prefixes, as well as labels that do not match
// the aforementioned filtering criteria.
func FilterLabels(lbls Labels) (identityLabels, informationLabels Labels) {
	return validLabelPrefixes.filterLabels(lbls)
}
<|endoftext|>


<commit_before>/*
Copyright 2015 All rights reserved.
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 main

import (
	"fmt"
	"os"
	"time"

	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/golang/glog"
)

var (
	// the aws client
	awsCli *awsClient
)

//
// Steps:
//  - grab the command line configuration
//  - retrieve the instance identity document from metadata service
//  - find the instances in the auto-scaling group
//  - create a etcd client from the instance and see if we can connect the cluster
//  - write out the environment file
//  - if in proxy mode we can exit here
//  - check if the instance id exists in the cluster and if not, try to add us
//  - list the members in the cluster and find the instance status, if terminated, try and remove
//

func main() {
	if err := getConfig(); err != nil {
		printUsage(fmt.Sprintf("invalid configuration, error: %s", err))
	}
	glog.Infof("starting %s version: %s, author: %s <%s>", program, version, author, email)

	// step: retrieve this instances identity
	identity, err := getInstanceIdentity()
	if err != nil {
		glog.Errorf("failed to get the instance identity, error: %s", err)
		os.Exit(1)
	}

	// step: create a aws client
	awsCli, err = newAwsClient(identity.Region)
	if err != nil {
		glog.Errorf("failed to create a aws client, error: %s", err)
		os.Exit(1)
	}

	// step: get a list of instances in the scaling group
	instances, err := getAutoScalingMembers(identity)
	if err != nil {
		glog.Errorf("failed to retrieve a list of instance from auto-scaling group, error: %s", err)
		os.Exit(1)
	}

	cluster_state := "new"

	// step: create an etcd client for us
	client, err := newEtcdClient(getEtcdEndpoints(instances))
	if err != nil {
		glog.Warningf("failed to create an etcd client, error: %s", err)
	} else {
		found, err := client.hasMember(identity.InstanceID)
		if err != nil {
			glog.Warningf("failed to check if we're in the cluster, error: %s", err)
		}
		if found {
			cluster_state = "existing"
		}
	}

	if config.proxyMode {
		cluster_state = "existing"
	}

	// step: write out the environment file
	glog.Infof("writing the environment variables to file: %s", config.environmentFile)
	if err := writeEnvironment(config.environmentFile, identity, instances, cluster_state, config.proxyMode); err != nil {
		glog.Errorf("failed to write the environment file, error: %s", err)
		os.Exit(1)
	}

	// step: create an etcd client from the members if NOT in proxy mode
	if !config.proxyMode {
		glog.Infof("attempting to add the member: %s into the cluster", identity.InstanceID)
		// step: update the etcd cluster
		if err := syncMembership(identity, getEtcdEndpoints(instances)); err != nil {
			glog.Errorf("failed to update the etcd cluster, error: %s", err)
			os.Exit(1)
		}
	}
}

// getAutoScalingMembers retrieve the members from the auto-scaling group
func getAutoScalingMembers(identity *awsIdentity) ([]*ec2.Instance, error) {
	autoScalingGroupName := config.groupName
	// step: are we in proxy mode?
	if autoScalingGroupName == "" {
		glog.Infof("etcd auto-scaling group not set, using instance id %s for search", identity.InstanceID)
		name, err := awsCli.getAutoScalingGroupWithInstanceID(identity.InstanceID)
		if err != nil {
			return nil, err
		}
		autoScalingGroupName = name
	}

	glog.Infof("retrieving the instances from the group: %s", autoScalingGroupName)
	// step: get a list of the instance
	instances, err := awsCli.getAutoScalingInstances(autoScalingGroupName)
	if err != nil {
		return nil, err
	}

	return instances, nil
}

// syncMembership is responsible for adding the new member into the cluster and cleaning up anyone
// that doesn't need to be there anymore
func syncMembership(identity *awsIdentity, instances []string) error {
	memberID := identity.InstanceID
	client, err := newEtcdClient(instances)
	if err != nil {
		return err
	}

	glog.Infof("retrievig a list of cluster members")
	members, err := client.listMembers()
	if err != nil {
		return err
	}

	glog.Infof("found %d members in the cluster", len(members))

	// step: if we are not in the cluster, attempt to add our self
	glog.V(3).Infof("checking the member %s is part of the cluster", memberID)
	if found, err := client.hasMember(memberID); err != nil {
		return err
	} else if !found {
		glog.Infof("member %s is not presently part of the cluster, adding now", memberID)

		nodeAddress := identity.PrivateDNSName
		if config.privateIPs {
			nodeAddress = identity.LocalIP
		}
		peerURL := getPeerURL(nodeAddress)

		glog.Infof("attempting to add the member, peerURL: %s", peerURL)

		if err := client.addMember(memberID, peerURL); err != nil {
			glog.Errorf("failed to add the member into the cluster, error: %s", err)
		}
		glog.Infof("successfully added the member: %s to cluster", memberID)
	} else {
		glog.Infof("member %s is already in the cluster, moving to cleanup", memberID)
	}

	// step: attempt to remove any boxes from the cluster which have terminated
	glog.Infof("checking if any cluster members can been cleaned out")

	// step: remove any members no longer required
	for _, i := range members {
		glog.V(10).Infof("checking if instance: %s, url: %s is still alive", i.Name, i.PeerURLs)
		terminated, err := awsCli.isInstanceTerminated(i.Name)
		if err != nil {
			glog.Warningf("failed to determine if member %s is running, error: %s", i.Name, err)
			continue
		}
		if terminated {
			glog.Infof("member %s has been terminated, removing from the cluster", i.Name)
			removed := false
			for j := 0; j < 3; j++ {
				if err := client.deleteMember(i.ID); err != nil {
					glog.Errorf("failed to remove the member %s, error: %s", i.Name, err)
					<-time.After(time.Duration(3) * time.Second)
				} else {
					glog.Infof("successfully remove the member: %s", i.Name)
					removed = true
					break
				}
			}

			if !removed {
				return fmt.Errorf("failed to remove the member: %s", i.Name)
			}
		}
	}

	return nil
}

func writeEnvironment(filename string, identity *awsIdentity, members []*ec2.Instance, state string, proxy bool) error {
	// step: generate the cluster url
	peersURL := getPeerURLs(members)
	mode := "off"
	if proxy {
		mode  = "on"
	}
	content := fmt.Sprintf(`
ETCD_INITIAL_CLUSTER_STATE=%s
ETCD_NAME=%s
ETCD_INITIAL_CLUSTER="%s"
ETCD_PROXY="%s"
`, state, identity.InstanceID, peersURL, mode)

	if err := writeFile(filename, content); err != nil {
		return err
	}

	return nil
}
<commit_msg>- fixing the logic for preexisting clusters<commit_after>/*
Copyright 2015 All rights reserved.
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 main

import (
	"fmt"
	"os"
	"time"

	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/golang/glog"
)

var (
	// the aws client
	awsCli *awsClient
)

//
// Steps:
//  - grab the command line configuration
//  - retrieve the instance identity document from metadata service
//  - find the instances in the auto-scaling group
//  - create a etcd client from the instance and see if we can connect the cluster
//  - write out the environment file
//  - if in proxy mode we can exit here
//  - check if the instance id exists in the cluster and if not, try to add us
//  - list the members in the cluster and find the instance status, if terminated, try and remove
//

func main() {
	if err := getConfig(); err != nil {
		printUsage(fmt.Sprintf("invalid configuration, error: %s", err))
	}
	glog.Infof("starting %s version: %s, author: %s <%s>", program, version, author, email)

	// step: retrieve this instances identity
	identity, err := getInstanceIdentity()
	if err != nil {
		glog.Errorf("failed to get the instance identity, error: %s", err)
		os.Exit(1)
	}

	// step: create a aws client
	awsCli, err = newAwsClient(identity.Region)
	if err != nil {
		glog.Errorf("failed to create a aws client, error: %s", err)
		os.Exit(1)
	}

	// step: get a list of instances in the scaling group
	instances, err := getAutoScalingMembers(identity)
	if err != nil {
		glog.Errorf("failed to retrieve a list of instance from auto-scaling group, error: %s", err)
		os.Exit(1)
	}

	cluster_state := "new"

	// step: create an etcd client for us and check if we can connect to cluster
	if _, err := newEtcdClient(getEtcdEndpoints(instances)); err == nil {
		cluster_state = "existing"
	}

	// step: if in proxy mode, its always existing
	if config.proxyMode {
		cluster_state = "existing"
	}

	// step: write out the environment file
	glog.Infof("writing the environment variables to file: %s", config.environmentFile)
	if err := writeEnvironment(config.environmentFile, identity, instances, cluster_state, config.proxyMode); err != nil {
		glog.Errorf("failed to write the environment file, error: %s", err)
		os.Exit(1)
	}

	// step: create an etcd client from the members if NOT in proxy mode
	if !config.proxyMode {
		glog.Infof("attempting to add the member: %s into the cluster", identity.InstanceID)
		// step: update the etcd cluster
		if err := syncMembership(identity, getEtcdEndpoints(instances)); err != nil {
			glog.Errorf("failed to update the etcd cluster, error: %s", err)
			os.Exit(1)
		}
	}
}

// getAutoScalingMembers retrieve the members from the auto-scaling group
func getAutoScalingMembers(identity *awsIdentity) ([]*ec2.Instance, error) {
	autoScalingGroupName := config.groupName
	// step: are we in proxy mode?
	if autoScalingGroupName == "" {
		glog.Infof("etcd auto-scaling group not set, using instance id %s for search", identity.InstanceID)
		name, err := awsCli.getAutoScalingGroupWithInstanceID(identity.InstanceID)
		if err != nil {
			return nil, err
		}
		autoScalingGroupName = name
	}

	glog.Infof("retrieving the instances from the group: %s", autoScalingGroupName)
	// step: get a list of the instance
	instances, err := awsCli.getAutoScalingInstances(autoScalingGroupName)
	if err != nil {
		return nil, err
	}

	return instances, nil
}

// syncMembership is responsible for adding the new member into the cluster and cleaning up anyone
// that doesn't need to be there anymore
func syncMembership(identity *awsIdentity, instances []string) error {
	memberID := identity.InstanceID
	client, err := newEtcdClient(instances)
	if err != nil {
		return err
	}

	glog.Infof("retrievig a list of cluster members")
	members, err := client.listMembers()
	if err != nil {
		return err
	}

	glog.Infof("found %d members in the cluster", len(members))

	// step: if we are not in the cluster, attempt to add our self
	glog.V(3).Infof("checking the member %s is part of the cluster", memberID)
	if found, err := client.hasMember(memberID); err != nil {
		return err
	} else if !found {
		glog.Infof("member %s is not presently part of the cluster, adding now", memberID)

		nodeAddress := identity.PrivateDNSName
		if config.privateIPs {
			nodeAddress = identity.LocalIP
		}
		peerURL := getPeerURL(nodeAddress)

		glog.Infof("attempting to add the member, peerURL: %s", peerURL)

		if err := client.addMember(memberID, peerURL); err != nil {
			glog.Errorf("failed to add the member into the cluster, error: %s", err)
		}
		glog.Infof("successfully added the member: %s to cluster", memberID)
	} else {
		glog.Infof("member %s is already in the cluster, moving to cleanup", memberID)
	}

	// step: attempt to remove any boxes from the cluster which have terminated
	glog.Infof("checking if any cluster members can been cleaned out")

	// step: remove any members no longer required
	for _, i := range members {
		glog.V(10).Infof("checking if instance: %s, url: %s is still alive", i.Name, i.PeerURLs)
		terminated, err := awsCli.isInstanceTerminated(i.Name)
		if err != nil {
			glog.Warningf("failed to determine if member %s is running, error: %s", i.Name, err)
			continue
		}
		if terminated {
			glog.Infof("member %s has been terminated, removing from the cluster", i.Name)
			removed := false
			for j := 0; j < 3; j++ {
				if err := client.deleteMember(i.ID); err != nil {
					glog.Errorf("failed to remove the member %s, error: %s", i.Name, err)
					<-time.After(time.Duration(3) * time.Second)
				} else {
					glog.Infof("successfully remove the member: %s", i.Name)
					removed = true
					break
				}
			}

			if !removed {
				return fmt.Errorf("failed to remove the member: %s", i.Name)
			}
		}
	}

	return nil
}

func writeEnvironment(filename string, identity *awsIdentity, members []*ec2.Instance, state string, proxy bool) error {
	// step: generate the cluster url
	peersURL := getPeerURLs(members)
	mode := "off"
	if proxy {
		mode  = "on"
	}
	content := fmt.Sprintf(`
ETCD_INITIAL_CLUSTER_STATE=%s
ETCD_NAME=%s
ETCD_INITIAL_CLUSTER="%s"
ETCD_PROXY="%s"
`, state, identity.InstanceID, peersURL, mode)

	if err := writeFile(filename, content); err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package logger

import (
	"io/ioutil"
	"strings"
	"sync"
	"time"

	"github.com/go-redis/redis/v7"
	"github.com/sirupsen/logrus"
)

const (
	debugRedisAddChannel = "add:log-debug"
	debugRedisRmvChannel = "rmv:log-debug"
	debugRedisPrefix     = "debug:"
)

var opts Options
var loggers = make(map[string]domainEntry)
var loggersMu sync.RWMutex

// Options contains the configuration values of the logger system
type Options struct {
	Syslog bool
	Level  string
	Redis  redis.UniversalClient
}

type domainEntry struct {
	log       *logrus.Logger
	expiredAt *time.Time
}

func (entry *domainEntry) Expired() bool {
	if entry.expiredAt == nil {
		return false
	}
	return entry.expiredAt.Before(time.Now())
}

// Init initializes the logger module with the specified options.
func Init(opt Options) error {
	level := opt.Level
	if level == "" {
		level = "info"
	}
	logLevel, err := logrus.ParseLevel(level)
	if err != nil {
		return err
	}
	logrus.SetLevel(logLevel)
	if opt.Syslog {
		hook, err := syslogHook()
		if err != nil {
			return err
		}
		logrus.AddHook(hook)
		logrus.SetOutput(ioutil.Discard)
	}
	if cli := opt.Redis; cli != nil {
		go subscribeLoggersDebug(cli)
		go loadDebug(cli)
	}
	opts = opt
	return nil
}

// Clone clones a logrus.Logger struct.
func Clone(in *logrus.Logger) *logrus.Logger {
	out := &logrus.Logger{
		Out:       in.Out,
		Hooks:     make(logrus.LevelHooks),
		Formatter: in.Formatter,
		Level:     in.Level,
	}
	for k, v := range in.Hooks {
		out.Hooks[k] = v
	}
	return out
}

// AddDebugDomain adds the specified domain to the debug list.
func AddDebugDomain(domain string, ttl time.Duration) error {
	if cli := opts.Redis; cli != nil {
		return publishDebug(cli, debugRedisAddChannel, domain, ttl)
	}
	addDebugDomain(domain, ttl)
	return nil
}

// RemoveDebugDomain removes the specified domain from the debug list.
func RemoveDebugDomain(domain string) error {
	if cli := opts.Redis; cli != nil {
		return publishDebug(cli, debugRedisRmvChannel, domain, 0)
	}
	removeDebugDomain(domain)
	return nil
}

// WithNamespace returns a logger with the specified nspace field.
func WithNamespace(nspace string) *logrus.Entry {
	return logrus.WithField("nspace", nspace)
}

// WithDomain returns a logger with the specified domain field.
func WithDomain(domain string) *logrus.Entry {
	loggersMu.RLock()
	entry, ok := loggers[domain]
	loggersMu.RUnlock()
	if ok {
		if !entry.Expired() {
			return entry.log.WithField("domain", domain)
		}
		removeDebugDomain(domain)
	}
	return logrus.WithField("domain", domain)
}

func addDebugDomain(domain string, ttl time.Duration) {
	loggersMu.Lock()
	defer loggersMu.Unlock()
	_, ok := loggers[domain]
	if ok {
		return
	}
	logger := logrus.New()
	logger.Level = logrus.DebugLevel
	if opts.Syslog {
		hook, err := syslogHook()
		if err == nil {
			logger.Hooks.Add(hook)
			logger.Out = ioutil.Discard
		}
	}
	expiredAt := time.Now().Add(ttl)
	loggers[domain] = domainEntry{logger, &expiredAt}
}

func removeDebugDomain(domain string) {
	loggersMu.Lock()
	defer loggersMu.Unlock()
	delete(loggers, domain)
}

func subscribeLoggersDebug(cli redis.UniversalClient) {
	sub := cli.Subscribe(debugRedisAddChannel, debugRedisRmvChannel)
	for msg := range sub.Channel() {
		parts := strings.Split(msg.Payload, "/")
		domain := parts[0]
		switch msg.Channel {
		case debugRedisAddChannel:
			var ttl time.Duration
			if len(parts) >= 2 {
				ttl, _ = time.ParseDuration(parts[1])
			}
			addDebugDomain(domain, ttl)
		case debugRedisRmvChannel:
			removeDebugDomain(domain)
		}
	}
}

func loadDebug(cli redis.UniversalClient) {
	keys, err := cli.Keys(debugRedisPrefix + "*").Result()
	if err != nil {
		return
	}
	for _, key := range keys {
		ttl, err := cli.TTL(key).Result()
		if err != nil {
			continue
		}
		domain := strings.TrimPrefix(key, debugRedisPrefix)
		addDebugDomain(domain, ttl)
	}
}

func publishDebug(cli redis.UniversalClient, channel, domain string, ttl time.Duration) error {
	err := cli.Publish(channel, domain+"/"+ttl.String()).Err()
	if err != nil {
		return err
	}
	key := debugRedisPrefix + domain
	if channel == debugRedisAddChannel {
		err = cli.Set(key, 0, ttl).Err()
	} else {
		err = cli.Del(key).Err()
	}
	return err
}

// DebugExpiration returns the expiration date for the debug mode for the
// instance logger of the given domain (or nil if the debug mode is not
// activated).
func DebugExpiration(domain string) *time.Time {
	loggersMu.RLock()
	entry, ok := loggers[domain]
	loggersMu.RUnlock()
	if !ok {
		return nil
	}
	return entry.expiredAt
}

// IsDebug returns whether or not the debug mode is activated.
func IsDebug(logger *logrus.Entry) bool {
	return logger.Logger.Level == logrus.DebugLevel
}
<commit_msg>Increase time precision for logs in debug<commit_after>package logger

import (
	"io/ioutil"
	"strings"
	"sync"
	"time"

	build "github.com/cozy/cozy-stack/pkg/config"
	"github.com/go-redis/redis/v7"
	"github.com/sirupsen/logrus"
)

const (
	debugRedisAddChannel = "add:log-debug"
	debugRedisRmvChannel = "rmv:log-debug"
	debugRedisPrefix     = "debug:"
)

var opts Options
var loggers = make(map[string]domainEntry)
var loggersMu sync.RWMutex

// Options contains the configuration values of the logger system
type Options struct {
	Syslog bool
	Level  string
	Redis  redis.UniversalClient
}

type domainEntry struct {
	log       *logrus.Logger
	expiredAt *time.Time
}

func (entry *domainEntry) Expired() bool {
	if entry.expiredAt == nil {
		return false
	}
	return entry.expiredAt.Before(time.Now())
}

// Init initializes the logger module with the specified options.
func Init(opt Options) error {
	level := opt.Level
	if level == "" {
		level = "info"
	}
	logLevel, err := logrus.ParseLevel(level)
	if err != nil {
		return err
	}
	logrus.SetLevel(logLevel)
	if opt.Syslog {
		hook, err := syslogHook()
		if err != nil {
			return err
		}
		logrus.AddHook(hook)
		logrus.SetOutput(ioutil.Discard)
	} else if build.IsDevRelease() && logLevel == logrus.DebugLevel {
		formatter := logrus.StandardLogger().Formatter.(*logrus.TextFormatter)
		formatter.TimestampFormat = time.RFC3339Nano
	}
	if cli := opt.Redis; cli != nil {
		go subscribeLoggersDebug(cli)
		go loadDebug(cli)
	}
	opts = opt
	return nil
}

// Clone clones a logrus.Logger struct.
func Clone(in *logrus.Logger) *logrus.Logger {
	out := &logrus.Logger{
		Out:       in.Out,
		Hooks:     make(logrus.LevelHooks),
		Formatter: in.Formatter,
		Level:     in.Level,
	}
	for k, v := range in.Hooks {
		out.Hooks[k] = v
	}
	return out
}

// AddDebugDomain adds the specified domain to the debug list.
func AddDebugDomain(domain string, ttl time.Duration) error {
	if cli := opts.Redis; cli != nil {
		return publishDebug(cli, debugRedisAddChannel, domain, ttl)
	}
	addDebugDomain(domain, ttl)
	return nil
}

// RemoveDebugDomain removes the specified domain from the debug list.
func RemoveDebugDomain(domain string) error {
	if cli := opts.Redis; cli != nil {
		return publishDebug(cli, debugRedisRmvChannel, domain, 0)
	}
	removeDebugDomain(domain)
	return nil
}

// WithNamespace returns a logger with the specified nspace field.
func WithNamespace(nspace string) *logrus.Entry {
	return logrus.WithField("nspace", nspace)
}

// WithDomain returns a logger with the specified domain field.
func WithDomain(domain string) *logrus.Entry {
	loggersMu.RLock()
	entry, ok := loggers[domain]
	loggersMu.RUnlock()
	if ok {
		if !entry.Expired() {
			return entry.log.WithField("domain", domain)
		}
		removeDebugDomain(domain)
	}
	return logrus.WithField("domain", domain)
}

func addDebugDomain(domain string, ttl time.Duration) {
	loggersMu.Lock()
	defer loggersMu.Unlock()
	_, ok := loggers[domain]
	if ok {
		return
	}
	logger := logrus.New()
	logger.Level = logrus.DebugLevel
	if opts.Syslog {
		hook, err := syslogHook()
		if err == nil {
			logger.Hooks.Add(hook)
			logger.Out = ioutil.Discard
		}
	}
	expiredAt := time.Now().Add(ttl)
	loggers[domain] = domainEntry{logger, &expiredAt}
}

func removeDebugDomain(domain string) {
	loggersMu.Lock()
	defer loggersMu.Unlock()
	delete(loggers, domain)
}

func subscribeLoggersDebug(cli redis.UniversalClient) {
	sub := cli.Subscribe(debugRedisAddChannel, debugRedisRmvChannel)
	for msg := range sub.Channel() {
		parts := strings.Split(msg.Payload, "/")
		domain := parts[0]
		switch msg.Channel {
		case debugRedisAddChannel:
			var ttl time.Duration
			if len(parts) >= 2 {
				ttl, _ = time.ParseDuration(parts[1])
			}
			addDebugDomain(domain, ttl)
		case debugRedisRmvChannel:
			removeDebugDomain(domain)
		}
	}
}

func loadDebug(cli redis.UniversalClient) {
	keys, err := cli.Keys(debugRedisPrefix + "*").Result()
	if err != nil {
		return
	}
	for _, key := range keys {
		ttl, err := cli.TTL(key).Result()
		if err != nil {
			continue
		}
		domain := strings.TrimPrefix(key, debugRedisPrefix)
		addDebugDomain(domain, ttl)
	}
}

func publishDebug(cli redis.UniversalClient, channel, domain string, ttl time.Duration) error {
	err := cli.Publish(channel, domain+"/"+ttl.String()).Err()
	if err != nil {
		return err
	}
	key := debugRedisPrefix + domain
	if channel == debugRedisAddChannel {
		err = cli.Set(key, 0, ttl).Err()
	} else {
		err = cli.Del(key).Err()
	}
	return err
}

// DebugExpiration returns the expiration date for the debug mode for the
// instance logger of the given domain (or nil if the debug mode is not
// activated).
func DebugExpiration(domain string) *time.Time {
	loggersMu.RLock()
	entry, ok := loggers[domain]
	loggersMu.RUnlock()
	if !ok {
		return nil
	}
	return entry.expiredAt
}

// IsDebug returns whether or not the debug mode is activated.
func IsDebug(logger *logrus.Entry) bool {
	return logger.Logger.Level == logrus.DebugLevel
}
<|endoftext|>