<commit_before>package repo_manager

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"html/template"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"path"
	"strconv"
	"strings"
	"time"

	"go.skia.org/infra/autoroll/go/strategy"
	"go.skia.org/infra/go/exec"
	"go.skia.org/infra/go/gerrit"
	"go.skia.org/infra/go/git"
	"go.skia.org/infra/go/sklog"
	"go.skia.org/infra/go/util"
)

/*
	Repo Manager which rolls infra asset versions.
*/

const (
	// Path to asset VERSION files within a given repo.
	ASSET_VERSION_TMPL = "infra/bots/assets/%s/VERSION"

	// Commit message for AssetRepoManager.
	// TODO(borenet): autoroll.FromGerritChangeInfo wants to parse the
	// revisions from the roll subject. In this case, it makes no sense
	// bit it's still required for the roller to operate. We should find
	// a way around it so that we can be more free with our commit message.
	TMPL_COMMIT_MESSAGE_ASSETS = `Roll {{.Asset}} {{.From}}..{{.To}}

The AutoRoll server is located here: {{.ServerURL}}

Documentation for the AutoRoller is here:
https://skia.googlesource.com/buildbot/+/master/autoroll/README.md

If the roll is causing failures, please contact the current sheriff, who should
be CC'd on the roll, and stop the roller if necessary.

TBR={{.Emails}}
`
)

var (
	commitMsgTmplAssets = template.Must(template.New("commitMsgAssets").Parse(TMPL_COMMIT_MESSAGE_ASSETS))
)

// AssetRepoManagerConfig provides configuration for the AssetRepoManager.
type AssetRepoManagerConfig struct {
	DepotToolsRepoManagerConfig

	// Which assets to roll.
	Asset string `json:"asset"`

	// URL of the child repo.
	ChildRepo string `json:"childRepo"`
}

// Validate the AssetRepoManagerConfig.
func (c *AssetRepoManagerConfig) Validate() error {
	if err := c.DepotToolsRepoManagerConfig.Validate(); err != nil {
		return err
	}
	if c.Asset == "" {
		return errors.New("Asset is required.")
	}
	if c.ChildRepo == "" {
		return errors.New("ChildRepo is required.")
	}
	return nil
}

// assetRepoManager is a RepoManager which rolls infra asset versions so that
// the "parent" repo's versions match the "child" repo's versions.
type assetRepoManager struct {
	*depotToolsRepoManager
	asset string
}

// NewAssetRepoManager returns a RepoManager instance which rolls infra asset
// versions so that the "parent" repo's versions match the "child" repo's
// versions.
func NewAssetRepoManager(ctx context.Context, c *AssetRepoManagerConfig, workdir string, g gerrit.GerritInterface, recipeCfgFile, serverURL string, client *http.Client) (RepoManager, error) {
	if err := c.Validate(); err != nil {
		return nil, err
	}
	drm, err := newDepotToolsRepoManager(ctx, c.DepotToolsRepoManagerConfig, workdir, recipeCfgFile, serverURL, g, client)
	if err != nil {
		return nil, err
	}
	rv := &assetRepoManager{
		depotToolsRepoManager: drm,
		asset:                 c.Asset,
	}
	// By default, depotToolsRepoManager assumes that the child repo is
	// nested inside the parent repo and is obtained via "gclient sync".
	// This is not the case for AssetRepoManager, so we have to sync it
	// ourselves.
	childRepo, err := git.NewCheckout(ctx, c.ChildRepo, workdir)
	if err != nil {
		return nil, err
	}
	rv.childRepo = childRepo
	rv.childDir = childRepo.Dir()
	return rv, nil
}

// See documentation for noCheckoutRepoManagerBuildCommitMessageFunc.
func (rm *assetRepoManager) buildCommitMessage(from, to, serverURL, cqExtraTrybots string, emails []string) (string, error) {
	rm.infoMtx.RLock()
	defer rm.infoMtx.RUnlock()

	// Build the commit message.
	data := struct {
		Asset     string
		Emails    string
		From      string
		ServerURL string
		To        string
	}{
		Asset:     rm.asset,
		Emails:    strings.Join(emails, ","),
		From:      from,
		ServerURL: serverURL,
		To:        to,
	}
	var buf bytes.Buffer
	if err := commitMsgTmplAssets.Execute(&buf, data); err != nil {
		return "", err
	}
	msg := buf.String()
	if cqExtraTrybots != "" {
		msg += fmt.Sprintf("\nCQ_INCLUDE_TRYBOTS=%s", cqExtraTrybots)
	}
	return buf.String(), nil
}

// Read the version of the given asset from the given repo at the given commit.
func readAssetVersion(ctx context.Context, repo git.GitDir, commit, asset string) (string, error) {
	return repo.GetFile(ctx, fmt.Sprintf(ASSET_VERSION_TMPL, asset), commit)
}

// See documentation for RepoManager interface.
func (rm *assetRepoManager) Update(ctx context.Context) error {
	rm.repoMtx.Lock()
	defer rm.repoMtx.Unlock()

	if err := rm.createAndSyncParent(ctx); err != nil {
		return fmt.Errorf("Could not create and sync parent repo: %s", err)
	}

	// By default, depotToolsRepoManager assumes that the child repo is
	// nested inside the parent repo and is obtained via "gclient sync".
	// This is not the case for AssetRepoManager, so we have to sync it
	// ourselves.
	if err := rm.childRepo.Update(ctx); err != nil {
		return err
	}

	// Obtain the hash of the current child HEAD.
	childHead, err := rm.childRepo.RevParse(ctx, rm.childBranch)
	if err != nil {
		return err
	}

	// Read the asset versions from both repos.
	lastRollRev, err := readAssetVersion(ctx, git.GitDir(rm.parentDir), rm.parentBranch, rm.asset)
	if err != nil {
		return err
	}
	lastRollRev = strings.TrimSpace(lastRollRev)
	nextRollRev, err := readAssetVersion(ctx, git.GitDir(rm.childDir), childHead, rm.asset)
	if err != nil {
		return err
	}
	nextRollRev = strings.TrimSpace(nextRollRev)

	// Subtract the last-rolled version number from the next version number
	// to obtain commitsNotRolled.
	lastInt, err := strconv.Atoi(lastRollRev)
	if err != nil {
		return err
	}
	nextInt, err := strconv.Atoi(nextRollRev)
	if err != nil {
		return err
	}
	commitsNotRolled := nextInt - lastInt

	rm.infoMtx.Lock()
	defer rm.infoMtx.Unlock()
	rm.lastRollRev = lastRollRev
	rm.nextRollRev = nextRollRev
	rm.commitsNotRolled = commitsNotRolled
	return nil
}

// See documentation for RepoManager interface.
func (rm *assetRepoManager) CreateNewRoll(ctx context.Context, from, to string, emails []string, cqExtraTrybots string, dryRun bool) (int64, error) {
	rm.repoMtx.Lock()
	defer rm.repoMtx.Unlock()

	// Clean the checkout, get onto a fresh branch.
	if err := rm.cleanParent(ctx); err != nil {
		return 0, err
	}
	if _, err := exec.RunCwd(ctx, rm.parentDir, "git", "checkout", "-b", ROLL_BRANCH, "-t", fmt.Sprintf("origin/%s", rm.parentBranch), "-f"); err != nil {
		return 0, err
	}

	// Defer some more cleanup.
	defer func() {
		util.LogErr(rm.cleanParent(ctx))
	}()

	// Create the roll CL.
	if _, err := exec.RunCwd(ctx, rm.parentDir, "git", "config", "user.name", getLocalPartOfEmailAddress(rm.user)); err != nil {
		return 0, err
	}
	if _, err := exec.RunCwd(ctx, rm.parentDir, "git", "config", "user.email", rm.user); err != nil {
		return 0, err
	}
	assetFile := path.Join(rm.parentDir, fmt.Sprintf(ASSET_VERSION_TMPL, rm.asset))
	if err := util.WithWriteFile(assetFile, func(w io.Writer) error {
		_, err := w.Write([]byte(to))
		return err
	}); err != nil {
		return 0, err
	}

	// Run the pre-upload steps.
	sklog.Infof("Running pre-upload steps.")
	for _, s := range rm.PreUploadSteps() {
		if err := s(ctx, rm.httpClient, rm.parentDir); err != nil {
			return 0, fmt.Errorf("Failed pre-upload step: %s", err)
		}
	}

	// Commit.
	commitMsg, err := rm.buildCommitMessage(from, to, rm.serverURL, cqExtraTrybots, emails)
	if err != nil {
		return 0, err
	}
	if _, err := exec.RunCwd(ctx, rm.parentDir, "git", "commit", "-a", "-m", commitMsg); err != nil {
		return 0, err
	}

	// Upload the CL.
	uploadCmd := &exec.Command{
		Dir:        rm.parentDir,
		Env:        rm.depotToolsEnv,
		InheritEnv: true,
		Name:       "git",
		Args:       []string{"cl", "upload", "--bypass-hooks", "-f", "-v", "-v"},
		Timeout:    2 * time.Minute,
	}
	if dryRun {
		uploadCmd.Args = append(uploadCmd.Args, "--cq-dry-run")
	} else {
		uploadCmd.Args = append(uploadCmd.Args, "--use-commit-queue")
	}
	uploadCmd.Args = append(uploadCmd.Args, "--gerrit")
	uploadCmd.Args = append(uploadCmd.Args, "-m", commitMsg)

	// Upload the CL.
	sklog.Infof("Running command: git %s", strings.Join(uploadCmd.Args, " "))
	if _, err := exec.RunCommand(ctx, uploadCmd); err != nil {
		return 0, err
	}

	// Obtain the issue number.
	sklog.Infof("Retrieving issue number of uploaded CL.")
	tmp, err := ioutil.TempDir("", "")
	if err != nil {
		return 0, err
	}
	defer util.RemoveAll(tmp)
	jsonFile := path.Join(tmp, "issue.json")
	if _, err := exec.RunCommand(ctx, &exec.Command{
		Dir:        rm.parentDir,
		Env:        rm.depotToolsEnv,
		InheritEnv: true,
		Name:       "git",
		Args:       []string{"cl", "issue", fmt.Sprintf("--json=%s", jsonFile)},
	}); err != nil {
		return 0, err
	}
	f, err := os.Open(jsonFile)
	if err != nil {
		return 0, err
	}
	var issue issueJson
	if err := json.NewDecoder(f).Decode(&issue); err != nil {
		return 0, err
	}
	return issue.Issue, nil
}

// See documentation for RepoManager interface.
func (rm *assetRepoManager) FullChildHash(ctx context.Context, version string) (string, error) {
	// We're rolling asset version numbers, not commit hashes, and we don't
	// shorten them. Return the given version.
	return version, nil
}

// See documentation for RepoManager interface.
func (r *assetRepoManager) CreateNextRollStrategy(ctx context.Context, s string) (strategy.NextRollStrategy, error) {
	return strategy.GetNextRollStrategy(ctx, s, r.childBranch, DEFAULT_REMOTE, "", []string{}, nil, nil)
}

// See documentation for RepoManager interface.
func (r *assetRepoManager) SetStrategy(s strategy.NextRollStrategy) {
	r.strategyMtx.Lock()
	defer r.strategyMtx.Unlock()
	r.strategy = s
}

// See documentation for RepoManager interface.
func (r *assetRepoManager) DefaultStrategy() string {
	return strategy.ROLL_STRATEGY_BATCH
}

// See documentation for RepoManager interface.
func (r *assetRepoManager) ValidStrategies() []string {
	return []string{
		strategy.ROLL_STRATEGY_BATCH,
	}
}

// See documentation for RepoManager interface.
func (rm *assetRepoManager) RolledPast(ctx context.Context, version string) (bool, error) {
	lastInt, err := strconv.Atoi(rm.lastRollRev)
	if err != nil {
		return false, err
	}
	versionInt, err := strconv.Atoi(version)
	if err != nil {
		return false, err
	}
	return lastInt >= versionInt, nil
}
<commit_msg>[autoroll] Fixes for AssetRepoManager<commit_after>package repo_manager

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"html/template"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"path"
	"strconv"
	"strings"
	"time"

	"go.skia.org/infra/autoroll/go/strategy"
	"go.skia.org/infra/go/exec"
	"go.skia.org/infra/go/gerrit"
	"go.skia.org/infra/go/git"
	"go.skia.org/infra/go/sklog"
	"go.skia.org/infra/go/util"
)

/*
	Repo Manager which rolls infra asset versions.
*/

const (
	// Path to asset VERSION files within a given repo.
	ASSET_VERSION_TMPL = "infra/bots/assets/%s/VERSION"

	// Commit message for AssetRepoManager.
	// TODO(borenet): autoroll.FromGerritChangeInfo wants to parse the
	// revisions from the roll subject. In this case, it makes no sense
	// bit it's still required for the roller to operate. We should find
	// a way around it so that we can be more free with our commit message.
	TMPL_COMMIT_MESSAGE_ASSETS = `Roll {{.Asset}} {{.From}}..{{.To}}

The AutoRoll server is located here: {{.ServerURL}}

Documentation for the AutoRoller is here:
https://skia.googlesource.com/buildbot/+/master/autoroll/README.md

If the roll is causing failures, please contact the current sheriff, who should
be CC'd on the roll, and stop the roller if necessary.

TBR={{.Emails}}
`
)

var (
	commitMsgTmplAssets = template.Must(template.New("commitMsgAssets").Parse(TMPL_COMMIT_MESSAGE_ASSETS))
)

// AssetRepoManagerConfig provides configuration for the AssetRepoManager.
type AssetRepoManagerConfig struct {
	DepotToolsRepoManagerConfig

	// Which assets to roll.
	Asset string `json:"asset"`

	// URL of the child repo.
	ChildRepo string `json:"childRepo"`
}

// Validate the AssetRepoManagerConfig.
func (c *AssetRepoManagerConfig) Validate() error {
	if err := c.DepotToolsRepoManagerConfig.Validate(); err != nil {
		return err
	}
	if c.Asset == "" {
		return errors.New("Asset is required.")
	}
	if c.ChildRepo == "" {
		return errors.New("ChildRepo is required.")
	}
	return nil
}

// assetRepoManager is a RepoManager which rolls infra asset versions so that
// the "parent" repo's versions match the "child" repo's versions.
type assetRepoManager struct {
	*depotToolsRepoManager
	asset string
}

// NewAssetRepoManager returns a RepoManager instance which rolls infra asset
// versions so that the "parent" repo's versions match the "child" repo's
// versions.
func NewAssetRepoManager(ctx context.Context, c *AssetRepoManagerConfig, workdir string, g gerrit.GerritInterface, recipeCfgFile, serverURL string, client *http.Client) (RepoManager, error) {
	if err := c.Validate(); err != nil {
		return nil, err
	}
	drm, err := newDepotToolsRepoManager(ctx, c.DepotToolsRepoManagerConfig, workdir, recipeCfgFile, serverURL, g, client)
	if err != nil {
		return nil, err
	}
	rv := &assetRepoManager{
		depotToolsRepoManager: drm,
		asset:                 c.Asset,
	}
	// By default, depotToolsRepoManager assumes that the child repo is
	// nested inside the parent repo and is obtained via "gclient sync".
	// This is not the case for AssetRepoManager, so we have to sync it
	// ourselves.
	childRepo, err := git.NewCheckout(ctx, c.ChildRepo, workdir)
	if err != nil {
		return nil, err
	}
	rv.childRepo = childRepo
	rv.childDir = childRepo.Dir()
	return rv, nil
}

// See documentation for noCheckoutRepoManagerBuildCommitMessageFunc.
func (rm *assetRepoManager) buildCommitMessage(from, to, serverURL, cqExtraTrybots string, emails []string) (string, error) {
	rm.infoMtx.RLock()
	defer rm.infoMtx.RUnlock()

	// Build the commit message.
	data := struct {
		Asset     string
		Emails    string
		From      string
		ServerURL string
		To        string
	}{
		Asset:     rm.asset,
		Emails:    strings.Join(emails, ","),
		From:      from,
		ServerURL: serverURL,
		To:        to,
	}
	var buf bytes.Buffer
	if err := commitMsgTmplAssets.Execute(&buf, data); err != nil {
		return "", err
	}
	msg := buf.String()
	if cqExtraTrybots != "" {
		msg += fmt.Sprintf("\nCQ_INCLUDE_TRYBOTS=%s", cqExtraTrybots)
	}
	return buf.String(), nil
}

// Read the version of the given asset from the given repo at the given commit.
func readAssetVersion(ctx context.Context, repo git.GitDir, commit, asset string) (string, error) {
	return repo.GetFile(ctx, fmt.Sprintf(ASSET_VERSION_TMPL, asset), commit)
}

// See documentation for RepoManager interface.
func (rm *assetRepoManager) Update(ctx context.Context) error {
	rm.repoMtx.Lock()
	defer rm.repoMtx.Unlock()

	if err := rm.createAndSyncParent(ctx); err != nil {
		return fmt.Errorf("Could not create and sync parent repo: %s", err)
	}

	// By default, depotToolsRepoManager assumes that the child repo is
	// nested inside the parent repo and is obtained via "gclient sync".
	// This is not the case for AssetRepoManager, so we have to sync it
	// ourselves.
	if err := rm.childRepo.Update(ctx); err != nil {
		return err
	}

	// Obtain the hash of the current parent and child HEADs.
	parentHead, err := git.GitDir(rm.parentDir).RevParse(ctx, "origin/"+rm.parentBranch)
	if err != nil {
		return err
	}
	childHead, err := rm.childRepo.RevParse(ctx, "origin/"+rm.childBranch)
	if err != nil {
		return err
	}

	// Read the asset versions from both repos.
	lastRollRev, err := readAssetVersion(ctx, git.GitDir(rm.parentDir), parentHead, rm.asset)
	if err != nil {
		return err
	}
	lastRollRev = strings.TrimSpace(lastRollRev)
	nextRollRev, err := readAssetVersion(ctx, git.GitDir(rm.childDir), childHead, rm.asset)
	if err != nil {
		return err
	}
	nextRollRev = strings.TrimSpace(nextRollRev)

	// Subtract the last-rolled version number from the next version number
	// to obtain commitsNotRolled.
	lastInt, err := strconv.Atoi(lastRollRev)
	if err != nil {
		return err
	}
	nextInt, err := strconv.Atoi(nextRollRev)
	if err != nil {
		return err
	}
	commitsNotRolled := nextInt - lastInt

	rm.infoMtx.Lock()
	defer rm.infoMtx.Unlock()
	rm.lastRollRev = lastRollRev
	rm.nextRollRev = nextRollRev
	rm.commitsNotRolled = commitsNotRolled
	return nil
}

// See documentation for RepoManager interface.
func (rm *assetRepoManager) CreateNewRoll(ctx context.Context, from, to string, emails []string, cqExtraTrybots string, dryRun bool) (int64, error) {
	rm.repoMtx.Lock()
	defer rm.repoMtx.Unlock()

	// Clean the checkout, get onto a fresh branch.
	if err := rm.cleanParent(ctx); err != nil {
		return 0, err
	}
	if _, err := exec.RunCwd(ctx, rm.parentDir, "git", "checkout", "-b", ROLL_BRANCH, "-t", fmt.Sprintf("origin/%s", rm.parentBranch), "-f"); err != nil {
		return 0, err
	}

	// Defer some more cleanup.
	defer func() {
		util.LogErr(rm.cleanParent(ctx))
	}()

	// Create the roll CL.
	if _, err := exec.RunCwd(ctx, rm.parentDir, "git", "config", "user.name", getLocalPartOfEmailAddress(rm.user)); err != nil {
		return 0, err
	}
	if _, err := exec.RunCwd(ctx, rm.parentDir, "git", "config", "user.email", rm.user); err != nil {
		return 0, err
	}
	assetFile := path.Join(rm.parentDir, fmt.Sprintf(ASSET_VERSION_TMPL, rm.asset))
	if err := util.WithWriteFile(assetFile, func(w io.Writer) error {
		_, err := w.Write([]byte(to))
		return err
	}); err != nil {
		return 0, err
	}

	// Run the pre-upload steps.
	sklog.Infof("Running pre-upload steps.")
	for _, s := range rm.PreUploadSteps() {
		if err := s(ctx, rm.httpClient, rm.parentDir); err != nil {
			return 0, fmt.Errorf("Failed pre-upload step: %s", err)
		}
	}

	// Commit.
	commitMsg, err := rm.buildCommitMessage(from, to, rm.serverURL, cqExtraTrybots, emails)
	if err != nil {
		return 0, err
	}
	if _, err := exec.RunCwd(ctx, rm.parentDir, "git", "commit", "-a", "-m", commitMsg); err != nil {
		return 0, err
	}

	// Upload the CL.
	uploadCmd := &exec.Command{
		Dir:        rm.parentDir,
		Env:        rm.depotToolsEnv,
		InheritEnv: true,
		Name:       "git",
		Args:       []string{"cl", "upload", "--bypass-hooks", "-f", "-v", "-v"},
		Timeout:    2 * time.Minute,
	}
	if dryRun {
		uploadCmd.Args = append(uploadCmd.Args, "--cq-dry-run")
	} else {
		uploadCmd.Args = append(uploadCmd.Args, "--use-commit-queue")
	}
	uploadCmd.Args = append(uploadCmd.Args, "--gerrit")
	uploadCmd.Args = append(uploadCmd.Args, "-m", commitMsg)

	// Upload the CL.
	sklog.Infof("Running command: git %s", strings.Join(uploadCmd.Args, " "))
	if _, err := exec.RunCommand(ctx, uploadCmd); err != nil {
		return 0, err
	}

	// Obtain the issue number.
	sklog.Infof("Retrieving issue number of uploaded CL.")
	tmp, err := ioutil.TempDir("", "")
	if err != nil {
		return 0, err
	}
	defer util.RemoveAll(tmp)
	jsonFile := path.Join(tmp, "issue.json")
	if _, err := exec.RunCommand(ctx, &exec.Command{
		Dir:        rm.parentDir,
		Env:        rm.depotToolsEnv,
		InheritEnv: true,
		Name:       "git",
		Args:       []string{"cl", "issue", fmt.Sprintf("--json=%s", jsonFile)},
	}); err != nil {
		return 0, err
	}
	f, err := os.Open(jsonFile)
	if err != nil {
		return 0, err
	}
	var issue issueJson
	if err := json.NewDecoder(f).Decode(&issue); err != nil {
		return 0, err
	}
	return issue.Issue, nil
}

// See documentation for RepoManager interface.
func (rm *assetRepoManager) FullChildHash(ctx context.Context, version string) (string, error) {
	// We're rolling asset version numbers, not commit hashes, and we don't
	// shorten them. Return the given version.
	return version, nil
}

// See documentation for RepoManager interface.
func (r *assetRepoManager) CreateNextRollStrategy(ctx context.Context, s string) (strategy.NextRollStrategy, error) {
	return strategy.GetNextRollStrategy(ctx, s, r.childBranch, DEFAULT_REMOTE, "", []string{}, nil, nil)
}

// See documentation for RepoManager interface.
func (r *assetRepoManager) SetStrategy(s strategy.NextRollStrategy) {
	r.strategyMtx.Lock()
	defer r.strategyMtx.Unlock()
	r.strategy = s
}

// See documentation for RepoManager interface.
func (r *assetRepoManager) DefaultStrategy() string {
	return strategy.ROLL_STRATEGY_BATCH
}

// See documentation for RepoManager interface.
func (r *assetRepoManager) ValidStrategies() []string {
	return []string{
		strategy.ROLL_STRATEGY_BATCH,
	}
}

// See documentation for RepoManager interface.
func (rm *assetRepoManager) RolledPast(ctx context.Context, version string) (bool, error) {
	lastInt, err := strconv.Atoi(rm.lastRollRev)
	if err != nil {
		return false, err
	}
	versionInt, err := strconv.Atoi(version)
	if err != nil {
		return false, err
	}
	return lastInt >= versionInt, nil
}
<|endoftext|>


<commit_before>package matchers

import (
	"testing"
)

func TestEqualTo(t *testing.T) {
	m := EqualTo{true}
	if !m.Match(true) {
		t.Error("true is not true")
	}
	if m.Match(false) {
		t.Error("true is false")
	}
}

func TestAssertThat(t *testing.T) {
	f := new(testing.T)
	AssertThat(Protect{f}, true, EqualTo{true})
	if f.Failed() {
		t.Error("true is not true")
	}
	AssertThat(Protect{f}, true, EqualTo{false})
	if !f.Failed() {
		t.Error("true is false")
	}
}

func TestIs(t *testing.T) {
	AssertThat(t, true, Is{true})
	AssertThat(t, true, Is{EqualTo{true}})
	AssertThat(t, Expect{true, Is{false}}, Fails{})
}

func TestNot(t *testing.T) {
	AssertThat(t, true, Not{false})
	AssertThat(t, Expect{true, Not{true}}, Fails{})
}

func TestAllOf(t *testing.T) {
	AssertThat(t, true, AllOf{Is{true}, Not{false}})
	AssertThat(t, Expect{true, AllOf{Is{false}, Not{true}}}, Fails{})
	AssertThat(t, Expect{true, AllOf{Is{true}, Not{true}}}, Fails{})
	AssertThat(t, Expect{true, AllOf{Is{false}, Not{true}}}, Fails{})
}

func TestAnyOf(t *testing.T) {
	AssertThat(t, true, AnyOf{Is{true}, Not{true}})
	AssertThat(t, true, AnyOf{Is{false}, Not{false}})
	AssertThat(t, true, AnyOf{Is{true}, Not{false}})
	AssertThat(t, Expect{true, AnyOf{Is{false}, Not{true}}}, Fails{})
}

func TestElementsAre(t *testing.T) {
	AssertThat(t, []int{1, 2, 3, 4, 5}, ElementsAre{5, 1, 4, 2, 3})
	AssertThat(t, Expect{[]int{1, 2}, ElementsAre{1, 2, 2}}, Fails{})
	AssertThat(t, Expect{[]int{1, 2}, ElementsAre{1, 3}}, Fails{})
}

func TestContains(t *testing.T) {
	AssertThat(t, []int{1, 2, 3, 4, 5}, Contains{5, 1})
	AssertThat(t, Expect{[]int{1, 2}, Contains{1, 3}}, Fails{})
	AssertThat(t, Expect{[]int{1, 2}, Contains{3}}, Fails{})
	AssertThat(t, []int{1, 2}, Contains{})
}

func TestTypeOf(t *testing.T) {
	AssertThat(t, "zzzzz", TypeOf{string("")})
	AssertThat(t, Expect{1, TypeOf{string("")}}, Fails{})
}

func TestFails(t *testing.T) {
	AssertThat(t, Expect{Expect{true, Is{true}}, Fails{}}, Fails{})
}

func TestFailsPanic(t *testing.T) {
	defer func() {
		e := recover()
		AssertThat(t, e, Is{Not{nil}})
	}()
	AssertThat(t, true, Fails{})
}
<commit_msg>Up code coverage to 100%<commit_after>package matchers

import (
	"testing"
)

func TestEqualTo(t *testing.T) {
	m := EqualTo{true}
	if !m.Match(true) {
		t.Error("true is not true")
	}
	if m.Match(false) {
		t.Error("true is false")
	}
}

func TestAssertThat(t *testing.T) {
	f := new(testing.T)
	AssertThat(Protect{f}, true, EqualTo{true})
	if f.Failed() {
		t.Error("true is not true")
	}
	AssertThat(Protect{f}, true, EqualTo{false})
	if !f.Failed() {
		t.Error("true is false")
	}
}

func TestIs(t *testing.T) {
	AssertThat(t, true, Is{true})
	AssertThat(t, true, Is{EqualTo{true}})
	AssertThat(t, Expect{true, Is{false}}, Fails{})
}

func TestNot(t *testing.T) {
	AssertThat(t, true, Not{false})
	AssertThat(t, Expect{true, Not{true}}, Fails{})
}

func TestAllOf(t *testing.T) {
	AssertThat(t, true, AllOf{Is{true}, Not{false}})
	AssertThat(t, Expect{true, AllOf{Is{false}, Not{true}}}, Fails{})
	AssertThat(t, Expect{true, AllOf{Is{true}, Not{true}}}, Fails{})
	AssertThat(t, Expect{true, AllOf{Is{false}, Not{true}}}, Fails{})
}

func TestAnyOf(t *testing.T) {
	AssertThat(t, true, AnyOf{Is{true}, Not{true}})
	AssertThat(t, true, AnyOf{Is{false}, Not{false}})
	AssertThat(t, true, AnyOf{Is{true}, Not{false}})
	AssertThat(t, Expect{true, AnyOf{Is{false}, Not{true}}}, Fails{})
}

func TestElementsAre(t *testing.T) {
	AssertThat(t, []int{1, 2, 3, 4, 5}, ElementsAre{5, 1, 4, 2, 3})
	AssertThat(t, Expect{[]int{1, 2}, ElementsAre{1, 2, 2}}, Fails{})
	AssertThat(t, Expect{[]int{1, 2}, ElementsAre{1, 3}}, Fails{})
	AssertThat(t, Expect{1, ElementsAre{1, 3}}, Fails{})
	v := []interface{}{}
	a, b := 1, ""
	v = append(v, a)
	v = append(v, b)
	AssertThat(t, v, ElementsAre{TypeOf{1}, TypeOf{""}})
}

func TestContains(t *testing.T) {
	AssertThat(t, []int{1, 2, 3, 4, 5}, Contains{5, 1})
	AssertThat(t, Expect{[]int{1, 2}, Contains{1, 3}}, Fails{})
	AssertThat(t, Expect{[]int{1, 2}, Contains{3}}, Fails{})
	AssertThat(t, []int{1, 2}, Contains{})
	AssertThat(t, Expect{1, Contains{}}, Fails{})
	v := []interface{}{}
	a, b := 1, ""
	v = append(v, a)
	v = append(v, b)
	AssertThat(t, v, Contains{TypeOf{1}})
}

func TestTypeOf(t *testing.T) {
	AssertThat(t, "zzzzz", TypeOf{string("")})
	AssertThat(t, Expect{1, TypeOf{string("")}}, Fails{})
}

func TestFails(t *testing.T) {
	AssertThat(t, Expect{Expect{true, Is{true}}, Fails{}}, Fails{})
}

func TestFailsPanic(t *testing.T) {
	defer func() {
		e := recover()
		AssertThat(t, e, Is{Not{nil}})
	}()
	AssertThat(t, true, Fails{})
}
<|endoftext|>


<commit_before>// Package jsonseq provides methods for reading and writing JSON text sequences
// (`application/json-seq`) as defined in RFC 7464 (https://tools.ietf.org/html/rfc7464).
package jsonseq

import (
	"bufio"
	"bytes"
	"encoding/json"
	"fmt"
	"io"
)

// ContentType is the MIME media type for JSON text sequences.
// See: https://tools.ietf.org/html/rfc7464#section-4
const ContentType = `application/json-seq`

const digitSet = "1234567980"

const (
	rs = 0x1E
	lf = 0x0A
	sp = 0x20
	tb = 0x09
	cr = 0x0D
)

// whitespace characters defined in https://tools.ietf.org/html/rfc7159#section-2.
var wsSet = []byte{sp, tb, lf, cr}

func wsByte(b byte) bool {
	return bytes.IndexByte(wsSet, b) >= 0
}

func wsRune(r rune) bool {
	return bytes.ContainsRune(wsSet, r)
}

// WriteRecord writes a JSON text sequence record with beginning
// (RS) and end (LF) marker bytes.
func WriteRecord(w io.Writer, json []byte) error {
	_, err := w.Write([]byte{rs})
	if err != nil {
		return err
	}
	_, err = w.Write(json)
	if err != nil {
		return err
	}
	_, err = w.Write([]byte{lf})
	return err
}

// A RecordWriter delimits the start of written records with a record separator.
//
// The standard library's json.Encoder calls Write just once for each value and
// always with a trailing line feed, so it can be adapted very simply to emit a
// JSON text sequence.
//
// 	encoder := json.NewEncoder(&jsonseq.RecordWriter{writer})
type RecordWriter struct {
	io.Writer
}

// Write prefixes every written record with an ASCII record separator. The caller
// is responsible for including a trailing line feed when necessary. Calls the
// underlying Writer exactly once.
func (w *RecordWriter) Write(record []byte) (int, error) {
	_, err := w.Writer.Write([]byte{rs})
	if err != nil {
		return -1, err
	}
	n, err := w.Writer.Write(record)
	return n + 1, err
}

// A Decoder reads and decodes JSON text sequence records from an input stream.
type Decoder struct {
	s *bufio.Scanner
}

// NewDecoder creates a Decoder.
func NewDecoder(r io.Reader) *Decoder {
	s := bufio.NewScanner(r)
	s.Split(ScanRecord)
	return &Decoder{s: s}
}

// Decode scans the next record, or returns an error. Extra trailing data is discarded.
// The Decoder remains valid until io.EOF is returned.
func (d *Decoder) Decode(v interface{}) error {
	if !d.s.Scan() {
		if err := d.s.Err(); err != nil {
			return err
		}
		return io.EOF
	}
	b := d.s.Bytes()

	b, ok := RecordValue(b)
	if !ok {
		return fmt.Errorf("invalid record: %q", string(b))
	}
	// Decode the first value, and discard any remaining data.
	return json.NewDecoder(bytes.NewReader(b)).Decode(v)
}

// RecordValue returns the *value* bytes from a JSON text sequence record and a flag
// indicating if the *record* is valid. This is *NOT* a validation of any contained JSON,
// which could itself be invalid or contain extra trailing values.
//
// See section 2.4: Top-Level Values: numbers, true, false, and null.
// https://tools.ietf.org/html/rfc7464#section-2.4
func RecordValue(b []byte) ([]byte, bool) {
	if len(b) < 2 {
		return b, false
	}
	if b[0] != rs {
		return b, false
	}
	// Drop rs and leading whitespace.
	b = bytes.TrimLeftFunc(b[1:], wsRune)

	// A number, true, false, or null value could be truncated if not
	// followed by whitespace.
	switch b[0] {
	case 'n':
		if bytes.HasPrefix(b, []byte("null")) {
			if wsByte(b[4]) {
				return b, true
			}
			return b, false
		}
	case 't':
		if bytes.HasPrefix(b, []byte("true")) {
			if wsByte(b[4]) {
				return b, true
			}
			return b, false
		}
	case 'f':
		if bytes.HasPrefix(b, []byte("false")) {
			if wsByte(b[5]) {
				return b, true
			}
			return b, false
		}
	case '-':
		if '0' <= b[1] && b[1] <= '9' {
			t := bytes.TrimLeft(b, digitSet)
			if len(t) > 0 && wsByte(t[0]) {
				return b, true
			}
			return b, false
		}
	case '0', '1', '2', '3', '4', '5', '6', '7', '8', '9':
		t := bytes.TrimLeft(b, digitSet)
		if len(t) > 0 && wsByte(t[0]) {
			return b, true
		}
		return b, false
	}

	return b, true
}

// ScanRecord is a bufio.SplitFunc which splits JSON text sequence records.
// Scanned bytes must be validated with the RecordValue function.
func ScanRecord(data []byte, atEOF bool) (advance int, token []byte, err error) {
	if atEOF && len(data) == 0 {
		return 0, nil, nil
	}
	// Find record start.
	switch i := bytes.IndexByte(data, rs); {
	case i < 0:
		if atEOF {
			// Partial record.
			return len(data), data, nil
		}
		// Request more data.
		return 0, nil, nil
	case i > 0:
		// Partial record.
		return i, data[0 : i-1], nil
	}
	// else i == 0

	// Drop consecutive leading rs's
	for len(data) > 1 && data[1] == rs {
		data = data[1:]
	}

	// Find end or next record.
	i := bytes.IndexByte(data[1:], rs)
	if i < 0 {
		if atEOF {
			return len(data), data, nil
		}
		// Request more data.
		return 0, nil, nil
	}
	return 1 + i, data[:1+i], nil
}
<commit_msg>add NewEncoder func; add NewDecoderFn func<commit_after>// Package jsonseq provides methods for reading and writing JSON text sequences
// (`application/json-seq`) as defined in RFC 7464 (https://tools.ietf.org/html/rfc7464).
package jsonseq

import (
	"bufio"
	"bytes"
	"encoding/json"
	"fmt"
	"io"
)

// ContentType is the MIME media type for JSON text sequences.
// See: https://tools.ietf.org/html/rfc7464#section-4
const ContentType = `application/json-seq`

const digitSet = "1234567980"

const (
	rs = 0x1E
	lf = 0x0A
	sp = 0x20
	tb = 0x09
	cr = 0x0D
)

// whitespace characters defined in https://tools.ietf.org/html/rfc7159#section-2.
var wsSet = []byte{sp, tb, lf, cr}

func wsByte(b byte) bool {
	return bytes.IndexByte(wsSet, b) >= 0
}

func wsRune(r rune) bool {
	return bytes.ContainsRune(wsSet, r)
}

// WriteRecord writes a JSON text sequence record with beginning
// (RS) and end (LF) marker bytes.
func WriteRecord(w io.Writer, json []byte) error {
	_, err := w.Write([]byte{rs})
	if err != nil {
		return err
	}
	_, err = w.Write(json)
	if err != nil {
		return err
	}
	_, err = w.Write([]byte{lf})
	return err
}

// A RecordWriter prefixes Write calls with a record separator.
//
// Callers must only call Write once for each value, and are responsible for
// including trailing line feeds when necessary or desired.
type RecordWriter struct {
	io.Writer
}

// Write prefixes every written record with an ASCII record separator.
func (w *RecordWriter) Write(record []byte) (int, error) {
	_, err := w.Writer.Write([]byte{rs})
	if err != nil {
		return -1, err
	}
	n, err := w.Writer.Write(record)
	return n + 1, err
}

// NewEncoder returns a standard library json.Encoder that writes a JSON text sequence to w.
//
// The Encoder calls Write just once for each value and always with a trailing line feed.
func NewEncoder(w io.Writer) *json.Encoder {
	return json.NewEncoder(&RecordWriter{w})
}

// Decode functions decode the JSON-encoded data and store the result in the value
// pointed to by v, or return an error if invalid.
// Note that the encoded data may have extra trailing data, which is perfectly
// valid. This disqualifies parsers which assume a single value (e.g. json.Unmarshal).
type Decode func(b []byte, v interface{}) error

// A Decoder reads and decodes JSON text sequence records from an input stream.
type Decoder struct {
	s  *bufio.Scanner
	fn Decode
}

// NewDecoder creates a new Decoder backed by the standard library's encoding/json
// Decoder. Any extra trailing data is discarded.
func NewDecoder(r io.Reader) *Decoder {
	return NewDecoderFn(r, func(b []byte, v interface{}) error {
		// Decode the first value, and discard any remaining data.
		return json.NewDecoder(bytes.NewReader(b)).Decode(v)
	})
}

// NewDecoderFn creates a new Decoder backed by a custom Decode function.
func NewDecoderFn(r io.Reader, fn Decode) *Decoder {
	s := bufio.NewScanner(r)
	s.Split(ScanRecord)
	return &Decoder{
		s:  s,
		fn: fn,
	}
}

// Decode scans the next record, or returns an error.
// The Decoder remains valid until io.EOF is returned.
func (d *Decoder) Decode(v interface{}) error {
	if !d.s.Scan() {
		if err := d.s.Err(); err != nil {
			return err
		}
		return io.EOF
	}
	b := d.s.Bytes()

	b, ok := RecordValue(b)
	if !ok {
		return fmt.Errorf("invalid record: %q", string(b))
	}
	return d.fn(b, v)
}

// RecordValue returns the *value* bytes from a JSON text sequence record and a flag
// indicating if the *record* is valid. This is *NOT* a validation of any contained JSON,
// which could itself be invalid or contain extra trailing values.
//
// See section 2.4: Top-Level Values: numbers, true, false, and null.
// https://tools.ietf.org/html/rfc7464#section-2.4
func RecordValue(b []byte) ([]byte, bool) {
	if len(b) < 2 {
		return b, false
	}
	if b[0] != rs {
		return b, false
	}
	// Drop rs and leading whitespace.
	b = bytes.TrimLeftFunc(b[1:], wsRune)

	// A number, true, false, or null value could be truncated if not
	// followed by whitespace.
	switch b[0] {
	case 'n':
		if bytes.HasPrefix(b, []byte("null")) {
			if wsByte(b[4]) {
				return b, true
			}
			return b, false
		}
	case 't':
		if bytes.HasPrefix(b, []byte("true")) {
			if wsByte(b[4]) {
				return b, true
			}
			return b, false
		}
	case 'f':
		if bytes.HasPrefix(b, []byte("false")) {
			if wsByte(b[5]) {
				return b, true
			}
			return b, false
		}
	case '-':
		if '0' <= b[1] && b[1] <= '9' {
			t := bytes.TrimLeft(b, digitSet)
			if len(t) > 0 && wsByte(t[0]) {
				return b, true
			}
			return b, false
		}
	case '0', '1', '2', '3', '4', '5', '6', '7', '8', '9':
		t := bytes.TrimLeft(b, digitSet)
		if len(t) > 0 && wsByte(t[0]) {
			return b, true
		}
		return b, false
	}

	return b, true
}

// ScanRecord is a bufio.SplitFunc which splits JSON text sequence records.
// Scanned bytes must be validated with the RecordValue function.
func ScanRecord(data []byte, atEOF bool) (advance int, token []byte, err error) {
	if atEOF && len(data) == 0 {
		return 0, nil, nil
	}
	// Find record start.
	switch i := bytes.IndexByte(data, rs); {
	case i < 0:
		if atEOF {
			// Partial record.
			return len(data), data, nil
		}
		// Request more data.
		return 0, nil, nil
	case i > 0:
		// Partial record.
		return i, data[0 : i-1], nil
	}
	// else i == 0

	// Drop consecutive leading rs's
	for len(data) > 1 && data[1] == rs {
		data = data[1:]
	}

	// Find end or next record.
	i := bytes.IndexByte(data[1:], rs)
	if i < 0 {
		if atEOF {
			return len(data), data, nil
		}
		// Request more data.
		return 0, nil, nil
	}
	return 1 + i, data[:1+i], nil
}
<|endoftext|>


<commit_before>package scheduler

import (
	"encoding/json"
	"io"
	"net/http"
	"sync"
	"time"

	"github.com/concourse/atc/builds"
	"github.com/concourse/atc/db"
	tbuilds "github.com/concourse/turbine/api/builds"
	"github.com/concourse/turbine/event"
	"github.com/concourse/turbine/routes"
	"github.com/pivotal-golang/lager"
	"github.com/tedsuo/rata"
	"github.com/vito/go-sse/sse"
)

type tracker struct {
	logger lager.Logger

	db TrackerDB

	trackingBuilds map[int]bool
	lock           *sync.Mutex
}

type TrackerDB interface {
	AppendBuildEvent(buildID int, event db.BuildEvent) error

	SaveBuildStartTime(buildID int, startTime time.Time) error
	SaveBuildEndTime(buildID int, startTime time.Time) error

	SaveBuildInput(buildID int, vr builds.VersionedResource) error
	SaveBuildOutput(buildID int, vr builds.VersionedResource) error

	SaveBuildStatus(buildID int, status builds.Status) error
}

func NewTracker(logger lager.Logger, db TrackerDB) BuildTracker {
	return &tracker{
		logger: logger,
		db:     db,

		trackingBuilds: make(map[int]bool),
		lock:           new(sync.Mutex),
	}
}

func (tracker *tracker) TrackBuild(build builds.Build) error {
	tLog := tracker.logger.Session("track-build", lager.Data{
		"buld": build.ID,
	})

	alreadyTracking := tracker.markTracking(build.ID)
	if alreadyTracking {
		tLog.Info("already-tracking")
		return nil
	}

	tLog.Info("tracking")

	defer func() {
		tLog.Info("done-tracking")
		tracker.unmarkTracking(build.ID)
	}()

	generator := rata.NewRequestGenerator(build.Endpoint, routes.Routes)

	events, err := generator.CreateRequest(
		routes.GetBuildEvents,
		rata.Params{"guid": build.Guid},
		nil,
	)
	if err != nil {
		tLog.Error("failed-to-create-events-request", err)
		return err
	}

	resp, err := http.DefaultClient.Do(events)
	if err != nil {
		tLog.Error("failed-to-get-events", err)
		return err
	}

	if resp.StatusCode == http.StatusNotFound {
		tLog.Info("saving-orphaned-build-as-errored")

		err := tracker.db.SaveBuildStatus(build.ID, builds.StatusErrored)
		if err != nil {
			tLog.Error("failed-to-save-build-as-errored", err)
			return err
		}

		return nil
	}

	reader := sse.NewReader(resp.Body)

	outputs := map[string]builds.VersionedResource{}

	var currentVersion string

	for {
		se, err := reader.Next()
		if err != nil {
			if err == io.EOF {
				tLog.Info("event-stream-completed")
				return nil
			}

			return err
		}

		err = tracker.db.AppendBuildEvent(build.ID, db.BuildEvent{
			Type:    se.Name,
			Payload: string(se.Data),
		})

		if se.Name == "version" {
			var version event.Version
			err := json.Unmarshal(se.Data, &version)
			if err != nil {
				tLog.Error("failed-to-unmarshal-version", err)
				return err
			}

			tLog.Info("event-stream-version", lager.Data{
				"version": string(version),
			})

			currentVersion = string(version)
			continue
		}

		if se.Name == "end" {
			tLog.Info("event-stream-ended")

			del, err := generator.CreateRequest(
				routes.DeleteBuild,
				rata.Params{"guid": build.Guid},
				nil,
			)
			if err != nil {
				tLog.Error("failed-to-create-delete-request", err)
				return err
			}

			resp, err := http.DefaultClient.Do(del)
			if err != nil {
				tLog.Error("failed-to-delete-build", err)
				return err
			}

			resp.Body.Close()
			continue
		}

		switch currentVersion {
		case "1.0":
			fallthrough
		case "1.1":
			switch se.Name {
			case "status":
				tLog.Info("processing-build-status", lager.Data{
					"event": string(se.Data),
				})

				var status event.Status
				err := json.Unmarshal(se.Data, &status)
				if err != nil {
					tLog.Error("failed-to-unmarshal-status", err)
					return err
				}

				if status.Status == tbuilds.StatusStarted {
					err = tracker.db.SaveBuildStartTime(build.ID, time.Unix(status.Time, 0))
					if err != nil {
						tLog.Error("failed-to-save-build-start-time", err)
						return err
					}
				} else {
					err = tracker.db.SaveBuildEndTime(build.ID, time.Unix(status.Time, 0))
					if err != nil {
						tLog.Error("failed-to-save-build-end-time", err)
						return err
					}
				}

				err = tracker.db.SaveBuildStatus(build.ID, builds.Status(status.Status))
				if err != nil {
					tLog.Error("failed-to-save-build-status", err)
					return err
				}

				if status.Status == tbuilds.StatusSucceeded {
					for _, output := range outputs {
						err := tracker.db.SaveBuildOutput(build.ID, output)
						if err != nil {
							tLog.Error("failed-to-save-build-output", err)
							return err
						}
					}
				}

			case "input":
				if build.JobName == "" {
					tLog.Info("ignoring-build-input-for-one-off")
					break
				}

				tLog.Info("processing-build-input", lager.Data{
					"event": string(se.Data),
				})

				var input event.Input
				err := json.Unmarshal(se.Data, &input)
				if err != nil {
					tLog.Error("failed-to-unarshal-input", err)
					return err
				}

				err = tracker.db.SaveBuildInput(build.ID, vrFromInput(input.Input))
				if err != nil {
					tLog.Error("failed-to-save-build-input", err)
					return err
				}

				// record implicit output
				outputs[input.Input.Name] = vrFromInput(input.Input)

			case "output":
				if build.JobName == "" {
					tLog.Info("ignoring-build-output-for-one-off")
					break
				}

				var output event.Output
				err := json.Unmarshal(se.Data, &output)
				if err != nil {
					tLog.Error("failed-to-unarshal-output", err)
					return err
				}

				outputs[output.Output.Name] = vrFromOutput(output.Output)
			}
		}
	}

	return nil
}

func (tracker *tracker) markTracking(buildID int) bool {
	tracker.lock.Lock()
	alreadyTracking, found := tracker.trackingBuilds[buildID]
	if !found {
		tracker.trackingBuilds[buildID] = true
	}
	tracker.lock.Unlock()

	return alreadyTracking
}

func (tracker *tracker) unmarkTracking(buildID int) {
	tracker.lock.Lock()
	delete(tracker.trackingBuilds, buildID)
	tracker.lock.Unlock()
}

func vrFromInput(input tbuilds.Input) builds.VersionedResource {
	metadata := make([]builds.MetadataField, len(input.Metadata))
	for i, md := range input.Metadata {
		metadata[i] = builds.MetadataField{
			Name:  md.Name,
			Value: md.Value,
		}
	}

	return builds.VersionedResource{
		Name:     input.Name,
		Type:     input.Type,
		Source:   builds.Source(input.Source),
		Version:  builds.Version(input.Version),
		Metadata: metadata,
	}
}

// same as input, but type is different.
//
// :(
func vrFromOutput(output tbuilds.Output) builds.VersionedResource {
	metadata := make([]builds.MetadataField, len(output.Metadata))
	for i, md := range output.Metadata {
		metadata[i] = builds.MetadataField{
			Name:  md.Name,
			Value: md.Value,
		}
	}

	return builds.VersionedResource{
		Name:     output.Name,
		Type:     output.Type,
		Source:   builds.Source(output.Source),
		Version:  builds.Version(output.Version),
		Metadata: metadata,
	}
}
<commit_msg>don't leak build events response body<commit_after>package scheduler

import (
	"encoding/json"
	"io"
	"net/http"
	"sync"
	"time"

	"github.com/concourse/atc/builds"
	"github.com/concourse/atc/db"
	tbuilds "github.com/concourse/turbine/api/builds"
	"github.com/concourse/turbine/event"
	"github.com/concourse/turbine/routes"
	"github.com/pivotal-golang/lager"
	"github.com/tedsuo/rata"
	"github.com/vito/go-sse/sse"
)

type tracker struct {
	logger lager.Logger

	db TrackerDB

	trackingBuilds map[int]bool
	lock           *sync.Mutex
}

type TrackerDB interface {
	AppendBuildEvent(buildID int, event db.BuildEvent) error

	SaveBuildStartTime(buildID int, startTime time.Time) error
	SaveBuildEndTime(buildID int, startTime time.Time) error

	SaveBuildInput(buildID int, vr builds.VersionedResource) error
	SaveBuildOutput(buildID int, vr builds.VersionedResource) error

	SaveBuildStatus(buildID int, status builds.Status) error
}

func NewTracker(logger lager.Logger, db TrackerDB) BuildTracker {
	return &tracker{
		logger: logger,
		db:     db,

		trackingBuilds: make(map[int]bool),
		lock:           new(sync.Mutex),
	}
}

func (tracker *tracker) TrackBuild(build builds.Build) error {
	tLog := tracker.logger.Session("track-build", lager.Data{
		"buld": build.ID,
	})

	alreadyTracking := tracker.markTracking(build.ID)
	if alreadyTracking {
		tLog.Info("already-tracking")
		return nil
	}

	tLog.Info("tracking")

	defer func() {
		tLog.Info("done-tracking")
		tracker.unmarkTracking(build.ID)
	}()

	generator := rata.NewRequestGenerator(build.Endpoint, routes.Routes)

	events, err := generator.CreateRequest(
		routes.GetBuildEvents,
		rata.Params{"guid": build.Guid},
		nil,
	)
	if err != nil {
		tLog.Error("failed-to-create-events-request", err)
		return err
	}

	resp, err := http.DefaultClient.Do(events)
	if err != nil {
		tLog.Error("failed-to-get-events", err)
		return err
	}

	defer resp.Body.Close()

	if resp.StatusCode == http.StatusNotFound {
		tLog.Info("saving-orphaned-build-as-errored")

		err := tracker.db.SaveBuildStatus(build.ID, builds.StatusErrored)
		if err != nil {
			tLog.Error("failed-to-save-build-as-errored", err)
			return err
		}

		return nil
	}

	reader := sse.NewReader(resp.Body)

	outputs := map[string]builds.VersionedResource{}

	var currentVersion string

	for {
		se, err := reader.Next()
		if err != nil {
			if err == io.EOF {
				tLog.Info("event-stream-completed")
				return nil
			}

			return err
		}

		err = tracker.db.AppendBuildEvent(build.ID, db.BuildEvent{
			Type:    se.Name,
			Payload: string(se.Data),
		})

		if se.Name == "version" {
			var version event.Version
			err := json.Unmarshal(se.Data, &version)
			if err != nil {
				tLog.Error("failed-to-unmarshal-version", err)
				return err
			}

			tLog.Info("event-stream-version", lager.Data{
				"version": string(version),
			})

			currentVersion = string(version)
			continue
		}

		if se.Name == "end" {
			tLog.Info("event-stream-ended")

			del, err := generator.CreateRequest(
				routes.DeleteBuild,
				rata.Params{"guid": build.Guid},
				nil,
			)
			if err != nil {
				tLog.Error("failed-to-create-delete-request", err)
				return err
			}

			resp, err := http.DefaultClient.Do(del)
			if err != nil {
				tLog.Error("failed-to-delete-build", err)
				return err
			}

			resp.Body.Close()
			continue
		}

		switch currentVersion {
		case "1.0":
			fallthrough
		case "1.1":
			switch se.Name {
			case "status":
				tLog.Info("processing-build-status", lager.Data{
					"event": string(se.Data),
				})

				var status event.Status
				err := json.Unmarshal(se.Data, &status)
				if err != nil {
					tLog.Error("failed-to-unmarshal-status", err)
					return err
				}

				if status.Status == tbuilds.StatusStarted {
					err = tracker.db.SaveBuildStartTime(build.ID, time.Unix(status.Time, 0))
					if err != nil {
						tLog.Error("failed-to-save-build-start-time", err)
						return err
					}
				} else {
					err = tracker.db.SaveBuildEndTime(build.ID, time.Unix(status.Time, 0))
					if err != nil {
						tLog.Error("failed-to-save-build-end-time", err)
						return err
					}
				}

				err = tracker.db.SaveBuildStatus(build.ID, builds.Status(status.Status))
				if err != nil {
					tLog.Error("failed-to-save-build-status", err)
					return err
				}

				if status.Status == tbuilds.StatusSucceeded {
					for _, output := range outputs {
						err := tracker.db.SaveBuildOutput(build.ID, output)
						if err != nil {
							tLog.Error("failed-to-save-build-output", err)
							return err
						}
					}
				}

			case "input":
				if build.JobName == "" {
					tLog.Info("ignoring-build-input-for-one-off")
					break
				}

				tLog.Info("processing-build-input", lager.Data{
					"event": string(se.Data),
				})

				var input event.Input
				err := json.Unmarshal(se.Data, &input)
				if err != nil {
					tLog.Error("failed-to-unarshal-input", err)
					return err
				}

				err = tracker.db.SaveBuildInput(build.ID, vrFromInput(input.Input))
				if err != nil {
					tLog.Error("failed-to-save-build-input", err)
					return err
				}

				// record implicit output
				outputs[input.Input.Name] = vrFromInput(input.Input)

			case "output":
				if build.JobName == "" {
					tLog.Info("ignoring-build-output-for-one-off")
					break
				}

				var output event.Output
				err := json.Unmarshal(se.Data, &output)
				if err != nil {
					tLog.Error("failed-to-unarshal-output", err)
					return err
				}

				outputs[output.Output.Name] = vrFromOutput(output.Output)
			}
		}
	}

	return nil
}

func (tracker *tracker) markTracking(buildID int) bool {
	tracker.lock.Lock()
	alreadyTracking, found := tracker.trackingBuilds[buildID]
	if !found {
		tracker.trackingBuilds[buildID] = true
	}
	tracker.lock.Unlock()

	return alreadyTracking
}

func (tracker *tracker) unmarkTracking(buildID int) {
	tracker.lock.Lock()
	delete(tracker.trackingBuilds, buildID)
	tracker.lock.Unlock()
}

func vrFromInput(input tbuilds.Input) builds.VersionedResource {
	metadata := make([]builds.MetadataField, len(input.Metadata))
	for i, md := range input.Metadata {
		metadata[i] = builds.MetadataField{
			Name:  md.Name,
			Value: md.Value,
		}
	}

	return builds.VersionedResource{
		Name:     input.Name,
		Type:     input.Type,
		Source:   builds.Source(input.Source),
		Version:  builds.Version(input.Version),
		Metadata: metadata,
	}
}

// same as input, but type is different.
//
// :(
func vrFromOutput(output tbuilds.Output) builds.VersionedResource {
	metadata := make([]builds.MetadataField, len(output.Metadata))
	for i, md := range output.Metadata {
		metadata[i] = builds.MetadataField{
			Name:  md.Name,
			Value: md.Value,
		}
	}

	return builds.VersionedResource{
		Name:     output.Name,
		Type:     output.Type,
		Source:   builds.Source(output.Source),
		Version:  builds.Version(output.Version),
		Metadata: metadata,
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"flag"
	"html/template"
	"io/ioutil"
	"log"
	"net/http"
	"os/exec"
	"strconv"
	"strings"
)

type Album struct {
	Title, Artist, Folder string
}

var albums []Album

var html *template.Template

func handler(w http.ResponseWriter, r *http.Request) {

	if r.URL.Path == "/favicon.ico" {
		http.NotFound(w, r)
		return
	}

	if r.FormValue("u") != "" {
		ioutil.WriteFile("/tmp/mplayer", []byte("loadlist "+r.FormValue("u")+"\n"), 0644)
	} else if r.FormValue("f") != "" {
		ioutil.WriteFile("/tmp/mplayer", []byte("loadfile '"+r.FormValue("f")+"'\n"), 0644)
	} else if r.FormValue("d") != "" {
		folder := r.FormValue("d")
		cmd := exec.Command("find", folder, "-type", "f")
		var out bytes.Buffer
		cmd.Stdout = &out
		cmd.Run()
		playlist, _ := ioutil.TempFile("", "jukebox")
		ioutil.WriteFile(playlist.Name(), []byte(out.String()), 0644)
		ioutil.WriteFile("/tmp/mplayer", []byte("loadlist '"+playlist.Name()+"'\n"), 0644)
	} else if r.FormValue("c") != "" {
		ioutil.WriteFile("/tmp/mplayer", []byte(r.FormValue("c")+"\n"), 0644)
	}

	html.Execute(w, albums)
}

func main() {
	port := flag.Int("port", 80, "port")
	root := flag.String("root", "", "root")
	flag.Parse()
	if *root == "" {
		panic("root required")
	}

	cmd := exec.Command("find", *root, "-mindepth", "2", "-maxdepth", "2", "-type", "d")
	var out bytes.Buffer
	cmd.Stdout = &out
	cmd.Run()
	lines := strings.Split(out.String(), "\n")

	albums = make([]Album, len(lines)-1)
	for i, line := range lines {
		if len(line) > 0 {
			parts := strings.Split(line, "/")
			albums[i] = Album{
				parts[len(parts)-1],
				parts[len(parts)-2],
				line,
			}
		}
	}

	const _html = `<!DOCTYPE html>
<title>jukebox</title><meta name="viewport" content="width=device-width, initial-scale=1.0">

Stop

Radio 4

Radio 5 live

Groove Salad
  {{range .}}
{{.Artist}} - {{.Title}}{{end}}

`
	html = template.Must(template.New("html").Parse(_html))

	http.HandleFunc("/", handler)
	err := http.ListenAndServe(":"+strconv.Itoa(*port), nil)
	if err != nil {
		log.Fatal("ListenAndServe: ", err)
	}
}
<commit_msg>Tidy up exec<commit_after>package main

import (
	"bytes"
	"flag"
	"html/template"
	"io/ioutil"
	"log"
	"net/http"
	"os/exec"
	"strconv"
	"strings"
)

type Album struct {
	Title, Artist, Folder string
}

var albums []Album

var html *template.Template

func handler(w http.ResponseWriter, r *http.Request) {
	if r.URL.Path == "/favicon.ico" {
		http.NotFound(w, r)
		return
	}

	if r.FormValue("u") != "" {
		ioutil.WriteFile("/tmp/mplayer", []byte("loadlist "+r.FormValue("u")+"\n"), 0644)
	} else if r.FormValue("f") != "" {
		ioutil.WriteFile("/tmp/mplayer", []byte("loadfile '"+r.FormValue("f")+"'\n"), 0644)
	} else if r.FormValue("d") != "" {
		folder := r.FormValue("d")
		cmd := exec.Command("find", folder, "-type", "f")
		var out bytes.Buffer
		cmd.Stdout = &out
		cmd.Run()
		playlist, _ := ioutil.TempFile("", "jukebox")
		ioutil.WriteFile(playlist.Name(), []byte(out.String()), 0644)
		ioutil.WriteFile("/tmp/mplayer", []byte("loadlist '"+playlist.Name()+"'\n"), 0644)
	} else if r.FormValue("c") != "" {
		ioutil.WriteFile("/tmp/mplayer", []byte(r.FormValue("c")+"\n"), 0644)
	}

	html.Execute(w, albums)
}

func buildTemplates() {
	const _html = `<!DOCTYPE html>
<title>jukebox</title><meta name="viewport" content="width=device-width, initial-scale=1.0">

Stop

Radio 4

Radio 5 live

Groove Salad
  {{range .}}
{{.Artist}} - {{.Title}}{{end}}

`
	html = template.Must(template.New("html").Parse(_html))
}

func findAlbums(root string) {
	cmd := exec.Command("find", root, "-mindepth", "2", "-maxdepth", "2", "-type", "d")
	output, err := cmd.Output()
	if err != nil {
		log.Fatal("findAlbums: ", err)
	}
	lines := strings.Split(string(output), "\n")

	albums = make([]Album, len(lines)-1)
	for i, line := range lines {
		if len(line) > 0 {
			parts := strings.Split(line, "/")
			albums[i] = Album{
				parts[len(parts)-1],
				parts[len(parts)-2],
				line,
			}
		}
	}
}

func main() {
	port := flag.Int("port", 80, "port")
	root := flag.String("root", "", "root")
	flag.Parse()
	if *root == "" {
		panic("root required")
	}

	buildTemplates()
	findAlbums(*root)

	http.HandleFunc("/", handler)
	err := http.ListenAndServe(":"+strconv.Itoa(*port), nil)
	if err != nil {
		log.Fatal("ListenAndServe: ", err)
	}
}
<|endoftext|>


<commit_before>package kafka_avro

import (
	"fmt"
	"log"
	"net"
	"os"
	"strconv"
	"strings"
	"time"

	avro "github.com/elodina/go-avro"
	kafkaavro "github.com/elodina/go-kafka-avro"
	"github.com/gliderlabs/logspout/router"
	"gopkg.in/Shopify/sarama.v1"
	"github.com/fsouza/go-dockerclient"
	"github.com/joeshaw/iso8601"
)

var messageSchema = `{
  "type": "record",
  "name": "LogLine",
  "fields": [
    {"name": "timestamp", "type": "string"},
    {"name": "container_name", "type": "string"},
    {"name": "host", "type": "string"},
    {"name": "source", "type": "string"},
    {"name": "line", "type": "string"}
  ]
}`

func init() {
	router.AdapterFactories.Register(NewKafkaAvroAdapter, "kafka_avro")
}

type KafkaAvroAdapter struct {
	route    *router.Route
	brokers  []string
	topic    string
	schema   avro.Schema
	registry *kafkaavro.KafkaAvroEncoder
	producer sarama.AsyncProducer
}

func NewKafkaAvroAdapter(route *router.Route) (router.LogAdapter, error) {
	brokers := readBrokers(route.Address)
	if len(brokers) == 0 {
		return nil, errorf("The Kafka broker host:port is missing. Did you specify it as a route address?")
	}

	topic := readTopic(route.Address, route.Options)
	if topic == "" {
		return nil, errorf("The Kafka topic is missing. Did you specify it as a route option?")
	}

	schemaUrl := readSchemaRegistryUrl(route.Options)
	if schemaUrl == "" {
		return nil, errorf("The schema registry url is missing. Did you specify it as a route option?")
	}

	registry := kafkaavro.NewKafkaAvroEncoder(schemaUrl)

	var schema avro.Schema
	schema, err := avro.ParseSchema(messageSchema)
	if err != nil {
		return nil, errorf("The schema could not be parsed")
	}

	if os.Getenv("DEBUG") != "" {
		log.Printf("Starting Kafka producer for address: %s, topic: %s.\n", brokers, topic)
	}

	var retries int
	retries, err = strconv.Atoi(os.Getenv("KAFKA_CONNECT_RETRIES"))
	if err != nil {
		retries = 3
	}
	var producer sarama.AsyncProducer
	for i := 0; i < retries; i++ {
		producer, err = sarama.NewAsyncProducer(brokers, newConfig())
		if err != nil {
			if os.Getenv("DEBUG") != "" {
				log.Println("Couldn't create Kafka producer. Retrying...", err)
			}
			if i == retries-1 {
				return nil, errorf("Couldn't create Kafka producer. %v", err)
			}
		} else {
			time.Sleep(1 * time.Second)
		}
	}

	return &KafkaAvroAdapter{
		route:    route,
		brokers:  brokers,
		topic:    topic,
		registry: registry,
		schema:   schema,
		producer: producer,
	}, nil
}

func (a *KafkaAvroAdapter) Stream(logstream chan *router.Message) {
	defer a.producer.Close()
	for rm := range logstream {
		message, err := a.formatMessage(rm)
		if err != nil {
			log.Println("kafka:", err)
			a.route.Close()
			break
		}

		a.producer.Input() <- message
	}
}

func newConfig() *sarama.Config {
	config := sarama.NewConfig()
	config.ClientID = "logspout"
	config.Producer.Return.Errors = false
	config.Producer.Return.Successes = false
	config.Producer.Flush.Frequency = 1 * time.Second
	config.Producer.RequiredAcks = sarama.WaitForLocal

	if opt := os.Getenv("KAFKA_COMPRESSION_CODEC"); opt != "" {
		switch opt {
		case "gzip":
			config.Producer.Compression = sarama.CompressionGZIP
		case "snappy":
			config.Producer.Compression = sarama.CompressionSnappy
		}
	}

	return config
}

func (a *KafkaAvroAdapter) formatMessage(message *router.Message) (*sarama.ProducerMessage, error) {
	var encoder sarama.Encoder
	var containerName string
	var env docker.Env = message.Container.Config.Env

	if env != nil {
		containerName = env.Get("MARATHON_APP_ID")
	}
	if containerName == "" {
		containerName = message.Container.Name
	}
	host := env.Get("HOST")
	if host == "" {
		host = getLocalIP()
	}

	record := avro.NewGenericRecord(a.schema)
	record.Set("timestamp", string(iso8601.Time(message.Time)))
	record.Set("container_name", containerName)
	record.Set("host", host)
	record.Set("source", message.Source)
	record.Set("line", message.Data)

	b, err := a.registry.Encode(record)
	if err != nil {
		return nil, err
	}
	debug(b)

	encoder = sarama.ByteEncoder(b)

	return &sarama.ProducerMessage{
		Topic: a.topic,
		Value: encoder,
	}, nil
}

func readBrokers(address string) []string {
	if strings.Contains(address, "/") {
		slash := strings.Index(address, "/")
		address = address[:slash]
	}

	return strings.Split(address, ",")
}

func readTopic(address string, options map[string]string) string {
	var topic string
	if !strings.Contains(address, "/") {
		topic = options["topic"]
	} else {
		slash := strings.Index(address, "/")
		topic = address[slash+1:]
	}

	return topic
}

func readSchemaRegistryUrl(options map[string]string) string {
	var url string
	if _, ok := options["schema_registry_url"]; ok {
		url = options["schema_registry_url"]
	}
	return url
}

func getLocalIP() string {
	addrs, err := net.InterfaceAddrs()
	if err != nil {
		return ""
	}
	for _, address := range addrs {
		// check the address type and if it is not a loopback the display it
		if ipnet, ok := address.(*net.IPNet); ok && !ipnet.IP.IsLoopback() {
			if ipnet.IP.To4() != nil {
				return ipnet.IP.String()
			}
		}
	}
	return ""
}

func errorf(format string, a ...interface{}) (err error) {
	err = fmt.Errorf(format, a...)
	if os.Getenv("DEBUG") != "" {
		fmt.Println(err.Error())
	}
	return
}

func debug(v ...interface{}) {
	if os.Getenv("DEBUG") != "" {
		log.Println(v...)
	}
}
<commit_msg>Fixing iso8601<commit_after>package kafka_avro

import (
	"fmt"
	"log"
	"net"
	"os"
	"strconv"
	"strings"
	"time"

	avro "github.com/elodina/go-avro"
	kafkaavro "github.com/elodina/go-kafka-avro"
	"github.com/gliderlabs/logspout/router"
	"gopkg.in/Shopify/sarama.v1"
	"github.com/fsouza/go-dockerclient"
	"github.com/joeshaw/iso8601"
)

var messageSchema = `{
  "type": "record",
  "name": "LogLine",
  "fields": [
    {"name": "timestamp", "type": "string"},
    {"name": "container_name", "type": "string"},
    {"name": "host", "type": "string"},
    {"name": "source", "type": "string"},
    {"name": "line", "type": "string"}
  ]
}`

func init() {
	router.AdapterFactories.Register(NewKafkaAvroAdapter, "kafka_avro")
}

type KafkaAvroAdapter struct {
	route    *router.Route
	brokers  []string
	topic    string
	schema   avro.Schema
	registry *kafkaavro.KafkaAvroEncoder
	producer sarama.AsyncProducer
}

func NewKafkaAvroAdapter(route *router.Route) (router.LogAdapter, error) {
	brokers := readBrokers(route.Address)
	if len(brokers) == 0 {
		return nil, errorf("The Kafka broker host:port is missing. Did you specify it as a route address?")
	}

	topic := readTopic(route.Address, route.Options)
	if topic == "" {
		return nil, errorf("The Kafka topic is missing. Did you specify it as a route option?")
	}

	schemaUrl := readSchemaRegistryUrl(route.Options)
	if schemaUrl == "" {
		return nil, errorf("The schema registry url is missing. Did you specify it as a route option?")
	}

	registry := kafkaavro.NewKafkaAvroEncoder(schemaUrl)

	var schema avro.Schema
	schema, err := avro.ParseSchema(messageSchema)
	if err != nil {
		return nil, errorf("The schema could not be parsed")
	}

	if os.Getenv("DEBUG") != "" {
		log.Printf("Starting Kafka producer for address: %s, topic: %s.\n", brokers, topic)
	}

	var retries int
	retries, err = strconv.Atoi(os.Getenv("KAFKA_CONNECT_RETRIES"))
	if err != nil {
		retries = 3
	}
	var producer sarama.AsyncProducer
	for i := 0; i < retries; i++ {
		producer, err = sarama.NewAsyncProducer(brokers, newConfig())
		if err != nil {
			if os.Getenv("DEBUG") != "" {
				log.Println("Couldn't create Kafka producer. Retrying...", err)
			}
			if i == retries-1 {
				return nil, errorf("Couldn't create Kafka producer. %v", err)
			}
		} else {
			time.Sleep(1 * time.Second)
		}
	}

	return &KafkaAvroAdapter{
		route:    route,
		brokers:  brokers,
		topic:    topic,
		registry: registry,
		schema:   schema,
		producer: producer,
	}, nil
}

func (a *KafkaAvroAdapter) Stream(logstream chan *router.Message) {
	defer a.producer.Close()
	for rm := range logstream {
		message, err := a.formatMessage(rm)
		if err != nil {
			log.Println("kafka:", err)
			a.route.Close()
			break
		}

		a.producer.Input() <- message
	}
}

func newConfig() *sarama.Config {
	config := sarama.NewConfig()
	config.ClientID = "logspout"
	config.Producer.Return.Errors = false
	config.Producer.Return.Successes = false
	config.Producer.Flush.Frequency = 1 * time.Second
	config.Producer.RequiredAcks = sarama.WaitForLocal

	if opt := os.Getenv("KAFKA_COMPRESSION_CODEC"); opt != "" {
		switch opt {
		case "gzip":
			config.Producer.Compression = sarama.CompressionGZIP
		case "snappy":
			config.Producer.Compression = sarama.CompressionSnappy
		}
	}

	return config
}

func (a *KafkaAvroAdapter) formatMessage(message *router.Message) (*sarama.ProducerMessage, error) {
	var encoder sarama.Encoder
	var containerName string
	var env docker.Env = message.Container.Config.Env

	if env != nil {
		containerName = env.Get("MARATHON_APP_ID")
	}
	if containerName == "" {
		containerName = message.Container.Name
	}
	host := env.Get("HOST")
	if host == "" {
		host = getLocalIP()
	}

	record := avro.NewGenericRecord(a.schema)
	record.Set("timestamp", iso8601.Time(message.Time).String())
	record.Set("container_name", containerName)
	record.Set("host", host)
	record.Set("source", message.Source)
	record.Set("line", message.Data)

	b, err := a.registry.Encode(record)
	if err != nil {
		return nil, err
	}
	debug(b)

	encoder = sarama.ByteEncoder(b)

	return &sarama.ProducerMessage{
		Topic: a.topic,
		Value: encoder,
	}, nil
}

func readBrokers(address string) []string {
	if strings.Contains(address, "/") {
		slash := strings.Index(address, "/")
		address = address[:slash]
	}

	return strings.Split(address, ",")
}

func readTopic(address string, options map[string]string) string {
	var topic string
	if !strings.Contains(address, "/") {
		topic = options["topic"]
	} else {
		slash := strings.Index(address, "/")
		topic = address[slash+1:]
	}

	return topic
}

func readSchemaRegistryUrl(options map[string]string) string {
	var url string
	if _, ok := options["schema_registry_url"]; ok {
		url = options["schema_registry_url"]
	}
	return url
}

func getLocalIP() string {
	addrs, err := net.InterfaceAddrs()
	if err != nil {
		return ""
	}
	for _, address := range addrs {
		// check the address type and if it is not a loopback the display it
		if ipnet, ok := address.(*net.IPNet); ok && !ipnet.IP.IsLoopback() {
			if ipnet.IP.To4() != nil {
				return ipnet.IP.String()
			}
		}
	}
	return ""
}

func errorf(format string, a ...interface{}) (err error) {
	err = fmt.Errorf(format, a...)
	if os.Getenv("DEBUG") != "" {
		fmt.Println(err.Error())
	}
	return
}

func debug(v ...interface{}) {
	if os.Getenv("DEBUG") != "" {
		log.Println(v...)
	}
}
<|endoftext|>


<commit_before>// Copyright 2017 syzkaller project authors. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

package targets

import (
	"os"
	"os/exec"
	"runtime"
	"strings"
	"sync"
)

type Target struct {
	init sync.Once
	osCommon
	OS               string
	Arch             string
	PtrSize          uint64
	PageSize         uint64
	NumPages         uint64
	DataOffset       uint64
	CArch            []string
	CFlags           []string
	CrossCFlags      []string
	CCompilerPrefix  string
	CCompiler        string
	KernelArch       string
	KernelHeaderArch string
	// NeedSyscallDefine is used by csource package to decide when to emit __NR_* defines.
	NeedSyscallDefine func(nr uint64) bool
}

type osCommon struct {
	// Does the OS use syscall numbers (e.g. Linux) or has interface based on functions (e.g. fuchsia).
	SyscallNumbers bool
	// E.g. "__NR_" or "SYS_".
	SyscallPrefix string
	// ipc<->executor communication tuning.
	// If ExecutorUsesShmem, programs and coverage are passed through shmem, otherwise via pipes.
	ExecutorUsesShmem bool
	// If ExecutorUsesForkServer, executor uses extended protocol with handshake.
	ExecutorUsesForkServer bool
	// Extension of executable files (notably, .exe for windows).
	ExeExtension string
}

func Get(OS, arch string) *Target {
	target := List[OS][arch]
	if target == nil {
		return nil
	}
	target.init.Do(func() {
		checkStaticBuild(target)
	})
	return target
}

// nolint: lll
var List = map[string]map[string]*Target{
	"test": map[string]*Target{
		"32": {
			PtrSize:  4,
			PageSize: 8 << 10,
		},
		"64": {
			PtrSize:  8,
			PageSize: 4 << 10,
		},
	},
	"linux": map[string]*Target{
		"amd64": {
			PtrSize:          8,
			PageSize:         4 << 10,
			CArch:            []string{"__x86_64__"},
			CFlags:           []string{"-m64"},
			CrossCFlags:      []string{"-m64", "-static"},
			CCompilerPrefix:  "x86_64-linux-gnu-",
			KernelArch:       "x86_64",
			KernelHeaderArch: "x86",
			NeedSyscallDefine: func(nr uint64) bool {
				// Only generate defines for new syscalls
				// (added after commit 8a1ab3155c2ac on 2012-10-04).
				return nr >= 313
			},
		},
		"386": {
			PtrSize:          4,
			PageSize:         4 << 10,
			CArch:            []string{"__i386__"},
			CFlags:           []string{"-m32"},
			CrossCFlags:      []string{"-m32", "-static"},
			CCompilerPrefix:  "x86_64-linux-gnu-",
			KernelArch:       "i386",
			KernelHeaderArch: "x86",
		},
		"arm64": {
			PtrSize:          8,
			PageSize:         4 << 10,
			CArch:            []string{"__aarch64__"},
			CrossCFlags:      []string{"-static"},
			CCompilerPrefix:  "aarch64-linux-gnu-",
			KernelArch:       "arm64",
			KernelHeaderArch: "arm64",
		},
		"arm": {
			PtrSize:          4,
			PageSize:         4 << 10,
			CArch:            []string{"__arm__"},
			CFlags:           []string{"-D__LINUX_ARM_ARCH__=6", "-m32", "-D__ARM_EABI__"},
			CrossCFlags:      []string{"-D__LINUX_ARM_ARCH__=6", "-march=armv6t2", "-static"},
			CCompilerPrefix:  "arm-linux-gnueabihf-",
			KernelArch:       "arm",
			KernelHeaderArch: "arm",
		},
		"ppc64le": {
			PtrSize:          8,
			PageSize:         4 << 10,
			CArch:            []string{"__ppc64__", "__PPC64__", "__powerpc64__"},
			CFlags:           []string{"-D__powerpc64__"},
			CrossCFlags:      []string{"-D__powerpc64__", "-static"},
			CCompilerPrefix:  "powerpc64le-linux-gnu-",
			KernelArch:       "powerpc",
			KernelHeaderArch: "powerpc",
		},
	},
	"freebsd": map[string]*Target{
		"amd64": {
			PtrSize:     8,
			PageSize:    4 << 10,
			CArch:       []string{"__x86_64__"},
			CFlags:      []string{"-m64"},
			CrossCFlags: []string{"-m64", "-static"},
		},
	},
	"netbsd": map[string]*Target{
		"amd64": {
			PtrSize:     8,
			PageSize:    4 << 10,
			CArch:       []string{"__x86_64__"},
			CFlags:      []string{"-m64"},
			CrossCFlags: []string{"-m64", "-static"},
		},
	},
	"fuchsia": map[string]*Target{
		"amd64": {
			PtrSize:          8,
			PageSize:         4 << 10,
			CArch:            []string{"__x86_64__"},
			KernelHeaderArch: "x64",
			CCompiler:        os.ExpandEnv("${SOURCEDIR}/buildtools/linux-x64/clang/bin/clang++"),
			CrossCFlags: []string{
				"-Wno-deprecated",
				"-Wno-error",
				"--target=x86_64-fuchsia",
				"-lfdio",
				"-lzircon",
				"--sysroot", os.ExpandEnv("${SOURCEDIR}/out/build-zircon/build-x64/sysroot"),
				"-L", "${SOURCEDIR}/out/x64/x64-shared",
				"-L", "${SOURCEDIR}/out/x64/sdks/zircon_sysroot/sysroot/lib",
			},
		},
		"arm64": {
			PtrSize:          8,
			PageSize:         4 << 10,
			CArch:            []string{"__aarch64__"},
			KernelHeaderArch: "arm64",
			CCompiler:        os.ExpandEnv("${SOURCEDIR}/buildtools/linux-x64/clang/bin/clang++"),
			CrossCFlags: []string{
				"-Wno-deprecated",
				"-Wno-error",
				"--target=aarch64-fuchsia",
				"-lfdio",
				"-lzircon",
				"--sysroot", os.ExpandEnv("${SOURCEDIR}/out/build-zircon/build-arm64/sysroot"),
			},
		},
	},
	"windows": map[string]*Target{
		"amd64": {
			PtrSize: 8,
			// TODO(dvyukov): what should we do about 4k vs 64k?
			PageSize: 4 << 10,
			CArch:    []string{"_M_X64"},
		},
	},
	"akaros": map[string]*Target{
		"amd64": {
			PtrSize:           8,
			PageSize:          4 << 10,
			CArch:             []string{"__x86_64__"},
			NeedSyscallDefine: dontNeedSyscallDefine,
			CCompiler:         os.ExpandEnv("${SOURCEDIR}/install/x86_64-ucb-akaros-gcc/bin/x86_64-ucb-akaros-g++"),
			// Most likely this is incorrect (why doesn't it know own sysroot?), but worked for me.
			CrossCFlags: []string{
				"-static",
				"-I", os.ExpandEnv("${SOURCEDIR}/tools/compilers/gcc-glibc/x86_64-ucb-akaros-gcc-stage3-builddir/x86_64-ucb-akaros/libstdc++-v3/include/x86_64-ucb-akaros"),
				"-I", os.ExpandEnv("${SOURCEDIR}/tools/compilers/gcc-glibc/x86_64-ucb-akaros-gcc-stage3-builddir/x86_64-ucb-akaros/libstdc++-v3/include"),
				"-I", os.ExpandEnv("${SOURCEDIR}/tools/compilers/gcc-glibc/gcc-4.9.2/libstdc++-v3/libsupc++"),
				"-L", os.ExpandEnv("${SOURCEDIR}/tools/compilers/gcc-glibc/x86_64-ucb-akaros-gcc-stage3-builddir/x86_64-ucb-akaros/libstdc++-v3/src/.libs"),
			},
		},
	},
}

var oses = map[string]osCommon{
	"linux": {
		SyscallNumbers:         true,
		SyscallPrefix:          "__NR_",
		ExecutorUsesShmem:      true,
		ExecutorUsesForkServer: true,
	},
	"freebsd": {
		SyscallNumbers:         true,
		SyscallPrefix:          "SYS_",
		ExecutorUsesShmem:      true,
		ExecutorUsesForkServer: true,
	},
	"netbsd": {
		SyscallNumbers:         true,
		SyscallPrefix:          "SYS_",
		ExecutorUsesShmem:      true,
		ExecutorUsesForkServer: true,
	},
	"fuchsia": {
		SyscallNumbers:         false,
		ExecutorUsesShmem:      false,
		ExecutorUsesForkServer: false,
	},
	"windows": {
		SyscallNumbers:         false,
		ExecutorUsesShmem:      false,
		ExecutorUsesForkServer: false,
		ExeExtension:           ".exe",
	},
	"akaros": {
		SyscallNumbers:         true,
		SyscallPrefix:          "SYS_",
		ExecutorUsesShmem:      false,
		ExecutorUsesForkServer: true,
	},
}

func init() {
	for OS, archs := range List {
		for arch, target := range archs {
			initTarget(target, OS, arch)
		}
	}
}

func initTarget(target *Target, OS, arch string) {
	target.osCommon = oses[OS]
	target.OS = OS
	target.Arch = arch
	if target.NeedSyscallDefine == nil {
		target.NeedSyscallDefine = needSyscallDefine
	}
	target.DataOffset = 512 << 20
	target.NumPages = (16 << 20) / target.PageSize
	if OS == runtime.GOOS && arch == runtime.GOARCH {
		// Don't use cross-compiler for native compilation, there are cases when this does not work:
		// https://github.com/google/syzkaller/pull/619
		// https://github.com/google/syzkaller/issues/387
		// https://github.com/google/syzkaller/commit/06db3cec94c54e1cf720cdd5db72761514569d56
		target.CCompilerPrefix = ""
	}
	if target.CCompiler == "" {
		target.CCompiler = target.CCompilerPrefix + "gcc"
	}
}

func checkStaticBuild(target *Target) {
	for i, flag := range target.CrossCFlags {
		if flag == "-static" {
			// Some distributions don't have static libraries.
			if !supportsStatic(target) {
				copy(target.CrossCFlags[i:], target.CrossCFlags[i+1:])
				target.CrossCFlags = target.CrossCFlags[:len(target.CrossCFlags)-1]
			}
			break
		}
	}
}

func supportsStatic(target *Target) bool {
	cmd := exec.Command(target.CCompiler, "-x", "c", "-", "-o", "/dev/null", "-static")
	cmd.Stdin = strings.NewReader("int main(){}")
	return cmd.Run() == nil
}

func needSyscallDefine(nr uint64) bool {
	return true
}
func dontNeedSyscallDefine(nr uint64) bool {
	return false
}
<commit_msg>sys/targets: fix env expansion for fuchsia<commit_after>// Copyright 2017 syzkaller project authors. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

package targets

import (
	"os"
	"os/exec"
	"runtime"
	"strings"
	"sync"
)

type Target struct {
	init sync.Once
	osCommon
	OS               string
	Arch             string
	PtrSize          uint64
	PageSize         uint64
	NumPages         uint64
	DataOffset       uint64
	CArch            []string
	CFlags           []string
	CrossCFlags      []string
	CCompilerPrefix  string
	CCompiler        string
	KernelArch       string
	KernelHeaderArch string
	// NeedSyscallDefine is used by csource package to decide when to emit __NR_* defines.
	NeedSyscallDefine func(nr uint64) bool
}

type osCommon struct {
	// Does the OS use syscall numbers (e.g. Linux) or has interface based on functions (e.g. fuchsia).
	SyscallNumbers bool
	// E.g. "__NR_" or "SYS_".
	SyscallPrefix string
	// ipc<->executor communication tuning.
	// If ExecutorUsesShmem, programs and coverage are passed through shmem, otherwise via pipes.
	ExecutorUsesShmem bool
	// If ExecutorUsesForkServer, executor uses extended protocol with handshake.
	ExecutorUsesForkServer bool
	// Extension of executable files (notably, .exe for windows).
	ExeExtension string
}

func Get(OS, arch string) *Target {
	target := List[OS][arch]
	if target == nil {
		return nil
	}
	target.init.Do(func() {
		checkStaticBuild(target)
	})
	return target
}

// nolint: lll
var List = map[string]map[string]*Target{
	"test": map[string]*Target{
		"32": {
			PtrSize:  4,
			PageSize: 8 << 10,
		},
		"64": {
			PtrSize:  8,
			PageSize: 4 << 10,
		},
	},
	"linux": map[string]*Target{
		"amd64": {
			PtrSize:          8,
			PageSize:         4 << 10,
			CArch:            []string{"__x86_64__"},
			CFlags:           []string{"-m64"},
			CrossCFlags:      []string{"-m64", "-static"},
			CCompilerPrefix:  "x86_64-linux-gnu-",
			KernelArch:       "x86_64",
			KernelHeaderArch: "x86",
			NeedSyscallDefine: func(nr uint64) bool {
				// Only generate defines for new syscalls
				// (added after commit 8a1ab3155c2ac on 2012-10-04).
				return nr >= 313
			},
		},
		"386": {
			PtrSize:          4,
			PageSize:         4 << 10,
			CArch:            []string{"__i386__"},
			CFlags:           []string{"-m32"},
			CrossCFlags:      []string{"-m32", "-static"},
			CCompilerPrefix:  "x86_64-linux-gnu-",
			KernelArch:       "i386",
			KernelHeaderArch: "x86",
		},
		"arm64": {
			PtrSize:          8,
			PageSize:         4 << 10,
			CArch:            []string{"__aarch64__"},
			CrossCFlags:      []string{"-static"},
			CCompilerPrefix:  "aarch64-linux-gnu-",
			KernelArch:       "arm64",
			KernelHeaderArch: "arm64",
		},
		"arm": {
			PtrSize:          4,
			PageSize:         4 << 10,
			CArch:            []string{"__arm__"},
			CFlags:           []string{"-D__LINUX_ARM_ARCH__=6", "-m32", "-D__ARM_EABI__"},
			CrossCFlags:      []string{"-D__LINUX_ARM_ARCH__=6", "-march=armv6t2", "-static"},
			CCompilerPrefix:  "arm-linux-gnueabihf-",
			KernelArch:       "arm",
			KernelHeaderArch: "arm",
		},
		"ppc64le": {
			PtrSize:          8,
			PageSize:         4 << 10,
			CArch:            []string{"__ppc64__", "__PPC64__", "__powerpc64__"},
			CFlags:           []string{"-D__powerpc64__"},
			CrossCFlags:      []string{"-D__powerpc64__", "-static"},
			CCompilerPrefix:  "powerpc64le-linux-gnu-",
			KernelArch:       "powerpc",
			KernelHeaderArch: "powerpc",
		},
	},
	"freebsd": map[string]*Target{
		"amd64": {
			PtrSize:     8,
			PageSize:    4 << 10,
			CArch:       []string{"__x86_64__"},
			CFlags:      []string{"-m64"},
			CrossCFlags: []string{"-m64", "-static"},
		},
	},
	"netbsd": map[string]*Target{
		"amd64": {
			PtrSize:     8,
			PageSize:    4 << 10,
			CArch:       []string{"__x86_64__"},
			CFlags:      []string{"-m64"},
			CrossCFlags: []string{"-m64", "-static"},
		},
	},
	"fuchsia": map[string]*Target{
		"amd64": {
			PtrSize:          8,
			PageSize:         4 << 10,
			CArch:            []string{"__x86_64__"},
			KernelHeaderArch: "x64",
			CCompiler:        os.ExpandEnv("${SOURCEDIR}/buildtools/linux-x64/clang/bin/clang++"),
			CrossCFlags: []string{
				"-Wno-deprecated",
				"-Wno-error",
				"--target=x86_64-fuchsia",
				"-lfdio",
				"-lzircon",
				"--sysroot", os.ExpandEnv("${SOURCEDIR}/out/build-zircon/build-x64/sysroot"),
				"-L", os.ExpandEnv("${SOURCEDIR}/out/x64/x64-shared"),
				"-L", os.ExpandEnv("${SOURCEDIR}/out/x64/sdks/zircon_sysroot/sysroot/lib"),
			},
		},
		"arm64": {
			PtrSize:          8,
			PageSize:         4 << 10,
			CArch:            []string{"__aarch64__"},
			KernelHeaderArch: "arm64",
			CCompiler:        os.ExpandEnv("${SOURCEDIR}/buildtools/linux-x64/clang/bin/clang++"),
			CrossCFlags: []string{
				"-Wno-deprecated",
				"-Wno-error",
				"--target=aarch64-fuchsia",
				"-lfdio",
				"-lzircon",
				"--sysroot", os.ExpandEnv("${SOURCEDIR}/out/build-zircon/build-arm64/sysroot"),
			},
		},
	},
	"windows": map[string]*Target{
		"amd64": {
			PtrSize: 8,
			// TODO(dvyukov): what should we do about 4k vs 64k?
			PageSize: 4 << 10,
			CArch:    []string{"_M_X64"},
		},
	},
	"akaros": map[string]*Target{
		"amd64": {
			PtrSize:           8,
			PageSize:          4 << 10,
			CArch:             []string{"__x86_64__"},
			NeedSyscallDefine: dontNeedSyscallDefine,
			CCompiler:         os.ExpandEnv("${SOURCEDIR}/install/x86_64-ucb-akaros-gcc/bin/x86_64-ucb-akaros-g++"),
			// Most likely this is incorrect (why doesn't it know own sysroot?), but worked for me.
			CrossCFlags: []string{
				"-static",
				"-I", os.ExpandEnv("${SOURCEDIR}/tools/compilers/gcc-glibc/x86_64-ucb-akaros-gcc-stage3-builddir/x86_64-ucb-akaros/libstdc++-v3/include/x86_64-ucb-akaros"),
				"-I", os.ExpandEnv("${SOURCEDIR}/tools/compilers/gcc-glibc/x86_64-ucb-akaros-gcc-stage3-builddir/x86_64-ucb-akaros/libstdc++-v3/include"),
				"-I", os.ExpandEnv("${SOURCEDIR}/tools/compilers/gcc-glibc/gcc-4.9.2/libstdc++-v3/libsupc++"),
				"-L", os.ExpandEnv("${SOURCEDIR}/tools/compilers/gcc-glibc/x86_64-ucb-akaros-gcc-stage3-builddir/x86_64-ucb-akaros/libstdc++-v3/src/.libs"),
			},
		},
	},
}

var oses = map[string]osCommon{
	"linux": {
		SyscallNumbers:         true,
		SyscallPrefix:          "__NR_",
		ExecutorUsesShmem:      true,
		ExecutorUsesForkServer: true,
	},
	"freebsd": {
		SyscallNumbers:         true,
		SyscallPrefix:          "SYS_",
		ExecutorUsesShmem:      true,
		ExecutorUsesForkServer: true,
	},
	"netbsd": {
		SyscallNumbers:         true,
		SyscallPrefix:          "SYS_",
		ExecutorUsesShmem:      true,
		ExecutorUsesForkServer: true,
	},
	"fuchsia": {
		SyscallNumbers:         false,
		ExecutorUsesShmem:      false,
		ExecutorUsesForkServer: false,
	},
	"windows": {
		SyscallNumbers:         false,
		ExecutorUsesShmem:      false,
		ExecutorUsesForkServer: false,
		ExeExtension:           ".exe",
	},
	"akaros": {
		SyscallNumbers:         true,
		SyscallPrefix:          "SYS_",
		ExecutorUsesShmem:      false,
		ExecutorUsesForkServer: true,
	},
}

func init() {
	for OS, archs := range List {
		for arch, target := range archs {
			initTarget(target, OS, arch)
		}
	}
}

func initTarget(target *Target, OS, arch string) {
	target.osCommon = oses[OS]
	target.OS = OS
	target.Arch = arch
	if target.NeedSyscallDefine == nil {
		target.NeedSyscallDefine = needSyscallDefine
	}
	target.DataOffset = 512 << 20
	target.NumPages = (16 << 20) / target.PageSize
	if OS == runtime.GOOS && arch == runtime.GOARCH {
		// Don't use cross-compiler for native compilation, there are cases when this does not work:
		// https://github.com/google/syzkaller/pull/619
		// https://github.com/google/syzkaller/issues/387
		// https://github.com/google/syzkaller/commit/06db3cec94c54e1cf720cdd5db72761514569d56
		target.CCompilerPrefix = ""
	}
	if target.CCompiler == "" {
		target.CCompiler = target.CCompilerPrefix + "gcc"
	}
}

func checkStaticBuild(target *Target) {
	for i, flag := range target.CrossCFlags {
		if flag == "-static" {
			// Some distributions don't have static libraries.
			if !supportsStatic(target) {
				copy(target.CrossCFlags[i:], target.CrossCFlags[i+1:])
				target.CrossCFlags = target.CrossCFlags[:len(target.CrossCFlags)-1]
			}
			break
		}
	}
}

func supportsStatic(target *Target) bool {
	cmd := exec.Command(target.CCompiler, "-x", "c", "-", "-o", "/dev/null", "-static")
	cmd.Stdin = strings.NewReader("int main(){}")
	return cmd.Run() == nil
}

func needSyscallDefine(nr uint64) bool {
	return true
}
func dontNeedSyscallDefine(nr uint64) bool {
	return false
}
<|endoftext|>


<commit_before>// Copyright (c) 2019 Arista Networks, Inc.
// Use of this source code is governed by the Apache License 2.0
// that can be found in the COPYING file.

package key

// An entry represents an entry in a map whose key is not normally hashable,
// and is therefore of type Hashable
// (that is, a Hash method has been defined for this entry's key, and we can index it)
type entry struct {
	k    Hashable
	v    interface{}
	next *entry
}

// Map represents a map of arbitrary keys and values.
type Map struct {
	normal map[interface{}]interface{}
	custom map[uint64]entry
	length int // length of the Map
}

// String will output the string representation of the map
func (m *Map) String() string {
	// TODO
	return ""
}

// true if two arbitrary maps are equal
func mapEqual(a, b map[interface{}]interface{}) bool {
	if len(a) != len(b) {
		return false
	}
	for k, av := range a {
		bv, ok := b[k]
		if !ok {
			return false
		}
		if !keyEqual(av, bv) {
			return false
		}
	}
	return true
}

// true if entry a in entry list b
func findEntry(a, b entry) bool {
	bn := &b
	for bn != nil {
		if a.k.Equal(bn.k) {
			return keyEqual(a.v, bn.v)
		}
		bn = bn.next
	}
	return false
}

// return true if all entries in list a can be found in b
func entryEqual(a, b entry) bool {
	an := &a
	for an != nil {
		if !findEntry(*an, b) {
			return false
		}
		an = an.next
	}
	return true
}

// Equal compares two Maps
func (m *Map) Equal(other interface{}) bool {
	o, ok := other.(*Map)
	if !ok {
		return false
	}
	if m.length != o.length {
		return false
	}
	if !mapEqual(m.normal, o.normal) {
		return false
	}
	for k, mv := range m.custom {
		if len(m.custom) != len(o.custom) {
			return false
		}
		if ov, ok := o.custom[k]; ok {
			if !entryEqual(mv, ov) {
				return false
			}
		} else {
			return false
		}
	}
	return true
}

// Hashable represents the key for an entry in a Map that cannot natively be hashed
type Hashable interface {
	Hash() uint64
	Equal(other interface{}) bool
}

// Set allows the indexing of entries with arbitrary key types, so long as the keys are
// either hashable natively or are of type Hashable
func (m *Map) Set(k, v interface{}) {
	if k == nil {
		return
	}
	if hkey, ok := k.(Hashable); ok {
		if m.custom == nil {
			m.custom = make(map[uint64]entry)
		}
		// get hash, add to custom if not present
		// if present, append to next of root entry
		h := hkey.Hash()
		rootentry, ok := m.custom[h]
		if !ok {
			rootentry = entry{k: hkey, v: v}
		} else {
			var prev *entry
			curr := &rootentry
			for curr != nil {
				if curr.k.Equal(hkey) {
					curr.v = v
					m.custom[h] = rootentry
					return
				}
				prev = curr
				curr = prev.next
			}
			prev.next = &entry{k: hkey, v: v}
		}
		// write the stack back
		m.custom[h] = rootentry
	} else {
		if m.normal == nil {
			m.normal = make(map[interface{}]interface{})
		}
		l := len(m.normal)
		m.normal[k] = v
		if l == len(m.normal) { // len hasn't changed
			return
		}
	}
	m.length++
}

// Get retrieves the value stored with key k from the Map
func (m *Map) Get(k interface{}) (interface{}, bool) {
	if hkey, ok := k.(Hashable); ok {
		h := hkey.Hash()
		hentry, ok := m.custom[h]
		if !ok {
			return nil, false
		}
		curr := &hentry
		for curr != nil {
			if curr.k.Equal(hkey) {
				return curr.v, true
			}
			curr = curr.next
		}
		return nil, false
	}
	v, ok := m.normal[k]
	return v, ok
}

// Del removes an entry with key k from the Map
func (m *Map) Del(k interface{}) {
	if hkey, ok := k.(Hashable); ok {
		if m.custom == nil {
			return
		}
		h := hkey.Hash()
		hentry, ok := m.custom[h]
		if !ok {
			return
		}
		var prev *entry
		curr := &hentry
		for curr != nil {
			if curr.k.Equal(hkey) { // del
				if prev == nil { // delete the head
					if curr.next == nil { // no more entries at this hash, remove it
						delete(m.custom, h)
					} else {
						m.custom[h] = *curr.next
					}
				} else {
					// delete a mid/tail entry node
					prev.next = curr.next
					m.custom[h] = hentry
				}
				m.length--
				return
			}
			prev = curr
			curr = curr.next
		}
		return
	}
	// not Hashable, check normal
	if m.normal == nil {
		return
	}
	l := len(m.normal)
	delete(m.normal, k)
	if l != len(m.normal) {
		m.length--
	}
}

// Iter applies func f to every key-value pair in the Map
func (m *Map) Iter(f func(k, v interface{}) error) error {
	for k, v := range m.normal {
		if err := f(k, v); err != nil {
			return err
		}
	}
	for _, e := range m.custom {
		curr := &e
		for curr != nil {
			if err := f(curr.k, curr.v); err != nil {
				return err
			}
			curr = curr.next
		}
	}
	return nil
}
<commit_msg>key: add Map.Len() to get length of map<commit_after>// Copyright (c) 2019 Arista Networks, Inc.
// Use of this source code is governed by the Apache License 2.0
// that can be found in the COPYING file.

package key

// An entry represents an entry in a map whose key is not normally hashable,
// and is therefore of type Hashable
// (that is, a Hash method has been defined for this entry's key, and we can index it)
type entry struct {
	k    Hashable
	v    interface{}
	next *entry
}

// Map represents a map of arbitrary keys and values.
type Map struct {
	normal map[interface{}]interface{}
	custom map[uint64]entry
	length int // length of the Map
}

// String will output the string representation of the map
func (m *Map) String() string {
	// TODO
	return ""
}

// Len returns the length of the Map
func (m *Map) Len() int {
	return m.length
}

// true if two arbitrary maps are equal
func mapEqual(a, b map[interface{}]interface{}) bool {
	if len(a) != len(b) {
		return false
	}
	for k, av := range a {
		bv, ok := b[k]
		if !ok {
			return false
		}
		if !keyEqual(av, bv) {
			return false
		}
	}
	return true
}

// true if entry a in entry list b
func findEntry(a, b entry) bool {
	bn := &b
	for bn != nil {
		if a.k.Equal(bn.k) {
			return keyEqual(a.v, bn.v)
		}
		bn = bn.next
	}
	return false
}

// return true if all entries in list a can be found in b
func entryEqual(a, b entry) bool {
	an := &a
	for an != nil {
		if !findEntry(*an, b) {
			return false
		}
		an = an.next
	}
	return true
}

// Equal compares two Maps
func (m *Map) Equal(other interface{}) bool {
	o, ok := other.(*Map)
	if !ok {
		return false
	}
	if m.length != o.length {
		return false
	}
	if !mapEqual(m.normal, o.normal) {
		return false
	}
	if len(m.custom) != len(o.custom) {
		return false
	}
	for k, mv := range m.custom {
		if ov, ok := o.custom[k]; ok {
			if !entryEqual(mv, ov) {
				return false
			}
		} else {
			return false
		}
	}
	return true
}

// Hashable represents the key for an entry in a Map that cannot natively be hashed
type Hashable interface {
	Hash() uint64
	Equal(other interface{}) bool
}

// Set allows the indexing of entries with arbitrary key types, so long as the keys are
// either hashable natively or are of type Hashable
func (m *Map) Set(k, v interface{}) {
	if k == nil {
		return
	}
	if hkey, ok := k.(Hashable); ok {
		if m.custom == nil {
			m.custom = make(map[uint64]entry)
		}
		// get hash, add to custom if not present
		// if present, append to next of root entry
		h := hkey.Hash()
		rootentry, ok := m.custom[h]
		if !ok {
			rootentry = entry{k: hkey, v: v}
		} else {
			var prev *entry
			curr := &rootentry
			for curr != nil {
				if curr.k.Equal(hkey) {
					curr.v = v
					m.custom[h] = rootentry
					return
				}
				prev = curr
				curr = prev.next
			}
			prev.next = &entry{k: hkey, v: v}
		}
		// write the stack back
		m.custom[h] = rootentry
	} else {
		if m.normal == nil {
			m.normal = make(map[interface{}]interface{})
		}
		l := len(m.normal)
		m.normal[k] = v
		if l == len(m.normal) { // len hasn't changed
			return
		}
	}
	m.length++
}

// Get retrieves the value stored with key k from the Map
func (m *Map) Get(k interface{}) (interface{}, bool) {
	if hkey, ok := k.(Hashable); ok {
		h := hkey.Hash()
		hentry, ok := m.custom[h]
		if !ok {
			return nil, false
		}
		curr := &hentry
		for curr != nil {
			if curr.k.Equal(hkey) {
				return curr.v, true
			}
			curr = curr.next
		}
		return nil, false
	}
	v, ok := m.normal[k]
	return v, ok
}

// Del removes an entry with key k from the Map
func (m *Map) Del(k interface{}) {
	if hkey, ok := k.(Hashable); ok {
		if m.custom == nil {
			return
		}
		h := hkey.Hash()
		hentry, ok := m.custom[h]
		if !ok {
			return
		}
		var prev *entry
		curr := &hentry
		for curr != nil {
			if curr.k.Equal(hkey) { // del
				if prev == nil { // delete the head
					if curr.next == nil { // no more entries at this hash, remove it
						delete(m.custom, h)
					} else {
						m.custom[h] = *curr.next
					}
				} else {
					// delete a mid/tail entry node
					prev.next = curr.next
					m.custom[h] = hentry
				}
				m.length--
				return
			}
			prev = curr
			curr = curr.next
		}
		return
	}
	// not Hashable, check normal
	if m.normal == nil {
		return
	}
	l := len(m.normal)
	delete(m.normal, k)
	if l != len(m.normal) {
		m.length--
	}
}

// Iter applies func f to every key-value pair in the Map
func (m *Map) Iter(f func(k, v interface{}) error) error {
	for k, v := range m.normal {
		if err := f(k, v); err != nil {
			return err
		}
	}
	for _, e := range m.custom {
		curr := &e
		for curr != nil {
			if err := f(curr.k, curr.v); err != nil {
				return err
			}
			curr = curr.next
		}
	}
	return nil
}
<|endoftext|>