<commit_before>package main

import (
	"encoding/binary"
	"errors"
	"fmt"
	libgit "github.com/driusan/git"
	"io"
	"sort"
	//   	"io/ioutil"
	"bytes"
	"crypto/sha1"
	"os"
)

var InvalidIndex error = errors.New("Invalid index")

// index file is defined as Network byte Order (Big Endian)

// 12 byte header:
// 4 bytes: D I R C (stands for "Dir cache")
// 4-byte version number (can be 2, 3 or 4)
// 32bit number of index entries
type fixedGitIndex struct {
	Signature          [4]byte // 4
	Version            uint32  // 8
	NumberIndexEntries uint32  // 12
}

type GitIndex struct {
	fixedGitIndex // 12
	Objects       []*GitIndexEntry
}
type GitIndexEntry struct {
	fixedIndexEntry

	PathName string
}

type fixedIndexEntry struct {
	Ctime     uint32 // 16
	Ctimenano uint32 // 20

	Mtime     uint32 // 24
	Mtimenano uint32 // 28

	Dev uint32 // 32
	Ino uint32 // 36

	Mode uint32 // 40

	Uid uint32 // 44
	Gid uint32 // 48

	Fsize uint32 // 52

	Sha1 [20]byte // 72

	Flags uint16 // 74
}

func ReadIndex(g *libgit.Repository) (*GitIndex, error) {
	file, err := os.Open(g.Path + "/index")
	if err != nil {
		return &GitIndex{
			fixedGitIndex{
				[4]byte{'D', 'I', 'R', 'C'},
				2, // version 2
				0, // no entries
			},
			make([]*GitIndexEntry, 0),
		}, err
	}
	defer file.Close()

	var i fixedGitIndex
	binary.Read(file, binary.BigEndian, &i)
	fmt.Println(i.Signature)
	if i.Signature != [4]byte{'D', 'I', 'R', 'C'} {
		return nil, InvalidIndex
	}
	if i.Version < 2 || i.Version > 4 {
		return nil, InvalidIndex
	}

	var idx uint32
	indexes := make([]*GitIndexEntry, i.NumberIndexEntries, i.NumberIndexEntries)
	for idx = 0; idx < i.NumberIndexEntries; idx += 1 {
		if index, err := ReadIndexEntry(file); err == nil {
			indexes[idx] = index
		}
	}
	return &GitIndex{i, indexes}, nil
}

func ReadIndexEntry(file *os.File) (*GitIndexEntry, error) {
	var f fixedIndexEntry
	var name []byte
	binary.Read(file, binary.BigEndian, &f)

	var nameLength uint16
	nameLength = f.Flags & 0x0FFF

	if nameLength&0xFFF != 0xFFF {
		name = make([]byte, nameLength, nameLength)
		n, err := file.Read(name)
		if err != nil {
			panic("I don't know what to do")
		}
		if n != int(nameLength) {
			panic("Error reading the name")
		}

		// I don't understand where this +4 comes from, but it seems to work
		// out with how the c git implementation calculates the padding..
		//
		// The definition of the index file format at:
		// https://github.com/git/git/blob/master/Documentation/technical/index-format.txt
		// claims that there should be "1-8 nul bytes as necessary to pad the entry to a multiple of eight
		// bytes while keeping the name NUL-terminated."
		//
		// The fixed size of the header is 82 bytes if you add up all the types.
		// the length of the name is nameLength bytes, so according to the spec
		// this *should* be 8 - ((82 + nameLength) % 8) bytes of padding.
		// But reading existant index files, there seems to be an extra 4 bytes
		// incorporated into the index size calculation.
		expectedOffset := 8 - ((82 + nameLength + 4) % 8)
		file.Seek(int64(expectedOffset), 1)
		/*
			This was used to verify that the offset is correct, but it causes problems if the data following
			the offset is empty..
			whitespace := make([]byte, 1, 1)
			var w uint16
			// Read all the whitespace that git uses for alignment.
			for _, _ = file.Read(whitespace); whitespace[0] == 0; _, _ = file.Read(whitespace) {
				w += 1
			}

			if w % 8 != expectedOffset {
				panic(fmt.Sprintf("Read incorrect number of whitespace characters %d vs %d", w, expectedOffset))
			}
			if w == 0 {
				panic("Name was not null terminated in index")
			}

			// Undo the last read, which wasn't whitespace..
			file.Seek(-1, 1)
		*/

	} else {
		panic("TODO: I can't handle such long names yet")
	}
	return &GitIndexEntry{f, string(name)}, nil
}

// Adds a file to the index, without writing it to disk.
// To write it to disk after calling this, use GitIndex.WriteIndex
//
// This will do the following:
// write git object blob of file contents to .git/objects
// normalize os.File name to path relative to gitRoot
// search GitIndex for normalized name
//	if GitIndexEntry found
//		update GitIndexEntry to point to the new object blob
// 	else
// 		add new GitIndexEntry if not found
//
func (g *GitIndex) AddFile(repo *libgit.Repository, file *os.File) {

	sha1, err := repo.StoreObjectLoose(libgit.ObjectBlob, file)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error storing object: %s", err)
	}
	fmt.Printf("Sha1: %s\n", sha1)
	fmt.Printf("Name is %s\n", file.Name())
	name := getNormalizedName(file)
	for _, entry := range g.Objects {
		if entry.PathName == name {
			entry.Sha1 = sha1
			return
		}
	}
	panic("TODO: Unimplemented: add new file to GitIndexEntry")
}

func getNormalizedName(file *os.File) string {
	return file.Name()
}

// This will write a new index file to w by doing the following:
// 1. Sort the objects in g.Index to ascending order based on name
// 2. Write g.fixedGitIndex to w
// 3. for each entry in g.Objects, write it to w.
// 4. Write the Sha1 of the contents of what was written
func (g GitIndex) WriteIndex(file io.Writer) {
	sort.Sort(ByPath(g.Objects))
	s := sha1.New()
	w := io.MultiWriter(file, s)
	binary.Write(w, binary.BigEndian, g.fixedGitIndex)
	for _, entry := range g.Objects {
		binary.Write(w, binary.BigEndian, entry.fixedIndexEntry)
		binary.Write(w, binary.BigEndian, []byte(entry.PathName))
		padding := 8 - ((82 + len(entry.PathName) + 4) % 8)
		p := make([]byte, padding)
		binary.Write(w, binary.BigEndian, p)
	}
	binary.Write(w, binary.BigEndian, s.Sum(nil))

}

// Implement the sort interface on *GitIndexEntry, so that
// it's easy to sort by name.
type ByPath []*GitIndexEntry

func (g ByPath) Len() int           { return len(g) }
func (g ByPath) Swap(i, j int)      { g[i], g[j] = g[j], g[i] }
func (g ByPath) Less(i, j int) bool { return g[i].PathName < g[j].PathName }

// Writes the current index to a tree object.
func (g GitIndex) WriteTree(repo *libgit.Repository) {
	content := bytes.NewBuffer(nil)
	// [mode] [file/folder name]\0[SHA-1 of referencing blob or tree as [20]byte]
	for _, obj := range g.Objects {
		fmt.Printf("%o %s\x00%x\n", obj.Mode, obj.PathName, obj.Sha1)
		fmt.Fprintf(content, "%o %s\x00", obj.Mode, obj.PathName)
		content.Write(obj.Sha1[:])
	}
	sha1, err := repo.StoreObjectLoose(libgit.ObjectTree, bytes.NewReader(content.Bytes()))
	fmt.Printf("% x err: %s %s", sha1, err, sha1)
}

func expandGitTreeIntoIndexes(repo *libgit.Repository, tree *libgit.Tree, prefix string) ([]*GitIndexEntry, error) {
	newEntries := make([]*GitIndexEntry, 0)

	for _, entry := range tree.ListEntries() {
		var dirname string
		if prefix == "" {
			dirname = entry.Name()
		} else {
			dirname = prefix + "/" + entry.Name()
		}
		switch entry.Type {
		case libgit.ObjectBlob:
			newEntry := GitIndexEntry{}
			newEntry.Sha1 = entry.Id
			// This isn't right.
			// This should be:
			// "32-bit mode, split into:
			//      4-bit object type: valid values in binary are
			//          1000 (regular file), 1010 (symbolic link), and
			//          1110 (gitlink)
			//      3-bit unused
			//      9-bit unix permission. Only 0755 and 0644 are valid
			//          for regular files, symbolic links have 0 in this
			//          field"

			//go-gits entry mode is an int, but it needs to be a uint32
			switch entry.EntryMode() {
			case libgit.ModeBlob:
				newEntry.Mode = 0100644
			case libgit.ModeExec:
				newEntry.Mode = 0100755
			case libgit.ModeSymlink:
				newEntry.Mode = 0120000
			case libgit.ModeCommit:
				newEntry.Mode = 0160000
			case libgit.ModeTree:
				newEntry.Mode = 0040000
			}
			newEntry.PathName = dirname
			newEntry.Fsize = uint32(entry.Size())

			modTime := entry.ModTime()
			newEntry.Mtime = uint32(modTime.Unix())
			newEntry.Mtimenano = uint32(modTime.Nanosecond())
			newEntry.Flags = uint16(len(dirname)) & 0xFFF
			/* I don't know how git can extract these values
			   from a tree. For now, leave them empty

			   Ctime     uint32
			   Ctimenano uint32
			   Dev uint32
			   Ino uint32
			   Uid uint32
			   Gid uint32
			*/
			newEntries = append(newEntries, &newEntry)
		case libgit.ObjectTree:

			subTree, err := tree.SubTree(entry.Name())
			if err != nil {
				panic(err)
			}
			subindexes, err := expandGitTreeIntoIndexes(repo, subTree, dirname)
			if err != nil {
				panic(err)
			}
			newEntries = append(newEntries, subindexes...)

		default:
			fmt.Fprintf(os.Stderr, "Unknown object type: %s\n", entry.Type)
		}
	}
	return newEntries, nil

}
func (g *GitIndex) ResetIndex(repo *libgit.Repository, tree *libgit.Tree) error {
	newEntries, err := expandGitTreeIntoIndexes(repo, tree, "")
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error resetting index: %s\n", err)
		return err
	}
	fmt.Printf("%s", newEntries)
	g.NumberIndexEntries = uint32(len(newEntries))
	g.Objects = newEntries
	return nil
}
<commit_msg>Implemented write-tree without directory support<commit_after>package main

import (
	"encoding/binary"
	"errors"
	"fmt"
	libgit "github.com/driusan/git"
	"io"
	"sort"
	//   	"io/ioutil"
	"bytes"
	"crypto/sha1"
	"os"
	"strings"
)

var InvalidIndex error = errors.New("Invalid index")

// index file is defined as Network byte Order (Big Endian)

// 12 byte header:
// 4 bytes: D I R C (stands for "Dir cache")
// 4-byte version number (can be 2, 3 or 4)
// 32bit number of index entries
type fixedGitIndex struct {
	Signature          [4]byte // 4
	Version            uint32  // 8
	NumberIndexEntries uint32  // 12
}

type GitIndex struct {
	fixedGitIndex // 12
	Objects       []*GitIndexEntry
}
type GitIndexEntry struct {
	fixedIndexEntry

	PathName string
}

type fixedIndexEntry struct {
	Ctime     uint32 // 16
	Ctimenano uint32 // 20

	Mtime     uint32 // 24
	Mtimenano uint32 // 28

	Dev uint32 // 32
	Ino uint32 // 36

	Mode uint32 // 40

	Uid uint32 // 44
	Gid uint32 // 48

	Fsize uint32 // 52

	Sha1 [20]byte // 72

	Flags uint16 // 74
}

func ReadIndex(g *libgit.Repository) (*GitIndex, error) {
	file, err := os.Open(g.Path + "/index")
	if err != nil {
		return &GitIndex{
			fixedGitIndex{
				[4]byte{'D', 'I', 'R', 'C'},
				2, // version 2
				0, // no entries
			},
			make([]*GitIndexEntry, 0),
		}, err
	}
	defer file.Close()

	var i fixedGitIndex
	binary.Read(file, binary.BigEndian, &i)
	fmt.Println(i.Signature)
	if i.Signature != [4]byte{'D', 'I', 'R', 'C'} {
		return nil, InvalidIndex
	}
	if i.Version < 2 || i.Version > 4 {
		return nil, InvalidIndex
	}

	var idx uint32
	indexes := make([]*GitIndexEntry, i.NumberIndexEntries, i.NumberIndexEntries)
	for idx = 0; idx < i.NumberIndexEntries; idx += 1 {
		if index, err := ReadIndexEntry(file); err == nil {
			indexes[idx] = index
		}
	}
	return &GitIndex{i, indexes}, nil
}

func ReadIndexEntry(file *os.File) (*GitIndexEntry, error) {
	var f fixedIndexEntry
	var name []byte
	binary.Read(file, binary.BigEndian, &f)

	var nameLength uint16
	nameLength = f.Flags & 0x0FFF

	if nameLength&0xFFF != 0xFFF {
		name = make([]byte, nameLength, nameLength)
		n, err := file.Read(name)
		if err != nil {
			panic("I don't know what to do")
		}
		if n != int(nameLength) {
			panic("Error reading the name")
		}

		// I don't understand where this +4 comes from, but it seems to work
		// out with how the c git implementation calculates the padding..
		//
		// The definition of the index file format at:
		// https://github.com/git/git/blob/master/Documentation/technical/index-format.txt
		// claims that there should be "1-8 nul bytes as necessary to pad the entry to a multiple of eight
		// bytes while keeping the name NUL-terminated."
		//
		// The fixed size of the header is 82 bytes if you add up all the types.
		// the length of the name is nameLength bytes, so according to the spec
		// this *should* be 8 - ((82 + nameLength) % 8) bytes of padding.
		// But reading existant index files, there seems to be an extra 4 bytes
		// incorporated into the index size calculation.
		expectedOffset := 8 - ((82 + nameLength + 4) % 8)
		file.Seek(int64(expectedOffset), 1)
		/*
			This was used to verify that the offset is correct, but it causes problems if the data following
			the offset is empty..
			whitespace := make([]byte, 1, 1)
			var w uint16
			// Read all the whitespace that git uses for alignment.
			for _, _ = file.Read(whitespace); whitespace[0] == 0; _, _ = file.Read(whitespace) {
				w += 1
			}

			if w % 8 != expectedOffset {
				panic(fmt.Sprintf("Read incorrect number of whitespace characters %d vs %d", w, expectedOffset))
			}
			if w == 0 {
				panic("Name was not null terminated in index")
			}

			// Undo the last read, which wasn't whitespace..
			file.Seek(-1, 1)
		*/

	} else {
		panic("TODO: I can't handle such long names yet")
	}
	return &GitIndexEntry{f, string(name)}, nil
}

// Adds a file to the index, without writing it to disk.
// To write it to disk after calling this, use GitIndex.WriteIndex
//
// This will do the following:
// write git object blob of file contents to .git/objects
// normalize os.File name to path relative to gitRoot
// search GitIndex for normalized name
//	if GitIndexEntry found
//		update GitIndexEntry to point to the new object blob
// 	else
// 		add new GitIndexEntry if not found
//
func (g *GitIndex) AddFile(repo *libgit.Repository, file *os.File) {

	sha1, err := repo.StoreObjectLoose(libgit.ObjectBlob, file)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error storing object: %s", err)
	}
	fmt.Printf("Sha1: %s\n", sha1)
	fmt.Printf("Name is %s\n", file.Name())
	name := getNormalizedName(file)
	for _, entry := range g.Objects {
		if entry.PathName == name {
			entry.Sha1 = sha1
			return
		}
	}
	panic("TODO: Unimplemented: add new file to GitIndexEntry")
}

func getNormalizedName(file *os.File) string {
	return file.Name()
}

// This will write a new index file to w by doing the following:
// 1. Sort the objects in g.Index to ascending order based on name
// 2. Write g.fixedGitIndex to w
// 3. for each entry in g.Objects, write it to w.
// 4. Write the Sha1 of the contents of what was written
func (g GitIndex) WriteIndex(file io.Writer) {
	sort.Sort(ByPath(g.Objects))
	s := sha1.New()
	w := io.MultiWriter(file, s)
	binary.Write(w, binary.BigEndian, g.fixedGitIndex)
	for _, entry := range g.Objects {
		binary.Write(w, binary.BigEndian, entry.fixedIndexEntry)
		binary.Write(w, binary.BigEndian, []byte(entry.PathName))
		padding := 8 - ((82 + len(entry.PathName) + 4) % 8)
		p := make([]byte, padding)
		binary.Write(w, binary.BigEndian, p)
	}
	binary.Write(w, binary.BigEndian, s.Sum(nil))

}

// Implement the sort interface on *GitIndexEntry, so that
// it's easy to sort by name.
type ByPath []*GitIndexEntry

func (g ByPath) Len() int           { return len(g) }
func (g ByPath) Swap(i, j int)      { g[i], g[j] = g[j], g[i] }
func (g ByPath) Less(i, j int) bool { return g[i].PathName < g[j].PathName }

func writeIndexSubtree(repo *libgit.Repository, name string, entries []*GitIndexEntry) ([20]byte, error) {
	var ret [20]byte
	return ret, nil
}
func writeIndexEntries(repo *libgit.Repository, prefix string, entries []*GitIndexEntry) ([20]byte, error) {
	content := bytes.NewBuffer(nil)
	// [mode] [file/folder name]\0[SHA-1 of referencing blob or tree as [20]byte]

	lastname := ""
	firstIdxForTree := -1

	for idx, obj := range entries {
		nameBits := strings.Split(obj.PathName, "/")

		// Either it's the last entry and we haven't written a tree yet, or it's not the last
		// entry but the directory changed
		if (nameBits[0] != lastname || idx == len(entries)-1) && lastname != "" {
			fmt.Printf("Should write tree for %s (entries %d to %d)\n", lastname, firstIdxForTree, idx)

			subsha1, err := writeIndexSubtree(repo, lastname, entries[firstIdxForTree:idx])
			if err != nil {
				panic(err)
			}

			fmt.Printf("%o %s %x", 0040000, lastname, subsha1)
			// Write the object
			fmt.Fprintf(content, "%o %s\x00", 0040000, lastname)
			content.Write(subsha1[:])

			// Reset the data keeping track of what this tree is.
			lastname = ""
			firstIdxForTree = -1
			continue
		}
		if len(nameBits) == 1 {
			fmt.Printf("%s\n", obj)
			//write the blob for the file portion
			fmt.Printf("%o %s\x00%x\n", obj.Mode, obj.PathName, obj.Sha1)
			fmt.Fprintf(content, "%o %s\x00", obj.Mode, obj.PathName)
			content.Write(obj.Sha1[:])
			lastname = ""
		} else {
			lastname = nameBits[0]
			if firstIdxForTree == -1 {
				firstIdxForTree = idx
			}
		}
	}

	sha1, err := repo.StoreObjectLoose(libgit.ObjectTree, bytes.NewReader(content.Bytes()))
	return [20]byte(sha1), err
}

// Writes the current index to a tree object.
func (g GitIndex) WriteTree(repo *libgit.Repository) {

	sha1, err := writeIndexEntries(repo, "", g.Objects)
	fmt.Printf("% x err: %s\n", sha1, err)
}

func expandGitTreeIntoIndexes(repo *libgit.Repository, tree *libgit.Tree, prefix string) ([]*GitIndexEntry, error) {
	newEntries := make([]*GitIndexEntry, 0)

	for _, entry := range tree.ListEntries() {
		var dirname string
		if prefix == "" {
			dirname = entry.Name()
		} else {
			dirname = prefix + "/" + entry.Name()
		}
		switch entry.Type {
		case libgit.ObjectBlob:
			newEntry := GitIndexEntry{}
			newEntry.Sha1 = entry.Id
			// This isn't right.
			// This should be:
			// "32-bit mode, split into:
			//      4-bit object type: valid values in binary are
			//          1000 (regular file), 1010 (symbolic link), and
			//          1110 (gitlink)
			//      3-bit unused
			//      9-bit unix permission. Only 0755 and 0644 are valid
			//          for regular files, symbolic links have 0 in this
			//          field"

			//go-gits entry mode is an int, but it needs to be a uint32
			switch entry.EntryMode() {
			case libgit.ModeBlob:
				newEntry.Mode = 0100644
			case libgit.ModeExec:
				newEntry.Mode = 0100755
			case libgit.ModeSymlink:
				newEntry.Mode = 0120000
			case libgit.ModeCommit:
				newEntry.Mode = 0160000
			case libgit.ModeTree:
				newEntry.Mode = 0040000
			}
			newEntry.PathName = dirname
			newEntry.Fsize = uint32(entry.Size())

			modTime := entry.ModTime()
			newEntry.Mtime = uint32(modTime.Unix())
			newEntry.Mtimenano = uint32(modTime.Nanosecond())
			newEntry.Flags = uint16(len(dirname)) & 0xFFF
			/* I don't know how git can extract these values
			   from a tree. For now, leave them empty

			   Ctime     uint32
			   Ctimenano uint32
			   Dev uint32
			   Ino uint32
			   Uid uint32
			   Gid uint32
			*/
			newEntries = append(newEntries, &newEntry)
		case libgit.ObjectTree:

			subTree, err := tree.SubTree(entry.Name())
			if err != nil {
				panic(err)
			}
			subindexes, err := expandGitTreeIntoIndexes(repo, subTree, dirname)
			if err != nil {
				panic(err)
			}
			newEntries = append(newEntries, subindexes...)

		default:
			fmt.Fprintf(os.Stderr, "Unknown object type: %s\n", entry.Type)
		}
	}
	return newEntries, nil

}
func (g *GitIndex) ResetIndex(repo *libgit.Repository, tree *libgit.Tree) error {
	newEntries, err := expandGitTreeIntoIndexes(repo, tree, "")
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error resetting index: %s\n", err)
		return err
	}
	fmt.Printf("%s", newEntries)
	g.NumberIndexEntries = uint32(len(newEntries))
	g.Objects = newEntries
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2014 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// TODO(adg): accept -a flag on 'create' and 'commit' (like git commit -a)
// TODO(adg): accept -r flag on 'upload' to nominate reviewer
// TODO(adg): support 'create' from non-master branches
// TODO(adg): check style of commit message
// TODO(adg): write doc comment
// TOOD(adg): print gerrit votes on 'pending'
// TODO(adg): add gofmt commit hook
// TODO(adg): 'upload' warn about uncommitted changes (maybe commit/create too?)

package main // import "golang.org/x/review"

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

var (
	hookFile = filepath.FromSlash(".git/hooks/commit-msg")
	verbose  = flag.Bool("v", false, "verbose output")
	noRun    = flag.Bool("n", false, "print but do not run commands")
)

const usage = `Usage: %s [-n] [-v] <command>
Type "%s help" for more information.
`

const help = `Usage: %s [-n] [-v] <command>

The review command is a wrapper for the git command that provides a simple
interface to the "single-commit feature branch" development model.

Available comands:

	create <name>
		Create a local branch with the provided name
		and commit the staged changes to it.

	commit
		Amend local branch HEAD commit with the staged changes.

	diff
		View differences between remote branch HEAD and
		the local branch HEAD.
		(The differences introduced by this change.)

	upload
		Upload HEAD commit to the code review server.

	sync
		Fetch changes from the remote repository and merge them to the
		current branch, rebasing the HEAD commit (if any) on top of
		them. If the HEAD commit has been submitted, switch back to the
		master branch and delete the feature branch.

	pending
		Show local branches and their head commits.

`

func main() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, usage, os.Args[0], os.Args[0])
		os.Exit(2)
	}
	flag.Parse()

	goToRepoRoot()
	installHook()

	command, args := flag.Arg(0), flag.Args()[1:]

	switch command {
	case "help":
		fmt.Fprintf(os.Stdout, help, os.Args[0])
	case "create", "cr":
		create(args)
	case "commit", "co":
		commit(args)
	case "diff", "d":
		diff(args)
	case "upload", "u":
		upload(args)
	case "sync", "s":
		doSync(args)
	case "pending", "p":
		pending(args)
	default:
		flag.Usage()
	}
}

func create(args []string) {
	if len(args) != 1 || args[0] == "" {
		flag.Usage()
	}
	name := args[0]
	if !hasStagedChanges() {
		dief("no staged changes.\nDid you forget to 'git add'?")
	}
	if b := currentBranch(); b != "master" {
		dief("must run 'create' from the master branch (now on %q).\n"+
			"(Try 'review sync' or 'git checkout master' first.)",
			b)
	}
	run("git", "checkout", "-q", "-b", name)
	if err := runErr("git", "commit", "-q"); err != nil {
		verbosef("Commit failed: %v\nSwitching back to master.\n", err)
		run("git", "checkout", "-q", "master")
		run("git", "branch", "-q", "-d", name)
	}
}

func commit(args []string) {
	if len(args) != 0 {
		flag.Usage()
	}
	if !hasStagedChanges() {
		dief("no staged changes. Did you forget to 'git add'?")
	}
	if currentBranch() == "master" {
		dief("can't commit to master branch.")
	}
	run("git", "commit", "-q", "--amend", "-C", "HEAD")
}

func diff(args []string) {
	if len(args) != 0 {
		flag.Usage()
	}
	run("git", "diff", "HEAD^", "HEAD")
}

func upload(args []string) {
	if len(args) != 0 {
		flag.Usage()
	}
	if currentBranch() == "master" {
		dief("can't upload from master branch.")
	}
	run("git", "push", "-q", "origin", "HEAD:refs/for/master")
}

func doSync(args []string) {
	if len(args) != 0 {
		flag.Usage()
	}
	run("git", "fetch", "-q")

	// If we're on master, just fast-forward.
	branch := currentBranch()
	if branch == "master" {
		run("git", "merge", "-q", "--ff-only", "origin/master")
		return
	}

	// Check that exactly this commit was submitted to master. If so,
	// switch back to master, fast-forward, delete the feature branch.
	if branchContains("origin/master", "HEAD") {
		run("git", "checkout", "-q", "master")
		run("git", "merge", "-q", "--ff-only", "origin/master")
		run("git", "branch", "-q", "-d", branch)
		fmt.Printf("Change on %q submitted; branch deleted.\n"+
			"On master branch.\n", branch)
		return
	}

	// Check whether a rebased version of this commit was submitted to
	// master. If so, switch back to master, fast-forward, and
	// provide instructions for deleting the feature branch.
	// (We're not 100% sure that the feature branch HEAD was submitted,
	// so be cautious.)
	if headSubmitted(branch) {
		run("git", "checkout", "-q", "master")
		run("git", "merge", "-q", "--ff-only", "origin/master")
		fmt.Fprintf(os.Stderr,
			"I think the change on %q has been submitted.\n"+
				"If you agree, and no longer need branch %q, "+
				"run:\n\tgit branch -D %v\n"+
				"On master branch.\n",
			branch, branch, branch)
		return
	}

	// Bump master HEAD to that of origin/master, just in case the user
	// switches back to master with "git checkout master" later.
	// TODO(adg): maybe we shouldn't do this at all?
	if !branchContains("origin/master", "master") {
		run("git", "branch", "-f", "master", "origin/master")
	}

	// We have un-submitted changes on this feature branch; rebase.
	run("git", "rebase", "-q", "origin/master")
}

func pending(args []string) {
	if len(args) != 0 {
		flag.Usage()
	}
	var (
		wg      sync.WaitGroup
		origin  = originURL()
		current = currentBranch()
	)
	if current == "master" {
		fmt.Println("On master branch.")
	}
	for _, branch := range localBranches() {
		if branch == "master" {
			continue
		}
		wg.Add(1)
		go func(branch string) {
			defer wg.Done()
			p := ""
			if branch == current {
				p = "* "
			}
			id := headChangeId(branch)
			c, err := getChange(origin, id)
			switch err {
			case notFound:
				// TODO(adg): read local commit msg
				var msg string
				fmt.Printf("%v%v:\n\t%v\n\t(not uploaded)\n",
					p, branch, msg)
			case nil:
				status := ""
				switch c.Status {
				case "MERGED":
					status += " [submitted]"
				case "ABANDONED":
					status += " [abandoned]"
				}
				fmt.Printf("%v%v%v:\n\t%v\n\t%v\n",
					p, branch, status, c.Subject, c.URL)
			default:
				fmt.Fprintf(os.Stderr, "fetching change for %q: %v\n", branch, err)
			}
		}(branch)
	}
	wg.Wait()
}

func originURL() string {
	out, err := exec.Command("git", "config", "remote.origin.url").CombinedOutput()
	if err != nil {
		dief("could not find URL for 'origin' remote.\n"+
			"Did you check out from the right place?\n"+
			"git config remote.origin.url: %v\n"+
			"%s", err, out)
	}
	return string(out)
}

func localBranches() (branches []string) {
	for _, s := range getLines("git", "branch", "-l", "-q") {
		branches = append(branches, strings.TrimPrefix(s, "* "))
	}
	return branches
}

func branchContains(branch, rev string) bool {
	for _, s := range getLines("git", "branch", "-r", "--contains", rev) {
		if s == branch {
			return true
		}
	}
	return false
}

var stagedRe = regexp.MustCompile(`^[ACDMR]  `)

func hasStagedChanges() bool {
	for _, s := range getLines("git", "status", "-b", "--porcelain") {
		if stagedRe.MatchString(s) {
			return true
		}
	}
	return false
}

func currentBranch() string {
	return strings.TrimSpace(getOutput("git", "rev-parse", "--abbrev-ref", "HEAD"))
}

func headSubmitted(branch string) bool {
	s := "Change-Id: " + headChangeId(branch)
	return len(getOutput("git", "log", "--grep", s, "origin/master")) > 0
}

func headChangeId(branch string) string {
	const (
		p = "Change-Id: "
		f = "--format=format:%b"
	)
	for _, s := range getLines("git", "log", "-n", "1", f, branch, "--") {
		if strings.HasPrefix(s, p) {
			return strings.TrimSpace(strings.TrimPrefix(s, p))
		}
	}
	dief("no Change-Id line found in HEAD commit on branch %s.", branch)
	panic("unreachable")
}

func goToRepoRoot() {
	prevDir, err := os.Getwd()
	if err != nil {
		dief("could not get current directory: %v", err)
	}
	for {
		if _, err := os.Stat(".git"); err == nil {
			return
		}
		if err := os.Chdir(".."); err != nil {
			dief("could not chdir: %v", err)
		}
		currentDir, err := os.Getwd()
		if err != nil {
			dief("could not get current directory: %v", err)
		}
		if currentDir == prevDir {
			dief("git root not found.\n" +
				"Run from within the Git tree please.")
		}
		prevDir = currentDir
	}
}

func installHook() {
	_, err := os.Stat(hookFile)
	if err == nil {
		return
	}
	if !os.IsNotExist(err) {
		dief("error checking for hook file: %v", err)
	}
	verbosef("Presubmit hook to add Change-Id to commit messages is missing.\n"+
		"Automatically creating it at %v.\n", hookFile)
	hookContent := []byte(commitMsgHook)
	if err := ioutil.WriteFile(hookFile, hookContent, 0700); err != nil {
		dief("error writing hook file: %v", err)
	}
}

func run(command string, args ...string) {
	if err := runErr(command, args...); err != nil {
		if !*verbose {
			// If we're not in verbose mode, print the command
			// before dying to give context to the failure.
			fmt.Fprintln(os.Stderr, commandString(command, args))
		}
		dief("%v", err)
	}
}

func runErr(command string, args ...string) error {
	if *verbose || *noRun {
		fmt.Fprintln(os.Stderr, commandString(command, args))
	}
	if *noRun {
		return nil
	}
	cmd := exec.Command(command, args...)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	return cmd.Run()
}

// getOutput runs the specified command and returns its combined standard
// output and standard error outputs.
// NOTE: It should only be used to run commands that return information,
// **not** commands that make any actual changes.
func getOutput(command string, args ...string) string {
	b, err := exec.Command(command, args...).CombinedOutput()
	if err != nil {
		fmt.Fprintf(os.Stderr, "%v\n%s\n", commandString(command, args), b)
		dief("%v", err)
	}
	return string(b)
}

// getLines is like getOutput but it returns non-empty output lines.
// NOTE: It should only be used to run commands that return information,
// **not** commands that make any actual changes.
func getLines(command string, args ...string) []string {
	var s []string
	for _, l := range strings.Split(getOutput(command, args...), "\n") {
		l = strings.TrimSpace(l)
		if l != "" {
			s = append(s, l)
		}
	}
	return s
}

func commandString(command string, args []string) string {
	return strings.Join(append([]string{command}, args...), " ")
}

func dief(format string, args ...interface{}) {
	fmt.Fprintf(os.Stderr, "review: "+format+"\n", args...)
	os.Exit(1)
}

func verbosef(format string, args ...interface{}) {
	if *verbose {
		fmt.Fprintf(os.Stderr, format, args...)
	}
}
<commit_msg>review: add todo<commit_after>// Copyright 2014 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// TODO(adg): accept -a flag on 'create' and 'commit' (like git commit -a)
// TODO(adg): accept -r flag on 'upload' to nominate reviewer
// TODO(adg): support 'create' from non-master branches
// TODO(adg): check style of commit message
// TODO(adg): write doc comment
// TOOD(adg): print gerrit votes on 'pending'
// TODO(adg): add gofmt commit hook
// TODO(adg): 'upload' warn about uncommitted changes (maybe commit/create too?)
// TODO(adg): ability to edit commit message

package main // import "golang.org/x/review"

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

var (
	hookFile = filepath.FromSlash(".git/hooks/commit-msg")
	verbose  = flag.Bool("v", false, "verbose output")
	noRun    = flag.Bool("n", false, "print but do not run commands")
)

const usage = `Usage: %s [-n] [-v] <command>
Type "%s help" for more information.
`

const help = `Usage: %s [-n] [-v] <command>

The review command is a wrapper for the git command that provides a simple
interface to the "single-commit feature branch" development model.

Available comands:

	create <name>
		Create a local branch with the provided name
		and commit the staged changes to it.

	commit
		Amend local branch HEAD commit with the staged changes.

	diff
		View differences between remote branch HEAD and
		the local branch HEAD.
		(The differences introduced by this change.)

	upload
		Upload HEAD commit to the code review server.

	sync
		Fetch changes from the remote repository and merge them to the
		current branch, rebasing the HEAD commit (if any) on top of
		them. If the HEAD commit has been submitted, switch back to the
		master branch and delete the feature branch.

	pending
		Show local branches and their head commits.

`

func main() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, usage, os.Args[0], os.Args[0])
		os.Exit(2)
	}
	flag.Parse()

	goToRepoRoot()
	installHook()

	command, args := flag.Arg(0), flag.Args()[1:]

	switch command {
	case "help":
		fmt.Fprintf(os.Stdout, help, os.Args[0])
	case "create", "cr":
		create(args)
	case "commit", "co":
		commit(args)
	case "diff", "d":
		diff(args)
	case "upload", "u":
		upload(args)
	case "sync", "s":
		doSync(args)
	case "pending", "p":
		pending(args)
	default:
		flag.Usage()
	}
}

func create(args []string) {
	if len(args) != 1 || args[0] == "" {
		flag.Usage()
	}
	name := args[0]
	if !hasStagedChanges() {
		dief("no staged changes.\nDid you forget to 'git add'?")
	}
	if b := currentBranch(); b != "master" {
		dief("must run 'create' from the master branch (now on %q).\n"+
			"(Try 'review sync' or 'git checkout master' first.)",
			b)
	}
	run("git", "checkout", "-q", "-b", name)
	if err := runErr("git", "commit", "-q"); err != nil {
		verbosef("Commit failed: %v\nSwitching back to master.\n", err)
		run("git", "checkout", "-q", "master")
		run("git", "branch", "-q", "-d", name)
	}
}

func commit(args []string) {
	if len(args) != 0 {
		flag.Usage()
	}
	if !hasStagedChanges() {
		dief("no staged changes. Did you forget to 'git add'?")
	}
	if currentBranch() == "master" {
		dief("can't commit to master branch.")
	}
	run("git", "commit", "-q", "--amend", "-C", "HEAD")
}

func diff(args []string) {
	if len(args) != 0 {
		flag.Usage()
	}
	run("git", "diff", "HEAD^", "HEAD")
}

func upload(args []string) {
	if len(args) != 0 {
		flag.Usage()
	}
	if currentBranch() == "master" {
		dief("can't upload from master branch.")
	}
	run("git", "push", "-q", "origin", "HEAD:refs/for/master")
}

func doSync(args []string) {
	if len(args) != 0 {
		flag.Usage()
	}
	run("git", "fetch", "-q")

	// If we're on master, just fast-forward.
	branch := currentBranch()
	if branch == "master" {
		run("git", "merge", "-q", "--ff-only", "origin/master")
		return
	}

	// Check that exactly this commit was submitted to master. If so,
	// switch back to master, fast-forward, delete the feature branch.
	if branchContains("origin/master", "HEAD") {
		run("git", "checkout", "-q", "master")
		run("git", "merge", "-q", "--ff-only", "origin/master")
		run("git", "branch", "-q", "-d", branch)
		fmt.Printf("Change on %q submitted; branch deleted.\n"+
			"On master branch.\n", branch)
		return
	}

	// Check whether a rebased version of this commit was submitted to
	// master. If so, switch back to master, fast-forward, and
	// provide instructions for deleting the feature branch.
	// (We're not 100% sure that the feature branch HEAD was submitted,
	// so be cautious.)
	if headSubmitted(branch) {
		run("git", "checkout", "-q", "master")
		run("git", "merge", "-q", "--ff-only", "origin/master")
		fmt.Fprintf(os.Stderr,
			"I think the change on %q has been submitted.\n"+
				"If you agree, and no longer need branch %q, "+
				"run:\n\tgit branch -D %v\n"+
				"On master branch.\n",
			branch, branch, branch)
		return
	}

	// Bump master HEAD to that of origin/master, just in case the user
	// switches back to master with "git checkout master" later.
	// TODO(adg): maybe we shouldn't do this at all?
	if !branchContains("origin/master", "master") {
		run("git", "branch", "-f", "master", "origin/master")
	}

	// We have un-submitted changes on this feature branch; rebase.
	run("git", "rebase", "-q", "origin/master")
}

func pending(args []string) {
	if len(args) != 0 {
		flag.Usage()
	}
	var (
		wg      sync.WaitGroup
		origin  = originURL()
		current = currentBranch()
	)
	if current == "master" {
		fmt.Println("On master branch.")
	}
	for _, branch := range localBranches() {
		if branch == "master" {
			continue
		}
		wg.Add(1)
		go func(branch string) {
			defer wg.Done()
			p := ""
			if branch == current {
				p = "* "
			}
			id := headChangeId(branch)
			c, err := getChange(origin, id)
			switch err {
			case notFound:
				// TODO(adg): read local commit msg
				var msg string
				fmt.Printf("%v%v:\n\t%v\n\t(not uploaded)\n",
					p, branch, msg)
			case nil:
				status := ""
				switch c.Status {
				case "MERGED":
					status += " [submitted]"
				case "ABANDONED":
					status += " [abandoned]"
				}
				fmt.Printf("%v%v%v:\n\t%v\n\t%v\n",
					p, branch, status, c.Subject, c.URL)
			default:
				fmt.Fprintf(os.Stderr, "fetching change for %q: %v\n", branch, err)
			}
		}(branch)
	}
	wg.Wait()
}

func originURL() string {
	out, err := exec.Command("git", "config", "remote.origin.url").CombinedOutput()
	if err != nil {
		dief("could not find URL for 'origin' remote.\n"+
			"Did you check out from the right place?\n"+
			"git config remote.origin.url: %v\n"+
			"%s", err, out)
	}
	return string(out)
}

func localBranches() (branches []string) {
	for _, s := range getLines("git", "branch", "-l", "-q") {
		branches = append(branches, strings.TrimPrefix(s, "* "))
	}
	return branches
}

func branchContains(branch, rev string) bool {
	for _, s := range getLines("git", "branch", "-r", "--contains", rev) {
		if s == branch {
			return true
		}
	}
	return false
}

var stagedRe = regexp.MustCompile(`^[ACDMR]  `)

func hasStagedChanges() bool {
	for _, s := range getLines("git", "status", "-b", "--porcelain") {
		if stagedRe.MatchString(s) {
			return true
		}
	}
	return false
}

func currentBranch() string {
	return strings.TrimSpace(getOutput("git", "rev-parse", "--abbrev-ref", "HEAD"))
}

func headSubmitted(branch string) bool {
	s := "Change-Id: " + headChangeId(branch)
	return len(getOutput("git", "log", "--grep", s, "origin/master")) > 0
}

func headChangeId(branch string) string {
	const (
		p = "Change-Id: "
		f = "--format=format:%b"
	)
	for _, s := range getLines("git", "log", "-n", "1", f, branch, "--") {
		if strings.HasPrefix(s, p) {
			return strings.TrimSpace(strings.TrimPrefix(s, p))
		}
	}
	dief("no Change-Id line found in HEAD commit on branch %s.", branch)
	panic("unreachable")
}

func goToRepoRoot() {
	prevDir, err := os.Getwd()
	if err != nil {
		dief("could not get current directory: %v", err)
	}
	for {
		if _, err := os.Stat(".git"); err == nil {
			return
		}
		if err := os.Chdir(".."); err != nil {
			dief("could not chdir: %v", err)
		}
		currentDir, err := os.Getwd()
		if err != nil {
			dief("could not get current directory: %v", err)
		}
		if currentDir == prevDir {
			dief("git root not found.\n" +
				"Run from within the Git tree please.")
		}
		prevDir = currentDir
	}
}

func installHook() {
	_, err := os.Stat(hookFile)
	if err == nil {
		return
	}
	if !os.IsNotExist(err) {
		dief("error checking for hook file: %v", err)
	}
	verbosef("Presubmit hook to add Change-Id to commit messages is missing.\n"+
		"Automatically creating it at %v.\n", hookFile)
	hookContent := []byte(commitMsgHook)
	if err := ioutil.WriteFile(hookFile, hookContent, 0700); err != nil {
		dief("error writing hook file: %v", err)
	}
}

func run(command string, args ...string) {
	if err := runErr(command, args...); err != nil {
		if !*verbose {
			// If we're not in verbose mode, print the command
			// before dying to give context to the failure.
			fmt.Fprintln(os.Stderr, commandString(command, args))
		}
		dief("%v", err)
	}
}

func runErr(command string, args ...string) error {
	if *verbose || *noRun {
		fmt.Fprintln(os.Stderr, commandString(command, args))
	}
	if *noRun {
		return nil
	}
	cmd := exec.Command(command, args...)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	return cmd.Run()
}

// getOutput runs the specified command and returns its combined standard
// output and standard error outputs.
// NOTE: It should only be used to run commands that return information,
// **not** commands that make any actual changes.
func getOutput(command string, args ...string) string {
	b, err := exec.Command(command, args...).CombinedOutput()
	if err != nil {
		fmt.Fprintf(os.Stderr, "%v\n%s\n", commandString(command, args), b)
		dief("%v", err)
	}
	return string(b)
}

// getLines is like getOutput but it returns non-empty output lines.
// NOTE: It should only be used to run commands that return information,
// **not** commands that make any actual changes.
func getLines(command string, args ...string) []string {
	var s []string
	for _, l := range strings.Split(getOutput(command, args...), "\n") {
		l = strings.TrimSpace(l)
		if l != "" {
			s = append(s, l)
		}
	}
	return s
}

func commandString(command string, args []string) string {
	return strings.Join(append([]string{command}, args...), " ")
}

func dief(format string, args ...interface{}) {
	fmt.Fprintf(os.Stderr, "review: "+format+"\n", args...)
	os.Exit(1)
}

func verbosef(format string, args ...interface{}) {
	if *verbose {
		fmt.Fprintf(os.Stderr, format, args...)
	}
}
<|endoftext|>


<commit_before>package amigo

import (
	"os"

	"github.com/pelletier/go-toml"
)

const EnvMapKey = "envmap"

type Config struct {
	env      map[string]string
	confFile *toml.TomlTree
}

// Associate a given config file key with an environment var
func (c *Config) Env(confKey string, envKey string) {
	if val := os.Getenv(envKey); val != "" {
		c.env[confKey] = val
	}
}

func (c *Config) Get(key string) interface{} {
	if val, ok := c.env[key]; ok {
		return val
	}
	return c.confFile.Get(key)
}

// Return a new configuration object for use by library consumers
func New(filepath string) (*Config, error) {
	file, err := toml.LoadFile(filepath)

	if err != nil {
		return nil, err
	}

	c := &Config{}
	c.env = make(map[string]string)
	c.confFile = file

	// if a envmap table is defined, associate the specified keys with the env
	// vars defined there
	if file.Has(EnvMapKey) {
		envmap := file.Get(EnvMapKey).(*toml.TomlTree)
		for _, confKey := range envmap.Keys() {
			envKey := envmap.GetPath([]string{confKey}).(string)
			c.Env(confKey, envKey)
		}
	}
	return c, nil
}
<commit_msg>fixing issue with quoted strings in envmap table<commit_after>package amigo

import (
	"os"
	"regexp"

	"github.com/pelletier/go-toml"
)

const EnvMapKey = "envmap"

type Config struct {
	env      map[string]string
	confFile *toml.TomlTree
}

// Associate a given config file key with an environment var
func (c *Config) Env(confKey string, envKey string) {
	if val := os.Getenv(envKey); val != "" {
		c.env[confKey] = val
	}
}

func (c *Config) Get(key string) interface{} {
	if val, ok := c.env[key]; ok {
		return val
	}
	return c.confFile.Get(key)
}

func normalizeKey(key string) string {
	re := regexp.MustCompile(`^"(.+)"$`)
	if re.MatchString(key) {
		return key[1 : len(key)-1]
	}
	return key
}

// Return a new configuration object for use by library consumers
func New(filepath string) (*Config, error) {
	file, err := toml.LoadFile(filepath)

	if err != nil {
		return nil, err
	}

	c := &Config{}
	c.env = make(map[string]string)
	c.confFile = file

	// if an envmap table is defined, associate the specified keys with the env
	// vars defined there
	if file.Has(EnvMapKey) {
		envmap := file.Get(EnvMapKey).(*toml.TomlTree)
		for _, confKey := range envmap.Keys() {
			envKey := envmap.GetPath([]string{confKey}).(string)
			c.Env(normalizeKey(confKey), envKey)
		}
	}
	return c, nil
}
<|endoftext|>


<commit_before>package nats

import (
	"fmt"
	"time"

	"github.com/apcera/nats"
)

const (
	subject = "test"

	// Maximum bytes we will get behind before we start slowing down publishing.
	maxBytesBehind = 1024 * 1024 // 1MB

	// Maximum msgs we will get behind before we start slowing down publishing.
	maxMsgsBehind = 65536 // 64k

	// Time to delay publishing when we are behind.
	delay = 1 * time.Millisecond
)

// Peer implements the peer interface for NATS.
type Peer struct {
	conn     *nats.Conn
	messages chan []byte
	send     chan []byte
	errors   chan error
	done     chan bool
}

// NewPeer creates and returns a new Peer for communicating with NATS.
func NewPeer(host string) (*Peer, error) {
	conn, err := nats.Connect(fmt.Sprintf("nats://%s", host))
	if err != nil {
		return nil, err
	}

	// We want to be alerted if we get disconnected, this will be due to Slow
	// Consumer.
	conn.Opts.AllowReconnect = false

	return &Peer{
		conn:     conn,
		messages: make(chan []byte, 10000),
		send:     make(chan []byte),
		errors:   make(chan error, 1),
		done:     make(chan bool),
	}, nil
}

// Subscribe prepares the peer to consume messages.
func (n *Peer) Subscribe() error {
	n.conn.Subscribe(subject, func(message *nats.Msg) {
		n.messages <- message.Data
	})
	return nil
}

// Recv returns a single message consumed by the peer. Subscribe must be called
// before this. It returns an error if the receive failed.
func (n *Peer) Recv() ([]byte, error) {
	return <-n.messages, nil
}

// Send returns a channel on which messages can be sent for publishing.
func (n *Peer) Send() chan<- []byte {
	return n.send
}

// Errors returns the channel on which the peer sends publish errors.
func (n *Peer) Errors() <-chan error {
	return n.errors
}

// Done signals to the peer that message publishing has completed.
func (n *Peer) Done() {
	n.done <- true
}

// Setup prepares the peer for testing.
func (n *Peer) Setup() {
	go func() {
		for {
			select {
			case msg := <-n.send:
				if err := n.sendMessage(msg); err != nil {
					n.errors <- err
				}
			case <-n.done:
				return
			}
		}
	}()
}

func (n *Peer) sendMessage(message []byte) error {
	// Check if we are behind by >= 1MB bytes.
	bytesDeltaOver := n.conn.OutBytes-n.conn.InBytes >= maxBytesBehind

	// Check if we are behind by >= 65k msgs.
	msgsDeltaOver := n.conn.OutMsgs-n.conn.InMsgs >= maxMsgsBehind

	// If we are behind on either condition, sleep a bit to catch up receiver.
	if bytesDeltaOver || msgsDeltaOver {
		time.Sleep(delay)
	}

	return n.conn.Publish(subject, message)
}

// Teardown performs any cleanup logic that needs to be performed after the
// test is complete.
func (n *Peer) Teardown() {
	n.conn.Close()
}
<commit_msg>Update nats import<commit_after>package nats

import (
	"fmt"
	"time"

	"github.com/nats-io/nats"
)

const (
	subject = "test"

	// Maximum bytes we will get behind before we start slowing down publishing.
	maxBytesBehind = 1024 * 1024 // 1MB

	// Maximum msgs we will get behind before we start slowing down publishing.
	maxMsgsBehind = 65536 // 64k

	// Time to delay publishing when we are behind.
	delay = 1 * time.Millisecond
)

// Peer implements the peer interface for NATS.
type Peer struct {
	conn     *nats.Conn
	messages chan []byte
	send     chan []byte
	errors   chan error
	done     chan bool
}

// NewPeer creates and returns a new Peer for communicating with NATS.
func NewPeer(host string) (*Peer, error) {
	conn, err := nats.Connect(fmt.Sprintf("nats://%s", host))
	if err != nil {
		return nil, err
	}

	// We want to be alerted if we get disconnected, this will be due to Slow
	// Consumer.
	conn.Opts.AllowReconnect = false

	return &Peer{
		conn:     conn,
		messages: make(chan []byte, 10000),
		send:     make(chan []byte),
		errors:   make(chan error, 1),
		done:     make(chan bool),
	}, nil
}

// Subscribe prepares the peer to consume messages.
func (n *Peer) Subscribe() error {
	n.conn.Subscribe(subject, func(message *nats.Msg) {
		n.messages <- message.Data
	})
	return nil
}

// Recv returns a single message consumed by the peer. Subscribe must be called
// before this. It returns an error if the receive failed.
func (n *Peer) Recv() ([]byte, error) {
	return <-n.messages, nil
}

// Send returns a channel on which messages can be sent for publishing.
func (n *Peer) Send() chan<- []byte {
	return n.send
}

// Errors returns the channel on which the peer sends publish errors.
func (n *Peer) Errors() <-chan error {
	return n.errors
}

// Done signals to the peer that message publishing has completed.
func (n *Peer) Done() {
	n.done <- true
}

// Setup prepares the peer for testing.
func (n *Peer) Setup() {
	go func() {
		for {
			select {
			case msg := <-n.send:
				if err := n.sendMessage(msg); err != nil {
					n.errors <- err
				}
			case <-n.done:
				return
			}
		}
	}()
}

func (n *Peer) sendMessage(message []byte) error {
	// Check if we are behind by >= 1MB bytes.
	bytesDeltaOver := n.conn.OutBytes-n.conn.InBytes >= maxBytesBehind

	// Check if we are behind by >= 65k msgs.
	msgsDeltaOver := n.conn.OutMsgs-n.conn.InMsgs >= maxMsgsBehind

	// If we are behind on either condition, sleep a bit to catch up receiver.
	if bytesDeltaOver || msgsDeltaOver {
		time.Sleep(delay)
	}

	return n.conn.Publish(subject, message)
}

// Teardown performs any cleanup logic that needs to be performed after the
// test is complete.
func (n *Peer) Teardown() {
	n.conn.Close()
}
<|endoftext|>


<commit_before>/*
   Copyright 2010-2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.

   This file is licensed under the Apache License, Version 2.0 (the "License").
   You may not use this file except in compliance with the License. A copy of
   the License is located at

    http://aws.amazon.com/apache2.0/

   This file 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"

    "github.com/aws/aws-sdk-go/aws/session"
    "github.com/aws/aws-sdk-go/service/sqs"
)

// Usage:
// go run sqs_listqueues.go
func main() {
    // Initialize a session in us-west-2 that the SDK will use to load
    // credentials from the shared credentials file ~/.aws/credentials.
    sess, err := session.NewSession(&aws.Config{
        Region: aws.String("us-west-2")},
    )

    // Create a SQS service client.
    svc := sqs.New(sess)

    // List the queues available in a given region.
    result, err := svc.ListQueues(nil)
    if err != nil {
        fmt.Println("Error", err)
        return
    }

    fmt.Println("Success")
    // As these are pointers, printing them out directly would not be useful.
    for i, urls := range result.QueueUrls {
        // Avoid dereferencing a nil pointer.
        if urls == nil {
            continue
        }
        fmt.Printf("%d: %s\n", i, *urls)
    }
}
<commit_msg>Adding missing import<commit_after>/*
   Copyright 2010-2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.

   This file is licensed under the Apache License, Version 2.0 (the "License").
   You may not use this file except in compliance with the License. A copy of
   the License is located at

    http://aws.amazon.com/apache2.0/

   This file 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"

    "github.com/aws/aws-sdk-go/aws"
    "github.com/aws/aws-sdk-go/aws/session"
    "github.com/aws/aws-sdk-go/service/sqs"
)

// Usage:
// go run sqs_listqueues.go
func main() {
    // Initialize a session in us-west-2 that the SDK will use to load
    // credentials from the shared credentials file ~/.aws/credentials.
    sess, err := session.NewSession(&aws.Config{
        Region: aws.String("us-west-2")},
    )

    // Create a SQS service client.
    svc := sqs.New(sess)

    // List the queues available in a given region.
    result, err := svc.ListQueues(nil)
    if err != nil {
        fmt.Println("Error", err)
        return
    }

    fmt.Println("Success")
    // As these are pointers, printing them out directly would not be useful.
    for i, urls := range result.QueueUrls {
        // Avoid dereferencing a nil pointer.
        if urls == nil {
            continue
        }
        fmt.Printf("%d: %s\n", i, *urls)
    }
}
<|endoftext|>


<commit_before>package libkb

func getUsernameIfProvisioned(m MetaContext, uc UserConfig) (ret NormalizedUsername, err error) {
	m.Debug("getUsernameIfProvisioned(%+v)", uc)
	did := uc.GetDeviceID()
	if did.IsNil() {
		m.Debug("- no valid username since nil deviceID")
		return ret, nil
	}
	err = checkDeviceValidForUID(m.Ctx(), m.G().GetUPAKLoader(), uc.GetUID(), did)
	switch err.(type) {
	case nil:
		m.Debug("- checks out")
		return uc.GetUsername(), nil
	case DeviceNotFoundError:
		m.Debug("- user was likely reset (%s)", err)
		return ret, nil
	case KeyRevokedError:
		m.Debug("- device was revoked (s)", err)
		return ret, nil
	case UserDeletedError:
		m.Debug(" - user was deleted (%s)", err)
		return ret, nil
	case NotFoundError:
		// This can happen in development if the dev db is nuked or a mobile
		// device is connected to dev servers.
		m.Debug(" - user wasn't found (%s)", err)
		return ret, nil
	default:
		m.Debug("- unexpected error; propagating (%s)", err)
		return ret, err
	}
}

// GetAllProvisionedUsernames looks into the current config.json file, and finds all usernames
// that are currently provisioned on this device. That is, it filters out those that are on revoked
// devices or have reset their accounts. It uses UPAK loading for verifying the current user/device
// statuses, so should be fast if everything is cached recently.
func GetAllProvisionedUsernames(m MetaContext) (current NormalizedUsername, all []NormalizedUsername, err error) {

	m = m.WithLogTag("GAPU")
	defer m.Trace("GetAllProvisionedUsernames", func() error { return err })()

	currentUC, allUCs, err := m.G().Env.GetConfig().GetAllUserConfigs()
	if err != nil {
		return current, nil, err
	}

	if currentUC != nil {
		current, err = getUsernameIfProvisioned(m, *currentUC)
		if err != nil {
			m.Error("Error while checking user %q uid=%q, `current` will be nil", currentUC.GetUsername(), currentUC.GetUID())
		}
	}

	for _, u := range allUCs {
		tmp, err := getUsernameIfProvisioned(m, u)
		if err != nil {
			m.Error("Error while checking user %q uid=%q, skipping", currentUC.GetUsername(), currentUC.GetUID())
			continue
		}
		if !tmp.IsNil() {
			all = append(all, tmp)
		}
	}

	return current, all, nil
}
<commit_msg>fix panic in GetAllProvisionedUsernames (#16811)<commit_after>package libkb

func getUsernameIfProvisioned(m MetaContext, uc UserConfig) (ret NormalizedUsername, err error) {
	m.Debug("getUsernameIfProvisioned(%+v)", uc)
	did := uc.GetDeviceID()
	if did.IsNil() {
		m.Debug("- no valid username since nil deviceID")
		return ret, nil
	}
	err = checkDeviceValidForUID(m.Ctx(), m.G().GetUPAKLoader(), uc.GetUID(), did)
	switch err.(type) {
	case nil:
		m.Debug("- checks out")
		return uc.GetUsername(), nil
	case DeviceNotFoundError:
		m.Debug("- user was likely reset (%s)", err)
		return ret, nil
	case KeyRevokedError:
		m.Debug("- device was revoked (s)", err)
		return ret, nil
	case UserDeletedError:
		m.Debug(" - user was deleted (%s)", err)
		return ret, nil
	case NotFoundError:
		// This can happen in development if the dev db is nuked or a mobile
		// device is connected to dev servers.
		m.Debug(" - user wasn't found (%s)", err)
		return ret, nil
	default:
		m.Debug("- unexpected error; propagating (%s)", err)
		return ret, err
	}
}

// GetAllProvisionedUsernames looks into the current config.json file, and finds all usernames
// that are currently provisioned on this device. That is, it filters out those that are on revoked
// devices or have reset their accounts. It uses UPAK loading for verifying the current user/device
// statuses, so should be fast if everything is cached recently.
func GetAllProvisionedUsernames(m MetaContext) (current NormalizedUsername, all []NormalizedUsername, err error) {

	m = m.WithLogTag("GAPU")
	defer m.Trace("GetAllProvisionedUsernames", func() error { return err })()

	currentUC, allUCs, err := m.G().Env.GetConfig().GetAllUserConfigs()
	if err != nil {
		return current, nil, err
	}

	if currentUC != nil {
		current, err = getUsernameIfProvisioned(m, *currentUC)
		if err != nil {
			m.Error("Error while checking user %q uid=%q, `current` will be nil", currentUC.GetUsername(), currentUC.GetUID())
		}
	}

	for _, uc := range allUCs {
		tmp, err := getUsernameIfProvisioned(m, uc)
		if err != nil {
			m.Error("Error while checking user %q uid=%q, skipping", uc.GetUsername(), uc.GetUID())
			continue
		}
		if !tmp.IsNil() {
			all = append(all, tmp)
		}
	}

	return current, all, nil
}
<|endoftext|>


<commit_before>// Copyright 2017 Istio Authors
//
// 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 framework

import (
	"cmd/pprof/internal/tempfile"
	"flag"
	"io/ioutil"
	"os"
	"sync"

	"github.com/golang/glog"
)

var (
	skipCleanup = flag.Bool("skip_cleanup", false, "Debug, skip clean up")
	logProvider = flag.String("log_provider", "", "Cluster log storage provider")
)

type testCleanup struct {
	Cleanables         []Cleanable
	CleanablesLock     sync.Mutex
	CleanupActions     []func() error
	CleanupActionsLock sync.Mutex
	skipCleanup        bool
}

// CommonConfig regroup all common test configuration.
type CommonConfig struct {
	// Test Cleanup registration
	Cleanup *testCleanup
	// Test Information
	Info *testInfo
	// Kubernetes and istio installation information
	Kube *KubeInfo
}

// Cleanable interfaces that need to be registered to CommonConfig
type Cleanable interface {
	Setup() error
	Teardown() error
}

// Runnable is used for Testing purposes.
type runnable interface {
	Run() int
}

// InitGlog sets the logging directory.
// Should be called right after flag.Parse().
func InitGlog() error {
	tmpDir, err := ioutil.TempDir(os.TempDir(), tmpPrefix)
	if err != nil {
		return err
	}
	f := flag.Lookup("log_dir")
	if err = f.Value.Set(tmpDir); err != nil {
		return err
	}
	glog.Info("Logging initialized")
	return nil
}

// NewCommonConfig creates a full config will all supported configs.
func NewCommonConfig(testID string) (*CommonConfig, error) {
	t, err := newTestInfo(testID)
	if err != nil {
		return nil, err
	}
	k, err := newKubeInfo(t.LogsPath, t.RunID)
	if err != nil {
		return nil, err
	}
	cl := new(testCleanup)
	cl.skipCleanup = *skipCleanup

	c := &CommonConfig{
		Info:    t,
		Kube:    k,
		Cleanup: cl,
	}
	c.Cleanup.RegisterCleanable(c.Info)
	c.Cleanup.RegisterCleanable(c.Kube)
	c.Cleanup.RegisterCleanable(c.Kube.Istioctl)
	c.Cleanup.RegisterCleanable(c.Kube.AppManager)
	return c, nil
}

func (t *testCleanup) RegisterCleanable(c Cleanable) {
	t.CleanablesLock.Lock()
	defer t.CleanablesLock.Unlock()
	t.Cleanables = append(t.Cleanables, c)
}

func (t *testCleanup) popCleanable() Cleanable {
	t.CleanablesLock.Lock()
	defer t.CleanablesLock.Unlock()
	if len(t.Cleanables) == 0 {
		return nil
	}
	c := t.Cleanables[0]
	t.Cleanables = t.Cleanables[1:]
	return c
}

func (t *testCleanup) addCleanupAction(fn func() error) {
	t.CleanupActionsLock.Lock()
	defer t.CleanupActionsLock.Unlock()
	t.CleanupActions = append(t.CleanupActions, fn)
}

func (t *testCleanup) popCleanupAction() func() error {
	t.CleanupActionsLock.Lock()
	defer t.CleanupActionsLock.Unlock()
	if len(t.CleanupActions) == 0 {
		return nil
	}
	fn := t.CleanupActions[len(t.CleanupActions)-1]
	t.CleanupActions = t.CleanupActions[:len(t.CleanupActions)-1]
	return fn
}

func (t *testCleanup) init() error {
	// Run setup on all cleanable
	glog.Info("Starting Initialization")
	c := t.popCleanable()
	for c != nil {
		err := c.Setup()
		t.addCleanupAction(c.Teardown)
		if err != nil {
			return err
		}
		c = t.popCleanable()
	}
	glog.Info("Initialization complete")
	return nil
}

func (t *testCleanup) cleanup() {
	defer tempfile.Cleanup()
	if t.skipCleanup {
		glog.Info("Debug model, skip cleanup")
		return
	}
	// Run tear down on all cleanable
	glog.Info("Starting Cleanup")
	fn := t.popCleanupAction()
	for fn != nil {
		if err := fn(); err != nil {
			glog.Errorf("Failed to cleanup. Error %s", err)
		}
		fn = t.popCleanupAction()
	}
	glog.Info("Cleanup complete")
}

// Save test logs to tmp dir
// Fetch and save cluster tracing logs if logProvider specified
// Logs are uploaded during test tear down
func (c *CommonConfig) saveLogs(r int) error {
	if c.Info == nil {
		glog.Warning("Skipping log saving as Info is not initialized")
		return nil
	}
	if c.Info.LogBucketPath == "" {
		return nil
	}
	glog.Info("Saving logs")
	if err := c.Info.Update(r); err != nil {
		glog.Errorf("Could not create status file. Error %s", err)
		return err
	}
	if r != 0 && *logProvider == "stackdriver" { // fetch logs only if tests failed
		if err := c.Info.FetchAndSaveClusterLogs(c.Kube.Namespace); err != nil {
			return err
		}
	}
	return nil
}

// RunTest sets up all registered cleanables in FIFO order
// Execute the runnable
// Call teardown on all the cleanables in LIFO order.
func (c *CommonConfig) RunTest(m runnable) int {
	var ret int
	if err := c.Cleanup.init(); err != nil {
		glog.Errorf("Failed to complete Init. Error %s", err)
		ret = 1
	} else {
		glog.Info("Running test")
		ret = m.Run()
	}
	if err := c.saveLogs(ret); err != nil {
		glog.Warningf("Log saving incomplete: %v", err)
	}
	c.Cleanup.cleanup()
	return ret
}
<commit_msg>Removed not need call of tempfile.Cleanup() (#752)<commit_after>// Copyright 2017 Istio Authors
//
// 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 framework

import (
	"flag"
	"io/ioutil"
	"os"
	"sync"

	"github.com/golang/glog"
)

var (
	skipCleanup = flag.Bool("skip_cleanup", false, "Debug, skip clean up")
	logProvider = flag.String("log_provider", "", "Cluster log storage provider")
)

type testCleanup struct {
	Cleanables         []Cleanable
	CleanablesLock     sync.Mutex
	CleanupActions     []func() error
	CleanupActionsLock sync.Mutex
	skipCleanup        bool
}

// CommonConfig regroup all common test configuration.
type CommonConfig struct {
	// Test Cleanup registration
	Cleanup *testCleanup
	// Test Information
	Info *testInfo
	// Kubernetes and istio installation information
	Kube *KubeInfo
}

// Cleanable interfaces that need to be registered to CommonConfig
type Cleanable interface {
	Setup() error
	Teardown() error
}

// Runnable is used for Testing purposes.
type runnable interface {
	Run() int
}

// InitGlog sets the logging directory.
// Should be called right after flag.Parse().
func InitGlog() error {
	tmpDir, err := ioutil.TempDir(os.TempDir(), tmpPrefix)
	if err != nil {
		return err
	}
	f := flag.Lookup("log_dir")
	if err = f.Value.Set(tmpDir); err != nil {
		return err
	}
	glog.Info("Logging initialized")
	return nil
}

// NewCommonConfig creates a full config will all supported configs.
func NewCommonConfig(testID string) (*CommonConfig, error) {
	t, err := newTestInfo(testID)
	if err != nil {
		return nil, err
	}
	k, err := newKubeInfo(t.LogsPath, t.RunID)
	if err != nil {
		return nil, err
	}
	cl := new(testCleanup)
	cl.skipCleanup = *skipCleanup

	c := &CommonConfig{
		Info:    t,
		Kube:    k,
		Cleanup: cl,
	}
	c.Cleanup.RegisterCleanable(c.Info)
	c.Cleanup.RegisterCleanable(c.Kube)
	c.Cleanup.RegisterCleanable(c.Kube.Istioctl)
	c.Cleanup.RegisterCleanable(c.Kube.AppManager)
	return c, nil
}

func (t *testCleanup) RegisterCleanable(c Cleanable) {
	t.CleanablesLock.Lock()
	defer t.CleanablesLock.Unlock()
	t.Cleanables = append(t.Cleanables, c)
}

func (t *testCleanup) popCleanable() Cleanable {
	t.CleanablesLock.Lock()
	defer t.CleanablesLock.Unlock()
	if len(t.Cleanables) == 0 {
		return nil
	}
	c := t.Cleanables[0]
	t.Cleanables = t.Cleanables[1:]
	return c
}

func (t *testCleanup) addCleanupAction(fn func() error) {
	t.CleanupActionsLock.Lock()
	defer t.CleanupActionsLock.Unlock()
	t.CleanupActions = append(t.CleanupActions, fn)
}

func (t *testCleanup) popCleanupAction() func() error {
	t.CleanupActionsLock.Lock()
	defer t.CleanupActionsLock.Unlock()
	if len(t.CleanupActions) == 0 {
		return nil
	}
	fn := t.CleanupActions[len(t.CleanupActions)-1]
	t.CleanupActions = t.CleanupActions[:len(t.CleanupActions)-1]
	return fn
}

func (t *testCleanup) init() error {
	// Run setup on all cleanable
	glog.Info("Starting Initialization")
	c := t.popCleanable()
	for c != nil {
		err := c.Setup()
		t.addCleanupAction(c.Teardown)
		if err != nil {
			return err
		}
		c = t.popCleanable()
	}
	glog.Info("Initialization complete")
	return nil
}

func (t *testCleanup) cleanup() {
	if t.skipCleanup {
		glog.Info("Debug model, skip cleanup")
		return
	}
	// Run tear down on all cleanable
	glog.Info("Starting Cleanup")
	fn := t.popCleanupAction()
	for fn != nil {
		if err := fn(); err != nil {
			glog.Errorf("Failed to cleanup. Error %s", err)
		}
		fn = t.popCleanupAction()
	}
	glog.Info("Cleanup complete")
}

// Save test logs to tmp dir
// Fetch and save cluster tracing logs if logProvider specified
// Logs are uploaded during test tear down
func (c *CommonConfig) saveLogs(r int) error {
	if c.Info == nil {
		glog.Warning("Skipping log saving as Info is not initialized")
		return nil
	}
	if c.Info.LogBucketPath == "" {
		return nil
	}
	glog.Info("Saving logs")
	if err := c.Info.Update(r); err != nil {
		glog.Errorf("Could not create status file. Error %s", err)
		return err
	}
	if r != 0 && *logProvider == "stackdriver" { // fetch logs only if tests failed
		if err := c.Info.FetchAndSaveClusterLogs(c.Kube.Namespace); err != nil {
			return err
		}
	}
	return nil
}

// RunTest sets up all registered cleanables in FIFO order
// Execute the runnable
// Call teardown on all the cleanables in LIFO order.
func (c *CommonConfig) RunTest(m runnable) int {
	var ret int
	if err := c.Cleanup.init(); err != nil {
		glog.Errorf("Failed to complete Init. Error %s", err)
		ret = 1
	} else {
		glog.Info("Running test")
		ret = m.Run()
	}
	if err := c.saveLogs(ret); err != nil {
		glog.Warningf("Log saving incomplete: %v", err)
	}
	c.Cleanup.cleanup()
	return ret
}
<|endoftext|>


<commit_before>package zkclient

import (
	"fmt"
	"strconv"
	"time"

	"github.com/fsouza/go-dockerclient"
	"github.com/samuel/go-zookeeper/zk"
)

type Client struct {
	dkrClient *docker.Client
	zkConn    *zk.Conn
}

var err error

func New() *Client {
	return &Client{}
}
func (c *Client) Init(zkIP string) {
	c.dkrClient, err = docker.NewClientFromEnv()
	handleError(err)
	c.zkConn, _, err = zk.Connect([]string{zkIP}, time.Second) //TODO :: need to get ip of zookeeper from docker
	handleError(err)
	time.Sleep(time.Second)
}

func (c *Client) Connect(containerName string, zkPath string) {

	prevcontainers := 0

	for {
		containers, err := c.dkrClient.ListContainers(docker.ListContainersOptions{All: false, Filters: map[string][]string{"name": {containerName}, "status": {"running"}}})
		handleError(err)
		if len(containers) == prevcontainers {
			continue
		} else {
			if len(containers) == 1 {
				_, err := c.zkConn.Create(zkPath, []byte("localhost:"+strconv.FormatInt(containers[0].Ports[0].PrivatePort, 10)), int32(zk.FlagEphemeral)|int32(zk.FlagSequence), zk.WorldACL(zk.PermAll))
				handleError(err)
				fmt.Println("[zookeeper] : Server registered")
			} else {
				err := c.zkConn.Delete(zkPath, -1)
				handleError(err)
				fmt.Println("[zookeeper] : Server deleted")
			}
		}
		prevcontainers = len(containers)
	}

}

//handleError stops the program execution
func handleError(err error) {
	if err != nil {
		panic(err)
	}
}
<commit_msg>Feature to delete Sequential node<commit_after>package zkclient

import (
	"fmt"
	"strconv"
	"time"

	"github.com/fsouza/go-dockerclient"
	"github.com/samuel/go-zookeeper/zk"
)

type Client struct {
	dkrClient *docker.Client
	zkConn    *zk.Conn
}

var err error

func New() *Client {
	return &Client{}
}
func (c *Client) Init(zkIP string) {
	c.dkrClient, err = docker.NewClientFromEnv()
	handleError(err)
	c.zkConn, _, err = zk.Connect([]string{zkIP}, time.Second) //TODO :: need to get ip of zookeeper from docker
	handleError(err)
	time.Sleep(time.Second)
}

func (c *Client) Connect(containerName string, zkPath string) {

	prevcontainers := 0
	var address string

	for {
		containers, err := c.dkrClient.ListContainers(docker.ListContainersOptions{All: false, Filters: map[string][]string{"name": {containerName}, "status": {"running"}}})
		handleError(err)
		if len(containers) == prevcontainers {
			continue
		} else {
			if len(containers) == 1 {
				address, err = c.zkConn.Create(zkPath, []byte("localhost:"+strconv.FormatInt(containers[0].Ports[0].PrivatePort, 10)), int32(zk.FlagEphemeral)|int32(zk.FlagSequence), zk.WorldACL(zk.PermAll))
				handleError(err)
				time.Sleep(time.Second)
				fmt.Println("[zookeeper] : Server registered")
			} else {
				deleteErr := c.zkConn.Delete(address, -1)
				handleError(deleteErr)
				time.Sleep(time.Second)
				fmt.Println("[zookeeper] : Server deleted")
			}
		}
		prevcontainers = len(containers)
	}

}

//handleError stops the program execution
func handleError(err error) {
	if err != nil {
		panic(err)
	}
}
<|endoftext|>


<commit_before>// Copyright 2014 CoreOS, Inc.
//
// 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"
	"os"
	"path/filepath"
	"syscall"
	"time"

	"github.com/coreos/rocket/pkg/lock"
)

const (
	defaultGracePeriod = 30 * time.Minute
)

var (
	flagGracePeriod time.Duration
	cmdGC           = &Command{
		Name:    "gc",
		Summary: "Garbage-collect rkt containers no longer in use",
		Usage:   "[--grace-period=duration]",
		Run:     runGC,
	}
)

func init() {
	commands = append(commands, cmdGC)
	cmdGC.Flags.DurationVar(&flagGracePeriod, "grace-period", defaultGracePeriod, "duration to wait before discarding inactive containers from garbage")
}

func runGC(args []string) (exit int) {
	if err := os.MkdirAll(garbageDir(), 0755); err != nil {
		fmt.Fprintf(os.Stderr, "Unable to create garbage dir: %v\n", err)
		return 1
	}

	cs, err := getContainers()
	if err != nil {
		fmt.Fprintf(os.Stderr, "Unable to get containers list: %v\n", err)
		return 1
	}
	for _, c := range cs {
		cp := filepath.Join(containersDir(), c)
		l, err := lock.TryExclusiveLock(cp)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Unable to open lock, ignoring %q: %v\n", c, err)
			continue
		}

		fmt.Printf("Moving container %q to garbage\n", c)
		err = os.Rename(cp, filepath.Join(garbageDir(), c))
		if err != nil {
			fmt.Println(err)
		}
		l.Close()
	}

	// clean up anything old in the garbage dir
	err = emptyGarbage(flagGracePeriod)
	if err != nil {
		fmt.Println(err)
		return 1
	}

	return
}

// getContainers returns a slice representing the containers in the given rocket directory
func getContainers() ([]string, error) {
	cdir := containersDir()
	ls, err := ioutil.ReadDir(cdir)
	if err != nil {
		return nil, fmt.Errorf("cannot read containers directory: %v", err)
	}
	var cs []string
	for _, dir := range ls {
		if !dir.IsDir() {
			fmt.Fprintf(os.Stderr, "Unrecognized file: %q, ignoring", dir)
			continue
		}
		cs = append(cs, dir.Name())
	}
	return cs, nil
}

// emptyGarbage discards sufficiently aged containers from garbageDir()
func emptyGarbage(gracePeriod time.Duration) error {
	g := garbageDir()

	ls, err := ioutil.ReadDir(g)
	if err != nil {
		return err
	}

	for _, dir := range ls {
		gp := filepath.Join(g, dir.Name())
		st := &syscall.Stat_t{}
		err := syscall.Lstat(gp, st)
		if err != nil {
			if err != syscall.ENOENT {
				fmt.Fprintf(os.Stderr, "Unable to stat %q, ignoring: %v", gp, err)
			}
			continue
		}

		expiration := time.Unix(st.Ctim.Unix()).Add(gracePeriod)
		if time.Now().After(expiration) {
			l, err := lock.ExclusiveLock(gp)
			if err != nil {
				continue
			}
			fmt.Printf("Garbage collecting container %q\n", dir.Name())
			if err = os.RemoveAll(gp); err != nil {
				fmt.Fprintf(os.Stderr, "Unable to remove container %q: %v\n", dir.Name(), err)
			}
			l.Close()
		}
	}
	return nil
}
<commit_msg>rkt: fix trailing newlines in errors<commit_after>// Copyright 2014 CoreOS, Inc.
//
// 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"
	"os"
	"path/filepath"
	"syscall"
	"time"

	"github.com/coreos/rocket/pkg/lock"
)

const (
	defaultGracePeriod = 30 * time.Minute
)

var (
	flagGracePeriod time.Duration
	cmdGC           = &Command{
		Name:    "gc",
		Summary: "Garbage-collect rkt containers no longer in use",
		Usage:   "[--grace-period=duration]",
		Run:     runGC,
	}
)

func init() {
	commands = append(commands, cmdGC)
	cmdGC.Flags.DurationVar(&flagGracePeriod, "grace-period", defaultGracePeriod, "duration to wait before discarding inactive containers from garbage")
}

func runGC(args []string) (exit int) {
	if err := os.MkdirAll(garbageDir(), 0755); err != nil {
		fmt.Fprintf(os.Stderr, "Unable to create garbage dir: %v\n", err)
		return 1
	}

	cs, err := getContainers()
	if err != nil {
		fmt.Fprintf(os.Stderr, "Unable to get containers list: %v\n", err)
		return 1
	}
	for _, c := range cs {
		cp := filepath.Join(containersDir(), c)
		l, err := lock.TryExclusiveLock(cp)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Unable to open lock, ignoring %q: %v\n", c, err)
			continue
		}

		fmt.Printf("Moving container %q to garbage\n", c)
		err = os.Rename(cp, filepath.Join(garbageDir(), c))
		if err != nil {
			fmt.Println(err)
		}
		l.Close()
	}

	// clean up anything old in the garbage dir
	err = emptyGarbage(flagGracePeriod)
	if err != nil {
		fmt.Println(err)
		return 1
	}

	return
}

// getContainers returns a slice representing the containers in the given rocket directory
func getContainers() ([]string, error) {
	cdir := containersDir()
	ls, err := ioutil.ReadDir(cdir)
	if err != nil {
		return nil, fmt.Errorf("cannot read containers directory: %v", err)
	}
	var cs []string
	for _, dir := range ls {
		if !dir.IsDir() {
			fmt.Fprintf(os.Stderr, "Unrecognized file: %q, ignoring\n", dir)
			continue
		}
		cs = append(cs, dir.Name())
	}
	return cs, nil
}

// emptyGarbage discards sufficiently aged containers from garbageDir()
func emptyGarbage(gracePeriod time.Duration) error {
	g := garbageDir()

	ls, err := ioutil.ReadDir(g)
	if err != nil {
		return err
	}

	for _, dir := range ls {
		gp := filepath.Join(g, dir.Name())
		st := &syscall.Stat_t{}
		err := syscall.Lstat(gp, st)
		if err != nil {
			if err != syscall.ENOENT {
				fmt.Fprintf(os.Stderr, "Unable to stat %q, ignoring: %v\n", gp, err)
			}
			continue
		}

		expiration := time.Unix(st.Ctim.Unix()).Add(gracePeriod)
		if time.Now().After(expiration) {
			l, err := lock.ExclusiveLock(gp)
			if err != nil {
				continue
			}
			fmt.Printf("Garbage collecting container %q\n", dir.Name())
			if err = os.RemoveAll(gp); err != nil {
				fmt.Fprintf(os.Stderr, "Unable to remove container %q: %v\n", dir.Name(), err)
			}
			l.Close()
		}
	}
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2014 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package files_test

import (
	"os"
	"path/filepath"

	jc "github.com/juju/testing/checkers"
	gc "gopkg.in/check.v1"

	"github.com/juju/juju/state/backups/files"
	"github.com/juju/juju/testing"
)

var _ = gc.Suite(&filesSuite{})

type filesSuite struct {
	testing.BaseSuite
	root string
}

func (s *filesSuite) SetUpTest(c *gc.C) {
	s.BaseSuite.SetUpTest(c)

	s.root = c.MkDir()

	// Prep the fake FS.
	mkdir := func(path string) string {
		dirname := filepath.Join(s.root, path)
		os.MkdirAll(dirname, 0777)
		return dirname
	}
	touch := func(dirname, name string) {
		path := filepath.Join(dirname, name)
		file, err := os.Create(path)
		c.Assert(err, gc.IsNil)
		file.Close()
	}

	dirname := mkdir("/var/lib/juju")
	touch(dirname, "system-identity")
	touch(dirname, "nonce.txt")
	touch(dirname, "server.pem")
	touch(dirname, "shared-secret")
	mkdir("/var/lib/juju/tools")

	dirname = mkdir("/var/lib/juju/agents")
	touch(dirname, "machine-0.conf")

	dirname = mkdir("/var/log/juju")
	touch(dirname, "all-machines.log")
	touch(dirname, "machine-0.log")

	dirname = mkdir("/etc/init")
	touch(dirname, "jujud-machine-0.conf")
	touch(dirname, "juju-db.conf")

	dirname = mkdir("/etc/rsyslog.d")
	touch(dirname, "spam-juju.conf")

	dirname = mkdir("/home/ubuntu/.ssh")
	touch(dirname, "authorized_keys")
}

func (s *filesSuite) TestGetFilesToBackUp(c *gc.C) {
	paths := files.Paths{DataDir: "/var/lib/juju"}
	files, err := files.GetFilesToBackUp(s.root, paths)
	c.Assert(err, gc.IsNil)

	c.Check(files, jc.SameContents, []string{
		filepath.Join(s.root, "/etc/init/juju-db.conf"),
		filepath.Join(s.root, "/etc/init/jujud-machine-0.conf"),
		filepath.Join(s.root, "/etc/rsyslog.d/spam-juju.conf"),
		filepath.Join(s.root, "/home/ubuntu/.ssh/authorized_keys"),
		filepath.Join(s.root, "/var/lib/juju/agents/machine-0.conf"),
		filepath.Join(s.root, "/var/lib/juju/nonce.txt"),
		filepath.Join(s.root, "/var/lib/juju/server.pem"),
		filepath.Join(s.root, "/var/lib/juju/shared-secret"),
		filepath.Join(s.root, "/var/lib/juju/system-identity"),
		filepath.Join(s.root, "/var/lib/juju/tools"),
		filepath.Join(s.root, "/var/log/juju/all-machines.log"),
		filepath.Join(s.root, "/var/log/juju/machine-0.log"),
	})
}
<commit_msg>Set LogsDir for test.<commit_after>// Copyright 2014 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package files_test

import (
	"os"
	"path/filepath"
	"sort"

	jc "github.com/juju/testing/checkers"
	"github.com/juju/utils/set"
	gc "gopkg.in/check.v1"

	"github.com/juju/juju/state/backups/files"
	"github.com/juju/juju/testing"
)

var _ = gc.Suite(&filesSuite{})

type filesSuite struct {
	testing.BaseSuite
	root string
}

func (s *filesSuite) SetUpTest(c *gc.C) {
	s.BaseSuite.SetUpTest(c)

	s.root = c.MkDir()

	// Prep the fake FS.
	mkdir := func(path string) string {
		dirname := filepath.Join(s.root, path)
		os.MkdirAll(dirname, 0777)
		return dirname
	}
	touch := func(dirname, name string) {
		path := filepath.Join(dirname, name)
		file, err := os.Create(path)
		c.Assert(err, gc.IsNil)
		file.Close()
	}

	dirname := mkdir("/var/lib/juju")
	touch(dirname, "system-identity")
	touch(dirname, "nonce.txt")
	touch(dirname, "server.pem")
	touch(dirname, "shared-secret")
	mkdir("/var/lib/juju/tools")

	dirname = mkdir("/var/lib/juju/agents")
	touch(dirname, "machine-0.conf")

	dirname = mkdir("/var/log/juju")
	touch(dirname, "all-machines.log")
	touch(dirname, "machine-0.log")

	dirname = mkdir("/etc/init")
	touch(dirname, "jujud-machine-0.conf")
	touch(dirname, "juju-db.conf")

	dirname = mkdir("/etc/rsyslog.d")
	touch(dirname, "spam-juju.conf")

	dirname = mkdir("/home/ubuntu/.ssh")
	touch(dirname, "authorized_keys")
}

func (s *filesSuite) checkSameStrings(c *gc.C, actual, expected []string) {
	sActual := set.NewStrings(actual...)
	sExpected := set.NewStrings(expected...)

	sActualOnly := sActual.Difference(sExpected)
	sExpectedOnly := sExpected.Difference(sActual)

	if !sActualOnly.IsEmpty() || !sExpectedOnly.IsEmpty() {
		c.Error("strings mismatch")
		onlyActual := sActualOnly.Values()
		onlyExpected := sExpectedOnly.Values()
		sort.Strings(onlyActual)
		sort.Strings(onlyExpected)

		if !sActualOnly.IsEmpty() {
			c.Log("...unexpected values:")
			for _, str := range onlyActual {
				c.Log(" " + str)
			}
		}
		if !sExpectedOnly.IsEmpty() {
			c.Log("...missing values:")
			for _, str := range onlyExpected {
				c.Log(" " + str)
			}
		}
	}
}

func (s *filesSuite) TestGetFilesToBackUp(c *gc.C) {
	paths := files.Paths{
		DataDir: "/var/lib/juju",
		LogsDir: "/var/log/juju",
	}
	files, err := files.GetFilesToBackUp(s.root, paths)
	c.Assert(err, gc.IsNil)

	expected := []string{
		filepath.Join(s.root, "/etc/init/juju-db.conf"),
		filepath.Join(s.root, "/etc/init/jujud-machine-0.conf"),
		filepath.Join(s.root, "/etc/rsyslog.d/spam-juju.conf"),
		filepath.Join(s.root, "/home/ubuntu/.ssh/authorized_keys"),
		filepath.Join(s.root, "/var/lib/juju/agents/machine-0.conf"),
		filepath.Join(s.root, "/var/lib/juju/nonce.txt"),
		filepath.Join(s.root, "/var/lib/juju/server.pem"),
		filepath.Join(s.root, "/var/lib/juju/shared-secret"),
		filepath.Join(s.root, "/var/lib/juju/system-identity"),
		filepath.Join(s.root, "/var/lib/juju/tools"),
		filepath.Join(s.root, "/var/log/juju/all-machines.log"),
		filepath.Join(s.root, "/var/log/juju/machine-0.log"),
	}
	c.Check(files, jc.SameContents, expected)
	s.checkSameStrings(c, files, expected)
}
<|endoftext|>


<commit_before>//  Copyright (c) 2014 Spotify AB.
//
//  Licensed to the Apache Software Foundation (ASF) under one
//  or more contributor license agreements.  See the NOTICE file
//  distributed with this work for additional information
//  regarding copyright ownership.  The ASF licenses this file
//  to you 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 (
	"bytes"
	"errors"
	"flag"
	"log"
	"os"
	"os/exec"
	"syscall"
)

type Syslogger struct {
	logger   *Writer
	stream   string
	buffer   *bytes.Buffer
	hostPort string
	priority Priority
	prefix   string
	logFlags int
        protocol string
}

func (s *Syslogger) Write(p []byte) (n int, err error) {
	if s.logger == nil {
		sl, err := Dial(s.protocol, s.hostPort, s.priority, s.prefix)
		if err != nil {
			// while syslog is down, dump the output
			return len(p), nil
		}
		s.logger = sl
	}
	for b := range p {
		s.buffer.WriteByte(p[b])
		if p[b] == 10 { // newline
			_, err := s.logger.Write(s.buffer.Bytes())
			if err != nil {
				s.logger = nil
			}
			s.buffer = bytes.NewBuffer([]byte{})
		}
	}
	return len(p), nil
}

func (s *Syslogger) Close() error {
	return nil
}

func NewSysLogger(stream, hostPort, prefix, protocol string) (*Syslogger, error) {
	var priority Priority
	if stream == "stderr" {
		priority = LOG_ERR | LOG_LOCAL0
	} else if stream == "stdout" {
		priority = LOG_INFO | LOG_LOCAL0
	} else {
		return nil, errors.New("cannot create syslogger for stream " + stream)
	}
	logFlags := 0

	return &Syslogger{nil, stream, bytes.NewBuffer([]byte{}), hostPort, priority, prefix, logFlags, protocol}, nil
}

func usage() {
	log.Printf("usage: %s -h syslog_host:port -n name -- executable [arg ...]", os.Args[0])
	flag.PrintDefaults()
	os.Exit(1)
}

func main() {
	log.SetFlags(0)

	flHostPort := flag.String("h", "", "Host port of where to connect to the syslog daemon")
	flLogName := flag.String("n", "", "Name to log as")
	flProtocol := flag.Bool("t", false, "use TCP instead of UDP (the default) for syslog communication")
	flag.Parse()

	if *flHostPort == "" {
		log.Printf("Must set the syslog host:port argument")
		usage()
	}

	if *flLogName == "" {
		log.Printf("Must set the syslog log name argument")
		usage()
	}

	protocol := "udp"
	if *flProtocol {
		protocol = "tcp"
	}

	//Example ./syslog-redirector -h 10.0.3.1:6514 -n test-ls-thingy -- \
	//            /bin/bash -c 'while true; do date; echo $SHELL; sleep 1; done'
	if len(os.Args) < 4 {
		log.Printf("at least 3 arguments required")
		usage()
	}
	hostPort := *flHostPort
	name := *flLogName

	if len(flag.Args()) == 0 {
		log.Printf("must supply a command")
		usage()
	}

	cmdName := flag.Args()[0]
	cmdArgs := flag.Args()[1:]

	var err error

	path, err := exec.LookPath(cmdName)
	if err != nil {
		log.Printf("Unable to locate %v", cmdName)
		os.Exit(127)
	}

	cmd := exec.Command(path, cmdArgs...)

	// TODO (dano): tolerate syslog downtime by reconnecting

	cmd.Stdout, err = NewSysLogger("stdout", hostPort, name, protocol)
	if err != nil {
		log.Printf("error creating syslog writer for stdout: %v", err)
	}

	cmd.Stderr, err = NewSysLogger("stderr", hostPort, name, protocol)
	if err != nil {
		log.Printf("error creating syslog writer for stderr: %v", err)
	}

	err = cmd.Run()
	if err != nil {
		if msg, ok := err.(*exec.ExitError); ok {
			os.Exit(msg.Sys().(syscall.WaitStatus).ExitStatus())
		} else {
			log.Printf("error running command: %v", err)
			os.Exit(1)
		}
	}

	os.Exit(0)
}
<commit_msg>Forward signals<commit_after>//  Copyright (c) 2014 Spotify AB.
//
//  Licensed to the Apache Software Foundation (ASF) under one
//  or more contributor license agreements.  See the NOTICE file
//  distributed with this work for additional information
//  regarding copyright ownership.  The ASF licenses this file
//  to you 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 (
	"bytes"
	"errors"
	"flag"
	"log"
	"os"
	"os/exec"
	"os/signal"
	"syscall"
)

type Syslogger struct {
	logger   *Writer
	stream   string
	buffer   *bytes.Buffer
	hostPort string
	priority Priority
	prefix   string
	logFlags int
	protocol string
}

func (s *Syslogger) Write(p []byte) (n int, err error) {
	if s.logger == nil {
		sl, err := Dial(s.protocol, s.hostPort, s.priority, s.prefix)
		if err != nil {
			// while syslog is down, dump the output
			return len(p), nil
		}
		s.logger = sl
	}
	for b := range p {
		s.buffer.WriteByte(p[b])
		if p[b] == 10 { // newline
			_, err := s.logger.Write(s.buffer.Bytes())
			if err != nil {
				s.logger = nil
			}
			s.buffer = bytes.NewBuffer([]byte{})
		}
	}
	return len(p), nil
}

func (s *Syslogger) Close() error {
	return nil
}

func NewSysLogger(stream, hostPort, prefix, protocol string) (*Syslogger, error) {
	var priority Priority
	if stream == "stderr" {
		priority = LOG_ERR | LOG_LOCAL0
	} else if stream == "stdout" {
		priority = LOG_INFO | LOG_LOCAL0
	} else {
		return nil, errors.New("cannot create syslogger for stream " + stream)
	}
	logFlags := 0

	return &Syslogger{nil, stream, bytes.NewBuffer([]byte{}), hostPort, priority, prefix, logFlags, protocol}, nil
}

func usage() {
	log.Printf("usage: %s -h syslog_host:port -n name -- executable [arg ...]", os.Args[0])
	flag.PrintDefaults()
	os.Exit(1)
}

func setupSignalHandlers(pid *exec.Cmd) {

	sigc := make(chan os.Signal, 1)

	sigarray := make([]os.Signal, 30)
	for i := 0; i < 30; i++ {
		sigarray[i] = syscall.Signal(i + 1)
	}

	signal.Notify(sigc, sigarray...)

	go func() {
		for {
			if pid.ProcessState != nil {
				break
			}
			s := <-sigc
			var ss syscall.Signal
			ss = s.(syscall.Signal)
			syscall.Kill(pid.Process.Pid, ss)
		}
	}()

}

func main() {
	log.SetFlags(0)

	flHostPort := flag.String("h", "", "Host port of where to connect to the syslog daemon")
	flLogName := flag.String("n", "", "Name to log as")
	flProtocol := flag.Bool("t", false, "use TCP instead of UDP (the default) for syslog communication")
	flag.Parse()

	if *flHostPort == "" {
		log.Printf("Must set the syslog host:port argument")
		usage()
	}

	if *flLogName == "" {
		log.Printf("Must set the syslog log name argument")
		usage()
	}

	protocol := "udp"
	if *flProtocol {
		protocol = "tcp"
	}

	//Example ./syslog-redirector -h 10.0.3.1:6514 -n test-ls-thingy -- \
	//            /bin/bash -c 'while true; do date; echo $SHELL; sleep 1; done'
	if len(os.Args) < 4 {
		log.Printf("at least 3 arguments required")
		usage()
	}
	hostPort := *flHostPort
	name := *flLogName

	if len(flag.Args()) == 0 {
		log.Printf("must supply a command")
		usage()
	}

	cmdName := flag.Args()[0]
	cmdArgs := flag.Args()[1:]

	var err error

	path, err := exec.LookPath(cmdName)
	if err != nil {
		log.Printf("Unable to locate %v", cmdName)
		os.Exit(127)
	}

	cmd := exec.Command(path, cmdArgs...)

	// TODO (dano): tolerate syslog downtime by reconnecting

	cmd.Stdout, err = NewSysLogger("stdout", hostPort, name, protocol)
	if err != nil {
		log.Printf("error creating syslog writer for stdout: %v", err)
	}

	cmd.Stderr, err = NewSysLogger("stderr", hostPort, name, protocol)
	if err != nil {
		log.Printf("error creating syslog writer for stderr: %v", err)
	}

	err = cmd.Start()

	if err != nil {
		if msg, ok := err.(*exec.ExitError); ok {
			os.Exit(msg.Sys().(syscall.WaitStatus).ExitStatus())
		} else {
			log.Printf("error running command: %v", err)
			os.Exit(1)
		}
	}

	setupSignalHandlers(cmd)
	err = cmd.Wait()

	os.Exit(0)
}
<|endoftext|>