<commit_before>package main

import (
	"encoding/json"
	"gopkg.in/libgit2/git2go.v22"
	"strings"
)

const (
	commitsPageSize        = 30 // The max page size for commits.
	branchesAllocationSize = 10 // The initial allocation size for branches.
	patchIndexLength       = 40 // The patch index length.
)

// GitRepository extends RepositoryInfo and is meant to represent a Git
// Repository.
type GitRepository struct {
	RepositoryInfo
}

// GetName is a Repository implementation that returns the name of the
// GitRepository.
func (repo *GitRepository) GetName() string {
	return repo.Name
}

// GetPath is a Repository implementation that returns the path of the
// GitRepository.
func (repo *GitRepository) GetPath() string {
	return repo.Path
}

// GetFile is a Repository implementation that returns the contents of a file
// in the GitRepository based on the file revision sha. On success, it returns
// the file contents in a byte array. On failure, the error will be returned.
func (repo *GitRepository) GetFile(id string) ([]byte, error) {
	gitRepo, err := git.OpenRepository(repo.Path)
	if err != nil {
		return nil, err
	}

	oid, err := git.NewOid(id)
	if err != nil {
		return nil, err
	}

	blob, err := gitRepo.LookupBlob(oid)
	if err != nil {
		return nil, err
	}

	return blob.Contents(), nil
}

// GetFileByCommit is a Repository implementation that returns the contents of
// a file in the GitRepository based on a commit sha and the file path. On
// success, it returns the file contents in a byte array. On failure, the error
// will be returned.
func (repo *GitRepository) GetFileByCommit(commit,
	filepath string) ([]byte, error) {
	gitRepo, err := git.OpenRepository(repo.Path)
	if err != nil {
		return nil, err
	}

	oid, err := git.NewOid(commit)
	if err != nil {
		return nil, err
	}

	c, err := gitRepo.LookupCommit(oid)
	if err != nil {
		return nil, err
	}

	tree, err := c.Tree()
	if err != nil {
		return nil, err
	}

	file, err := tree.EntryByPath(filepath)
	if err != nil {
		return nil, err
	}

	blob, err := gitRepo.LookupBlob(file.Id)
	if err != nil {
		return nil, err
	}

	return blob.Contents(), nil
}

// FileExists is a Repository implementation that returns whether a file exists
// in the GitRepository based on the file revision sha. It returns true if the
// file exists, false otherwise. On failure, the error will also be returned.
func (repo *GitRepository) FileExists(id string) (bool, error) {
	gitRepo, err := git.OpenRepository(repo.Path)
	if err != nil {
		return false, err
	}

	oid, err := git.NewOid(id)
	if err != nil {
		return false, err
	}

	if _, err := gitRepo.Lookup(oid); err != nil {
		return false, nil
	}

	return true, nil
}

// FileExistsByCommit is a Repository implementation that returns whether a
// file exists in the GitRepository based on a commit sha and the file path.
// It returns true if the file exists, false otherwise. On failure, the error
// will also be returned.
func (repo *GitRepository) FileExistsByCommit(commit,
	filepath string) (bool, error) {
	gitRepo, err := git.OpenRepository(repo.Path)
	if err != nil {
		return false, err
	}

	oid, err := git.NewOid(commit)
	if err != nil {
		return false, err
	}

	c, err := gitRepo.LookupCommit(oid)
	if err != nil {
		return false, err
	}

	tree, err := c.Tree()
	if err != nil {
		return false, err
	}

	file, err := tree.EntryByPath(filepath)
	if err != nil {
		return false, err
	}

	if _, err := gitRepo.Lookup(file.Id); err != nil {
		return false, nil
	}

	return true, nil
}

// GetBranches is a Repository implementation that returns all the branches in
// the repository. It returns a JSON representation of the branches containing
// the branch name and sha id. On failure, the error will also be returned.
//
// The JSON returned has the following format:
// [
//   {
//     "name": master,
//     "id": "1b6f00c0fe975dd12251431ed2ea561e0acc6d44"
//   }
// ]
func (repo *GitRepository) GetBranches() ([]byte, error) {
	type GitBranch struct {
		Name string `json:"name"`
		Id   string `json:"id"`
	}

	var branches []GitBranch = make([]GitBranch, 0, branchesAllocationSize)

	gitRepo, err := git.OpenRepository(repo.Path)
	if err != nil {
		return nil, err
	}

	iter, err := gitRepo.NewReferenceIterator()
	if err != nil {
		return nil, err
	}

	ref, err := iter.Next()
	for err == nil {
		if ref.IsBranch() {
			name := strings.Split(ref.Name(), "refs/heads/")[1]
			id := ref.Target().String()

			branches = append(branches, GitBranch{name, id})
		}
		ref, err = iter.Next()
	}

	json, err := json.Marshal(branches)
	if err != nil {
		return nil, err
	}

	return json, nil
}

// GetCommits is a Repository implementation that returns all the commits in
// the repository for the specified branch. It also takes an optional start
// commit sha, which will return all commits starting from the start commit
// sha instead. The returned JSON representation of the commits contains the
// author's name, the sha id, the date, the commit message, and the parent sha.
// On failure, the error will also be returned.
//
// The JSON returned has the following format:
// [
//   {
//     "author": "John Doe",
//     "id": "1b6f00c0fe975dd12251431ed2ea561e0acc6d44",
//     "date": "2015-06-27T05:51:39-07:00",
//     "message": "Add README.md",
//     "parent_id": "bfdde95432b3af879af969bd2377dc3e55ee46e6"
//   }
// ]
func (repo *GitRepository) GetCommits(branch string,
	start string) ([]byte, error) {
	type GitCommit struct {
		Author   string `json:"author"`
		Id       string `json:"id"`
		Date     string `json:"date"`
		Message  string `json:"message"`
		ParentId string `json:"parent_id"`
	}

	var commits []GitCommit = make([]GitCommit, 0, commitsPageSize)

	gitRepo, err := git.OpenRepository(repo.Path)
	if err != nil {
		return nil, err
	}

	revWalk, err := gitRepo.Walk()
	if err != nil {
		return nil, err
	}

	revWalk.Sorting(git.SortTopological | git.SortTime)

	// First try to look up the branch by its sha. If this fails, attempt to
	// get the branch by name.
	oid, err := git.NewOid(branch)
	if err != nil {
		gitBranch, err := gitRepo.LookupBranch(branch, git.BranchLocal)
		if err != nil {
			return nil, err
		}
		oid = gitBranch.Target()
		branch = gitBranch.Target().String()
	}

	if len(start) == 0 {
		start = branch
	}

	startOid, err := git.NewOid(start)
	if err != nil {
		return nil, err
	}
	revWalk.Push(startOid)

	err = revWalk.HideGlob("tags/*")
	if err != nil {
		return nil, err
	}

	for revWalk.Next(oid) == nil {
		commit, err := gitRepo.LookupCommit(oid)
		if err != nil {
			return nil, err
		}

		var parent string
		if commit.ParentCount() > 0 {
			parent = commit.Parent(0).Id().String()
		}

		gitCommit := GitCommit{
			commit.Author().Name,
			commit.Id().String(),
			commit.Author().When.String(),
			commit.Message(),
			parent,
		}

		commits = append(commits, gitCommit)

		if len(commits) == commitsPageSize {
			// We only want to return at max one page of commits.
			break
		}
	}

	json, err := json.Marshal(commits)
	if err != nil {
		return nil, err
	}

	return json, nil
}

// GetCommit is a Repository implementation that returns the commit information
// in the repository for the specified commit id. The returned JSON
// representation of the commit contains the author's name, the sha id, the
// date, the commit message, the parent sha, and the diff. On failure, the
// error will also be returned.
//
// The JSON returned has the following format:
// {
//   "author": "John Doe",
//   "id": "1b6f00c0fe975dd12251431ed2ea561e0acc6d44",
//   "date": "2015-06-27T05:51:39-07:00",
//   "message": "Add README.md",
//   "parent_id": "bfdde95432b3af879af969bd2377dc3e55ee46e6",
//   "diff": "diff --git a/test b/test
//            index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..044f599c9a720fe1a7d02e694a8dab492cbda8f0 100644
//            --- a/test
//            +++ b/test
//            @@ -1 +1,3 @@
//            test
//            +
//            +test"
// }
func (repo *GitRepository) GetCommit(commitId string) ([]byte, error) {
	type GitCommit struct {
		Author   string `json:"author"`
		Id       string `json:"id"`
		Date     string `json:"date"`
		Message  string `json:"message"`
		ParentId string `json:"parent_id"`
		Diff     string `json:"diff"`
	}

	gitRepo, err := git.OpenRepository(repo.Path)
	if err != nil {
		return nil, err
	}

	commitOid, err := git.NewOid(commitId)
	if err != nil {
		return nil, err
	}

	commit, err := gitRepo.LookupCommit(commitOid)
	if err != nil {
		return nil, err
	}

	var parent string
	var diff string

	commitTree, err := commit.Tree()
	if err != nil {
		return nil, err
	}

	options, err := git.DefaultDiffOptions()
	if err != nil {
		return nil, err
	}

	// Specifying full patch indices.
	options.IdAbbrev = patchIndexLength

	var parentTree *git.Tree
	if commit.ParentCount() > 0 {
		parent = commit.Parent(0).Id().String()
		parentTree, err = commit.Parent(0).Tree()
		if err != nil {
			return nil, err
		}
	}

	gitDiff, err := gitRepo.DiffTreeToTree(parentTree, commitTree, &options)
	if err != nil {
		return nil, err
	}

	deltas, err := gitDiff.NumDeltas()
	if err != nil {
		return nil, err
	}

	if deltas > 0 {
		patch, err := gitDiff.Patch(0)
		if err != nil {
			return nil, err
		}

		patchString, err := patch.String()
		if err != nil {
			return nil, err
		}

		diff = patchString
		patch.Free()
	}

	gitCommit := GitCommit{
		commit.Author().Name,
		commit.Id().String(),
		commit.Author().When.String(),
		commit.Message(),
		parent,
		diff,
	}

	json, err := json.Marshal(gitCommit)
	if err != nil {
		return nil, err
	}

	return json, nil
}
<commit_msg>Retrieve a maximum of 20 commits from Git repositories<commit_after>package main

import (
	"encoding/json"
	"gopkg.in/libgit2/git2go.v22"
	"strings"
)

const (
	commitsPageSize        = 20 // The max page size for commits.
	branchesAllocationSize = 10 // The initial allocation size for branches.
	patchIndexLength       = 40 // The patch index length.
)

// GitRepository extends RepositoryInfo and is meant to represent a Git
// Repository.
type GitRepository struct {
	RepositoryInfo
}

// GetName is a Repository implementation that returns the name of the
// GitRepository.
func (repo *GitRepository) GetName() string {
	return repo.Name
}

// GetPath is a Repository implementation that returns the path of the
// GitRepository.
func (repo *GitRepository) GetPath() string {
	return repo.Path
}

// GetFile is a Repository implementation that returns the contents of a file
// in the GitRepository based on the file revision sha. On success, it returns
// the file contents in a byte array. On failure, the error will be returned.
func (repo *GitRepository) GetFile(id string) ([]byte, error) {
	gitRepo, err := git.OpenRepository(repo.Path)
	if err != nil {
		return nil, err
	}

	oid, err := git.NewOid(id)
	if err != nil {
		return nil, err
	}

	blob, err := gitRepo.LookupBlob(oid)
	if err != nil {
		return nil, err
	}

	return blob.Contents(), nil
}

// GetFileByCommit is a Repository implementation that returns the contents of
// a file in the GitRepository based on a commit sha and the file path. On
// success, it returns the file contents in a byte array. On failure, the error
// will be returned.
func (repo *GitRepository) GetFileByCommit(commit,
	filepath string) ([]byte, error) {
	gitRepo, err := git.OpenRepository(repo.Path)
	if err != nil {
		return nil, err
	}

	oid, err := git.NewOid(commit)
	if err != nil {
		return nil, err
	}

	c, err := gitRepo.LookupCommit(oid)
	if err != nil {
		return nil, err
	}

	tree, err := c.Tree()
	if err != nil {
		return nil, err
	}

	file, err := tree.EntryByPath(filepath)
	if err != nil {
		return nil, err
	}

	blob, err := gitRepo.LookupBlob(file.Id)
	if err != nil {
		return nil, err
	}

	return blob.Contents(), nil
}

// FileExists is a Repository implementation that returns whether a file exists
// in the GitRepository based on the file revision sha. It returns true if the
// file exists, false otherwise. On failure, the error will also be returned.
func (repo *GitRepository) FileExists(id string) (bool, error) {
	gitRepo, err := git.OpenRepository(repo.Path)
	if err != nil {
		return false, err
	}

	oid, err := git.NewOid(id)
	if err != nil {
		return false, err
	}

	if _, err := gitRepo.Lookup(oid); err != nil {
		return false, nil
	}

	return true, nil
}

// FileExistsByCommit is a Repository implementation that returns whether a
// file exists in the GitRepository based on a commit sha and the file path.
// It returns true if the file exists, false otherwise. On failure, the error
// will also be returned.
func (repo *GitRepository) FileExistsByCommit(commit,
	filepath string) (bool, error) {
	gitRepo, err := git.OpenRepository(repo.Path)
	if err != nil {
		return false, err
	}

	oid, err := git.NewOid(commit)
	if err != nil {
		return false, err
	}

	c, err := gitRepo.LookupCommit(oid)
	if err != nil {
		return false, err
	}

	tree, err := c.Tree()
	if err != nil {
		return false, err
	}

	file, err := tree.EntryByPath(filepath)
	if err != nil {
		return false, err
	}

	if _, err := gitRepo.Lookup(file.Id); err != nil {
		return false, nil
	}

	return true, nil
}

// GetBranches is a Repository implementation that returns all the branches in
// the repository. It returns a JSON representation of the branches containing
// the branch name and sha id. On failure, the error will also be returned.
//
// The JSON returned has the following format:
// [
//   {
//     "name": master,
//     "id": "1b6f00c0fe975dd12251431ed2ea561e0acc6d44"
//   }
// ]
func (repo *GitRepository) GetBranches() ([]byte, error) {
	type GitBranch struct {
		Name string `json:"name"`
		Id   string `json:"id"`
	}

	var branches []GitBranch = make([]GitBranch, 0, branchesAllocationSize)

	gitRepo, err := git.OpenRepository(repo.Path)
	if err != nil {
		return nil, err
	}

	iter, err := gitRepo.NewReferenceIterator()
	if err != nil {
		return nil, err
	}

	ref, err := iter.Next()
	for err == nil {
		if ref.IsBranch() {
			name := strings.Split(ref.Name(), "refs/heads/")[1]
			id := ref.Target().String()

			branches = append(branches, GitBranch{name, id})
		}
		ref, err = iter.Next()
	}

	json, err := json.Marshal(branches)
	if err != nil {
		return nil, err
	}

	return json, nil
}

// GetCommits is a Repository implementation that returns all the commits in
// the repository for the specified branch. It also takes an optional start
// commit sha, which will return all commits starting from the start commit
// sha instead. The returned JSON representation of the commits contains the
// author's name, the sha id, the date, the commit message, and the parent sha.
// On failure, the error will also be returned.
//
// The JSON returned has the following format:
// [
//   {
//     "author": "John Doe",
//     "id": "1b6f00c0fe975dd12251431ed2ea561e0acc6d44",
//     "date": "2015-06-27T05:51:39-07:00",
//     "message": "Add README.md",
//     "parent_id": "bfdde95432b3af879af969bd2377dc3e55ee46e6"
//   }
// ]
func (repo *GitRepository) GetCommits(branch string,
	start string) ([]byte, error) {
	type GitCommit struct {
		Author   string `json:"author"`
		Id       string `json:"id"`
		Date     string `json:"date"`
		Message  string `json:"message"`
		ParentId string `json:"parent_id"`
	}

	var commits []GitCommit = make([]GitCommit, 0, commitsPageSize)

	gitRepo, err := git.OpenRepository(repo.Path)
	if err != nil {
		return nil, err
	}

	revWalk, err := gitRepo.Walk()
	if err != nil {
		return nil, err
	}

	revWalk.Sorting(git.SortTopological | git.SortTime)

	// First try to look up the branch by its sha. If this fails, attempt to
	// get the branch by name.
	oid, err := git.NewOid(branch)
	if err != nil {
		gitBranch, err := gitRepo.LookupBranch(branch, git.BranchLocal)
		if err != nil {
			return nil, err
		}
		oid = gitBranch.Target()
		branch = gitBranch.Target().String()
	}

	if len(start) == 0 {
		start = branch
	}

	startOid, err := git.NewOid(start)
	if err != nil {
		return nil, err
	}
	revWalk.Push(startOid)

	err = revWalk.HideGlob("tags/*")
	if err != nil {
		return nil, err
	}

	for revWalk.Next(oid) == nil {
		commit, err := gitRepo.LookupCommit(oid)
		if err != nil {
			return nil, err
		}

		var parent string
		if commit.ParentCount() > 0 {
			parent = commit.Parent(0).Id().String()
		}

		gitCommit := GitCommit{
			commit.Author().Name,
			commit.Id().String(),
			commit.Author().When.String(),
			commit.Message(),
			parent,
		}

		commits = append(commits, gitCommit)

		if len(commits) == commitsPageSize {
			// We only want to return at max one page of commits.
			break
		}
	}

	json, err := json.Marshal(commits)
	if err != nil {
		return nil, err
	}

	return json, nil
}

// GetCommit is a Repository implementation that returns the commit information
// in the repository for the specified commit id. The returned JSON
// representation of the commit contains the author's name, the sha id, the
// date, the commit message, the parent sha, and the diff. On failure, the
// error will also be returned.
//
// The JSON returned has the following format:
// {
//   "author": "John Doe",
//   "id": "1b6f00c0fe975dd12251431ed2ea561e0acc6d44",
//   "date": "2015-06-27T05:51:39-07:00",
//   "message": "Add README.md",
//   "parent_id": "bfdde95432b3af879af969bd2377dc3e55ee46e6",
//   "diff": "diff --git a/test b/test
//            index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..044f599c9a720fe1a7d02e694a8dab492cbda8f0 100644
//            --- a/test
//            +++ b/test
//            @@ -1 +1,3 @@
//            test
//            +
//            +test"
// }
func (repo *GitRepository) GetCommit(commitId string) ([]byte, error) {
	type GitCommit struct {
		Author   string `json:"author"`
		Id       string `json:"id"`
		Date     string `json:"date"`
		Message  string `json:"message"`
		ParentId string `json:"parent_id"`
		Diff     string `json:"diff"`
	}

	gitRepo, err := git.OpenRepository(repo.Path)
	if err != nil {
		return nil, err
	}

	commitOid, err := git.NewOid(commitId)
	if err != nil {
		return nil, err
	}

	commit, err := gitRepo.LookupCommit(commitOid)
	if err != nil {
		return nil, err
	}

	var parent string
	var diff string

	commitTree, err := commit.Tree()
	if err != nil {
		return nil, err
	}

	options, err := git.DefaultDiffOptions()
	if err != nil {
		return nil, err
	}

	// Specifying full patch indices.
	options.IdAbbrev = patchIndexLength

	var parentTree *git.Tree
	if commit.ParentCount() > 0 {
		parent = commit.Parent(0).Id().String()
		parentTree, err = commit.Parent(0).Tree()
		if err != nil {
			return nil, err
		}
	}

	gitDiff, err := gitRepo.DiffTreeToTree(parentTree, commitTree, &options)
	if err != nil {
		return nil, err
	}

	deltas, err := gitDiff.NumDeltas()
	if err != nil {
		return nil, err
	}

	if deltas > 0 {
		patch, err := gitDiff.Patch(0)
		if err != nil {
			return nil, err
		}

		patchString, err := patch.String()
		if err != nil {
			return nil, err
		}

		diff = patchString
		patch.Free()
	}

	gitCommit := GitCommit{
		commit.Author().Name,
		commit.Id().String(),
		commit.Author().When.String(),
		commit.Message(),
		parent,
		diff,
	}

	json, err := json.Marshal(gitCommit)
	if err != nil {
		return nil, err
	}

	return json, nil
}
<|endoftext|>


<commit_before>package provider_test

import (
	"code.cloudfoundry.org/lager/lagertest"
	"github.com/concourse/atc/auth"
	"github.com/concourse/atc/auth/genericoauth"
	"github.com/concourse/atc/auth/github"
	"github.com/concourse/atc/auth/provider"
	"github.com/concourse/atc/auth/uaa"
	"github.com/concourse/atc/db"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("OAuthFactory", func() {
	var oauthFactory provider.OAuthFactory

	BeforeEach(func() {
		oauthFactory = provider.NewOAuthFactory(
			lagertest.NewTestLogger("test"),
			"http://foo.bar",
			auth.OAuthRoutes,
			auth.OAuthCallback,
		)
	})

	Describe("GetProvider", func() {
		Context("when asking for github provider", func() {
			Context("when github provider is setup", func() {
				It("returns back GitHub's auth provider", func() {
					provider, found, err := oauthFactory.GetProvider(db.SavedTeam{
						Team: db.Team{
							Name: "some-team",
							GitHubAuth: &db.GitHubAuth{
								ClientID:     "user1",
								ClientSecret: "password1",
								Users:        []string{"thecandyman"},
							},
						},
					}, github.ProviderName)
					Expect(err).NotTo(HaveOccurred())
					Expect(found).To(BeTrue())
					Expect(provider).NotTo(BeNil())
				})
			})

			Context("when github provider is not setup", func() {
				It("returns false", func() {
					_, found, err := oauthFactory.GetProvider(db.SavedTeam{
						Team: db.Team{
							Name: "some-team",
						},
					}, github.ProviderName)
					Expect(err).NotTo(HaveOccurred())
					Expect(found).To(BeFalse())
				})
			})
		})

		Context("when asking for uaa provider", func() {
			Context("when UAA provider is setup", func() {
				It("returns back UAA's auth provider", func() {
					provider, found, err := oauthFactory.GetProvider(db.SavedTeam{
						Team: db.Team{
							Name: "some-team",
							UAAAuth: &db.UAAAuth{
								ClientID:     "user1",
								ClientSecret: "password1",
							},
						},
					}, uaa.ProviderName)
					Expect(err).NotTo(HaveOccurred())
					Expect(found).To(BeTrue())
					Expect(provider).NotTo(BeNil())
				})
			})

			Context("when uaa provider is not setup", func() {
				It("returns false", func() {
					_, found, err := oauthFactory.GetProvider(db.SavedTeam{
						Team: db.Team{
							Name: "some-team",
						},
					}, uaa.ProviderName)
					Expect(err).NotTo(HaveOccurred())
					Expect(found).To(BeFalse())
				})
			})
		})

		Context("when asking for goa provider", func() {
			Context("when Generic OAuth provider is setup", func() {
				It("returns back GOA's auth provider", func() {
					provider, found, err := oauthFactory.GetProvider(db.SavedTeam{
						Team: db.Team{
							Name: "some-team",
							GenericOAuth: &db.GenericOAuth{
								ClientID:     "user1",
								ClientSecret: "password1",
							},
						},
					}, genericoauth.ProviderName)
					Expect(err).NotTo(HaveOccurred())
					Expect(found).To(BeTrue())
					Expect(provider).NotTo(BeNil())
				})
			})

			Context("when Generic OAuth provider is not setup", func() {
				It("returns false", func() {
					_, found, err := oauthFactory.GetProvider(db.SavedTeam{
						Team: db.Team{
							Name: "some-team",
						},
					}, genericoauth.ProviderName)
					Expect(err).NotTo(HaveOccurred())
					Expect(found).To(BeFalse())
				})
			})
		})

		Context("when asking for unknown provider", func() {
			It("returns false", func() {
				_, found, err := oauthFactory.GetProvider(db.SavedTeam{
					Team: db.Team{
						Name: "some-team",
					},
				}, "bogus")
				Expect(err).NotTo(HaveOccurred())
				Expect(found).To(BeFalse())
			})
		})
	})
})
<commit_msg>Added missing api setter and tests<commit_after>package provider_test

import (
	"code.cloudfoundry.org/lager/lagertest"
	"github.com/concourse/atc/auth"
	"github.com/concourse/atc/auth/genericoauth"
	"github.com/concourse/atc/auth/github"
	"github.com/concourse/atc/auth/provider"
	"github.com/concourse/atc/auth/uaa"
	"github.com/concourse/atc/db"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("OAuthFactory", func() {
	var oauthFactory provider.OAuthFactory

	BeforeEach(func() {
		oauthFactory = provider.NewOAuthFactory(
			lagertest.NewTestLogger("test"),
			"http://foo.bar",
			auth.OAuthRoutes,
			auth.OAuthCallback,
		)
	})

	Describe("GetProvider", func() {
		Context("when asking for github provider", func() {
			Context("when github provider is setup", func() {
				It("returns back GitHub's auth provider", func() {
					provider, found, err := oauthFactory.GetProvider(db.SavedTeam{
						Team: db.Team{
							Name: "some-team",
							GitHubAuth: &db.GitHubAuth{
								ClientID:     "user1",
								ClientSecret: "password1",
								Users:        []string{"thecandyman"},
							},
						},
					}, github.ProviderName)
					Expect(err).NotTo(HaveOccurred())
					Expect(found).To(BeTrue())
					Expect(provider).NotTo(BeNil())
				})
			})

			Context("when github provider is not setup", func() {
				It("returns false", func() {
					_, found, err := oauthFactory.GetProvider(db.SavedTeam{
						Team: db.Team{
							Name: "some-team",
						},
					}, github.ProviderName)
					Expect(err).NotTo(HaveOccurred())
					Expect(found).To(BeFalse())
				})
			})
		})

		Context("when asking for uaa provider", func() {
			Context("when UAA provider is setup", func() {
				It("returns back UAA's auth provider", func() {
					provider, found, err := oauthFactory.GetProvider(db.SavedTeam{
						Team: db.Team{
							Name: "some-team",
							UAAAuth: &db.UAAAuth{
								ClientID:     "user1",
								ClientSecret: "password1",
							},
						},
					}, uaa.ProviderName)
					Expect(err).NotTo(HaveOccurred())
					Expect(found).To(BeTrue())
					Expect(provider).NotTo(BeNil())
				})
			})

			Context("when uaa provider is not setup", func() {
				It("returns false", func() {
					_, found, err := oauthFactory.GetProvider(db.SavedTeam{
						Team: db.Team{
							Name: "some-team",
						},
					}, uaa.ProviderName)
					Expect(err).NotTo(HaveOccurred())
					Expect(found).To(BeFalse())
				})
			})
		})

		Context("when asking for goa provider", func() {
			Context("when Generic OAuth provider is setup", func() {
				It("returns back GOA's auth provider", func() {
					provider, found, err := oauthFactory.GetProvider(db.SavedTeam{
						Team: db.Team{
							Name: "some-team",
							GenericOAuth: &db.GenericOAuth{
								ClientID:     "user1",
								ClientSecret: "password1",
							},
						},
					}, genericoauth.ProviderName)
					Expect(err).NotTo(HaveOccurred())
					Expect(found).To(BeTrue())
					Expect(provider).NotTo(BeNil())
				})
			})

			Context("when asking for generic oauth", func() {
				Context("when Generic OAuth provider is setup", func() {
					It("returns back GOA's auth provider", func() {
						providers, err := oauthFactory.GetProviders(db.SavedTeam{
							Team: db.Team{
								Name: "some-team",
								GenericOAuth: &db.GenericOAuth{
									ClientID:     "user1",
									ClientSecret: "password1",
								},
							},
						})
						Expect(err).NotTo(HaveOccurred())
						Expect(providers).To(HaveLen(1))
						Expect(providers[genericoauth.ProviderName]).NotTo(BeNil())
					})
				})

				Context("when Generic OAuth provider is not setup", func() {
					It("returns false", func() {
						_, found, err := oauthFactory.GetProvider(db.SavedTeam{
							Team: db.Team{
								Name: "some-team",
							},
						}, genericoauth.ProviderName)
						Expect(err).NotTo(HaveOccurred())
						Expect(found).To(BeFalse())
					})
				})
			})
		})

		Context("when asking for unknown provider", func() {
			It("returns false", func() {
				_, found, err := oauthFactory.GetProvider(db.SavedTeam{
					Team: db.Team{
						Name: "some-team",
					},
				}, "bogus")
				Expect(err).NotTo(HaveOccurred())
				Expect(found).To(BeFalse())
			})
		})
	})
})
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"net/http"
	"os"
	"os/exec"
	"strings"
	"syscall"
	"time"
)

type Service struct {
	Name       string `json:"name"`
	BridgeName string `json:"bridge-name"`
	BridgeIP   string `json:"bridge-ip"`
	ServicePid int
}

type Container struct {
	Name        string `json:"name"`
	ServiceName string `json:"service-name"`
	Command     string `json:"command"`
	Pid         int
	IP          string
	StartTime   time.Time
}

var services map[string]Service
var containers []Containter

func init() {
	services = make(map[string]Service)
}

func main() {
	http.HandleFunc("/api/v1/service/add", service_add)
	http.HandleFunc("/api/v1/container/run", container_run)
	err := http.ListenAndServe(":3000", nil)
	if err != nil {
		fmt.Println(err)
	}
}

func service_add(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" {
		http.Error(w, "Invalid Request!", http.StatusMethodNotAllowed)
		return
	}

	var s Service
	if err := json.NewDecoder(r.Body).Decode(&s); err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	if _, ok := services[s.Name]; ok {
		http.Error(w, "Service already exists", http.StatusInternalServerError)
		return
	}

	if err := service_create_network(s); err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	w.WriteHeader(http.StatusCreated)
}

func container_run(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" {
		http.Error(w, "Invalid Request!", http.StatusMethodNotAllowed)
		return
	}

	var c Container
	if err := json.NewDecoder(r.Body).Decode(&c); err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	if _, ok := services[c.ServiceName]; ok == false {
		http.Error(w, "Service does not exists", http.StatusInternalServerError)
		return
	}

	go run(c)

	w.WriteHeader(http.StatusCreated)
}

func service_create_network(s Service) error {
	create_bridge := strings.Split(fmt.Sprintf("/sbin/ip link add name %s type bridge", s.BridgeName), " ")
	set_bridge_up := strings.Split(fmt.Sprintf("/sbin/ip link set %s up", s.BridgeName), " ")
	set_bridge_ip := strings.Split(fmt.Sprintf("/sbin/ifconfig %s %s", s.BridgeName, s.BridgeIP), " ")

	cmd1 := exec.Command(create_bridge[0], create_bridge[1:]...)
	err := cmd1.Run()
	if err != nil {
		return err
	}

	cmd2 := exec.Command(set_bridge_up[0], set_bridge_up[1:]...)
	err = cmd2.Run()
	if err != nil {
		return err
	}

	cmd3 := exec.Command(set_bridge_ip[0], set_bridge_ip[1:]...)
	err = cmd3.Run()
	if err != nil {
		return err
	}

	services[s.Name] = s
	return nil
}

func run(c Container) {
	fmt.Println("running parent")
	runcmd := "/home/yup/p/containers/brocker-run/brocker-run"

	cmd := &exec.Cmd{
		Path: runcmd,
		Args: append([]string{runcmd}, c.Command),
	}
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	cmd.SysProcAttr = &syscall.SysProcAttr{
		Cloneflags: syscall.CLONE_NEWPID |
			syscall.CLONE_NEWNET,
	}

	if err := cmd.Start(); err != nil {
		fmt.Println(err)
	}

	c.Pid = cmd.Process.Pid
	c.StartTime = time.Now()
	containers = append(containers, c)
	fmt.Println(cmd.Process.Pid)

	cmd.Wait()
}
<commit_msg>Fixed spelling mistake<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"net/http"
	"os"
	"os/exec"
	"strings"
	"syscall"
	"time"
)

type Service struct {
	Name       string `json:"name"`
	BridgeName string `json:"bridge-name"`
	BridgeIP   string `json:"bridge-ip"`
	ServicePid int
}

type Container struct {
	Name        string `json:"name"`
	ServiceName string `json:"service-name"`
	Command     string `json:"command"`
	Pid         int
	IP          string
	StartTime   time.Time
}

var services map[string]Service
var containers []Container

func init() {
	services = make(map[string]Service)
}

func main() {
	http.HandleFunc("/api/v1/service/add", service_add)
	http.HandleFunc("/api/v1/container/run", container_run)
	err := http.ListenAndServe(":3000", nil)
	if err != nil {
		fmt.Println(err)
	}
}

func service_add(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" {
		http.Error(w, "Invalid Request!", http.StatusMethodNotAllowed)
		return
	}

	var s Service
	if err := json.NewDecoder(r.Body).Decode(&s); err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	if _, ok := services[s.Name]; ok {
		http.Error(w, "Service already exists", http.StatusInternalServerError)
		return
	}

	if err := service_create_network(s); err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	w.WriteHeader(http.StatusCreated)
}

func container_run(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" {
		http.Error(w, "Invalid Request!", http.StatusMethodNotAllowed)
		return
	}

	var c Container
	if err := json.NewDecoder(r.Body).Decode(&c); err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	if _, ok := services[c.ServiceName]; ok == false {
		http.Error(w, "Service does not exists", http.StatusInternalServerError)
		return
	}

	go run(c)

	w.WriteHeader(http.StatusCreated)
}

func service_create_network(s Service) error {
	create_bridge := strings.Split(fmt.Sprintf("/sbin/ip link add name %s type bridge", s.BridgeName), " ")
	set_bridge_up := strings.Split(fmt.Sprintf("/sbin/ip link set %s up", s.BridgeName), " ")
	set_bridge_ip := strings.Split(fmt.Sprintf("/sbin/ifconfig %s %s", s.BridgeName, s.BridgeIP), " ")

	cmd1 := exec.Command(create_bridge[0], create_bridge[1:]...)
	err := cmd1.Run()
	if err != nil {
		return err
	}

	cmd2 := exec.Command(set_bridge_up[0], set_bridge_up[1:]...)
	err = cmd2.Run()
	if err != nil {
		return err
	}

	cmd3 := exec.Command(set_bridge_ip[0], set_bridge_ip[1:]...)
	err = cmd3.Run()
	if err != nil {
		return err
	}

	services[s.Name] = s
	return nil
}

func run(c Container) {
	fmt.Println("running parent")
	runcmd := "/home/yup/p/containers/brocker-run/brocker-run"

	cmd := &exec.Cmd{
		Path: runcmd,
		Args: append([]string{runcmd}, c.Command),
	}
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	cmd.SysProcAttr = &syscall.SysProcAttr{
		Cloneflags: syscall.CLONE_NEWPID |
			syscall.CLONE_NEWNET,
	}

	if err := cmd.Start(); err != nil {
		fmt.Println(err)
	}

	c.Pid = cmd.Process.Pid
	c.StartTime = time.Now()
	containers = append(containers, c)
	fmt.Println(cmd.Process.Pid)

	cmd.Wait()
}
<|endoftext|>


<commit_before>/*
   Copyright 2016 GitHub Inc.
	 See https://github.com/github/gh-ost/blob/master/LICENSE
*/

package logic

import (
	"fmt"
	"os"
	"os/exec"
	"path/filepath"

	"github.com/github/gh-ost/go/base"
	"github.com/openark/golib/log"
)

const (
	onStartup            = "gh-ost-on-startup"
	onValidated          = "gh-ost-on-validated"
	onBeforeRowCopy      = "gh-ost-on-before-row-copy"
	onRowCopyComplete    = "gh-ost-on-row-copy-complete"
	onBeginPostponed     = "gh-ost-on-begin-postponed"
	onBeforeCutOver      = "gh-ost-on-before-cut-over"
	onInteractiveCommand = "gh-ost-on-interactive-command"
	onSuccess            = "gh-ost-on-success"
	onFailure            = "gh-ost-on-failure"
	onStatus             = "gh-ost-on-status"
	onStopReplication    = "gh-ost-on-stop-replication"
)

type HooksExecutor struct {
	migrationContext *base.MigrationContext
}

func NewHooksExecutor() *HooksExecutor {
	return &HooksExecutor{
		migrationContext: base.GetMigrationContext(),
	}
}

func (this *HooksExecutor) initHooks() error {
	return nil
}

func (this *HooksExecutor) applyEnvironmentVairables(extraVariables ...string) []string {
	env := os.Environ()
	env = append(env, fmt.Sprintf("GH_OST_DATABASE_NAME=%s", this.migrationContext.DatabaseName))
	env = append(env, fmt.Sprintf("GH_OST_TABLE_NAME=%s", this.migrationContext.OriginalTableName))
	env = append(env, fmt.Sprintf("GH_OST_GHOST_TABLE_NAME=%s", this.migrationContext.GetGhostTableName()))
	env = append(env, fmt.Sprintf("GH_OST_OLD_TABLE_NAME=%s", this.migrationContext.GetOldTableName()))
	env = append(env, fmt.Sprintf("GH_OST_DDL=%s", this.migrationContext.AlterStatement))
	env = append(env, fmt.Sprintf("GH_OST_ELAPSED_SECONDS=%f", this.migrationContext.ElapsedTime().Seconds()))
	env = append(env, fmt.Sprintf("GH_OST_MIGRATED_HOST=%s", this.migrationContext.ApplierConnectionConfig.ImpliedKey.Hostname))
	env = append(env, fmt.Sprintf("GH_OST_INSPECTED_HOST=%s", this.migrationContext.InspectorConnectionConfig.ImpliedKey.Hostname))
	env = append(env, fmt.Sprintf("GH_OST_EXECUTING_HOST=%s", this.migrationContext.Hostname))

	for _, variable := range extraVariables {
		env = append(env, variable)
	}
	return env
}

// executeHook executes a command, and sets relevant environment variables
func (this *HooksExecutor) executeHook(hook string, extraVariables ...string) error {
	cmd := exec.Command(hook)
	cmd.Env = this.applyEnvironmentVairables(extraVariables...)

	if err := cmd.Run(); err != nil {
		return log.Errore(err)
	}
	return nil
}

func (this *HooksExecutor) detectHooks(baseName string) (hooks []string, err error) {
	if this.migrationContext.HooksPath == "" {
		return hooks, err
	}
	pattern := fmt.Sprintf("%s/%s*", this.migrationContext.HooksPath, baseName)
	hooks, err = filepath.Glob(pattern)
	return hooks, err
}

func (this *HooksExecutor) executeHooks(baseName string, extraVariables ...string) error {
	hooks, err := this.detectHooks(baseName)
	if err != nil {
		return err
	}
	for _, hook := range hooks {
		if err := this.executeHook(hook, extraVariables...); err != nil {
			return err
		}
	}
	return nil
}

func (this *HooksExecutor) onStartup() error {
	return this.executeHooks(onStartup)
}

func (this *HooksExecutor) onValidated() error {
	return this.executeHooks(onValidated)
}

func (this *HooksExecutor) onBeforeRowCopy() error {
	return this.executeHooks(onBeforeRowCopy)
}

func (this *HooksExecutor) onRowCopyComplete() error {
	return this.executeHooks(onRowCopyComplete)
}

func (this *HooksExecutor) onBeginPostponed() error {
	return this.executeHooks(onBeginPostponed)
}

func (this *HooksExecutor) onBeforeCutOver() error {
	return this.executeHooks(onBeforeCutOver)
}

func (this *HooksExecutor) onInteractiveCommand(command string) error {
	v := fmt.Sprintf("GH_OST_COMMAND='%s'", command)
	return this.executeHooks(onInteractiveCommand, v)
}

func (this *HooksExecutor) onSuccess() error {
	return this.executeHooks(onSuccess)
}

func (this *HooksExecutor) onFailure() error {
	return this.executeHooks(onFailure)
}

func (this *HooksExecutor) onStatus(statusMessage string) error {
	v := fmt.Sprintf("GH_OST_STATUS='%s'", statusMessage)
	return this.executeHooks(onStatus, v)
}

func (this *HooksExecutor) onStopReplication() error {
	return this.executeHooks(onStopReplication)
}
<commit_msg>added onRowCountComplete(); supporting elapsedSeconds in on-status<commit_after>/*
   Copyright 2016 GitHub Inc.
	 See https://github.com/github/gh-ost/blob/master/LICENSE
*/

package logic

import (
	"fmt"
	"os"
	"os/exec"
	"path/filepath"

	"github.com/github/gh-ost/go/base"
	"github.com/openark/golib/log"
)

const (
	onStartup            = "gh-ost-on-startup"
	onValidated          = "gh-ost-on-validated"
	onRowCountComplete   = "gh-ost-on-rowcount-complete"
	onBeforeRowCopy      = "gh-ost-on-before-row-copy"
	onRowCopyComplete    = "gh-ost-on-row-copy-complete"
	onBeginPostponed     = "gh-ost-on-begin-postponed"
	onBeforeCutOver      = "gh-ost-on-before-cut-over"
	onInteractiveCommand = "gh-ost-on-interactive-command"
	onSuccess            = "gh-ost-on-success"
	onFailure            = "gh-ost-on-failure"
	onStatus             = "gh-ost-on-status"
	onStopReplication    = "gh-ost-on-stop-replication"
)

type HooksExecutor struct {
	migrationContext *base.MigrationContext
}

func NewHooksExecutor() *HooksExecutor {
	return &HooksExecutor{
		migrationContext: base.GetMigrationContext(),
	}
}

func (this *HooksExecutor) initHooks() error {
	return nil
}

func (this *HooksExecutor) applyEnvironmentVairables(extraVariables ...string) []string {
	env := os.Environ()
	env = append(env, fmt.Sprintf("GH_OST_DATABASE_NAME=%s", this.migrationContext.DatabaseName))
	env = append(env, fmt.Sprintf("GH_OST_TABLE_NAME=%s", this.migrationContext.OriginalTableName))
	env = append(env, fmt.Sprintf("GH_OST_GHOST_TABLE_NAME=%s", this.migrationContext.GetGhostTableName()))
	env = append(env, fmt.Sprintf("GH_OST_OLD_TABLE_NAME=%s", this.migrationContext.GetOldTableName()))
	env = append(env, fmt.Sprintf("GH_OST_DDL=%s", this.migrationContext.AlterStatement))
	env = append(env, fmt.Sprintf("GH_OST_ELAPSED_SECONDS=%f", this.migrationContext.ElapsedTime().Seconds()))
	env = append(env, fmt.Sprintf("GH_OST_MIGRATED_HOST=%s", this.migrationContext.ApplierConnectionConfig.ImpliedKey.Hostname))
	env = append(env, fmt.Sprintf("GH_OST_INSPECTED_HOST=%s", this.migrationContext.InspectorConnectionConfig.ImpliedKey.Hostname))
	env = append(env, fmt.Sprintf("GH_OST_EXECUTING_HOST=%s", this.migrationContext.Hostname))
	env = append(env, fmt.Sprintf("GH_OST_HOOKS_HINT=%s", this.migrationContext.HooksHintMessage))

	for _, variable := range extraVariables {
		env = append(env, variable)
	}
	return env
}

// executeHook executes a command, and sets relevant environment variables
func (this *HooksExecutor) executeHook(hook string, extraVariables ...string) error {
	cmd := exec.Command(hook)
	cmd.Env = this.applyEnvironmentVairables(extraVariables...)

	if err := cmd.Run(); err != nil {
		return log.Errore(err)
	}
	return nil
}

func (this *HooksExecutor) detectHooks(baseName string) (hooks []string, err error) {
	if this.migrationContext.HooksPath == "" {
		return hooks, err
	}
	pattern := fmt.Sprintf("%s/%s*", this.migrationContext.HooksPath, baseName)
	hooks, err = filepath.Glob(pattern)
	return hooks, err
}

func (this *HooksExecutor) executeHooks(baseName string, extraVariables ...string) error {
	hooks, err := this.detectHooks(baseName)
	if err != nil {
		return err
	}
	for _, hook := range hooks {
		if err := this.executeHook(hook, extraVariables...); err != nil {
			return err
		}
	}
	return nil
}

func (this *HooksExecutor) onStartup() error {
	return this.executeHooks(onStartup)
}

func (this *HooksExecutor) onValidated() error {
	return this.executeHooks(onValidated)
}

func (this *HooksExecutor) onRowCountComplete() error {
	return this.executeHooks(onRowCountComplete)
}
func (this *HooksExecutor) onBeforeRowCopy() error {
	return this.executeHooks(onBeforeRowCopy)
}

func (this *HooksExecutor) onRowCopyComplete() error {
	return this.executeHooks(onRowCopyComplete)
}

func (this *HooksExecutor) onBeginPostponed() error {
	return this.executeHooks(onBeginPostponed)
}

func (this *HooksExecutor) onBeforeCutOver() error {
	return this.executeHooks(onBeforeCutOver)
}

func (this *HooksExecutor) onInteractiveCommand(command string) error {
	v := fmt.Sprintf("GH_OST_COMMAND='%s'", command)
	return this.executeHooks(onInteractiveCommand, v)
}

func (this *HooksExecutor) onSuccess() error {
	return this.executeHooks(onSuccess)
}

func (this *HooksExecutor) onFailure() error {
	return this.executeHooks(onFailure)
}

func (this *HooksExecutor) onStatus(statusMessage string, elapsedSeconds int64) error {
	v0 := fmt.Sprintf("GH_OST_STATUS='%s'", statusMessage)
	v1 := fmt.Sprintf("GH_OST_ELAPSED_SECONDS='%d'", elapsedSeconds)
	return this.executeHooks(onStatus, v0, v1)
}

func (this *HooksExecutor) onStopReplication() error {
	return this.executeHooks(onStopReplication)
}
<|endoftext|>


<commit_before>package govsphereutils

import (
	"github.com/vmware/govmomi"
	"github.com/vmware/govmomi/vim25/mo"
	"github.com/vmware/govmomi/vim25/types"
)

type Searcher interface {
	Compare(*govmomi.Client, govmomi.Reference) bool
}

type nameSearcher struct {
	Name string
	Type string
}

func (v nameSearcher) Compare(c *govmomi.Client, ref govmomi.Reference) bool {
	// Check for VM Type
	if ref.Reference().Type != v.Type {
		return false
	}

	switch ref.Reference().Type {
	case "VirtualMachine":
		// Check for name
		var vm mo.VirtualMachine
		c.Properties(ref.Reference(), nil, &vm)

		if vm.Name == v.Name {
			return true
		}
	case "HostSystem":
		var h mo.HostSystem
		c.Properties(ref.Reference(), nil, &h)
		if h.Name == v.Name {
			return true
		}
	}

	return false
}

func FindVMByName(c *govmomi.Client, name string) (govmomi.Reference, error) {
	s := nameSearcher{name, "VirtualMachine"}

	return RecursiveSearch(c, s)
}

func FindHostByName(c *govmomi.Client, name string) (govmomi.Reference, error) {
	s := nameSearcher{name, "HostSystem"}

	return RecursiveSearch(c, s)
}

func RecursiveSearch(c *govmomi.Client, s Searcher) (govmomi.Reference, error) {
	// Take root vSphere folder and loop through each datacenter
	children, err := c.RootFolder().Children(c)
	if err != nil {
		return nil, err
	}

	// Loop through all datacenters and recursively search their trees
	for _, v := range children {
		ref, err := recursion(c, v, s)
		if err != nil {
			return nil, err
		}

		if ref != nil {
			return ref, nil
		}
	}

	return nil, nil
}

// recursion is meant to take a datacenter and recursively move down the tree
func recursion(c *govmomi.Client, ref govmomi.Reference, s Searcher) (govmomi.Reference, error) {
	// First thing is to check if we have found our reference
	if s.Compare(c, ref) {
		return ref, nil
	}

	// We have not found our reference so keep moving down the tree
	switch ref.Reference().Type {
	case "Datacenter":
		dc := govmomi.NewDatacenter(ref.Reference())
		dcFolders, err := dc.Folders(c)
		if err != nil {
			return nil, err
		}

		// Walk the VM folder
		found, err := recursion(c, dcFolders.VmFolder, s)
		if err != nil {
			return nil, err
		}
		if found != nil {
			return found, nil
		}

		// Walk the Host folder
		found, err = recursion(c, dcFolders.HostFolder, s)
		if err != nil {
			return nil, err
		}
		if found != nil {
			return found, nil
		}

		// Walk the Datastore Folder
		found, err = recursion(c, dcFolders.DatastoreFolder, s)
		if err != nil {
			return nil, err
		}
		if found != nil {
			return found, nil
		}

		// walk the Network folder
		found, err = recursion(c, dcFolders.NetworkFolder, s)
		if err != nil {
			return nil, err
		}
		if found != nil {
			return found, nil
		}

	case "Folder":
		folder := govmomi.NewFolder(ref.Reference())
		children, err := folder.Children(c)
		if err != nil {
			return nil, err
		}

		for _, v := range children {
			found, err := recursion(c, v, s)
			if err != nil {
				return nil, err
			}
			if found != nil {
				return found, nil
			}
		}
	case "StoragePod":
		var sp mo.StoragePod
		c.Properties(ref.Reference(), nil, &sp)

		for _, v := range sp.ChildEntity {
			newRef := newReference(v)
			found, err := recursion(c, newRef, s)
			if err != nil {
				return nil, err
			}
			if found != nil {
				return found, nil
			}
		}
	case "ClusterComputeResource", "ComputeResource":
		var ccr mo.ClusterComputeResource
		c.Properties(ref.Reference(), nil, &ccr)

		for _, v := range ccr.Host {
			newRef := newReference(v)
			found, err := recursion(c, newRef, s)
			if err != nil {
				return nil, err
			}
			if found != nil {
				return found, nil
			}
		}
	case "VmwareDistributedVirtualSwitch", "DistributedVirtualSwitch":
		var vDVS mo.VmwareDistributedVirtualSwitch
		c.Properties(ref.Reference(), nil, &vDVS)

		for _, v := range vDVS.Portgroup {
			newRef := newReference(v)
			found, err := recursion(c, newRef, s)
			if err != nil {
				return nil, err
			}
			if found != nil {
				return found, nil
			}
		}
	}

	return nil, nil
}

func newReference(e types.ManagedObjectReference) govmomi.Reference {
	switch e.Type {
	case "Folder":
		return &govmomi.Folder{ManagedObjectReference: e}
	case "StoragePod":
		return &govmomi.StoragePod{
			govmomi.Folder{ManagedObjectReference: e},
		}
	case "Datacenter":
		return &govmomi.Datacenter{ManagedObjectReference: e}
	case "VirtualMachine":
		return &govmomi.VirtualMachine{ManagedObjectReference: e}
	case "VirtualApp":
		return &govmomi.VirtualApp{
			govmomi.ResourcePool{ManagedObjectReference: e},
		}
	case "ComputeResource":
		return &govmomi.ComputeResource{ManagedObjectReference: e}
	case "ClusterComputeResource":
		return &govmomi.ClusterComputeResource{
			govmomi.ComputeResource{ManagedObjectReference: e},
		}
	case "HostSystem":
		return &govmomi.HostSystem{ManagedObjectReference: e}
	case "Network":
		return &govmomi.Network{ManagedObjectReference: e}
	case "ResourcePool":
		return &govmomi.ResourcePool{ManagedObjectReference: e}
	case "DistributedVirtualSwitch":
		return &govmomi.DistributedVirtualSwitch{ManagedObjectReference: e}
	case "VmwareDistributedVirtualSwitch":
		return &govmomi.VmwareDistributedVirtualSwitch{
			govmomi.DistributedVirtualSwitch{ManagedObjectReference: e},
		}
	case "DistributedVirtualPortgroup":
		return &govmomi.DistributedVirtualPortgroup{ManagedObjectReference: e}
	case "Datastore":
		return &govmomi.Datastore{ManagedObjectReference: e}
	default:
		panic("Unknown managed entity: " + e.Type)
	}
}
<commit_msg>Changed for govmomi problem gist<commit_after>package govsphereutils

import (
	"github.com/vmware/govmomi"
	"github.com/vmware/govmomi/vim25/mo"
	"github.com/vmware/govmomi/vim25/types"
)

func InventoryMap(c *govmomi.Client) (map[string][]govmomi.Reference, error) {
	// map to return
	m := map[string][]govmomi.Reference{}

	// Get the root Folder
	root := c.RootFolder()

	rootChildren, err := root.Children()
	if err != nil {
		return nil, err
	}

	// Loop through datacenters
	for _, v := range rootChildren {
		found, err := recursion(c, v)
		if err != nil {
			return nil, err
		}

		// Append found to map
		for t, v := range found {
			m[t] = append(m[t], v...)
		}
	}

	return m, nil
}

func recursion(c *govmomi.Client, ref govmomi.Reference) (map[string][]govmomi.Reference, error) {
	m := map[string][]govmomi.Reference{}

	switch ref.Reference().Type {
	default:
		// We should have a reference with no children so just return it
		m[ref.Reference().Type] = append(m[ref.Reference().Type], ref)
	case "Datacenter":
		dc := govmomi.NewDatacenter(c, ref.Reference())
		dcFolders, err := dc.Folders()
		if err != nil {
			return nil, err
		}

		// Walk the VM folder
		vmFound, err := recursion(c, dcFolders.VmFolder)
		if err != nil {
			return nil, err
		}

		// Walk the Host folder
		hFound, err := recursion(c, dcFolders.HostFolder)
		if err != nil {
			return nil, err
		}

		// Walk the Datastore Folder
		dFound, err := recursion(c, dcFolders.DatastoreFolder)
		if err != nil {
			return nil, err
		}

		// walk the Network folder
		nFound, err := recursion(c, dcFolders.NetworkFolder)
		if err != nil {
			return nil, err
		}

		// Add references to returned reference
		for t, v := range vmFound {
			m[t] = append(m[t], v...)
		}

		for t, v := range hFound {
			m[t] = append(m[t], v...)
		}

		for t, v := range dFound {
			m[t] = append(m[t], v...)
		}

		for t, v := range nFound {
			m[t] = append(m[t], v...)
		}

	case "Folder":
		folder := govmomi.NewFolder(c, ref.Reference())
		children, err := folder.Children()
		if err != nil {
			return nil, err
		}

		for _, v := range children {
			found, err := recursion(c, v)
			if err != nil {
				return nil, err
			}

			for t, v := range found {
				m[t] = append(m[t], v...)
			}
		}

	case "StoragePod":
		var sp mo.StoragePod
		c.Properties(ref.Reference(), nil, &sp)

		for _, v := range sp.ChildEntity {
			newRef := newReference(v)
			found, err := recursion(c, newRef)
			if err != nil {
				return nil, err
			}

			for t, v := range found {
				m[t] = append(m[t], v...)
			}
		}

	case "ClusterComputeResource", "ComputeResource":
		var ccr mo.ClusterComputeResource
		c.Properties(ref.Reference(), nil, &ccr)

		for _, v := range ccr.Host {
			newRef := newReference(v)
			found, err := recursion(c, newRef)
			if err != nil {
				return nil, err
			}

			for t, v := range found {
				m[t] = append(m[t], v...)
			}
		}

	case "VmwareDistributedVirtualSwitch", "DistributedVirtualSwitch":
		var vDVS mo.VmwareDistributedVirtualSwitch
		c.Properties(ref.Reference(), nil, &vDVS)

		for _, v := range vDVS.Portgroup {
			newRef := newReference(v)
			found, err := recursion(c, newRef)
			if err != nil {
				return nil, err
			}

			for t, v := range found {
				m[t] = append(m[t], v...)
			}
		}
	}

	return m, nil
}

func newReference(e types.ManagedObjectReference) govmomi.Reference {
	switch e.Type {
	case "Folder":
		return &govmomi.Folder{ManagedObjectReference: e}
	case "StoragePod":
		return &govmomi.StoragePod{
			&govmomi.Folder{ManagedObjectReference: e},
		}
	case "Datacenter":
		return &govmomi.Datacenter{ManagedObjectReference: e}
	case "VirtualMachine":
		return &govmomi.VirtualMachine{ManagedObjectReference: e}
	case "VirtualApp":
		return &govmomi.VirtualApp{
			&govmomi.ResourcePool{ManagedObjectReference: e},
		}
	case "ComputeResource":
		return &govmomi.ComputeResource{ManagedObjectReference: e}
	case "ClusterComputeResource":
		return &govmomi.ClusterComputeResource{
			govmomi.ComputeResource{ManagedObjectReference: e},
		}
	case "HostSystem":
		return &govmomi.HostSystem{ManagedObjectReference: e}
	case "Network":
		return &govmomi.Network{ManagedObjectReference: e}
	case "ResourcePool":
		return &govmomi.ResourcePool{ManagedObjectReference: e}
	case "DistributedVirtualSwitch":
		return &govmomi.DistributedVirtualSwitch{ManagedObjectReference: e}
	case "VmwareDistributedVirtualSwitch":
		return &govmomi.VmwareDistributedVirtualSwitch{
			govmomi.DistributedVirtualSwitch{ManagedObjectReference: e},
		}
	case "DistributedVirtualPortgroup":
		return &govmomi.DistributedVirtualPortgroup{ManagedObjectReference: e}
	case "Datastore":
		return &govmomi.Datastore{ManagedObjectReference: e}
	default:
		panic("Unknown managed entity: " + e.Type)
	}
}
<|endoftext|>


<commit_before>package grifts

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

	"github.com/markbates/grift/grift"
	"github.com/pkg/errors"
)

var _ = grift.Desc("release", "Generates a CHANGELOG and creates a new GitHub release based on what is in the version.go file.")
var _ = grift.Add("release", func(c *grift.Context) error {
	v, err := findVersion()
	if err != nil {
		return err
	}

	err = installBin()
	if err != nil {
		return err
	}

	err = localTest()
	if err != nil {
		return err
	}

	err = dockerTest()
	if err != nil {
		return err
	}

	grift.Run("shoulders", c)

	if err := push(); err != nil {
		return errors.WithStack(err)
	}

	err = tagRelease(v)
	if err != nil {
		return err
	}

	return runReleaser(v)
})

func installBin() error {
	cmd := exec.Command("go", "install", "-v", "./buffalo")
	cmd.Stdin = os.Stdin
	cmd.Stderr = os.Stderr
	cmd.Stdout = os.Stdout
	return cmd.Run()
}

func localTest() error {
	cmd := exec.Command("go", "test", "-v", "-race", "./...")
	cmd.Stdin = os.Stdin
	cmd.Stderr = os.Stderr
	cmd.Stdout = os.Stdout
	return cmd.Run()
}

func dockerTest() error {
	cmd := exec.Command("docker", "build", ".")
	cmd.Stdin = os.Stdin
	cmd.Stderr = os.Stderr
	cmd.Stdout = os.Stdout
	return cmd.Run()
}

func tagRelease(v string) error {
	cmd := exec.Command("git", "tag", v)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		return err
	}

	cmd = exec.Command("git", "push", "origin", "--tags")
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	return cmd.Run()
}

func runReleaser(v string) error {
	cmd := exec.Command("goreleaser", "--rm-dist")
	cmd.Stdin = os.Stdin
	cmd.Stderr = os.Stderr
	cmd.Stdout = os.Stdout
	return cmd.Run()
}

func push() error {
	cmd := exec.Command("git", "push", "origin", "master")
	cmd.Stdin = os.Stdin
	cmd.Stderr = os.Stderr
	cmd.Stdout = os.Stdout
	return cmd.Run()
}

func findVersion() (string, error) {
	pwd, err := os.Getwd()
	if err != nil {
		return "", err
	}
	vfile, err := ioutil.ReadFile(filepath.Join(pwd, "buffalo/cmd/version.go"))
	if err != nil {
		return "", err
	}

	//var Version = "v0.4.0"
	re := regexp.MustCompile(`const Version = "(.+)"`)
	matches := re.FindStringSubmatch(string(vfile))
	if len(matches) < 2 {
		return "", errors.New("failed to find the version")
	}
	v := matches[1]
	if strings.Contains(v, "dev") {
		return "", errors.Errorf("version can not be a dev version %s", v)
	}
	if !strings.HasPrefix(v, "v") {
		return "", errors.Errorf("version must match format `v0.0.0`: %s", v)
	}
	return v, nil
}
<commit_msg>make sure release tag uses -tags sqlite for tests<commit_after>package grifts

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

	"github.com/markbates/grift/grift"
	"github.com/pkg/errors"
)

var _ = grift.Desc("release", "Generates a CHANGELOG and creates a new GitHub release based on what is in the version.go file.")
var _ = grift.Add("release", func(c *grift.Context) error {
	v, err := findVersion()
	if err != nil {
		return err
	}

	err = installBin()
	if err != nil {
		return err
	}

	err = localTest()
	if err != nil {
		return err
	}

	err = dockerTest()
	if err != nil {
		return err
	}

	grift.Run("shoulders", c)

	if err := push(); err != nil {
		return errors.WithStack(err)
	}

	err = tagRelease(v)
	if err != nil {
		return err
	}

	return runReleaser(v)
})

func installBin() error {
	cmd := exec.Command("go", "install", "-v", "./buffalo")
	cmd.Stdin = os.Stdin
	cmd.Stderr = os.Stderr
	cmd.Stdout = os.Stdout
	return cmd.Run()
}

func localTest() error {
	cmd := exec.Command("go", "test", "-tags", "sqlite", "-v", "-race", "./...")
	cmd.Stdin = os.Stdin
	cmd.Stderr = os.Stderr
	cmd.Stdout = os.Stdout
	return cmd.Run()
}

func dockerTest() error {
	cmd := exec.Command("docker", "build", ".")
	cmd.Stdin = os.Stdin
	cmd.Stderr = os.Stderr
	cmd.Stdout = os.Stdout
	return cmd.Run()
}

func tagRelease(v string) error {
	cmd := exec.Command("git", "tag", v)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		return err
	}

	cmd = exec.Command("git", "push", "origin", "--tags")
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	return cmd.Run()
}

func runReleaser(v string) error {
	cmd := exec.Command("goreleaser", "--rm-dist")
	cmd.Stdin = os.Stdin
	cmd.Stderr = os.Stderr
	cmd.Stdout = os.Stdout
	return cmd.Run()
}

func push() error {
	cmd := exec.Command("git", "push", "origin", "master")
	cmd.Stdin = os.Stdin
	cmd.Stderr = os.Stderr
	cmd.Stdout = os.Stdout
	return cmd.Run()
}

func findVersion() (string, error) {
	pwd, err := os.Getwd()
	if err != nil {
		return "", err
	}
	vfile, err := ioutil.ReadFile(filepath.Join(pwd, "buffalo/cmd/version.go"))
	if err != nil {
		return "", err
	}

	//var Version = "v0.4.0"
	re := regexp.MustCompile(`const Version = "(.+)"`)
	matches := re.FindStringSubmatch(string(vfile))
	if len(matches) < 2 {
		return "", errors.New("failed to find the version")
	}
	v := matches[1]
	if strings.Contains(v, "dev") {
		return "", errors.Errorf("version can not be a dev version %s", v)
	}
	if !strings.HasPrefix(v, "v") {
		return "", errors.Errorf("version must match format `v0.0.0`: %s", v)
	}
	return v, nil
}
<|endoftext|>


<commit_before>package shh

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"time"

	"github.com/heroku/slog"
)

type LibratoMetric struct {
	Name       string             `json:"name"`
	Value      interface{}        `json:"value"`
	When       int64              `json:"measure_time"`
	Source     string             `json:"source,omitempty"`
	Attributes LibratoMetricAttrs `json:"attributes,omitempty"`
}

type LibratoMetricAttrs struct {
	UnitName string `json:"display_units_long,omitempty"`
	UnitAbbr string `json:"display_units_short,omitempty"`
}

type LibratoPostBody struct {
	Gauges   []LibratoMetric `json:"gauges,omitempty"`
	Counters []LibratoMetric `json:"counters,omitempty"`
}

const (
	LibratoBacklog         = 8 // No more than N pending batches in-flight
	LibratoMaxAttempts     = 4 // Max attempts before dropping batch
	LibratoStartingBackoff = 500 * time.Millisecond
)

type Librato struct {
	Timeout      time.Duration
	BatchSize    int
	User         string
	Token        string
	Url          string
	measurements <-chan Measurement
	batches      chan []Measurement
	prefix       string
	source       string
	client       *http.Client
	userAgent    string
	interval     time.Duration
	round        bool
}

func NewLibratoOutputter(measurements <-chan Measurement, config Config) *Librato {
	var user string
	var token string

	if config.LibratoUrl.User != nil {
		user = config.LibratoUrl.User.Username()
		token, _ = config.LibratoUrl.User.Password()
		config.LibratoUrl.User = nil
	}

	// override settings in URL if they were present
	if config.LibratoUser != "" {
		user = config.LibratoUser
	}
	if config.LibratoToken != "" {
		token = config.LibratoToken
	}

	return &Librato{
		measurements: measurements,
		prefix:       config.Prefix,
		source:       config.Source,
		batches:      make(chan []Measurement, LibratoBacklog),
		Timeout:      config.LibratoBatchTimeout,
		BatchSize:    config.LibratoBatchSize,
		User:         user,
		Token:        token,
		Url:          config.LibratoUrl.String(),
		interval:     config.Interval,
		round:        config.LibratoRound,
		userAgent:    config.UserAgent,
		client:       &http.Client{Timeout: config.NetworkTimeout},
	}
}

func (out *Librato) Start() {
	go out.deliver()
	go out.batch()
}

// Returns a batch that is ready to be submitted to Librato, either because it timed out
// after receiving it's first measurement or it is full.
func (out *Librato) readyBatch() []Measurement {
	batch := make([]Measurement, 0, out.BatchSize)
	timer := new(time.Timer) // "empty" timer so we don't timeout before we have any measurements
	for {
		select {
		case measurement := <-out.measurements:
			batch = append(batch, measurement)
			if len(batch) == 1 { // We got a measurement, so we want to start the timer.
				timer = time.NewTimer(out.Timeout)
				defer timer.Stop()
			}
			if len(batch) == cap(batch) {
				return batch
			}
		case <-timer.C:
			return batch
		}
	}
}

// Continuously batch measurments into the batch channel
func (out *Librato) batch() {
	ctx := slog.Context{"fn": "batch", "outputter": "librato"}
	for {
		batch := out.readyBatch()

		select {
		case out.batches <- batch:
		default:
			LogError(ctx, nil, "Batches backlogged, dropping")
		}
	}
}

func (out *Librato) appendLibratoMetric(counters, gauges []LibratoMetric, mm Measurement) ([]LibratoMetric, []LibratoMetric) {
	var t int64
	attrs := LibratoMetricAttrs{UnitName: mm.Unit().Name(), UnitAbbr: mm.Unit().Abbr()}

	if out.round {
		t = mm.Time().Round(out.interval).Unix()
	} else {
		t = mm.Time().Unix()
	}

	libratoMetric := LibratoMetric{mm.Name(out.prefix), mm.Value(), t, out.source, attrs}

	switch mm.Type() {
	case CounterType:
		counters = append(counters, libratoMetric)
	case GaugeType, FloatGaugeType:
		gauges = append(gauges, libratoMetric)
	}
	return counters, gauges
}

func (out *Librato) deliver() {
	ctx := slog.Context{"fn": "prepare", "outputter": "librato"}
	for batch := range out.batches {
		gauges := make([]LibratoMetric, 0)
		counters := make([]LibratoMetric, 0)
		for _, mm := range batch {
			counters, gauges = out.appendLibratoMetric(counters, gauges, mm)
		}

		counters, gauges = out.appendLibratoMetric(
			counters,
			gauges,
			GaugeMeasurement{time.Now(), "librato-outlet", []string{"batch", "guage", "size"}, uint64(len(gauges) + 2), Metrics},
		)
		counters, gauges = out.appendLibratoMetric(
			counters,
			gauges,
			GaugeMeasurement{time.Now(), "librato-outlet", []string{"batch", "counter", "size"}, uint64(len(counters)), Metrics},
		)

		payload := LibratoPostBody{gauges, counters}
		j, err := json.Marshal(payload)
		if err != nil {
			FatalError(ctx, err, "marshaling json")
		}

		out.sendWithBackoff(j)
	}
}

func (out *Librato) sendWithBackoff(payload []byte) bool {
	ctx := slog.Context{"fn": "sendWithBackoff", "outputter": "librato", "backoff": LibratoStartingBackoff, "attempts": 0}

	for ctx["attempts"].(int) < LibratoMaxAttempts {
		retry, err := out.send(payload)
		if retry {
			LogError(ctx, err, "backing off")
			ctx["backoff"] = backoff(ctx["backoff"].(time.Duration))
		} else {
			if err != nil {
				LogError(ctx, err, "error sending, no retry")
				return false
			} else {
				return true
			}
		}
		ctx["attempts"] = ctx["attempts"].(int) + 1
	}
	return false
}

// Attempts to send the payload and signals retries on errors
func (out *Librato) send(payload []byte) (bool, error) {
	body := bytes.NewReader(payload)
	req, err := http.NewRequest("POST", out.Url, body)
	if err != nil {
		return false, err
	}

	req.Header.Add("Content-Type", "application/json")
	req.Header.Add("User-Agent", out.userAgent)
	req.SetBasicAuth(out.User, out.Token)

	resp, err := out.client.Do(req)
	if err != nil {
		if err, ok := err.(*url.Error); ok && err.Err == io.EOF {
			return true, err
		}
		return false, err
	} else {
		defer resp.Body.Close()

		if resp.StatusCode >= 300 {
			b, _ := ioutil.ReadAll(resp.Body)

			if resp.StatusCode >= 500 {
				err = fmt.Errorf("server error: %d, body: %+q", resp.StatusCode, string(b))
				return true, err
			} else {
				err = fmt.Errorf("client error: %d, body: %+q", resp.StatusCode, string(b))
				return false, err
			}

		}
	}

	return false, nil
}

// Sleeps `bo` and then returns double
func backoff(bo time.Duration) time.Duration {
	time.Sleep(bo)
	return bo * 2
}
<commit_msg>Just retry all low level errors<commit_after>package shh

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"time"

	"github.com/heroku/slog"
)

type LibratoMetric struct {
	Name       string             `json:"name"`
	Value      interface{}        `json:"value"`
	When       int64              `json:"measure_time"`
	Source     string             `json:"source,omitempty"`
	Attributes LibratoMetricAttrs `json:"attributes,omitempty"`
}

type LibratoMetricAttrs struct {
	UnitName string `json:"display_units_long,omitempty"`
	UnitAbbr string `json:"display_units_short,omitempty"`
}

type LibratoPostBody struct {
	Gauges   []LibratoMetric `json:"gauges,omitempty"`
	Counters []LibratoMetric `json:"counters,omitempty"`
}

const (
	LibratoBacklog         = 8 // No more than N pending batches in-flight
	LibratoMaxAttempts     = 4 // Max attempts before dropping batch
	LibratoStartingBackoff = 500 * time.Millisecond
)

type Librato struct {
	Timeout      time.Duration
	BatchSize    int
	User         string
	Token        string
	Url          string
	measurements <-chan Measurement
	batches      chan []Measurement
	prefix       string
	source       string
	client       *http.Client
	userAgent    string
	interval     time.Duration
	round        bool
}

func NewLibratoOutputter(measurements <-chan Measurement, config Config) *Librato {
	var user string
	var token string

	if config.LibratoUrl.User != nil {
		user = config.LibratoUrl.User.Username()
		token, _ = config.LibratoUrl.User.Password()
		config.LibratoUrl.User = nil
	}

	// override settings in URL if they were present
	if config.LibratoUser != "" {
		user = config.LibratoUser
	}
	if config.LibratoToken != "" {
		token = config.LibratoToken
	}

	return &Librato{
		measurements: measurements,
		prefix:       config.Prefix,
		source:       config.Source,
		batches:      make(chan []Measurement, LibratoBacklog),
		Timeout:      config.LibratoBatchTimeout,
		BatchSize:    config.LibratoBatchSize,
		User:         user,
		Token:        token,
		Url:          config.LibratoUrl.String(),
		interval:     config.Interval,
		round:        config.LibratoRound,
		userAgent:    config.UserAgent,
		client:       &http.Client{Timeout: config.NetworkTimeout},
	}
}

func (out *Librato) Start() {
	go out.deliver()
	go out.batch()
}

// Returns a batch that is ready to be submitted to Librato, either because it timed out
// after receiving it's first measurement or it is full.
func (out *Librato) readyBatch() []Measurement {
	batch := make([]Measurement, 0, out.BatchSize)
	timer := new(time.Timer) // "empty" timer so we don't timeout before we have any measurements
	for {
		select {
		case measurement := <-out.measurements:
			batch = append(batch, measurement)
			if len(batch) == 1 { // We got a measurement, so we want to start the timer.
				timer = time.NewTimer(out.Timeout)
				defer timer.Stop()
			}
			if len(batch) == cap(batch) {
				return batch
			}
		case <-timer.C:
			return batch
		}
	}
}

// Continuously batch measurments into the batch channel
func (out *Librato) batch() {
	ctx := slog.Context{"fn": "batch", "outputter": "librato"}
	for {
		batch := out.readyBatch()

		select {
		case out.batches <- batch:
		default:
			LogError(ctx, nil, "Batches backlogged, dropping")
		}
	}
}

func (out *Librato) appendLibratoMetric(counters, gauges []LibratoMetric, mm Measurement) ([]LibratoMetric, []LibratoMetric) {
	var t int64
	attrs := LibratoMetricAttrs{UnitName: mm.Unit().Name(), UnitAbbr: mm.Unit().Abbr()}

	if out.round {
		t = mm.Time().Round(out.interval).Unix()
	} else {
		t = mm.Time().Unix()
	}

	libratoMetric := LibratoMetric{mm.Name(out.prefix), mm.Value(), t, out.source, attrs}

	switch mm.Type() {
	case CounterType:
		counters = append(counters, libratoMetric)
	case GaugeType, FloatGaugeType:
		gauges = append(gauges, libratoMetric)
	}
	return counters, gauges
}

func (out *Librato) deliver() {
	ctx := slog.Context{"fn": "prepare", "outputter": "librato"}
	for batch := range out.batches {
		gauges := make([]LibratoMetric, 0)
		counters := make([]LibratoMetric, 0)
		for _, mm := range batch {
			counters, gauges = out.appendLibratoMetric(counters, gauges, mm)
		}

		counters, gauges = out.appendLibratoMetric(
			counters,
			gauges,
			GaugeMeasurement{time.Now(), "librato-outlet", []string{"batch", "guage", "size"}, uint64(len(gauges) + 2), Metrics},
		)
		counters, gauges = out.appendLibratoMetric(
			counters,
			gauges,
			GaugeMeasurement{time.Now(), "librato-outlet", []string{"batch", "counter", "size"}, uint64(len(counters)), Metrics},
		)

		payload := LibratoPostBody{gauges, counters}
		j, err := json.Marshal(payload)
		if err != nil {
			FatalError(ctx, err, "marshaling json")
		}

		out.sendWithBackoff(j)
	}
}

func (out *Librato) sendWithBackoff(payload []byte) bool {
	ctx := slog.Context{"fn": "sendWithBackoff", "outputter": "librato", "backoff": LibratoStartingBackoff, "attempts": 0}

	for ctx["attempts"].(int) < LibratoMaxAttempts {
		retry, err := out.send(payload)
		if retry {
			LogError(ctx, err, "backing off")
			ctx["backoff"] = backoff(ctx["backoff"].(time.Duration))
		} else {
			if err != nil {
				LogError(ctx, err, "error sending, no retry")
				return false
			} else {
				return true
			}
		}
		ctx["attempts"] = ctx["attempts"].(int) + 1
	}
	return false
}

// Attempts to send the payload and signals retries on errors
func (out *Librato) send(payload []byte) (bool, error) {
	body := bytes.NewReader(payload)
	req, err := http.NewRequest("POST", out.Url, body)
	if err != nil {
		return false, err
	}

	req.Header.Add("Content-Type", "application/json")
	req.Header.Add("User-Agent", out.userAgent)
	req.SetBasicAuth(out.User, out.Token)

	resp, err := out.client.Do(req)
	if err != nil {
		return true, err
	} else {
		defer resp.Body.Close()

		if resp.StatusCode >= 300 {
			b, _ := ioutil.ReadAll(resp.Body)

			if resp.StatusCode >= 500 {
				err = fmt.Errorf("server error: %d, body: %+q", resp.StatusCode, string(b))
				return true, err
			} else {
				err = fmt.Errorf("client error: %d, body: %+q", resp.StatusCode, string(b))
				return false, err
			}

		}
	}

	return false, nil
}

// Sleeps `bo` and then returns double
func backoff(bo time.Duration) time.Duration {
	time.Sleep(bo)
	return bo * 2
}
<|endoftext|>


<commit_before>package hosts_updater

import (
	"bufio"
	"fmt"
	"os"
	"strings"
	"testing"

	"github.com/damonkelley/hostsup/host"
)

func TestAddHostsEntry(t *testing.T) {
	hostname, ip := "dev.dev", "192.168.0.1"
	host := host.NewHost(hostname, ip)

	h := NewHostsfile("testdata/hosts")
	h.AddHostsEntry(host)

	f, _ := os.Open("testdata/hosts")
	defer f.Close()

	scanner := bufio.NewScanner(f)

	for scanner.Scan() {
		if strings.Contains(scanner.Text(), "dev.dev") {
			return
		}
	}

	t.Error(fmt.Sprintf("Expected to find %s in testdata/hosts", hostname))
}

func TestRemoveHostsEntry(t *testing.T) {
	hostname, ip := "dev.dev", "192.168.0.1"
	host := host.NewHost(hostname, ip)

	h := NewHostsfile("testdata/hosts")
	h.RemoveHostsEntry(host)

	f, _ := os.Open("testdata/hosts")
	defer f.Close()

	scanner := bufio.NewScanner(f)

	for scanner.Scan() {
		if !strings.Contains(scanner.Text(), "dev.dev") {
			return
		}
	}

	t.Error(fmt.Sprintf("Expected to find %s in testdata/hosts", hostname))
}

func TestListHostsEntries(t *testing.T) {
	hostname1, ip1 := "dev1.dev", "192.168.0.1"
	host1 := host.NewHost(hostname1, ip1)

	hostname2, ip2 := "dev2.dev", "192.168.0.2"
	host2 := host.NewHost(hostname2, ip2)

	h := NewHostsfile("testdata/hosts")
	h.AddHostsEntry(host1)
	h.AddHostsEntry(host2)

	hosts := h.ListHostsEntries()

	if len(hosts) != 2 {
		t.Error(fmt.Sprintf("Expected to find 2 host entry. Found %d instead.", len(hosts)))
	}
}
<commit_msg>Use a temporary hosts file to test against.<commit_after>package hosts_updater

import (
	"bufio"
	"fmt"
	"io/ioutil"
	"os"
	"strings"
	"testing"
)

func createTestHostsFile() (*os.File, error) {
	f, err := ioutil.TempFile("testdata", "hosts-")

	if err != nil {
		return nil, err
	}

	contents, _ := ioutil.ReadFile("testdata/hosts")

	f.Write(contents)
	f.Seek(0, 0)

	return f, nil
}

func remove(f *os.File) error {
	return os.Remove(f.Name())
}

func TestAddHostsEntry(t *testing.T) {
	hostname, ip := "dev.dev", "192.168.0.1"
	host := NewHost(hostname, ip)

	f, _ := createTestHostsFile()
	defer remove(f)

	h := Hostsfile{f.Name(), f}
	h.AddHostsEntry(host)

	// Reset the offset after AddHostsEntry changes it.
	scanner := bufio.NewScanner(f)

	for scanner.Scan() {
		if strings.Contains(scanner.Text(), "dev.dev") {
			return
		}
	}

	t.Error(fmt.Sprintf("Expected to find %s in testdata/hosts", hostname))
}

func TestRemoveHostsEntry(t *testing.T) {
	hostname, ip := "dev.dev", "192.168.0.1"
	host := NewHost(hostname, ip)

	f, _ := createTestHostsFile()
	defer remove(f)

	h := Hostsfile{f.Name(), f}
	h.RemoveHostsEntry(host)

	scanner := bufio.NewScanner(f)

	for scanner.Scan() {
		if !strings.Contains(scanner.Text(), "dev.dev") {
			return
		}
	}

	t.Error(fmt.Sprintf("Expected to find %s in testdata/hosts", hostname))
}

func TestListHostsEntries(t *testing.T) {
	hostname1, ip1 := "dev1.dev", "192.168.0.1"
	host1 := NewHost(hostname1, ip1)

	hostname2, ip2 := "dev2.dev", "192.168.0.2"
	host2 := NewHost(hostname2, ip2)

	f, _ := createTestHostsFile()
	defer remove(f)

	h := Hostsfile{f.Name(), f}

	h.AddHostsEntry(host1)
	h.AddHostsEntry(host2)

	hosts := h.ListHostsEntries()

	if len(hosts) != 2 {
		t.Error(fmt.Sprintf("Expected to find 2 host entry. Found %d instead.", len(hosts)))
	}
}
<|endoftext|>


<commit_before>// Code derived from https:// github.com/btcsuite/btcd/blob/master/wire/message.go
package lnwire

import (
	"bytes"
	"fmt"
	"io"

	"github.com/roasbeef/btcd/wire"
)

// 4-byte network + 4-byte message id + payload-length 4-byte
const MessageHeaderSize = 12

const MaxMessagePayload = 1024 * 1024 * 32 //  32MB

const (
	// Funding channel open
	CmdFundingRequest      = uint32(200)
	CmdFundingResponse     = uint32(210)
	CmdFundingSignAccept   = uint32(220)
	CmdFundingSignComplete = uint32(230)

	// Close channel
	CmdCloseRequest  = uint32(300)
	CmdCloseComplete = uint32(310)

	// TODO Renumber to 1100
	// HTLC payment
	CmdHTLCAddRequest = uint32(1000)
	CmdHTLCAddAccept  = uint32(1010)
	CmdHTLCAddReject  = uint32(1020)

	// TODO Renumber to 1200
	// HTLC settlement
	CmdHTLCSettleRequest = uint32(1100)
	CmdHTLCSettleAccept  = uint32(1110)

	// HTLC timeout
	CmdHTLCTimeoutRequest = uint32(1300)
	CmdHTLCTimeoutAccept  = uint32(1310)

	// Commitments
	CmdCommitSignature  = uint32(2000)
	CmdCommitRevocation = uint32(2010)

	// Error
	CmdErrorGeneric = uint32(4000)
)

// Every message has these functions:
type Message interface {
	Decode(io.Reader, uint32) error // (io, protocol version)
	Encode(io.Writer, uint32) error // (io, protocol version)
	Command() uint32                // returns ID of the message
	MaxPayloadLength(uint32) uint32 // (version) maxpayloadsize
	Validate() error                // Validates the data struct
	String() string
}

func makeEmptyMessage(command uint32) (Message, error) {
	var msg Message

	switch command {
	case CmdFundingRequest:
		msg = &FundingRequest{}
	case CmdFundingResponse:
		msg = &FundingResponse{}
	case CmdFundingSignAccept:
		msg = &FundingSignAccept{}
	case CmdFundingSignComplete:
		msg = &FundingSignComplete{}
	case CmdCloseRequest:
		msg = &CloseRequest{}
	case CmdCloseComplete:
		msg = &CloseComplete{}
	case CmdHTLCAddRequest:
		msg = &HTLCAddRequest{}
	case CmdHTLCAddAccept:
		msg = &HTLCAddAccept{}
	case CmdHTLCAddReject:
		msg = &HTLCAddReject{}
	case CmdHTLCSettleRequest:
		msg = &HTLCSettleRequest{}
	case CmdHTLCSettleAccept:
		msg = &HTLCSettleAccept{}
	case CmdHTLCTimeoutRequest:
		msg = &HTLCTimeoutRequest{}
	case CmdHTLCTimeoutAccept:
		msg = &HTLCTimeoutAccept{}
	case CmdCommitSignature:
		msg = &CommitSignature{}
	case CmdCommitRevocation:
		msg = &CommitRevocation{}
	case CmdErrorGeneric:
		msg = &ErrorGeneric{}
	default:
		return nil, fmt.Errorf("unhandled command [%d]", command)
	}

	return msg, nil
}

type messageHeader struct {
	// NOTE(j): We don't need to worry about the magic overlapping with
	// bitcoin since this is inside encrypted comms anyway, but maybe we
	// should use the XOR (^wire.TestNet3) just in case???
	magic   wire.BitcoinNet // which Blockchain Technology(TM) to use
	command uint32
	length  uint32
}

func readMessageHeader(r io.Reader) (int, *messageHeader, error) {
	var headerBytes [MessageHeaderSize]byte
	n, err := io.ReadFull(r, headerBytes[:])
	if err != nil {
		return n, nil, err
	}
	hr := bytes.NewReader(headerBytes[:])

	hdr := messageHeader{}

	err = readElements(hr,
		&hdr.magic,
		&hdr.command,
		&hdr.length)
	if err != nil {
		return n, nil, err
	}

	return n, &hdr, nil
}

//  discardInput reads n bytes from reader r in chunks and discards the read
//  bytes.  This is used to skip payloads when various errors occur and helps
//  prevent rogue nodes from causing massive memory allocation through forging
//  header length.
func discardInput(r io.Reader, n uint32) {
	maxSize := uint32(10 * 1024) // 10k at a time
	numReads := n / maxSize
	bytesRemaining := n % maxSize
	if n > 0 {
		buf := make([]byte, maxSize)
		for i := uint32(0); i < numReads; i++ {
			io.ReadFull(r, buf)
		}
	}
	if bytesRemaining > 0 {
		buf := make([]byte, bytesRemaining)
		io.ReadFull(r, buf)
	}
}

func WriteMessage(w io.Writer, msg Message, pver uint32, btcnet wire.BitcoinNet) (int, error) {
	totalBytes := 0

	cmd := msg.Command()

	// Encode the message payload
	var bw bytes.Buffer
	err := msg.Encode(&bw, pver)
	if err != nil {
		return totalBytes, err
	}
	payload := bw.Bytes()
	lenp := len(payload)

	// Enforce maximum overall message payload
	if lenp > MaxMessagePayload {
		return totalBytes, fmt.Errorf("message payload is too large - encoded %d bytes, but maximum message payload is %d bytes", lenp, MaxMessagePayload)
	}

	// Enforce maximum message payload on the message type
	mpl := msg.MaxPayloadLength(pver)
	if uint32(lenp) > mpl {
		return totalBytes, fmt.Errorf("message payload is too large - encoded %d bytes, but maximum message payload of type %x is %d bytes", lenp, cmd, mpl)
	}

	// Create header for the message
	hdr := messageHeader{}
	hdr.magic = btcnet
	hdr.command = cmd
	hdr.length = uint32(lenp)

	//  Encode the header for the message.  This is done to a buffer
	//  rather than directly to the writer since writeElements doesn't
	//  return the number of bytes written.
	hw := bytes.NewBuffer(make([]byte, 0, MessageHeaderSize))
	writeElements(hw, hdr.magic, hdr.command, hdr.length)

	// Write header
	n, err := w.Write(hw.Bytes())
	totalBytes += n
	if err != nil {
		return totalBytes, err
	}

	// Write payload
	n, err = w.Write(payload)
	totalBytes += n
	if err != nil {
		return totalBytes, err
	}

	return totalBytes, nil
}

func ReadMessage(r io.Reader, pver uint32, btcnet wire.BitcoinNet) (int, Message, []byte, error) {
	totalBytes := 0
	n, hdr, err := readMessageHeader(r)
	totalBytes += n
	if err != nil {
		return totalBytes, nil, nil, err
	}

	// Enforce maximum message payload
	if hdr.length > MaxMessagePayload {
		return totalBytes, nil, nil, fmt.Errorf("message payload is too large - header indicates %d bytes, but max message payload is %d bytes.", hdr.length, MaxMessagePayload)
	}

	// Check for messages in the wrong bitcoin network
	if hdr.magic != btcnet {
		discardInput(r, hdr.length)
		return totalBytes, nil, nil, fmt.Errorf("message from other network [%v]", hdr.magic)
	}

	// Create struct of appropriate message type based on the command
	command := hdr.command
	msg, err := makeEmptyMessage(command)
	if err != nil {
		discardInput(r, hdr.length)
		return totalBytes, nil, nil, fmt.Errorf("ReadMessage %s", err.Error())
	}

	// Check for maximum length based on the message type
	mpl := msg.MaxPayloadLength(pver)
	if hdr.length > mpl {
		discardInput(r, hdr.length)
		return totalBytes, nil, nil, fmt.Errorf("payload exceeds max length. indicates %v bytes, but max of message type %v is %v.", hdr.length, command, mpl)
	}

	// Read payload
	payload := make([]byte, hdr.length)
	n, err = io.ReadFull(r, payload)
	totalBytes += n
	if err != nil {
		return totalBytes, nil, nil, err
	}

	// Unmarshal message
	pr := bytes.NewBuffer(payload)
	err = msg.Decode(pr, pver)
	if err != nil {
		return totalBytes, nil, nil, err
	}

	// Validate the data
	err = msg.Validate()
	if err != nil {
		return totalBytes, nil, nil, err
	}

	// We're good!
	return totalBytes, msg, payload, nil
}
<commit_msg>lnwire: update and document message.go<commit_after>// Code derived from https:// github.com/btcsuite/btcd/blob/master/wire/message.go
package lnwire

import (
	"bytes"
	"fmt"
	"io"

	"github.com/roasbeef/btcd/wire"
)

// MessageHeaderSize is the number of bytes in a lightning message header.
// The bytes are allocated as follows: network magic 4 bytes + command 4
// bytes + payload length 4 bytes. Note that a checksum is omitted as lightning
// messages are assumed to be transmitted over an AEAD secured connection which
// provides integrity over the entire message.
const MessageHeaderSize = 12

// MaxMessagePayload is the maximum bytes a message can be regardless of other
// individual limits imposed by messages themselves.
const MaxMessagePayload = 1024 * 1024 * 32 //  32MB

// Commands used in lightning message headers which detail the type of message.
const (
	// Commands for opening a channel funded by both parties (dual funder).
	CmdFundingRequest      = uint32(200)
	CmdFundingResponse     = uint32(210)
	CmdFundingSignAccept   = uint32(220)
	CmdFundingSignComplete = uint32(230)

	// Commands for opening a channel funded by one party (single funder).
	CmdSingleFundingRequest      = uint32(100)
	CmdSingleFundingResponse     = uint32(110)
	CmdSingleFundingComplete     = uint32(120)
	CmdSingleFundingSignComplete = uint32(130)
	CmdSingleFundingOpenProof    = uint32(140)

	// Commands for the workflow of cooperatively closing an active channel.
	CmdCloseRequest  = uint32(300)
	CmdCloseComplete = uint32(310)

	// Commands for negotiating HTLCs.
	CmdHTLCAddRequest     = uint32(1000)
	CmdHTLCAddAccept      = uint32(1010)
	CmdHTLCAddReject      = uint32(1020)
	CmdHTLCSettleRequest  = uint32(1100)
	CmdHTLCTimeoutRequest = uint32(1300)

	// Commands for modifying commitment transactions.
	CmdCommitSignature  = uint32(2000)
	CmdCommitRevocation = uint32(2010)

	// Commands for reporting protocol errors.
	CmdErrorGeneric = uint32(4000)
)

// Message is an interface that defines a lightning wire protocol message. The
// interface is general in order to allow implementing types full control over
// the representation of its data.
type Message interface {
	Decode(io.Reader, uint32) error
	Encode(io.Writer, uint32) error
	Command() uint32
	MaxPayloadLength(uint32) uint32
	Validate() error
	String() string
}

// makeEmptyMessage creates a new empty message of the proper concrete type
// based on the command ID.
func makeEmptyMessage(command uint32) (Message, error) {
	var msg Message

	switch command {
	case CmdFundingRequest:
		msg = &FundingRequest{}
	case CmdFundingResponse:
		msg = &FundingResponse{}
	case CmdFundingSignAccept:
		msg = &FundingSignAccept{}
	case CmdFundingSignComplete:
		msg = &FundingSignComplete{}
	case CmdSingleFundingRequest:
		msg = &SingleFundingRequest{}
	case CmdSingleFundingResponse:
		msg = &SingleFundingResponse{}
	case CmdSingleFundingComplete:
		msg = &SingleFundingComplete{}
	case CmdSingleFundingSignComplete:
		msg = &SingleFundingSignComplete{}
	case CmdSingleFundingOpenProof:
		msg = &SingleFundingOpenProof{}
	case CmdCloseRequest:
		msg = &CloseRequest{}
	case CmdCloseComplete:
		msg = &CloseComplete{}
	case CmdHTLCAddRequest:
		msg = &HTLCAddRequest{}
	case CmdHTLCAddReject:
		msg = &HTLCAddReject{}
	case CmdHTLCSettleRequest:
		msg = &HTLCSettleRequest{}
	case CmdHTLCTimeoutRequest:
		msg = &HTLCTimeoutRequest{}
	case CmdCommitSignature:
		msg = &CommitSignature{}
	case CmdCommitRevocation:
		msg = &CommitRevocation{}
	case CmdErrorGeneric:
		msg = &ErrorGeneric{}
	default:
		return nil, fmt.Errorf("unhandled command [%d]", command)
	}

	return msg, nil
}

// messageHeader represents the header structure for all lightning protocol
// messages.
type messageHeader struct {
	// magic represents Which Blockchain Technology(TM) to use.
	// NOTE(j): We don't need to worry about the magic overlapping with
	// bitcoin since this is inside encrypted comms anyway, but maybe we
	// should use the XOR (^wire.TestNet3) just in case???
	magic   wire.BitcoinNet // 4 bytes
	command uint32          // 4 bytes
	length  uint32          // 4 bytes
}

// readMessageHeader reads a lightning protocol message header from r.
func readMessageHeader(r io.Reader) (int, *messageHeader, error) {
	// As the message header is a fixed size structure, read bytes for the
	// entire header at once.
	var headerBytes [MessageHeaderSize]byte
	n, err := io.ReadFull(r, headerBytes[:])
	if err != nil {
		return n, nil, err
	}
	hr := bytes.NewReader(headerBytes[:])

	// Create and populate the message header from the raw header bytes.
	hdr := messageHeader{}
	err = readElements(hr,
		&hdr.magic,
		&hdr.command,
		&hdr.length)
	if err != nil {
		return n, nil, err
	}

	return n, &hdr, nil
}

// discardInput reads n bytes from reader r in chunks and discards the read
// bytes. This is used to skip payloads when various errors occur and helps
// prevent rogue nodes from causing massive memory allocation through forging
// header length.
func discardInput(r io.Reader, n uint32) {
	maxSize := uint32(10 * 1024) // 10k at a time
	numReads := n / maxSize
	bytesRemaining := n % maxSize
	if n > 0 {
		buf := make([]byte, maxSize)
		for i := uint32(0); i < numReads; i++ {
			io.ReadFull(r, buf)
		}
	}
	if bytesRemaining > 0 {
		buf := make([]byte, bytesRemaining)
		io.ReadFull(r, buf)
	}
}

// WriteMessage writes a lightning Message to w including the necessary header
// information and returns the number of bytes written.
func WriteMessage(w io.Writer, msg Message, pver uint32,
	btcnet wire.BitcoinNet) (int, error) {

	totalBytes := 0

	cmd := msg.Command()

	// Encode the message payload
	var bw bytes.Buffer
	err := msg.Encode(&bw, pver)
	if err != nil {
		return totalBytes, err
	}
	payload := bw.Bytes()
	lenp := len(payload)

	// Enforce maximum overall message payload
	if lenp > MaxMessagePayload {
		return totalBytes, fmt.Errorf("message payload is too large - "+
			"encoded %d bytes, but maximum message payload is %d bytes",
			lenp, MaxMessagePayload)
	}

	// Enforce maximum message payload on the message type
	mpl := msg.MaxPayloadLength(pver)
	if uint32(lenp) > mpl {
		return totalBytes, fmt.Errorf("message payload is too large - "+
			"encoded %d bytes, but maximum message payload of "+
			"type %x is %d bytes", lenp, cmd, mpl)
	}

	// Create header for the message.
	hdr := messageHeader{magic: btcnet, command: cmd, length: uint32(lenp)}

	// Encode the header for the message. This is done to a buffer
	// rather than directly to the writer since writeElements doesn't
	// return the number of bytes written.
	hw := bytes.NewBuffer(make([]byte, 0, MessageHeaderSize))
	if err := writeElements(hw, hdr.magic, hdr.command, hdr.length); err != nil {
		return 0, nil
	}

	// Write the head first.
	n, err := w.Write(hw.Bytes())
	totalBytes += n
	if err != nil {
		return totalBytes, err
	}

	// Write payload the payload itself after the header.
	n, err = w.Write(payload)
	totalBytes += n
	if err != nil {
		return totalBytes, err
	}

	return totalBytes, nil
}

// ReadMessageN reads, validates, and parses the next bitcoin Message from r for
// the provided protocol version and bitcoin network.  It returns the number of
// bytes read in addition to the parsed Message and raw bytes which comprise the
// message.  This function is the same as ReadMessage except it also returns the
// number of bytes read.
func ReadMessage(r io.Reader, pver uint32, btcnet wire.BitcoinNet) (int, Message, []byte, error) {
	totalBytes := 0
	n, hdr, err := readMessageHeader(r)
	totalBytes += n
	if err != nil {
		return totalBytes, nil, nil, err
	}

	// Enforce maximum message payload
	if hdr.length > MaxMessagePayload {
		return totalBytes, nil, nil, fmt.Errorf("message payload is "+
			"too large - header indicates %d bytes, but max "+
			"message payload is %d bytes.", hdr.length,
			MaxMessagePayload)
	}

	// Check for messages in the wrong network.
	if hdr.magic != btcnet {
		discardInput(r, hdr.length)
		return totalBytes, nil, nil, fmt.Errorf("message from other "+
			"network [%v]", hdr.magic)
	}

	// Create struct of appropriate message type based on the command.
	command := hdr.command
	msg, err := makeEmptyMessage(command)
	if err != nil {
		discardInput(r, hdr.length)
		return totalBytes, nil, nil, fmt.Errorf("ReadMessage %s",
			err.Error())
	}

	// Check for maximum length based on the message type.
	mpl := msg.MaxPayloadLength(pver)
	if hdr.length > mpl {
		discardInput(r, hdr.length)
		return totalBytes, nil, nil, fmt.Errorf("payload exceeds max "+
			"length. indicates %v bytes, but max of message type %v is %v.",
			hdr.length, command, mpl)
	}

	// Read payload.
	payload := make([]byte, hdr.length)
	n, err = io.ReadFull(r, payload)
	totalBytes += n
	if err != nil {
		return totalBytes, nil, nil, err
	}

	// Unmarshal message.
	pr := bytes.NewBuffer(payload)
	if err = msg.Decode(pr, pver); err != nil {
		return totalBytes, nil, nil, err
	}

	// Validate the data.
	if err = msg.Validate(); err != nil {
		return totalBytes, nil, nil, err
	}

	return totalBytes, msg, payload, nil
}
<|endoftext|>


<commit_before>package mem

import (
	"github.com/RoaringBitmap/roaring"
	"github.com/blevesearch/bleve/index/scorch/segment"
)

// _id field is always guaranteed to have fieldID of 0
const idFieldID uint16 = 0

// KNOWN ISSUES
// - LIMITATION - we decided whether or not to store term vectors for a field
//                at the segment level, based on the first definition of a
//                field we see.  in normal bleve usage this is fine, all
//                instances of a field definition will be the same.  however,
//                advanced users may violate this and provide unique field
//                definitions with each document.  this segment does not
//                support this usage.

// TODO
// - need better testing of multiple docs, iterating freqs, locations and
//   and verifying the correct results are returned
// - need tests for term dictionary iteration

// Segment is an in memory implementation of scorch.Segment
type Segment struct {

	// FieldsMap name -> id+1
	FieldsMap map[string]uint16
	// fields id -> name
	FieldsInv []string
	// field id -> has location info
	FieldsLoc []bool

	// term dictionary
	//  field id -> term -> posting id + 1
	Dicts []map[string]uint64

	// term dictionary keys
	//  field id -> []dictionary keys
	DictKeys [][]string

	// Postings list
	//  Postings list id -> Postings bitmap
	Postings []*roaring.Bitmap

	// term frequencies
	//  postings list id -> Freqs (one for each hit in bitmap)
	Freqs [][]uint64

	// field Norms
	//  postings list id -> Norms (one for each hit in bitmap)
	Norms [][]float32

	// field/start/end/pos/locarraypos
	//  postings list id -> start/end/pos/locarraypos (one for each freq)
	Locfields   [][]uint16
	Locstarts   [][]uint64
	Locends     [][]uint64
	Locpos      [][]uint64
	Locarraypos [][][]uint64

	// Stored field values
	//  docNum -> field id -> slice of values (each value []byte)
	Stored []map[uint16][][]byte

	// stored field types
	//  docNum -> field id -> slice of types (each type byte)
	StoredTypes []map[uint16][]byte

	// stored field array positions
	//  docNum -> field id -> slice of array positions (each is []uint64)
	StoredPos []map[uint16][][]uint64
}

// New builds a new empty Segment
func New() *Segment {
	return &Segment{
		FieldsMap: map[string]uint16{},
	}
}

// Fields returns the field names used in this segment
func (s *Segment) Fields() []string {
	return s.FieldsInv
}

// VisitDocument invokes the DocFieldValueVistor for each stored field
// for the specified doc number
func (s *Segment) VisitDocument(num uint64, visitor segment.DocumentFieldValueVisitor) error {
	// ensure document number exists
	if int(num) > len(s.Stored)-1 {
		return nil
	}
	docFields := s.Stored[int(num)]
	for field, values := range docFields {
		for i, value := range values {
			keepGoing := visitor(s.FieldsInv[field], s.StoredTypes[int(num)][field][i], value, s.StoredPos[int(num)][field][i])
			if !keepGoing {
				return nil
			}
		}
	}
	return nil
}

// Dictionary returns the term dictionary for the specified field
func (s *Segment) Dictionary(field string) segment.TermDictionary {
	return &Dictionary{
		segment: s,
		field:   field,
		fieldID: uint16(s.getOrDefineField(field, false)),
	}
}

// Count returns the number of documents in this segment
// (this has no notion of deleted docs)
func (s *Segment) Count() uint64 {
	return uint64(len(s.Stored))
}

// DocNumbers returns a bitset corresponding to the doc numbers of all the
// provided _id strings
func (s *Segment) DocNumbers(ids []string) *roaring.Bitmap {

	idDictionary := s.Dicts[idFieldID]
	rv := roaring.New()
	for _, id := range ids {
		postingID := idDictionary[id]
		if postingID > 0 {
			rv.Or(s.Postings[postingID-1])
		}
	}
	return rv
}
<commit_msg>fix crash in DocNumbers when segment is empty<commit_after>package mem

import (
	"github.com/RoaringBitmap/roaring"
	"github.com/blevesearch/bleve/index/scorch/segment"
)

// _id field is always guaranteed to have fieldID of 0
const idFieldID uint16 = 0

// KNOWN ISSUES
// - LIMITATION - we decided whether or not to store term vectors for a field
//                at the segment level, based on the first definition of a
//                field we see.  in normal bleve usage this is fine, all
//                instances of a field definition will be the same.  however,
//                advanced users may violate this and provide unique field
//                definitions with each document.  this segment does not
//                support this usage.

// TODO
// - need better testing of multiple docs, iterating freqs, locations and
//   and verifying the correct results are returned
// - need tests for term dictionary iteration

// Segment is an in memory implementation of scorch.Segment
type Segment struct {

	// FieldsMap name -> id+1
	FieldsMap map[string]uint16
	// fields id -> name
	FieldsInv []string
	// field id -> has location info
	FieldsLoc []bool

	// term dictionary
	//  field id -> term -> posting id + 1
	Dicts []map[string]uint64

	// term dictionary keys
	//  field id -> []dictionary keys
	DictKeys [][]string

	// Postings list
	//  Postings list id -> Postings bitmap
	Postings []*roaring.Bitmap

	// term frequencies
	//  postings list id -> Freqs (one for each hit in bitmap)
	Freqs [][]uint64

	// field Norms
	//  postings list id -> Norms (one for each hit in bitmap)
	Norms [][]float32

	// field/start/end/pos/locarraypos
	//  postings list id -> start/end/pos/locarraypos (one for each freq)
	Locfields   [][]uint16
	Locstarts   [][]uint64
	Locends     [][]uint64
	Locpos      [][]uint64
	Locarraypos [][][]uint64

	// Stored field values
	//  docNum -> field id -> slice of values (each value []byte)
	Stored []map[uint16][][]byte

	// stored field types
	//  docNum -> field id -> slice of types (each type byte)
	StoredTypes []map[uint16][]byte

	// stored field array positions
	//  docNum -> field id -> slice of array positions (each is []uint64)
	StoredPos []map[uint16][][]uint64
}

// New builds a new empty Segment
func New() *Segment {
	return &Segment{
		FieldsMap: map[string]uint16{},
	}
}

// Fields returns the field names used in this segment
func (s *Segment) Fields() []string {
	return s.FieldsInv
}

// VisitDocument invokes the DocFieldValueVistor for each stored field
// for the specified doc number
func (s *Segment) VisitDocument(num uint64, visitor segment.DocumentFieldValueVisitor) error {
	// ensure document number exists
	if int(num) > len(s.Stored)-1 {
		return nil
	}
	docFields := s.Stored[int(num)]
	for field, values := range docFields {
		for i, value := range values {
			keepGoing := visitor(s.FieldsInv[field], s.StoredTypes[int(num)][field][i], value, s.StoredPos[int(num)][field][i])
			if !keepGoing {
				return nil
			}
		}
	}
	return nil
}

// Dictionary returns the term dictionary for the specified field
func (s *Segment) Dictionary(field string) segment.TermDictionary {
	return &Dictionary{
		segment: s,
		field:   field,
		fieldID: uint16(s.getOrDefineField(field, false)),
	}
}

// Count returns the number of documents in this segment
// (this has no notion of deleted docs)
func (s *Segment) Count() uint64 {
	return uint64(len(s.Stored))
}

// DocNumbers returns a bitset corresponding to the doc numbers of all the
// provided _id strings
func (s *Segment) DocNumbers(ids []string) *roaring.Bitmap {
	rv := roaring.New()

	// guard against empty segment
	if len(s.FieldsMap) > 0 {
		idDictionary := s.Dicts[idFieldID]

		for _, id := range ids {
			postingID := idDictionary[id]
			if postingID > 0 {
				rv.Or(s.Postings[postingID-1])
			}
		}
	}
	return rv
}
<|endoftext|>