<commit_before>package main

import (
	"archive/zip"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/tormoder/fit/cmd/fitgen/internal/profile"
)

const (
	workbookNameXLS  = "Profile.xls"
	workbookNameXLSX = "Profile.xlsx"
)

func main() {
	sdkOverride := flag.String(
		"sdk",
		"",
		"provide or override SDK version printed in generated code",
	)
	timestamp := flag.Bool(
		"timestamp",
		false,
		"add generation timestamp to generated code",
	)
	runTests := flag.Bool(
		"test",
		false,
		"run all tests in fit repository after code has been generated",
	)
	verbose := flag.Bool(
		"verbose",
		false,
		"print verbose debugging output for profile parsing and code generation",
	)

	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "usage: fitgen [flags] [path to sdk zip, xls or xlsx file] [output directory]\n")
		flag.PrintDefaults()
	}

	flag.Parse()
	if flag.NArg() != 2 {
		flag.Usage()
		os.Exit(2)
	}

	l := log.New(os.Stdout, "fitgen:\t", 0)

	fitSrcDir := flag.Arg(1)
	l.Println("fit source output directory:", fitSrcDir)

	var (
		messagesOut    = filepath.Join(fitSrcDir, "messages.go")
		typesOut       = filepath.Join(fitSrcDir, "types.go")
		profileOut     = filepath.Join(fitSrcDir, "profile.go")
		typesStringOut = filepath.Join(fitSrcDir, "types_string.go")
		stringerPath   = filepath.Join(fitSrcDir, "cmd/stringer/stringer.go")
	)

	var (
		inputData []byte
		input     = flag.Arg(0)
		inputExt  = filepath.Ext(input)
		err       error
	)

	switch inputExt {
	case ".zip":
		inputData, err = readDataFromZIP(input)
	case ".xls", ".xlsx":
		inputData, err = readDataFromXLSX(input)
		if *sdkOverride == "" {
			log.Fatal("-sdk flag required if input is .xls(x)")
		}
	default:
		l.Fatalln("input file must be of type [.zip | .xls | .xlsx], got:", inputExt)
	}
	if err != nil {
		l.Fatal(err)
	}

	genOptions := []profile.GeneratorOption{
		profile.WithGenerationTimestamp(*timestamp),
		profile.WithLogger(l),
	}
	if *verbose {
		genOptions = append(genOptions, profile.WithDebugOutput())
	}

	var sdkString string
	if *sdkOverride != "" {
		sdkString = *sdkOverride
	} else {
		sdkString = parseSDKVersionStringFromZipFilePath(input)
	}

	sdkMaj, sdkMin, err := parseMajorAndMinorSDKVersion(sdkString)
	if err != nil {
		l.Fatalln("error parsing sdk version:", err)
	}

	generator, err := profile.NewGenerator(sdkMaj, sdkMin, inputData, genOptions...)
	if err != nil {
		l.Fatal(err)
	}

	fitProfile, err := generator.GenerateProfile()
	if err != nil {
		l.Fatal(err)
	}

	if err = ioutil.WriteFile(typesOut, fitProfile.TypesSource, 0644); err != nil {
		l.Fatalf("typegen: error writing types output file: %v", err)
	}

	if err = ioutil.WriteFile(messagesOut, fitProfile.MessagesSource, 0644); err != nil {
		l.Fatalf("typegen: error writing messages output file: %v", err)
	}

	if err = ioutil.WriteFile(profileOut, fitProfile.ProfileSource, 0644); err != nil {
		l.Fatalf("typegen: error writing profile output file: %v", err)
	}

	l.Println("running stringer")
	err = runStringerOnTypes(stringerPath, fitSrcDir, typesStringOut, fitProfile.StringerInput)
	if err != nil {
		l.Fatal(err)
	}
	l.Println("stringer: types done")

	logMesgNumVsMessages(fitProfile.MesgNumsWithoutMessage, l)

	if *runTests {
		err = runAllTests(fitSrcDir)
		if err != nil {
			l.Fatal(err)
		}
		l.Println("go test: pass")
	}

	l.Println("done")
}

func runStringerOnTypes(stringerPath, fitSrcDir, typesStringOut, fitTypes string) error {
	tmpFile, err := ioutil.TempFile("", "stringer")
	if err != nil {
		return fmt.Errorf("stringer: error getting temporary file: %v", err)
	}
	defer func() {
		_ = tmpFile.Close()
		err = os.Remove(tmpFile.Name())
		if err != nil {
			fmt.Printf("stringer: error removing temporary file: %v\n", err)
		}
	}()

	stringerCmd := exec.Command(
		"go",
		"run",
		stringerPath,
		"-trimprefix",
		"-type", fitTypes,
		"-output",
		tmpFile.Name(),
		fitSrcDir,
	)

	output, err := stringerCmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("stringer: error running on types: %v\n%s", err, output)
	}

	outFile, err := os.Create(typesStringOut)
	if err != nil {
		return fmt.Errorf("stringer: error creating output file: %v", err)
	}
	defer func() { _ = outFile.Close() }()

	_, err = outFile.WriteString("//lint:file-ignore SA4003 Ignore checks of unsigned types >= 0. stringer generates these.\n")
	if err != nil {
		return fmt.Errorf("stringer: error creating output file: %v", err)
	}

	_, err = tmpFile.Seek(0, 0)
	if err != nil {
		return fmt.Errorf("stringer: error seeking output file: %v", err)
	}

	_, err = io.Copy(outFile, tmpFile)
	if err != nil {
		return fmt.Errorf("stringer: error creating output file: %v", err)
	}

	return nil
}

func runAllTests(pkgDir string) error {
	listCmd := exec.Command("go", "list", pkgDir+"/...")
	output, err := listCmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("go list: fail: %v\n%s", err, output)
	}

	splitted := strings.Split(string(output), "\n")
	var goTestArgs []string
	// Command
	goTestArgs = append(goTestArgs, "test")
	// Packages
	for _, s := range splitted {
		if strings.Contains(s, "/vendor/") {
			continue
		}
		goTestArgs = append(goTestArgs, s)
	}

	testCmd := exec.Command("go", goTestArgs...)
	output, err = testCmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("go test: fail: %v\n%s", err, output)
	}

	return nil
}

func logMesgNumVsMessages(msgs []string, l *log.Logger) {
	if len(msgs) == 0 {
		return
	}
	l.Println("mesgnum-vs-msgs: implementation detail below, this may be automated in the future")
	l.Println("mesgnum-vs-msgs: #mesgnum values != #generated messages, diff:", len(msgs))
	l.Println("mesgnum-vs-msgs: remember to add/verify map entries for sdk in sdk.go for the following message(s):")
	for _, msg := range msgs {
		l.Printf("mesgnum-vs-msgs: ----> mesgnum %q has no corresponding message\n", msg)
	}
}

func readDataFromZIP(path string) ([]byte, error) {
	r, err := zip.OpenReader(path)
	if err != nil {
		return nil, fmt.Errorf("error opening sdk zip file: %v", err)
	}
	defer r.Close()

	var wfile *zip.File
	for _, f := range r.File {
		if f.Name == workbookNameXLS {
			wfile = f
			break
		}
		if f.Name == workbookNameXLSX {
			wfile = f
			break
		}
	}
	if wfile == nil {
		return nil, fmt.Errorf(
			"no file named %q or %q found in zip archive",
			workbookNameXLS, workbookNameXLSX,
		)
	}

	rc, err := wfile.Open()
	if err != nil {
		return nil, fmt.Errorf("error opening zip archive: %v", err)
	}
	defer rc.Close()

	data, err := ioutil.ReadAll(rc)
	if err != nil {
		return nil, fmt.Errorf("error reading %q from archive: %v", wfile.Name, err)
	}

	return data, nil
}

func readDataFromXLSX(path string) ([]byte, error) {
	return ioutil.ReadFile(path)
}

func parseSDKVersionStringFromZipFilePath(path string) string {
	_, file := filepath.Split(path)
	ver := strings.TrimSuffix(file, ".zip")
	return strings.TrimPrefix(ver, "FitSDKRelease_")
}

func parseMajorAndMinorSDKVersion(sdkString string) (int, int, error) {
	splitted := strings.Split(sdkString, ".")
	if len(splitted) < 2 {
		return 0, 0, fmt.Errorf("could not parse major/minor version from input: %q", sdkString)
	}

	maj, err := strconv.Atoi(splitted[0])
	if err != nil {
		return 0, 0, fmt.Errorf("could not parse major version from input: %q", splitted[0])
	}

	min, err := strconv.Atoi(splitted[1])
	if err != nil {
		return 0, 0, fmt.Errorf("could not parse minor version from input: %q", splitted[1])
	}

	return maj, min, nil
}
<commit_msg>cmd/fitgen: don't filter out (nonexistent) vendor dir for tests<commit_after>package main

import (
	"archive/zip"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/tormoder/fit/cmd/fitgen/internal/profile"
)

const (
	workbookNameXLS  = "Profile.xls"
	workbookNameXLSX = "Profile.xlsx"
)

func main() {
	sdkOverride := flag.String(
		"sdk",
		"",
		"provide or override SDK version printed in generated code",
	)
	timestamp := flag.Bool(
		"timestamp",
		false,
		"add generation timestamp to generated code",
	)
	runTests := flag.Bool(
		"test",
		false,
		"run all tests in fit repository after code has been generated",
	)
	verbose := flag.Bool(
		"verbose",
		false,
		"print verbose debugging output for profile parsing and code generation",
	)

	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "usage: fitgen [flags] [path to sdk zip, xls or xlsx file] [output directory]\n")
		flag.PrintDefaults()
	}

	flag.Parse()
	if flag.NArg() != 2 {
		flag.Usage()
		os.Exit(2)
	}

	l := log.New(os.Stdout, "fitgen:\t", 0)

	fitSrcDir := flag.Arg(1)
	l.Println("fit source output directory:", fitSrcDir)

	var (
		messagesOut    = filepath.Join(fitSrcDir, "messages.go")
		typesOut       = filepath.Join(fitSrcDir, "types.go")
		profileOut     = filepath.Join(fitSrcDir, "profile.go")
		typesStringOut = filepath.Join(fitSrcDir, "types_string.go")
		stringerPath   = filepath.Join(fitSrcDir, "cmd/stringer/stringer.go")
	)

	var (
		inputData []byte
		input     = flag.Arg(0)
		inputExt  = filepath.Ext(input)
		err       error
	)

	switch inputExt {
	case ".zip":
		inputData, err = readDataFromZIP(input)
	case ".xls", ".xlsx":
		inputData, err = readDataFromXLSX(input)
		if *sdkOverride == "" {
			log.Fatal("-sdk flag required if input is .xls(x)")
		}
	default:
		l.Fatalln("input file must be of type [.zip | .xls | .xlsx], got:", inputExt)
	}
	if err != nil {
		l.Fatal(err)
	}

	genOptions := []profile.GeneratorOption{
		profile.WithGenerationTimestamp(*timestamp),
		profile.WithLogger(l),
	}
	if *verbose {
		genOptions = append(genOptions, profile.WithDebugOutput())
	}

	var sdkString string
	if *sdkOverride != "" {
		sdkString = *sdkOverride
	} else {
		sdkString = parseSDKVersionStringFromZipFilePath(input)
	}

	sdkMaj, sdkMin, err := parseMajorAndMinorSDKVersion(sdkString)
	if err != nil {
		l.Fatalln("error parsing sdk version:", err)
	}

	generator, err := profile.NewGenerator(sdkMaj, sdkMin, inputData, genOptions...)
	if err != nil {
		l.Fatal(err)
	}

	fitProfile, err := generator.GenerateProfile()
	if err != nil {
		l.Fatal(err)
	}

	if err = ioutil.WriteFile(typesOut, fitProfile.TypesSource, 0644); err != nil {
		l.Fatalf("typegen: error writing types output file: %v", err)
	}

	if err = ioutil.WriteFile(messagesOut, fitProfile.MessagesSource, 0644); err != nil {
		l.Fatalf("typegen: error writing messages output file: %v", err)
	}

	if err = ioutil.WriteFile(profileOut, fitProfile.ProfileSource, 0644); err != nil {
		l.Fatalf("typegen: error writing profile output file: %v", err)
	}

	l.Println("running stringer")
	err = runStringerOnTypes(stringerPath, fitSrcDir, typesStringOut, fitProfile.StringerInput)
	if err != nil {
		l.Fatal(err)
	}
	l.Println("stringer: types done")

	logMesgNumVsMessages(fitProfile.MesgNumsWithoutMessage, l)

	if *runTests {
		err = runAllTests(fitSrcDir)
		if err != nil {
			l.Fatal(err)
		}
		l.Println("go test: pass")
	}

	l.Println("done")
}

func runStringerOnTypes(stringerPath, fitSrcDir, typesStringOut, fitTypes string) error {
	tmpFile, err := ioutil.TempFile("", "stringer")
	if err != nil {
		return fmt.Errorf("stringer: error getting temporary file: %v", err)
	}
	defer func() {
		_ = tmpFile.Close()
		err = os.Remove(tmpFile.Name())
		if err != nil {
			fmt.Printf("stringer: error removing temporary file: %v\n", err)
		}
	}()

	stringerCmd := exec.Command(
		"go",
		"run",
		stringerPath,
		"-trimprefix",
		"-type", fitTypes,
		"-output",
		tmpFile.Name(),
		fitSrcDir,
	)

	output, err := stringerCmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("stringer: error running on types: %v\n%s", err, output)
	}

	outFile, err := os.Create(typesStringOut)
	if err != nil {
		return fmt.Errorf("stringer: error creating output file: %v", err)
	}
	defer func() { _ = outFile.Close() }()

	_, err = outFile.WriteString("//lint:file-ignore SA4003 Ignore checks of unsigned types >= 0. stringer generates these.\n")
	if err != nil {
		return fmt.Errorf("stringer: error creating output file: %v", err)
	}

	_, err = tmpFile.Seek(0, 0)
	if err != nil {
		return fmt.Errorf("stringer: error seeking output file: %v", err)
	}

	_, err = io.Copy(outFile, tmpFile)
	if err != nil {
		return fmt.Errorf("stringer: error creating output file: %v", err)
	}

	return nil
}

func runAllTests(pkgDir string) error {
	goTestArgs := []string{"test", pkgDir}
	testCmd := exec.Command("go", goTestArgs...)
	output, err := testCmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("go test: fail: %v\n%s", err, output)
	}

	return nil
}

func logMesgNumVsMessages(msgs []string, l *log.Logger) {
	if len(msgs) == 0 {
		return
	}
	l.Println("mesgnum-vs-msgs: implementation detail below, this may be automated in the future")
	l.Println("mesgnum-vs-msgs: #mesgnum values != #generated messages, diff:", len(msgs))
	l.Println("mesgnum-vs-msgs: remember to add/verify map entries for sdk in sdk.go for the following message(s):")
	for _, msg := range msgs {
		l.Printf("mesgnum-vs-msgs: ----> mesgnum %q has no corresponding message\n", msg)
	}
}

func readDataFromZIP(path string) ([]byte, error) {
	r, err := zip.OpenReader(path)
	if err != nil {
		return nil, fmt.Errorf("error opening sdk zip file: %v", err)
	}
	defer r.Close()

	var wfile *zip.File
	for _, f := range r.File {
		if f.Name == workbookNameXLS {
			wfile = f
			break
		}
		if f.Name == workbookNameXLSX {
			wfile = f
			break
		}
	}
	if wfile == nil {
		return nil, fmt.Errorf(
			"no file named %q or %q found in zip archive",
			workbookNameXLS, workbookNameXLSX,
		)
	}

	rc, err := wfile.Open()
	if err != nil {
		return nil, fmt.Errorf("error opening zip archive: %v", err)
	}
	defer rc.Close()

	data, err := ioutil.ReadAll(rc)
	if err != nil {
		return nil, fmt.Errorf("error reading %q from archive: %v", wfile.Name, err)
	}

	return data, nil
}

func readDataFromXLSX(path string) ([]byte, error) {
	return ioutil.ReadFile(path)
}

func parseSDKVersionStringFromZipFilePath(path string) string {
	_, file := filepath.Split(path)
	ver := strings.TrimSuffix(file, ".zip")
	return strings.TrimPrefix(ver, "FitSDKRelease_")
}

func parseMajorAndMinorSDKVersion(sdkString string) (int, int, error) {
	splitted := strings.Split(sdkString, ".")
	if len(splitted) < 2 {
		return 0, 0, fmt.Errorf("could not parse major/minor version from input: %q", sdkString)
	}

	maj, err := strconv.Atoi(splitted[0])
	if err != nil {
		return 0, 0, fmt.Errorf("could not parse major version from input: %q", splitted[0])
	}

	min, err := strconv.Atoi(splitted[1])
	if err != nil {
		return 0, 0, fmt.Errorf("could not parse minor version from input: %q", splitted[1])
	}

	return maj, min, nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"strings"

	"github.com/chzyer/readline"
)

const (
	initPrompt      = "!> "
	multilinePrompt = "!>>> "
)

type CmdShell struct {
	cmdQueryBase
}

func (c *CmdShell) Execute(args []string) error {
	if err := c.validate(); err != nil {
		return err
	}

	if err := c.buildDatabase(); err != nil {
		return err
	}

	rl, err := readline.NewEx(&readline.Config{
		Prompt:                 initPrompt,
		HistoryFile:            "/tmp/gitql-history",
		DisableAutoSaveHistory: true,
	})
	if err != nil {
		return err
	}

	rl.Terminal.Print(fmt.Sprint(`
           gitQL SHELL
           -----------
You must end your queries with ';'

`))

	var cmds []string
	for {
		line, err := rl.Readline()
		if err != nil {
			break
		}
		line = strings.TrimSpace(line)
		if len(line) == 0 {
			continue
		}
		cmds = append(cmds, line)
		if !strings.HasSuffix(line, ";") {
			rl.SetPrompt(multilinePrompt)
			continue
		}

		query := strings.Join(cmds, " ")
		cmds = cmds[:0]
		rl.SetPrompt(initPrompt)
		rl.SaveHistory(query)

		rl.Terminal.Print(fmt.Sprintf("\n--> Executing query: %s\n\n", query))

		schema, rowIter, err := c.executeQuery(query)
		if err != nil {
			rl.Terminal.Print(fmt.Sprintf("ERROR: %v\n\n", err))
			continue
		}

		if err := c.printQuery(schema, rowIter, "pretty"); err != nil {
			rl.Terminal.Print(fmt.Sprintf("ERROR: %v\n\n", err))
			continue
		}
	}

	return rl.Close()
}
<commit_msg>cmd: add colors to shell (#101)<commit_after>package main

import (
	"fmt"
	"strings"

	"github.com/chzyer/readline"
	"github.com/fatih/color"
)

const (
	initPrompt      = "!> "
	multilinePrompt = "!>>> "
)

type CmdShell struct {
	cmdQueryBase
}

func (c *CmdShell) Execute(args []string) error {
	if err := c.validate(); err != nil {
		return err
	}

	if err := c.buildDatabase(); err != nil {
		return err
	}

	blue := color.New(color.FgHiBlue).SprintfFunc()
	white := color.New(color.FgWhite).SprintfFunc()
	red := color.New(color.FgHiRed).PrintfFunc()

	prompt := blue(initPrompt)
	mlPrompt := blue(multilinePrompt)

	rl, err := readline.NewEx(&readline.Config{
		Prompt:                 prompt,
		HistoryFile:            "/tmp/gitql-history",
		DisableAutoSaveHistory: true,
	})
	if err != nil {
		return err
	}

	fmt.Println("          ", white("git")+blue("QL"), "SHELL")
	fmt.Println("           -----------")
	fmt.Println("You must end your queries with ';'")
	fmt.Println("")

	var cmds []string
	for {
		line, err := rl.Readline()
		if err != nil {
			break
		}
		line = strings.TrimSpace(line)
		if len(line) == 0 {
			continue
		}
		cmds = append(cmds, line)
		if !strings.HasSuffix(line, ";") {
			rl.SetPrompt(mlPrompt)
			continue
		}

		query := strings.Join(cmds, " ")
		cmds = cmds[:0]
		rl.SetPrompt(prompt)
		rl.SaveHistory(query)

		rl.Terminal.Print(fmt.Sprintf("\n--> Executing query: %s\n\n", query))

		schema, rowIter, err := c.executeQuery(query)
		if err != nil {
			red("ERROR: %v\n\n", err)
			continue
		}

		if err := c.printQuery(schema, rowIter, "pretty"); err != nil {
			red("ERROR: %v\n\n", err)
			continue
		}
	}

	return rl.Close()
}
<|endoftext|>


<commit_before>// match is a regex matching program similar to 'grep'. However, unlike 'grep',
// it's designed to be used in tests. If its input doesn't match its search
// expression, it prints an error to that effect on stderr. If the input *does*
// match, then it's sent to stdout unchanged, so that multiple 'match' commands
// can be chained.
//
// Note that to make the best use of this behavior, 'match' should be used in a
// bash shell after 'set -e -o pipefail', so that a command chaining several
// 'match' commands exits on the first failure and prints its error message

package main

import (
	"bufio"
	"flag"
	"fmt"
	"os"
	"regexp"
)

var invert = flag.Bool("v", false, "If set, match returns an error if the "+
	"given regex does not match any text on stdin")

func die(f string, args ...interface{}) {
	fmt.Fprintf(os.Stderr, f, args...)
	os.Exit(1)
}

func main() {
	flag.Parse()
	if flag.NArg() == 0 {
		die("Must provide a regex to match")
	} else if flag.NArg() > 1 {
		die("Currently, match can only accept a single regex")
	}
	re := regexp.MustCompile(flag.Args()[0])

	// Read through input looking for match
	s := bufio.NewScanner(os.Stdin)
	exit := -1
	for s.Scan() {
		os.Stdout.Write(s.Bytes())
		os.Stdout.Write([]byte{'\n'})
		if exit >= 0 {
			continue
		}
		switch {
		case !re.Match(s.Bytes()): // No match, continue
		case *invert: // Failure -- matched with -v
			fmt.Fprintf(
				os.Stderr,
				"did not expect to find\n  \"%s\"\nbut matched\n  \"%s\"\nin the text:\n",
				flag.Args()[0],
				s.Text(),
			)
			exit = 1
		default: // Success, match -- copy input to output
			exit = 0
		}
	}
	if exit >= 0 {
		os.Exit(0)
	}
	if !*invert {
		// Failure -- no match
		fmt.Fprintf(
			os.Stderr,
			"failed to find\n  \"%s\"\nin the text:\n",
			flag.Args()[0],
		)
		os.Exit(1)
	}
	// Success -- no match with -v
}
<commit_msg>Fix match s.t. missing match now prints input text<commit_after>// match is a regex matching program similar to 'grep'. However, unlike 'grep',
// it's designed to be used in tests. If its input doesn't match its search
// expression, it prints an error to that effect on stderr. If the input *does*
// match, then it's sent to stdout unchanged, so that multiple 'match' commands
// can be chained.
//
// Note that to make the best use of this behavior, 'match' should be used in a
// bash shell after 'set -e -o pipefail', so that a command chaining several
// 'match' commands exits on the first failure and prints its error message

package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"os"
	"regexp"
	"strings"
)

var invert = flag.Bool("v", false, "If set, match returns an error if the "+
	"given regex does not match any text on stdin")

func die(f string, args ...interface{}) {
	fmt.Fprintf(os.Stderr, f, args...)
	os.Exit(1)
}

func main() {
	flag.Parse()
	if flag.NArg() == 0 {
		die("Must provide a regex to match")
	} else if flag.NArg() > 1 {
		die("Currently, match can only accept a single regex")
	}
	re := regexp.MustCompile(flag.Args()[0])

	// Read through input looking for match && copying input to inputBuf (which is
	// written to stdout at the end)
	inputBuf := &strings.Builder{}
	matchFound := false
	s := bufio.NewScanner(io.TeeReader(os.Stdin, inputBuf))
	for s.Scan() {
		switch {
		case matchFound:
			continue // match was found in the past--just copy input to inputBuf
		case !re.Match(s.Bytes()):
			continue // no match--keep copying and looking
		case *invert:
			// Failure: matched with -v
			fmt.Fprintf(
				os.Stderr,
				"match for forbidden regex\n  \"%s\"\nin the text:\n  \"%s\"\n",
				flag.Args()[0],
				s.Text(),
			)
			os.Exit(1)
		default:
			matchFound = true // match found on this line
		}
	}
	if matchFound || *invert {
		// Success--dump inputBuf (with guaranteed newline @ end)
		io.Copy(os.Stdout, strings.NewReader(strings.TrimSuffix(inputBuf.String(), "\n")))
		os.Stdout.Write([]byte{'\n'})
		os.Exit(0)
	} else if !*invert {
		// Failure: match expected, no match found
		fmt.Fprintf(
			os.Stderr,
			"failed to find\n  \"%s\"\nin the text:\n  \"\"\"\n%s\n  \"\"\"\n",
			flag.Args()[0],
			strings.TrimSuffix(inputBuf.String(), "\n"),
		)
		os.Exit(1)
	}
}
<|endoftext|>


<commit_before>//
// Copyright (C) 2014-2017 Nippon Telegraph and Telephone Corporation.
//
// 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"
	"io/ioutil"
	"net/http"
	_ "net/http/pprof"
	"os"
	"os/signal"
	"runtime"
	"syscall"

	"github.com/coreos/go-systemd/daemon"
	"github.com/jessevdk/go-flags"
	"github.com/kr/pretty"
	log "github.com/sirupsen/logrus"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"

	api "github.com/osrg/gobgp/api"
	"github.com/osrg/gobgp/internal/pkg/version"
	"github.com/osrg/gobgp/pkg/config"
	"github.com/osrg/gobgp/pkg/server"
)

func main() {
	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, syscall.SIGTERM)

	var opts struct {
		ConfigFile      string `short:"f" long:"config-file" description:"specifying a config file"`
		ConfigType      string `short:"t" long:"config-type" description:"specifying config type (toml, yaml, json)" default:"toml"`
		LogLevel        string `short:"l" long:"log-level" description:"specifying log level"`
		LogPlain        bool   `short:"p" long:"log-plain" description:"use plain format for logging (json by default)"`
		UseSyslog       string `short:"s" long:"syslog" description:"use syslogd"`
		Facility        string `long:"syslog-facility" description:"specify syslog facility"`
		DisableStdlog   bool   `long:"disable-stdlog" description:"disable standard logging"`
		CPUs            int    `long:"cpus" description:"specify the number of CPUs to be used"`
		GrpcHosts       string `long:"api-hosts" description:"specify the hosts that gobgpd listens on" default:":50051"`
		GracefulRestart bool   `short:"r" long:"graceful-restart" description:"flag restart-state in graceful-restart capability"`
		Dry             bool   `short:"d" long:"dry-run" description:"check configuration"`
		PProfHost       string `long:"pprof-host" description:"specify the host that gobgpd listens on for pprof" default:"localhost:6060"`
		PProfDisable    bool   `long:"pprof-disable" description:"disable pprof profiling"`
		UseSdNotify     bool   `long:"sdnotify" description:"use sd_notify protocol"`
		TLS             bool   `long:"tls" description:"enable TLS authentication for gRPC API"`
		TLSCertFile     string `long:"tls-cert-file" description:"The TLS cert file"`
		TLSKeyFile      string `long:"tls-key-file" description:"The TLS key file"`
		Version         bool   `long:"version" description:"show version number"`
	}
	_, err := flags.Parse(&opts)
	if err != nil {
		os.Exit(1)
	}

	if opts.Version {
		fmt.Println("gobgpd version", version.Version())
		os.Exit(0)
	}

	if opts.CPUs == 0 {
		runtime.GOMAXPROCS(runtime.NumCPU())
	} else {
		if runtime.NumCPU() < opts.CPUs {
			log.Errorf("Only %d CPUs are available but %d is specified", runtime.NumCPU(), opts.CPUs)
			os.Exit(1)
		}
		runtime.GOMAXPROCS(opts.CPUs)
	}

	if !opts.PProfDisable {
		go func() {
			log.Println(http.ListenAndServe(opts.PProfHost, nil))
		}()
	}

	switch opts.LogLevel {
	case "debug":
		log.SetLevel(log.DebugLevel)
	case "info":
		log.SetLevel(log.InfoLevel)
	default:
		log.SetLevel(log.InfoLevel)
	}

	if opts.DisableStdlog {
		log.SetOutput(ioutil.Discard)
	} else {
		log.SetOutput(os.Stdout)
	}

	if opts.UseSyslog != "" {
		if err := addSyslogHook(opts.UseSyslog, opts.Facility); err != nil {
			log.Error("Unable to connect to syslog daemon, ", opts.UseSyslog)
		}
	}

	if opts.LogPlain {
		if opts.DisableStdlog {
			log.SetFormatter(&log.TextFormatter{
				DisableColors: true,
			})
		}
	} else {
		log.SetFormatter(&log.JSONFormatter{})
	}

	if opts.Dry {
		c, err := config.ReadConfigFile(opts.ConfigFile, opts.ConfigType)
		if err != nil {
			log.WithFields(log.Fields{
				"Topic": "Config",
				"Error": err,
			}).Fatalf("Can't read config file %s", opts.ConfigFile)
		}
		log.WithFields(log.Fields{
			"Topic": "Config",
		}).Info("Finished reading the config file")
		if opts.LogLevel == "debug" {
			pretty.Println(c)
		}
		os.Exit(0)
	}

	maxSize := 256 << 20
	grpcOpts := []grpc.ServerOption{grpc.MaxRecvMsgSize(maxSize), grpc.MaxSendMsgSize(maxSize)}
	if opts.TLS {
		creds, err := credentials.NewServerTLSFromFile(opts.TLSCertFile, opts.TLSKeyFile)
		if err != nil {
			log.Fatalf("Failed to generate credentials: %v", err)
		}
		grpcOpts = append(grpcOpts, grpc.Creds(creds))
	}

	log.Info("gobgpd started")
	bgpServer := server.NewBgpServer(server.GrpcListenAddress(opts.GrpcHosts), server.GrpcOption(grpcOpts))
	go bgpServer.Serve()

	if opts.UseSdNotify {
		if status, err := daemon.SdNotify(false, daemon.SdNotifyReady); !status {
			if err != nil {
				log.Warnf("Failed to send notification via sd_notify(): %s", err)
			} else {
				log.Warnf("The socket sd_notify() isn't available")
			}
		}
	}

	if opts.ConfigFile == "" {
		<-sigCh
		stopServer(bgpServer, opts.UseSdNotify)
		return
	}

	signal.Notify(sigCh, syscall.SIGHUP)

	initialConfig, err := config.ReadConfigFile(opts.ConfigFile, opts.ConfigType)
	if err != nil {
		log.WithFields(log.Fields{
			"Topic": "Config",
			"Error": err,
		}).Fatalf("Can't read config file %s", opts.ConfigFile)
	}
	log.WithFields(log.Fields{
		"Topic": "Config",
	}).Info("Finished reading the config file")

	c := config.ApplyInitialConfig(bgpServer, initialConfig, opts.GracefulRestart)
	for sig := range sigCh {
		if sig != syscall.SIGHUP {
			stopServer(bgpServer, opts.UseSdNotify)
			return
		}

		log.WithFields(log.Fields{
			"Topic": "Config",
		}).Info("Reload the config file")
		newConfig, err := config.ReadConfigFile(opts.ConfigFile, opts.ConfigType)
		if err != nil {
			log.WithFields(log.Fields{
				"Topic": "Config",
				"Error": err,
			}).Warningf("Can't read config file %s", opts.ConfigFile)
			continue
		}

		c = config.UpdateConfig(bgpServer, c, newConfig)
	}
}

func stopServer(bgpServer *server.BgpServer, useSdNotify bool) {
	bgpServer.StopBgp(context.Background(), &api.StopBgpRequest{})
	if useSdNotify {
		daemon.SdNotify(false, daemon.SdNotifyStopping)
	}
}
<commit_msg>Rename c to currentConfig to follow naming pattern<commit_after>//
// Copyright (C) 2014-2017 Nippon Telegraph and Telephone Corporation.
//
// 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"
	"io/ioutil"
	"net/http"
	_ "net/http/pprof"
	"os"
	"os/signal"
	"runtime"
	"syscall"

	"github.com/coreos/go-systemd/daemon"
	"github.com/jessevdk/go-flags"
	"github.com/kr/pretty"
	log "github.com/sirupsen/logrus"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"

	api "github.com/osrg/gobgp/api"
	"github.com/osrg/gobgp/internal/pkg/version"
	"github.com/osrg/gobgp/pkg/config"
	"github.com/osrg/gobgp/pkg/server"
)

func main() {
	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, syscall.SIGTERM)

	var opts struct {
		ConfigFile      string `short:"f" long:"config-file" description:"specifying a config file"`
		ConfigType      string `short:"t" long:"config-type" description:"specifying config type (toml, yaml, json)" default:"toml"`
		LogLevel        string `short:"l" long:"log-level" description:"specifying log level"`
		LogPlain        bool   `short:"p" long:"log-plain" description:"use plain format for logging (json by default)"`
		UseSyslog       string `short:"s" long:"syslog" description:"use syslogd"`
		Facility        string `long:"syslog-facility" description:"specify syslog facility"`
		DisableStdlog   bool   `long:"disable-stdlog" description:"disable standard logging"`
		CPUs            int    `long:"cpus" description:"specify the number of CPUs to be used"`
		GrpcHosts       string `long:"api-hosts" description:"specify the hosts that gobgpd listens on" default:":50051"`
		GracefulRestart bool   `short:"r" long:"graceful-restart" description:"flag restart-state in graceful-restart capability"`
		Dry             bool   `short:"d" long:"dry-run" description:"check configuration"`
		PProfHost       string `long:"pprof-host" description:"specify the host that gobgpd listens on for pprof" default:"localhost:6060"`
		PProfDisable    bool   `long:"pprof-disable" description:"disable pprof profiling"`
		UseSdNotify     bool   `long:"sdnotify" description:"use sd_notify protocol"`
		TLS             bool   `long:"tls" description:"enable TLS authentication for gRPC API"`
		TLSCertFile     string `long:"tls-cert-file" description:"The TLS cert file"`
		TLSKeyFile      string `long:"tls-key-file" description:"The TLS key file"`
		Version         bool   `long:"version" description:"show version number"`
	}
	_, err := flags.Parse(&opts)
	if err != nil {
		os.Exit(1)
	}

	if opts.Version {
		fmt.Println("gobgpd version", version.Version())
		os.Exit(0)
	}

	if opts.CPUs == 0 {
		runtime.GOMAXPROCS(runtime.NumCPU())
	} else {
		if runtime.NumCPU() < opts.CPUs {
			log.Errorf("Only %d CPUs are available but %d is specified", runtime.NumCPU(), opts.CPUs)
			os.Exit(1)
		}
		runtime.GOMAXPROCS(opts.CPUs)
	}

	if !opts.PProfDisable {
		go func() {
			log.Println(http.ListenAndServe(opts.PProfHost, nil))
		}()
	}

	switch opts.LogLevel {
	case "debug":
		log.SetLevel(log.DebugLevel)
	case "info":
		log.SetLevel(log.InfoLevel)
	default:
		log.SetLevel(log.InfoLevel)
	}

	if opts.DisableStdlog {
		log.SetOutput(ioutil.Discard)
	} else {
		log.SetOutput(os.Stdout)
	}

	if opts.UseSyslog != "" {
		if err := addSyslogHook(opts.UseSyslog, opts.Facility); err != nil {
			log.Error("Unable to connect to syslog daemon, ", opts.UseSyslog)
		}
	}

	if opts.LogPlain {
		if opts.DisableStdlog {
			log.SetFormatter(&log.TextFormatter{
				DisableColors: true,
			})
		}
	} else {
		log.SetFormatter(&log.JSONFormatter{})
	}

	if opts.Dry {
		c, err := config.ReadConfigFile(opts.ConfigFile, opts.ConfigType)
		if err != nil {
			log.WithFields(log.Fields{
				"Topic": "Config",
				"Error": err,
			}).Fatalf("Can't read config file %s", opts.ConfigFile)
		}
		log.WithFields(log.Fields{
			"Topic": "Config",
		}).Info("Finished reading the config file")
		if opts.LogLevel == "debug" {
			pretty.Println(c)
		}
		os.Exit(0)
	}

	maxSize := 256 << 20
	grpcOpts := []grpc.ServerOption{grpc.MaxRecvMsgSize(maxSize), grpc.MaxSendMsgSize(maxSize)}
	if opts.TLS {
		creds, err := credentials.NewServerTLSFromFile(opts.TLSCertFile, opts.TLSKeyFile)
		if err != nil {
			log.Fatalf("Failed to generate credentials: %v", err)
		}
		grpcOpts = append(grpcOpts, grpc.Creds(creds))
	}

	log.Info("gobgpd started")
	bgpServer := server.NewBgpServer(server.GrpcListenAddress(opts.GrpcHosts), server.GrpcOption(grpcOpts))
	go bgpServer.Serve()

	if opts.UseSdNotify {
		if status, err := daemon.SdNotify(false, daemon.SdNotifyReady); !status {
			if err != nil {
				log.Warnf("Failed to send notification via sd_notify(): %s", err)
			} else {
				log.Warnf("The socket sd_notify() isn't available")
			}
		}
	}

	if opts.ConfigFile == "" {
		<-sigCh
		stopServer(bgpServer, opts.UseSdNotify)
		return
	}

	signal.Notify(sigCh, syscall.SIGHUP)

	initialConfig, err := config.ReadConfigFile(opts.ConfigFile, opts.ConfigType)
	if err != nil {
		log.WithFields(log.Fields{
			"Topic": "Config",
			"Error": err,
		}).Fatalf("Can't read config file %s", opts.ConfigFile)
	}
	log.WithFields(log.Fields{
		"Topic": "Config",
	}).Info("Finished reading the config file")

	currentConfig := config.ApplyInitialConfig(bgpServer, initialConfig, opts.GracefulRestart)
	for sig := range sigCh {
		if sig != syscall.SIGHUP {
			stopServer(bgpServer, opts.UseSdNotify)
			return
		}

		log.WithFields(log.Fields{
			"Topic": "Config",
		}).Info("Reload the config file")
		newConfig, err := config.ReadConfigFile(opts.ConfigFile, opts.ConfigType)
		if err != nil {
			log.WithFields(log.Fields{
				"Topic": "Config",
				"Error": err,
			}).Warningf("Can't read config file %s", opts.ConfigFile)
			continue
		}

		currentConfig = config.UpdateConfig(bgpServer, currentConfig, newConfig)
	}
}

func stopServer(bgpServer *server.BgpServer, useSdNotify bool) {
	bgpServer.StopBgp(context.Background(), &api.StopBgpRequest{})
	if useSdNotify {
		daemon.SdNotify(false, daemon.SdNotifyStopping)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"net/http"

	"github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/camlistore/lock"

	cmds "github.com/jbenet/go-ipfs/commands"
	cmdsHttp "github.com/jbenet/go-ipfs/commands/http"
	"github.com/jbenet/go-ipfs/config"
)

// DaemonLockFile is the filename of the daemon lock, relative to config dir
const DaemonLockFile = "daemon.lock"

var Daemon = &cmds.Command{
	Options:     []cmds.Option{},
	Help:        "TODO",
	Subcommands: map[string]*cmds.Command{},
	Run:         daemonFunc,
}

func daemonFunc(req cmds.Request, res cmds.Response) {
	configPath, err := getConfigPath(req)
	if err != nil {
		res.SetError(err, cmds.ErrNormal)
		return
	}

	lockPath, err := config.Path(configPath, DaemonLockFile)
	if err != nil {
		res.SetError(err, cmds.ErrNormal)
		return
	}

	lk, err := lock.Lock(lockPath)
	if err != nil {
		res.SetError(fmt.Errorf("Couldn't obtain lock. Is another daemon already running?"), cmds.ErrNormal)
		return
	}
	defer lk.Close()

	handler := cmdsHttp.Handler{}
	http.Handle(cmdsHttp.ApiPath+"/", handler)
	// TODO: load listen address/port from config/options
	err = http.ListenAndServe(":8080", nil)
	if err != nil {
		res.SetError(err, cmds.ErrNormal)
		return
	}
	// TODO: log to indicate that we are now listening

}
<commit_msg>cmd/ipfs: Load API address from config<commit_after>package main

import (
	"fmt"
	"net/http"

	"github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/camlistore/lock"
	ma "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-multiaddr"
	manet "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-multiaddr/net"

	cmds "github.com/jbenet/go-ipfs/commands"
	cmdsHttp "github.com/jbenet/go-ipfs/commands/http"
	"github.com/jbenet/go-ipfs/config"
)

// DaemonLockFile is the filename of the daemon lock, relative to config dir
const DaemonLockFile = "daemon.lock"

var Daemon = &cmds.Command{
	Options:     []cmds.Option{},
	Help:        "TODO",
	Subcommands: map[string]*cmds.Command{},
	Run:         daemonFunc,
}

func daemonFunc(req cmds.Request, res cmds.Response) {
	configPath, err := getConfigPath(req)
	if err != nil {
		res.SetError(err, cmds.ErrNormal)
		return
	}

	lockPath, err := config.Path(configPath, DaemonLockFile)
	if err != nil {
		res.SetError(err, cmds.ErrNormal)
		return
	}

	lk, err := lock.Lock(lockPath)
	if err != nil {
		res.SetError(fmt.Errorf("Couldn't obtain lock. Is another daemon already running?"), cmds.ErrNormal)
		return
	}
	defer lk.Close()

	configFile, err := config.Filename(configPath)
	if err != nil {
		res.SetError(err, cmds.ErrNormal)
		return
	}

	config, err := config.Load(configFile)
	if err != nil {
		res.SetError(err, cmds.ErrNormal)
		return
	}

	addr, err := ma.NewMultiaddr(config.Addresses.API)
	if err != nil {
		res.SetError(err, cmds.ErrNormal)
		return
	}

	_, host, err := manet.DialArgs(addr)
	if err != nil {
		res.SetError(err, cmds.ErrNormal)
		return
	}

	handler := cmdsHttp.Handler{}
	http.Handle(cmdsHttp.ApiPath+"/", handler)
	err = http.ListenAndServe(host, nil)
	if err != nil {
		res.SetError(err, cmds.ErrNormal)
		return
	}
	// TODO: log to indicate that we are now listening

}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"launchpad.net/gnuflag"
	"launchpad.net/juju-core/cmd"
	"launchpad.net/juju-core/environs"
	"launchpad.net/juju-core/juju"
	"launchpad.net/juju-core/state"
	"launchpad.net/juju-core/state/api/params"
)

type StatusCommand struct {
	EnvCommandBase
	out cmd.Output
}

var statusDoc = "This command will report on the runtime state of various system entities."

func (c *StatusCommand) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "status",
		Purpose: "output status information about an environment",
		Doc:     statusDoc,
		Aliases: []string{"stat"},
	}
}

func (c *StatusCommand) SetFlags(f *gnuflag.FlagSet) {
	c.EnvCommandBase.SetFlags(f)
	c.out.AddFlags(f, "yaml", map[string]cmd.Formatter{
		"yaml": cmd.FormatYaml,
		"json": cmd.FormatJson,
	})
}

func (c *StatusCommand) Run(ctx *cmd.Context) error {
	conn, err := juju.NewConnFromName(c.EnvName)
	if err != nil {
		return err
	}
	defer conn.Close()

	instances, err := fetchAllInstances(conn.Environ)
	if err != nil {
		return err
	}

	machines, err := fetchAllMachines(conn.State)
	if err != nil {
		return err
	}

	services, err := fetchAllServices(conn.State)
	if err != nil {
		return err
	}

	result := map[string]interface{}{
		"machines": checkError(processMachines(machines, instances)),
		"services": checkError(processServices(services)),
	}

	return c.out.Write(ctx, result)
}

// fetchAllInstances returns a map[string]environs.Instance representing
// a mapping of instance ids to their respective instance.
func fetchAllInstances(env environs.Environ) (map[state.InstanceId]environs.Instance, error) {
	m := make(map[state.InstanceId]environs.Instance)
	insts, err := env.AllInstances()
	if err != nil {
		return nil, err
	}
	for _, i := range insts {
		m[i.Id()] = i
	}
	return m, nil
}

// fetchAllMachines returns a map[string]*state.Machine representing
// a mapping of machine ids to machines.
func fetchAllMachines(st *state.State) (map[string]*state.Machine, error) {
	v := make(map[string]*state.Machine)
	machines, err := st.AllMachines()
	if err != nil {
		return nil, err
	}
	for _, m := range machines {
		v[m.Id()] = m
	}
	return v, nil
}

// fetchAllServices returns a map representing a mapping of service
// names to services.
func fetchAllServices(st *state.State) (map[string]*state.Service, error) {
	v := make(map[string]*state.Service)
	services, err := st.AllServices()
	if err != nil {
		return nil, err
	}
	for _, s := range services {
		v[s.Name()] = s
	}
	return v, nil
}

// processMachines gathers information about machines.
func processMachines(machines map[string]*state.Machine, instances map[state.InstanceId]environs.Instance) (map[string]interface{}, error) {
	r := make(map[string]interface{})
	for _, m := range machines {
		instid, ok := m.InstanceId()
		if !ok {
			r[m.Id()] = map[string]interface{}{
				"instance-id": "pending",
			}
		} else {
			instance, ok := instances[instid]
			if !ok {
				// Double plus ungood. There is an instance id recorded for this machine in the state,
				// yet the environ cannot find that id.
				return nil, fmt.Errorf("instance %s for machine %s not found", instid, m.Id())
			}
			r[m.Id()] = checkError(processMachine(m, instance))
		}
	}
	return r, nil
}

func processMachine(machine *state.Machine, instance environs.Instance) (map[string]interface{}, error) {
	r := m()
	r["instance-id"] = instance.Id()

	if dnsname, err := instance.DNSName(); err == nil {
		r["dns-name"] = dnsname
	}

	processVersion(r, machine)

	// The following logic is ported from Python juju status command.
	agentAlive, err := machine.AgentAlive()
	if err != nil {
		return nil, err
	}
	machineDead := machine.Life() == state.Dead
	status, info, err := machine.Status()
	if err != nil {
		return nil, err
	}
	r["agent-state"] = "running"
	if !agentAlive && !machineDead {
		// Agent should be running but it's not.
		r["agent-state"] = "down"
	} else if !agentAlive || status == params.MachineError {
		// Either the provisioning failed, or the agent is still starting.
		r["agent-state"] = "not-started"
	}

	r["instance-state"] = status
	if len(info) > 0 {
		r["instance-state-info"] = info
	}

	return r, nil
}

// processServices gathers information about services.
func processServices(services map[string]*state.Service) (map[string]interface{}, error) {
	r := m()
	for _, s := range services {
		r[s.Name()] = checkError(processService(s))
	}
	return r, nil
}

func processService(service *state.Service) (map[string]interface{}, error) {
	r := m()
	ch, _, err := service.Charm()
	if err != nil {
		return nil, err
	}
	r["charm"] = ch.String()
	r["exposed"] = service.IsExposed()

	// TODO(dfc) service.IsSubordinate() ?

	units, err := service.AllUnits()
	if err != nil {
		return nil, err
	}

	u := checkError(processUnits(units))
	if len(u) > 0 {
		r["units"] = u
	}

	// TODO(dfc) process relations
	return r, nil
}

func processUnits(units []*state.Unit) (map[string]interface{}, error) {
	r := m()
	for _, unit := range units {
		r[unit.Name()] = checkError(processUnit(unit))
	}
	return r, nil
}

func processUnit(unit *state.Unit) (map[string]interface{}, error) {
	r := m()

	if addr, ok := unit.PublicAddress(); ok {
		r["public-address"] = addr
	}

	if id, err := unit.AssignedMachineId(); err == nil {
		// TODO(dfc) we could make this nicer, ie machine/0
		r["machine"] = id
	}

	processVersion(r, unit)

	agentAlive, err := unit.AgentAlive()
	if err != nil {
		return nil, err
	}
	unitDead := unit.Life() == state.Dead
	status, info, err := unit.Status()
	if err != nil {
		return nil, err
	}
	if status != params.UnitPending {
		if !agentAlive && !unitDead {
			// Agent should be running but it's not.
			status = params.UnitDown
		}
	}
	r["agent-state"] = status
	if len(info) > 0 {
		r["agent-state-info"] = info
	}
	return r, nil
}

type versioned interface {
	AgentTools() (*state.Tools, error)
}

func processVersion(r map[string]interface{}, v versioned) {
	if t, err := v.AgentTools(); err == nil {
		r["agent-version"] = t.Binary.Number.String()
	}
}

func m() map[string]interface{} { return make(map[string]interface{}) }

func checkError(m map[string]interface{}, err error) map[string]interface{} {
	if err != nil {
		return map[string]interface{}{"status-error": err.Error()}
	}
	return m
}
<commit_msg>Changes after review<commit_after>package main

import (
	"fmt"
	"launchpad.net/gnuflag"
	"launchpad.net/juju-core/cmd"
	"launchpad.net/juju-core/environs"
	"launchpad.net/juju-core/juju"
	"launchpad.net/juju-core/state"
	"launchpad.net/juju-core/state/api/params"
)

type StatusCommand struct {
	EnvCommandBase
	out cmd.Output
}

var statusDoc = "This command will report on the runtime state of various system entities."

func (c *StatusCommand) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "status",
		Purpose: "output status information about an environment",
		Doc:     statusDoc,
		Aliases: []string{"stat"},
	}
}

func (c *StatusCommand) SetFlags(f *gnuflag.FlagSet) {
	c.EnvCommandBase.SetFlags(f)
	c.out.AddFlags(f, "yaml", map[string]cmd.Formatter{
		"yaml": cmd.FormatYaml,
		"json": cmd.FormatJson,
	})
}

func (c *StatusCommand) Run(ctx *cmd.Context) error {
	conn, err := juju.NewConnFromName(c.EnvName)
	if err != nil {
		return err
	}
	defer conn.Close()

	instances, err := fetchAllInstances(conn.Environ)
	if err != nil {
		return err
	}

	machines, err := fetchAllMachines(conn.State)
	if err != nil {
		return err
	}

	services, err := fetchAllServices(conn.State)
	if err != nil {
		return err
	}

	result := map[string]interface{}{
		"machines": checkError(processMachines(machines, instances)),
		"services": checkError(processServices(services)),
	}

	return c.out.Write(ctx, result)
}

// fetchAllInstances returns a map[string]environs.Instance representing
// a mapping of instance ids to their respective instance.
func fetchAllInstances(env environs.Environ) (map[state.InstanceId]environs.Instance, error) {
	m := make(map[state.InstanceId]environs.Instance)
	insts, err := env.AllInstances()
	if err != nil {
		return nil, err
	}
	for _, i := range insts {
		m[i.Id()] = i
	}
	return m, nil
}

// fetchAllMachines returns a map[string]*state.Machine representing
// a mapping of machine ids to machines.
func fetchAllMachines(st *state.State) (map[string]*state.Machine, error) {
	v := make(map[string]*state.Machine)
	machines, err := st.AllMachines()
	if err != nil {
		return nil, err
	}
	for _, m := range machines {
		v[m.Id()] = m
	}
	return v, nil
}

// fetchAllServices returns a map representing a mapping of service
// names to services.
func fetchAllServices(st *state.State) (map[string]*state.Service, error) {
	v := make(map[string]*state.Service)
	services, err := st.AllServices()
	if err != nil {
		return nil, err
	}
	for _, s := range services {
		v[s.Name()] = s
	}
	return v, nil
}

// processMachines gathers information about machines.
func processMachines(machines map[string]*state.Machine, instances map[state.InstanceId]environs.Instance) (map[string]interface{}, error) {
	r := make(map[string]interface{})
	for _, m := range machines {
		instid, ok := m.InstanceId()
		if !ok {
			r[m.Id()] = map[string]interface{}{
				"instance-id": "pending",
			}
		} else {
			instance, ok := instances[instid]
			if !ok {
				// Double plus ungood. There is an instance id recorded for this machine in the state,
				// yet the environ cannot find that id.
				return nil, fmt.Errorf("instance %s for machine %s not found", instid, m.Id())
			}
			r[m.Id()] = checkError(processMachine(m, instance))
		}
	}
	return r, nil
}

func processMachine(machine *state.Machine, instance environs.Instance) (map[string]interface{}, error) {
	r := m()
	r["instance-id"] = instance.Id()

	if dnsname, err := instance.DNSName(); err == nil {
		r["dns-name"] = dnsname
	}

	processVersion(r, machine)

	// The following logic is ported from Python juju status command.
	agentAlive, err := machine.AgentAlive()
	if err != nil {
		return nil, err
	}
	machineDead := machine.Life() == state.Dead
	status, info, err := machine.Status()
	if err != nil {
		return nil, err
	}
	r["agent-state"] = "running"
	if !agentAlive && !machineDead {
		// Agent should be running but it's not.
		r["agent-state"] = "down"
	} else if !agentAlive || status == params.MachineError {
		// Either the provisioning failed, or the agent is still starting.
		r["agent-state"] = "not-started"
	}

	r["instance-state"] = status
	if info != "" {
		r["instance-state-info"] = info
	}

	return r, nil
}

// processServices gathers information about services.
func processServices(services map[string]*state.Service) (map[string]interface{}, error) {
	r := m()
	for _, s := range services {
		r[s.Name()] = checkError(processService(s))
	}
	return r, nil
}

func processService(service *state.Service) (map[string]interface{}, error) {
	r := m()
	ch, _, err := service.Charm()
	if err != nil {
		return nil, err
	}
	r["charm"] = ch.String()
	r["exposed"] = service.IsExposed()

	// TODO(dfc) service.IsSubordinate() ?

	units, err := service.AllUnits()
	if err != nil {
		return nil, err
	}

	u := checkError(processUnits(units))
	if len(u) > 0 {
		r["units"] = u
	}

	// TODO(dfc) process relations
	return r, nil
}

func processUnits(units []*state.Unit) (map[string]interface{}, error) {
	r := m()
	for _, unit := range units {
		r[unit.Name()] = checkError(processUnit(unit))
	}
	return r, nil
}

func processUnit(unit *state.Unit) (map[string]interface{}, error) {
	r := m()

	if addr, ok := unit.PublicAddress(); ok {
		r["public-address"] = addr
	}

	if id, err := unit.AssignedMachineId(); err == nil {
		// TODO(dfc) we could make this nicer, ie machine/0
		r["machine"] = id
	}

	processVersion(r, unit)

	agentAlive, err := unit.AgentAlive()
	if err != nil {
		return nil, err
	}
	unitDead := unit.Life() == state.Dead
	status, info, err := unit.Status()
	if err != nil {
		return nil, err
	}
	if status != params.UnitPending {
		if !agentAlive && !unitDead {
			// Agent should be running but it's not.
			status = params.UnitDown
		}
	}
	r["agent-state"] = status
	if len(info) > 0 {
		r["agent-state-info"] = info
	}
	return r, nil
}

type versioned interface {
	AgentTools() (*state.Tools, error)
}

func processVersion(r map[string]interface{}, v versioned) {
	if t, err := v.AgentTools(); err == nil {
		r["agent-version"] = t.Binary.Number.String()
	}
}

func m() map[string]interface{} { return make(map[string]interface{}) }

func checkError(m map[string]interface{}, err error) map[string]interface{} {
	if err != nil {
		return map[string]interface{}{"status-error": err.Error()}
	}
	return m
}
<|endoftext|>


<commit_before>/*
Copyright (c) 2014 Ashley Jeffs

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.
*/

package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"net/http"
	"os"
	"os/signal"
	"syscall"
	"time"
	gopath "path"
	"golang.org/x/net/websocket"

	"github.com/jeffail/leaps/lib/acl"
	"github.com/jeffail/leaps/lib/curator"
	leaphttp "github.com/jeffail/leaps/lib/http"
	"github.com/jeffail/leaps/lib/store"
	"github.com/jeffail/util/log"
	"github.com/jeffail/util/metrics"
)

//--------------------------------------------------------------------------------------------------

// Flags
var (
	httpAddress *string
	showHidden  *bool
	debugWWWDir *string
	logLevel    *string
	subdirPath  *string
)

func init() {
	httpAddress = flag.String("address", ":8080", "The HTTP address to bind to")
	showHidden = flag.Bool("all", false, "Display all files, including hidden")
	debugWWWDir = flag.String("use_www", "", "Serve alternative web files from this dir")
	logLevel = flag.String("log_level", "INFO", "Log level (NONE, ERROR, WARM, INFO, DEBUG, TRACE)")
	subdirPath = flag.String("path", "/", "Subdirectory (when running leaps in a webserver subdirectory as example.com/myleaps)")
}

//--------------------------------------------------------------------------------------------------

var endpoints = []interface{}{}

func handle(path, description string, handler http.HandlerFunc) {
	path = gopath.Join(*subdirPath, path)
	http.HandleFunc(path, handler)
	endpoints = append(endpoints, struct {
		Path string `json:"path"`
		Desc string `json:"description"`
	}{
		Path: path,
		Desc: description,
	})
}

//--------------------------------------------------------------------------------------------------

func main() {
	var (
		err       error
		closeChan = make(chan bool)
	)

	flag.Usage = func() {
		fmt.Println(`Usage: leaps [flags...] [path/to/share]

If a path is not specified the current directory is shared instead.
`)
		flag.PrintDefaults()
	}

	flag.Parse()

	targetPath := "."
	if flag.NArg() == 1 {
		targetPath = flag.Arg(0)
	}

	// Logging and metrics aggregation
	logConf := log.NewLoggerConfig()
	logConf.Prefix = "leaps"
	logConf.LogLevel = *logLevel

	logger := log.NewLogger(os.Stdout, logConf)

	statConf := metrics.NewConfig()
	statConf.Prefix = "leaps"

	stats, err := metrics.New(statConf)
	if err != nil {
		fmt.Fprintln(os.Stderr, fmt.Sprintf("Metrics init error: %v\n", err))
		return
	}
	defer stats.Close()

	logger.Infoln("Launching a leaps instance, use CTRL+C to close.")

	// Document storage engine
	docStore, err := store.NewFile(".")
	if err != nil {
		fmt.Fprintln(os.Stderr, fmt.Sprintf("Document store error: %v\n", err))
		return
	}

	// Authenticator
	storeConf := acl.NewFileExistsConfig()
	storeConf.Path = targetPath
	storeConf.ShowHidden = *showHidden

	authenticator := acl.NewFileExists(storeConf, logger)

	// Curator of documents
	curatorConf := curator.NewConfig()
	curator, err := curator.New(curatorConf, logger, stats, authenticator, docStore)
	if err != nil {
		fmt.Fprintln(os.Stderr, fmt.Sprintf("Curator error: %v\n", err))
		return
	}
	defer curator.Close()

	handle("/endpoints", "Lists all available endpoints (including this one).",
		func(w http.ResponseWriter, r *http.Request) {
			data, reqErr := json.Marshal(endpoints)
			if reqErr != nil {
				logger.Errorf("Failed to serve endpoints: %v\n", reqErr)
				http.Error(w, reqErr.Error(), http.StatusInternalServerError)
				return
			}
			w.Header().Add("Content-Type", "application/json")
			w.Write(data)
		})

	handle("/files", "Returns a list of available files and a map of users per document.",
		func(w http.ResponseWriter, r *http.Request) {
			var reqErr error
			var users map[string][]string
			if users, reqErr = curator.GetUsers(time.Second); reqErr == nil {
				var data []byte
				data, reqErr = json.Marshal(struct {
					Paths []string            `json:"paths"`
					Users map[string][]string `json:"users"`
				}{
					Paths: authenticator.GetPaths(),
					Users: users,
				})
				if reqErr == nil {
					w.Write(data)
				}
			}
			if reqErr != nil {
				http.Error(w, reqErr.Error(), http.StatusInternalServerError)
				logger.Errorf("Failed to serve users: %v\n", reqErr)
				return
			}
			w.Header().Add("Content-Type", "application/json")
		})

	handle("/stats", "Lists all aggregated metrics as a json blob.", stats.JSONHandler())
	var myPath = gopath.Join("/", *subdirPath) + "/"

	if len(*debugWWWDir) > 0 {
		logger.Warnf("Serving web files from alternative www dir: %v\n", *debugWWWDir)
		http.Handle(myPath, http.StripPrefix(myPath, http.FileServer(http.Dir(*debugWWWDir))))
	} else {
		http.Handle(myPath, http.StripPrefix(myPath, http.FileServer(assetFS())))
	}
	http.Handle(gopath.Join("/", *subdirPath) + "/leaps/ws",
		websocket.Handler(leaphttp.WebsocketHandler(curator, time.Second, logger, stats)))

	go func() {
		logger.Infof("Serving HTTP requests at: %v%v\n", *httpAddress, *subdirPath)
		if httperr := http.ListenAndServe(*httpAddress, nil); httperr != nil {
			fmt.Fprintln(os.Stderr, fmt.Sprintf("HTTP listen error: %v\n", httperr))
		}
		closeChan <- true
	}()

	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)

	// Wait for termination signal
	select {
	case <-sigChan:
	case <-closeChan:
	}
}

//--------------------------------------------------------------------------------------------------
<commit_msg>Fix strip prefix for empty path<commit_after>/*
Copyright (c) 2014 Ashley Jeffs

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.
*/

package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"net/http"
	"os"
	"os/signal"
	gopath "path"
	"syscall"
	"time"

	"golang.org/x/net/websocket"

	"github.com/jeffail/leaps/lib/acl"
	"github.com/jeffail/leaps/lib/curator"
	leaphttp "github.com/jeffail/leaps/lib/http"
	"github.com/jeffail/leaps/lib/store"
	"github.com/jeffail/util/log"
	"github.com/jeffail/util/metrics"
)

//--------------------------------------------------------------------------------------------------

// Flags
var (
	httpAddress *string
	showHidden  *bool
	debugWWWDir *string
	logLevel    *string
	subdirPath  *string
)

func init() {
	httpAddress = flag.String("address", ":8080", "The HTTP address to bind to")
	showHidden = flag.Bool("all", false, "Display all files, including hidden")
	debugWWWDir = flag.String("use_www", "", "Serve alternative web files from this dir")
	logLevel = flag.String("log_level", "INFO", "Log level (NONE, ERROR, WARM, INFO, DEBUG, TRACE)")
	subdirPath = flag.String("path", "/", "Subdirectory (when running leaps in a webserver subdirectory as example.com/myleaps)")
}

//--------------------------------------------------------------------------------------------------

var endpoints = []interface{}{}

func handle(path, description string, handler http.HandlerFunc) {
	path = gopath.Join("/", *subdirPath, path)
	http.HandleFunc(path, handler)
	endpoints = append(endpoints, struct {
		Path string `json:"path"`
		Desc string `json:"description"`
	}{
		Path: path,
		Desc: description,
	})
}

//--------------------------------------------------------------------------------------------------

func main() {
	var (
		err       error
		closeChan = make(chan bool)
	)

	flag.Usage = func() {
		fmt.Println(`Usage: leaps [flags...] [path/to/share]

If a path is not specified the current directory is shared instead.
`)
		flag.PrintDefaults()
	}

	flag.Parse()

	targetPath := "."
	if flag.NArg() == 1 {
		targetPath = flag.Arg(0)
	}

	// Logging and metrics aggregation
	logConf := log.NewLoggerConfig()
	logConf.Prefix = "leaps"
	logConf.LogLevel = *logLevel

	logger := log.NewLogger(os.Stdout, logConf)

	statConf := metrics.NewConfig()
	statConf.Prefix = "leaps"

	stats, err := metrics.New(statConf)
	if err != nil {
		fmt.Fprintln(os.Stderr, fmt.Sprintf("Metrics init error: %v\n", err))
		return
	}
	defer stats.Close()

	logger.Infoln("Launching a leaps instance, use CTRL+C to close.")

	// Document storage engine
	docStore, err := store.NewFile(".")
	if err != nil {
		fmt.Fprintln(os.Stderr, fmt.Sprintf("Document store error: %v\n", err))
		return
	}

	// Authenticator
	storeConf := acl.NewFileExistsConfig()
	storeConf.Path = targetPath
	storeConf.ShowHidden = *showHidden

	authenticator := acl.NewFileExists(storeConf, logger)

	// Curator of documents
	curatorConf := curator.NewConfig()
	curator, err := curator.New(curatorConf, logger, stats, authenticator, docStore)
	if err != nil {
		fmt.Fprintln(os.Stderr, fmt.Sprintf("Curator error: %v\n", err))
		return
	}
	defer curator.Close()

	handle("/endpoints", "Lists all available endpoints (including this one).",
		func(w http.ResponseWriter, r *http.Request) {
			data, reqErr := json.Marshal(endpoints)
			if reqErr != nil {
				logger.Errorf("Failed to serve endpoints: %v\n", reqErr)
				http.Error(w, reqErr.Error(), http.StatusInternalServerError)
				return
			}
			w.Header().Add("Content-Type", "application/json")
			w.Write(data)
		})

	handle("/files", "Returns a list of available files and a map of users per document.",
		func(w http.ResponseWriter, r *http.Request) {
			var reqErr error
			var users map[string][]string
			if users, reqErr = curator.GetUsers(time.Second); reqErr == nil {
				var data []byte
				data, reqErr = json.Marshal(struct {
					Paths []string            `json:"paths"`
					Users map[string][]string `json:"users"`
				}{
					Paths: authenticator.GetPaths(),
					Users: users,
				})
				if reqErr == nil {
					w.Write(data)
				}
			}
			if reqErr != nil {
				http.Error(w, reqErr.Error(), http.StatusInternalServerError)
				logger.Errorf("Failed to serve users: %v\n", reqErr)
				return
			}
			w.Header().Add("Content-Type", "application/json")
		})

	handle("/stats", "Lists all aggregated metrics as a json blob.", stats.JSONHandler())

	wwwPath := gopath.Join("/", *subdirPath)
	stripPath := ""
	if wwwPath != "/" {
		wwwPath = wwwPath + "/"
		stripPath = wwwPath
	}
	if len(*debugWWWDir) > 0 {
		logger.Warnf("Serving web files from alternative www dir: %v\n", *debugWWWDir)
		http.Handle(wwwPath, http.StripPrefix(stripPath, http.FileServer(http.Dir(*debugWWWDir))))
	} else {
		http.Handle(wwwPath, http.StripPrefix(stripPath, http.FileServer(assetFS())))
	}
	http.Handle(gopath.Join("/", *subdirPath, "/leaps/ws"),
		websocket.Handler(leaphttp.WebsocketHandler(curator, time.Second, logger, stats)))

	go func() {
		logger.Infof("Serving HTTP requests at: %v%v\n", *httpAddress, *subdirPath)
		if httperr := http.ListenAndServe(*httpAddress, nil); httperr != nil {
			fmt.Fprintln(os.Stderr, fmt.Sprintf("HTTP listen error: %v\n", httperr))
		}
		closeChan <- true
	}()

	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)

	// Wait for termination signal
	select {
	case <-sigChan:
	case <-closeChan:
	}
}

//--------------------------------------------------------------------------------------------------
<|endoftext|>


<commit_before>package main

import (
	"encoding/gob"
	"github.com/Symantec/Dominator/lib/mdb"
	"github.com/Symantec/Dominator/lib/srpc"
	"github.com/Symantec/Dominator/proto/mdbserver"
	"io"
	"log"
	"sync"
)

type rpcType struct {
	currentMdb *mdb.Mdb
	logger     *log.Logger
	rwMutex    sync.RWMutex
	// Protected by lock.
	updateChannels map[*srpc.Conn]chan<- mdbserver.MdbUpdate
}

func startRpcd(logger *log.Logger) func(old, new *mdb.Mdb) {
	rpcObj := &rpcType{
		currentMdb:     &mdb.Mdb{},
		logger:         logger,
		updateChannels: make(map[*srpc.Conn]chan<- mdbserver.MdbUpdate),
	}
	srpc.RegisterName("MdbServer", rpcObj)
	return rpcObj.pushUpdateToAll
}

func (t *rpcType) GetMdbUpdates(conn *srpc.Conn) error {
	encoder := gob.NewEncoder(conn)
	updateChannel := make(chan mdbserver.MdbUpdate, 1)
	t.rwMutex.Lock()
	t.updateChannels[conn] = updateChannel
	t.rwMutex.Unlock()
	mdbUpdate := mdbserver.MdbUpdate{MachinesToAdd: t.currentMdb.Machines}
	if err := encoder.Encode(mdbUpdate); err != nil {
		return nil
	}
	if err := conn.Flush(); err != nil {
		return nil
	}
	closeChannel := getCloseNotifier(conn)
	for {
		var err error
		select {
		case mdbUpdate := <-updateChannel:
			if err = encoder.Encode(mdbUpdate); err != nil {
				break
			}
			if err = conn.Flush(); err != nil {
				break
			}
		case err = <-closeChannel:
			break
		}
		if err != nil {
			t.rwMutex.Lock()
			delete(t.updateChannels, conn)
			t.rwMutex.Unlock()
			if err != io.EOF {
				t.logger.Println(err)
				return err
			} else {
				return nil
			}
		}
	}
}

func (t *rpcType) pushUpdateToAll(old, new *mdb.Mdb) {
	t.currentMdb = new
	t.rwMutex.RLock()
	numConnections := len(t.updateChannels)
	t.rwMutex.RUnlock()
	if numConnections < 1 {
		return
	}
	mdbUpdate := mdbserver.MdbUpdate{}
	oldMachines := make(map[string]mdb.Machine, len(old.Machines))
	for _, machine := range old.Machines {
		oldMachines[machine.Hostname] = machine
	}
	for _, newMachine := range new.Machines {
		if oldMachine, ok := oldMachines[newMachine.Hostname]; ok {
			if newMachine != oldMachine {
				mdbUpdate.MachinesToUpdate = append(mdbUpdate.MachinesToUpdate,
					newMachine)
			}
		} else {
			mdbUpdate.MachinesToAdd = append(mdbUpdate.MachinesToAdd,
				newMachine)
		}
	}
	for _, machine := range new.Machines {
		delete(oldMachines, machine.Hostname)
	}
	for name := range oldMachines {
		mdbUpdate.MachinesToDelete = append(mdbUpdate.MachinesToDelete, name)
	}
	t.rwMutex.RLock()
	defer t.rwMutex.RUnlock()
	for _, channel := range t.updateChannels {
		channel <- mdbUpdate
	}
}

func getCloseNotifier(conn *srpc.Conn) <-chan error {
	closeChannel := make(chan error)
	go func() {
		for {
			buf := make([]byte, 1)
			if _, err := conn.Read(buf); err != nil {
				closeChannel <- err
				return
			}
		}
	}()
	return closeChannel
}
<commit_msg>Fix mdbd panic when first MDB data are pushed to clients.<commit_after>package main

import (
	"encoding/gob"
	"github.com/Symantec/Dominator/lib/mdb"
	"github.com/Symantec/Dominator/lib/srpc"
	"github.com/Symantec/Dominator/proto/mdbserver"
	"io"
	"log"
	"sync"
)

type rpcType struct {
	currentMdb *mdb.Mdb
	logger     *log.Logger
	rwMutex    sync.RWMutex
	// Protected by lock.
	updateChannels map[*srpc.Conn]chan<- mdbserver.MdbUpdate
}

func startRpcd(logger *log.Logger) func(old, new *mdb.Mdb) {
	rpcObj := &rpcType{
		currentMdb:     &mdb.Mdb{},
		logger:         logger,
		updateChannels: make(map[*srpc.Conn]chan<- mdbserver.MdbUpdate),
	}
	srpc.RegisterName("MdbServer", rpcObj)
	return rpcObj.pushUpdateToAll
}

func (t *rpcType) GetMdbUpdates(conn *srpc.Conn) error {
	encoder := gob.NewEncoder(conn)
	updateChannel := make(chan mdbserver.MdbUpdate, 1)
	t.rwMutex.Lock()
	t.updateChannels[conn] = updateChannel
	t.rwMutex.Unlock()
	mdbUpdate := mdbserver.MdbUpdate{MachinesToAdd: t.currentMdb.Machines}
	if err := encoder.Encode(mdbUpdate); err != nil {
		return nil
	}
	if err := conn.Flush(); err != nil {
		return nil
	}
	closeChannel := getCloseNotifier(conn)
	for {
		var err error
		select {
		case mdbUpdate := <-updateChannel:
			if err = encoder.Encode(mdbUpdate); err != nil {
				break
			}
			if err = conn.Flush(); err != nil {
				break
			}
		case err = <-closeChannel:
			break
		}
		if err != nil {
			t.rwMutex.Lock()
			delete(t.updateChannels, conn)
			t.rwMutex.Unlock()
			if err != io.EOF {
				t.logger.Println(err)
				return err
			} else {
				return nil
			}
		}
	}
}

func (t *rpcType) pushUpdateToAll(old, new *mdb.Mdb) {
	t.currentMdb = new
	t.rwMutex.RLock()
	numConnections := len(t.updateChannels)
	t.rwMutex.RUnlock()
	if numConnections < 1 {
		return
	}
	mdbUpdate := mdbserver.MdbUpdate{}
	if old == nil {
		old = &mdb.Mdb{}
	}
	oldMachines := make(map[string]mdb.Machine, len(old.Machines))
	for _, machine := range old.Machines {
		oldMachines[machine.Hostname] = machine
	}
	for _, newMachine := range new.Machines {
		if oldMachine, ok := oldMachines[newMachine.Hostname]; ok {
			if newMachine != oldMachine {
				mdbUpdate.MachinesToUpdate = append(mdbUpdate.MachinesToUpdate,
					newMachine)
			}
		} else {
			mdbUpdate.MachinesToAdd = append(mdbUpdate.MachinesToAdd,
				newMachine)
		}
	}
	for _, machine := range new.Machines {
		delete(oldMachines, machine.Hostname)
	}
	for name := range oldMachines {
		mdbUpdate.MachinesToDelete = append(mdbUpdate.MachinesToDelete, name)
	}
	t.rwMutex.RLock()
	defer t.rwMutex.RUnlock()
	for _, channel := range t.updateChannels {
		channel <- mdbUpdate
	}
}

func getCloseNotifier(conn *srpc.Conn) <-chan error {
	closeChannel := make(chan error)
	go func() {
		for {
			buf := make([]byte, 1)
			if _, err := conn.Read(buf); err != nil {
				closeChannel <- err
				return
			}
		}
	}()
	return closeChannel
}
<|endoftext|>


<commit_before>// Copyright 2017 The nem-toolchain project authors. All rights reserved.
// Use of this source code is governed by a MIT license that can be found in the LICENSE file.

package main

import (
	"bufio"
	"encoding/hex"
	"fmt"
	"math"
	"os"
	"runtime"
	"strings"
	"time"

	"github.com/nem-toolchain/nem-toolchain/pkg/core"
	"github.com/nem-toolchain/nem-toolchain/pkg/keypair"
	"github.com/nem-toolchain/nem-toolchain/pkg/vanity"
	"github.com/urfave/cli"
)

func accountCommand() cli.Command {
	return cli.Command{
		Name:  "account",
		Usage: "Account related bundle of actions",
		Subcommands: []cli.Command{
			{
				Name:   "generate",
				Usage:  "Generate a new account",
				Action: generateAction,
				Flags: []cli.Flag{
					cli.UintFlag{
						Name:  "number, n",
						Usage: "Number of generated accounts",
						Value: 1,
					},
					cli.BoolFlag{
						Name:  "strip, s",
						Usage: "Strip output to private key only",
					},
				},
			},
			{
				Name:   "info",
				Usage:  "Show info for given account",
				Action: infoAction,
				Flags: []cli.Flag{
					cli.BoolFlag{
						Name:  "address",
						Usage: "Show public address only for given private key",
					},
					cli.BoolFlag{
						Name:  "public",
						Usage: "Show public key only for given private key",
					},
				},
			},
			{
				Name:   "vanity",
				Usage:  "Find vanity address by a given list of prefixes",
				Action: vanityAction,
				Flags: []cli.Flag{
					cli.UintFlag{
						Name:  "number, n",
						Usage: "Number of generated accounts",
						Value: 1,
					},
					cli.BoolFlag{
						Name:  "strip, s",
						Usage: "Strip output to private key only",
					},
					cli.UintFlag{
						Name:  "workers, w",
						Usage: "Number of workers for generation",
						Value: uint(runtime.NumCPU()),
					},
					cli.BoolFlag{
						Name:  "show-complexity",
						Usage: "Show additionally the specified search complexity",
					},
					cli.BoolFlag{
						Name:  "skip-estimate",
						Usage: "Skip the step to calculate estimation times to search",
					},
					cli.BoolFlag{
						Name:  "no-digits",
						Usage: "Digits in address are disallow ('234567')",
					},
					cli.StringFlag{
						Name:  "exclude",
						Usage: "Characters that must not be in the address",
					},
				},
			},
		},
	}
}

func generateAction(c *cli.Context) error {
	ch, err := chainGlobalOption(c)
	if err != nil {
		return cli.NewExitError(err.Error(), 1)
	}

	num := c.Uint("number")
	for i := uint(0); i < num; i++ {
		pair := keypair.Gen()
		if c.Bool("strip") {
			printlnPrivateKey(pair, true)
		} else {
			printAccountDetails(ch, pair)
		}
	}

	return nil
}

func infoAction(c *cli.Context) error {
	ch, err := chainGlobalOption(c)
	if err != nil {
		return cli.NewExitError(err.Error(), 1)
	}

	fi, err := os.Stdin.Stat()
	if err != nil {
		return cli.NewExitError(err.Error(), 1)
	}

	if (fi.Mode()&os.ModeCharDevice) != 0 && fi.Size() == 0 {
		fmt.Print("Enter private key: ")
	}

	rd := bufio.NewReader(os.Stdin)
	s, err := rd.ReadString('\n')
	if err != nil {
		return cli.NewExitError(err.Error(), 1)
	}

	pk, err := hex.DecodeString(strings.TrimSpace(s))
	if err != nil {
		return cli.NewExitError(err.Error(), 1)
	}

	pair, err := keypair.FromSeed(pk)
	if err != nil {
		return cli.NewExitError(err.Error(), 1)
	}

	if c.Bool("address") {
		printlnAddress(ch, pair, true)
	} else if c.Bool("public") {
		printlnPublicKey(pair, true)
	} else {
		printAccountDetails(ch, pair)
	}

	return nil
}

func vanityAction(c *cli.Context) error {
	ch, err := chainGlobalOption(c)
	if err != nil {
		return cli.NewExitError(err.Error(), 1)
	}

	num := c.Uint("number")
	if num == 0 {
		return nil
	}

	workers := c.Uint("workers")
	if m := uint(runtime.NumCPU()); workers == 0 || workers > m {
		workers = m
	}

	var excludeSel vanity.Selector = vanity.TrueSelector{}
	if c.IsSet("exclude") {
		excludeSel, err = vanity.NewExcludeSelector(c.String("exclude"))
		if err != nil {
			return cli.NewExitError(err.Error(), 1)
		}
	}

	var noDigitsSel vanity.Selector = vanity.TrueSelector{}
	if c.Bool("no-digits") {
		noDigitsSel, _ = vanity.NewExcludeSelector("234567")
	}

	prefixes := make([]vanity.Selector, len(c.Args()))
	for i, pr := range c.Args() {
		sel, err := vanity.NewPrefixSelector(ch, strings.ToUpper(pr))
		if err != nil {
			return cli.NewExitError(err.Error(), 1)
		}
		prefixes[i] = sel
	}

	sel := vanity.AndSelector(
		excludeSel, noDigitsSel, vanity.OrSelector(prefixes...))

	if !c.Bool("strip") && !c.Bool("skip-estimate") {
		printEstimate(workers,
			vanity.Probability(sel)/float64(num), c.Bool("show-complexity"))
	}

	rs := make(chan keypair.KeyPair)
	for i := uint(0); i < workers; i++ {
		go vanity.StartSearch(ch, sel, rs)
	}

	for i := uint(0); i < num; i++ {
		pair := <-rs
		if c.Bool("strip") {
			printlnPrivateKey(pair, true)
		} else {
			printAccountDetails(ch, pair)
		}
		go vanity.StartSearch(ch, sel, rs)
	}

	return nil
}

// printEstimate prints vanity account search time estimate
func printEstimate(workers uint, pbty float64, cplx bool) {
	fmt.Print("Calculate accounts rate")
	ticker := time.NewTicker(time.Second)
	go func() {
		for range ticker.C {
			fmt.Print(".")
		}
	}()
	rate := countActualRate(workers)
	ticker.Stop()
	fmt.Printf(" %v accounts/sec\n", math.Trunc(rate))
	printEstimateDetails(pbty, rate, cplx)
}

// countActualRate counts total number of generated keypairs per second
func countActualRate(workers uint) float64 {
	res := make(chan int, workers)
	for i := 0; i < cap(res); i++ {
		go countKeyPairs(3200, res)
	}
	rate := float64(0)
	for i := 0; i < cap(res); i++ {
		rate += float64(<-res) / 3.2
	}
	return rate
}

// countKeyPairs counts number of generated keypairs for specified interval
func countKeyPairs(milliseconds time.Duration, res chan int) {
	timeout := time.After(time.Millisecond * milliseconds)
	for count := 0; ; count++ {
		keypair.Gen().Address(core.Mainnet)
		select {
		case <-timeout:
			res <- count
			return
		default:
			continue
		}
	}
}

// printEstimateDetails prints estimate search time details
func printEstimateDetails(pbty, rate float64, cplx bool) {
	if cplx {
		fmt.Printf("Specified search complexity: %v\n", math.Trunc(1.0/pbty))
	}
	fmt.Printf("Estimate search times: %v (50%%), %v (80%%), %v (99.9%%)\n",
		timeInSeconds(vanity.NumberOfAttempts(pbty, 0.5)/rate),
		timeInSeconds(vanity.NumberOfAttempts(pbty, 0.8)/rate),
		timeInSeconds(vanity.NumberOfAttempts(pbty, 0.99)/rate))
}

// timeInSeconds formats estimated time
func timeInSeconds(val float64) string {
	val = 1e9 * math.Trunc(val)
	if val >= math.MaxInt64 || math.IsInf(val, 0) {
		return "Inf"
	}
	return time.Duration(val).String()
}

// printAccountDetails prints account details in pretty user-oriented multi-line format
func printAccountDetails(chain core.Chain, pair keypair.KeyPair) {
	fmt.Println("----")
	printlnAddress(chain, pair, false)
	printlnPublicKey(pair, false)
	printlnPrivateKey(pair, false)
}

func printlnPrivateKey(pair keypair.KeyPair, strip bool) {
	printlnCustom("Private key: ", hex.EncodeToString(pair.Private), strip)
}

func printlnPublicKey(pair keypair.KeyPair, strip bool) {
	printlnCustom("Public key: ", hex.EncodeToString(pair.Public), strip)
}

func printlnAddress(chain core.Chain, pair keypair.KeyPair, strip bool) {
	printlnCustom("Address: ", pair.Address(chain).PrettyString(), strip)
}

func printlnCustom(title, value string, strip bool) {
	if !strip {
		fmt.Print(title)
	}
	fmt.Println(value)
}
<commit_msg>Lack of handling of stdin/out piping for `info` command Added<commit_after>// Copyright 2017 The nem-toolchain project authors. All rights reserved.
// Use of this source code is governed by a MIT license that can be found in the LICENSE file.

package main

import (
	"bufio"
	"encoding/hex"
	"fmt"
	"io/ioutil"
	"math"
	"os"
	"runtime"
	"strings"
	"time"

	"github.com/nem-toolchain/nem-toolchain/pkg/core"
	"github.com/nem-toolchain/nem-toolchain/pkg/keypair"
	"github.com/nem-toolchain/nem-toolchain/pkg/vanity"
	"github.com/urfave/cli"
)

func accountCommand() cli.Command {
	return cli.Command{
		Name:  "account",
		Usage: "Account related bundle of actions",
		Subcommands: []cli.Command{
			{
				Name:   "generate",
				Usage:  "Generate a new account",
				Action: generateAction,
				Flags: []cli.Flag{
					cli.UintFlag{
						Name:  "number, n",
						Usage: "Number of generated accounts",
						Value: 1,
					},
					cli.BoolFlag{
						Name:  "strip, s",
						Usage: "Strip output to private key only",
					},
				},
			},
			{
				Name:   "info",
				Usage:  "Show info for given account",
				Action: infoAction,
				Flags: []cli.Flag{
					cli.BoolFlag{
						Name:  "address",
						Usage: "Show public address only for given private key",
					},
					cli.BoolFlag{
						Name:  "public",
						Usage: "Show public key only for given private key",
					},
				},
			},
			{
				Name:   "vanity",
				Usage:  "Find vanity address by a given list of prefixes",
				Action: vanityAction,
				Flags: []cli.Flag{
					cli.UintFlag{
						Name:  "number, n",
						Usage: "Number of generated accounts",
						Value: 1,
					},
					cli.BoolFlag{
						Name:  "strip, s",
						Usage: "Strip output to private key only",
					},
					cli.UintFlag{
						Name:  "workers, w",
						Usage: "Number of workers for generation",
						Value: uint(runtime.NumCPU()),
					},
					cli.BoolFlag{
						Name:  "show-complexity",
						Usage: "Show additionally the specified search complexity",
					},
					cli.BoolFlag{
						Name:  "skip-estimate",
						Usage: "Skip the step to calculate estimation times to search",
					},
					cli.BoolFlag{
						Name:  "no-digits",
						Usage: "Digits in address are disallow ('234567')",
					},
					cli.StringFlag{
						Name:  "exclude",
						Usage: "Characters that must not be in the address",
					},
				},
			},
		},
	}
}

func generateAction(c *cli.Context) error {
	ch, err := chainGlobalOption(c)
	if err != nil {
		return cli.NewExitError(err.Error(), 1)
	}

	num := c.Uint("number")
	for i := uint(0); i < num; i++ {
		pair := keypair.Gen()
		if c.Bool("strip") {
			printlnPrivateKey(pair, true)
		} else {
			printAccountDetails(ch, pair)
		}
	}

	return nil
}

func infoAction(c *cli.Context) error {
	ch, err := chainGlobalOption(c)
	if err != nil {
		return cli.NewExitError(err.Error(), 1)
	}

	fi, err := os.Stdin.Stat()
	if err != nil {
		return cli.NewExitError(err.Error(), 1)
	}

	var privKeys []string
	if (fi.Mode()&os.ModeCharDevice) != 0 && fi.Size() == 0 {
		fmt.Print("Enter private key: ")
		rd := bufio.NewReader(os.Stdin)
		val, err := rd.ReadString('\n')
		if err != nil {
			return cli.NewExitError(err.Error(), 1)
		}
		privKeys = []string{val}
	} else {
		data, err := ioutil.ReadAll(os.Stdin)
		if err != nil {
			return cli.NewExitError(err.Error(), 1)
		}
		privKeys = strings.Split(string(data), "\n")
		privKeys = privKeys[:len(privKeys)-1]
	}

	for _, s := range privKeys {
		pk, err := hex.DecodeString(strings.TrimSpace(s))
		if err != nil {
			return cli.NewExitError(err.Error(), 1)
		}

		pair, err := keypair.FromSeed(pk)
		if err != nil {
			return cli.NewExitError(err.Error(), 1)
		}

		if c.Bool("address") {
			printlnAddress(ch, pair, true)
		} else if c.Bool("public") {
			printlnPublicKey(pair, true)
		} else {
			printAccountDetails(ch, pair)
		}
	}

	return nil
}

func vanityAction(c *cli.Context) error {
	ch, err := chainGlobalOption(c)
	if err != nil {
		return cli.NewExitError(err.Error(), 1)
	}

	num := c.Uint("number")
	if num == 0 {
		return nil
	}

	workers := c.Uint("workers")
	if m := uint(runtime.NumCPU()); workers == 0 || workers > m {
		workers = m
	}

	var excludeSel vanity.Selector = vanity.TrueSelector{}
	if c.IsSet("exclude") {
		excludeSel, err = vanity.NewExcludeSelector(c.String("exclude"))
		if err != nil {
			return cli.NewExitError(err.Error(), 1)
		}
	}

	var noDigitsSel vanity.Selector = vanity.TrueSelector{}
	if c.Bool("no-digits") {
		noDigitsSel, _ = vanity.NewExcludeSelector("234567")
	}

	prefixes := make([]vanity.Selector, len(c.Args()))
	for i, pr := range c.Args() {
		sel, err := vanity.NewPrefixSelector(ch, strings.ToUpper(pr))
		if err != nil {
			return cli.NewExitError(err.Error(), 1)
		}
		prefixes[i] = sel
	}

	sel := vanity.AndSelector(
		excludeSel, noDigitsSel, vanity.OrSelector(prefixes...))

	if !c.Bool("strip") && !c.Bool("skip-estimate") {
		printEstimate(workers,
			vanity.Probability(sel)/float64(num), c.Bool("show-complexity"))
	}

	rs := make(chan keypair.KeyPair)
	for i := uint(0); i < workers; i++ {
		go vanity.StartSearch(ch, sel, rs)
	}

	for i := uint(0); i < num; i++ {
		pair := <-rs
		if c.Bool("strip") {
			printlnPrivateKey(pair, true)
		} else {
			printAccountDetails(ch, pair)
		}
		go vanity.StartSearch(ch, sel, rs)
	}

	return nil
}

// printEstimate prints vanity account search time estimate
func printEstimate(workers uint, pbty float64, cplx bool) {
	fmt.Print("Calculate accounts rate")
	ticker := time.NewTicker(time.Second)
	go func() {
		for range ticker.C {
			fmt.Print(".")
		}
	}()
	rate := countActualRate(workers)
	ticker.Stop()
	fmt.Printf(" %v accounts/sec\n", math.Trunc(rate))
	printEstimateDetails(pbty, rate, cplx)
}

// countActualRate counts total number of generated keypairs per second
func countActualRate(workers uint) float64 {
	res := make(chan int, workers)
	for i := 0; i < cap(res); i++ {
		go countKeyPairs(3200, res)
	}
	rate := float64(0)
	for i := 0; i < cap(res); i++ {
		rate += float64(<-res) / 3.2
	}
	return rate
}

// countKeyPairs counts number of generated keypairs for specified interval
func countKeyPairs(milliseconds time.Duration, res chan int) {
	timeout := time.After(time.Millisecond * milliseconds)
	for count := 0; ; count++ {
		keypair.Gen().Address(core.Mainnet)
		select {
		case <-timeout:
			res <- count
			return
		default:
			continue
		}
	}
}

// printEstimateDetails prints estimate search time details
func printEstimateDetails(pbty, rate float64, cplx bool) {
	if cplx {
		fmt.Printf("Specified search complexity: %v\n", math.Trunc(1.0/pbty))
	}
	fmt.Printf("Estimate search times: %v (50%%), %v (80%%), %v (99.9%%)\n",
		timeInSeconds(vanity.NumberOfAttempts(pbty, 0.5)/rate),
		timeInSeconds(vanity.NumberOfAttempts(pbty, 0.8)/rate),
		timeInSeconds(vanity.NumberOfAttempts(pbty, 0.99)/rate))
}

// timeInSeconds formats estimated time
func timeInSeconds(val float64) string {
	val = 1e9 * math.Trunc(val)
	if val >= math.MaxInt64 || math.IsInf(val, 0) {
		return "Inf"
	}
	return time.Duration(val).String()
}

// printAccountDetails prints account details in pretty user-oriented multi-line format
func printAccountDetails(chain core.Chain, pair keypair.KeyPair) {
	fmt.Println("----")
	printlnAddress(chain, pair, false)
	printlnPublicKey(pair, false)
	printlnPrivateKey(pair, false)
}

func printlnPrivateKey(pair keypair.KeyPair, strip bool) {
	printlnCustom("Private key: ", hex.EncodeToString(pair.Private), strip)
}

func printlnPublicKey(pair keypair.KeyPair, strip bool) {
	printlnCustom("Public key: ", hex.EncodeToString(pair.Public), strip)
}

func printlnAddress(chain core.Chain, pair keypair.KeyPair, strip bool) {
	printlnCustom("Address: ", pair.Address(chain).PrettyString(), strip)
}

func printlnCustom(title, value string, strip bool) {
	if !strip {
		fmt.Print(title)
	}
	fmt.Println(value)
}
<|endoftext|>