<commit_before>package server

import (
	"bufio"
	"bytes"
	"fmt"
	"hash/fnv"
	"io"
	"math/rand"
	"os"
	"path"
	"reflect"
	"strings"
	"sync"
	"testing"

	"github.com/pachyderm/pachyderm/src/client"
	"github.com/pachyderm/pachyderm/src/client/pfs"
	"github.com/pachyderm/pachyderm/src/client/pkg/require"
	"github.com/pachyderm/pachyderm/src/client/pps"
	versionlib "github.com/pachyderm/pachyderm/src/client/version"
	"github.com/pachyderm/pachyderm/src/server/pkg/backoff"
	tu "github.com/pachyderm/pachyderm/src/server/pkg/testutil"
	"github.com/pachyderm/pachyderm/src/server/pkg/workload"

	"github.com/golang/snappy"
)

const (
	KB = 1024
	MB = 1024 * KB
)

var pachClient *client.APIClient
var getPachClientOnce sync.Once

func getPachClient(t testing.TB) *client.APIClient {
	getPachClientOnce.Do(func() {
		var err error
		if addr := os.Getenv("PACHD_PORT_650_TCP_ADDR"); addr != "" {
			pachClient, err = client.NewInCluster()
		} else {
			pachClient, err = client.NewOnUserMachine(false, false, "user")
		}
		require.NoError(t, err)
	})
	return pachClient
}

func collectCommitInfos(t testing.TB, commitInfoIter client.CommitInfoIterator) []*pfs.CommitInfo {
	var commitInfos []*pfs.CommitInfo
	for {
		commitInfo, err := commitInfoIter.Next()
		if err == io.EOF {
			return commitInfos
		}
		require.NoError(t, err)
		commitInfos = append(commitInfos, commitInfo)
	}
}

// TODO(msteffen) equivalent to funciton in src/server/auth/server/admin_test.go.
// These should be unified.
func RepoInfoToName(repoInfo interface{}) interface{} {
	return repoInfo.(*pfs.RepoInfo).Repo.Name
}

// testExtractRestored effectively implements both TestExtractRestoreObjects
// TestExtractRestoreObjects, and their logic is mostly the same
func testExtractRestore(t *testing.T, testObjects bool) {
	if testing.Short() {
		t.Skip("Skipping integration tests in short mode")
	}

	c := getPachClient(t)
	require.NoError(t, c.DeleteAll())

	dataRepo := tu.UniqueString("TestExtractRestoreObjects-in-")
	require.NoError(t, c.CreateRepo(dataRepo))

	// Create input data
	nCommits := 2
	r := rand.New(rand.NewSource(45))
	fileHashes := make([]string, 0, nCommits)
	for i := 0; i < nCommits; i++ {
		hash := fnv.New64a()
		fileContent := workload.RandString(r, 40*MB)
		_, err := c.PutFile(dataRepo, "master", fmt.Sprintf("file-%d", i),
			io.TeeReader(strings.NewReader(fileContent), hash))
		require.NoError(t, err)
		fileHashes = append(fileHashes, string(hash.Sum(nil)))
	}

	// Create test pipelines
	numPipelines := 3
	var input, pipeline string
	input = dataRepo
	for i := 0; i < numPipelines; i++ {
		pipeline = tu.UniqueString(fmt.Sprintf("TestExtractRestoreObjects-P%d-", i))
		require.NoError(t, c.CreatePipeline(
			pipeline,
			"",
			[]string{"bash"},
			[]string{
				fmt.Sprintf("cp /pfs/%s/* /pfs/out/", input),
			},
			&pps.ParallelismSpec{
				Constant: 1,
			},
			client.NewPFSInput(input, "/*"),
			"",
			false,
		))
		input = pipeline
	}

	// Wait for pipelines to process input data
	commitIter, err := c.FlushCommit([]*pfs.Commit{client.NewCommit(dataRepo, "master")}, nil)
	require.NoError(t, err)
	commitInfos := collectCommitInfos(t, commitIter)
	require.Equal(t, numPipelines, len(commitInfos))

	// Extract existing cluster state
	ops, err := c.ExtractAll(testObjects)
	require.NoError(t, err)

	// Delete existing metadata
	require.NoError(t, c.DeleteAll())

	if testObjects {
		// Delete existing objects
		require.NoError(t, c.GarbageCollect(10000))
	}

	// Restore metadata and possibly objects
	require.NoError(t, c.Restore(ops))

	// Wait for re-created pipelines to process recreated input data
	commitIter, err = c.FlushCommit([]*pfs.Commit{client.NewCommit(dataRepo, "master")}, nil)
	require.NoError(t, err)
	commitInfos = collectCommitInfos(t, commitIter)
	require.Equal(t, numPipelines, len(commitInfos))

	// Make sure recreated jobs all succeeded
	jis, err := c.ListJob("", nil, nil) // make sure jobs all succeeded
	require.NoError(t, err)
	for _, ji := range jis {
		require.Equal(t, pps.JobState_JOB_SUCCESS, ji.State)
	}

	// Make sure all branches were recreated
	bis, err := c.ListBranch(dataRepo)
	require.NoError(t, err)
	require.Equal(t, 1, len(bis))

	// Check input data
	// This check uses a backoff because sometimes GetFile causes pachd to OOM
	var restoredFileHashes []string
	require.NoError(t, backoff.Retry(func() error {
		restoredFileHashes = make([]string, 0, nCommits)
		for i := 0; i < nCommits; i++ {
			hash := fnv.New64a()
			err := c.GetFile(dataRepo, "master", fmt.Sprintf("file-%d", i), 0, 0, hash)
			if err != nil {
				return err
			}
			restoredFileHashes = append(restoredFileHashes, string(hash.Sum(nil)))
		}
		return nil
	}, backoff.NewTestingBackOff()))
	require.ElementsEqual(t, fileHashes, restoredFileHashes)

	// Check output data
	// This check uses a backoff because sometimes GetFile causes pachd to OOM
	require.NoError(t, backoff.Retry(func() error {
		restoredFileHashes = make([]string, 0, nCommits)
		for i := 0; i < nCommits; i++ {
			hash := fnv.New64a()
			err := c.GetFile(pipeline, "master", fmt.Sprintf("file-%d", i), 0, 0, hash)
			if err != nil {
				return err
			}
			restoredFileHashes = append(restoredFileHashes, string(hash.Sum(nil)))
		}
		return nil
	}, backoff.NewTestingBackOff()))
	require.ElementsEqual(t, fileHashes, restoredFileHashes)
}

// TestExtractRestoreNoObjects tests extraction and restoration in the case
// where existing objects are re-used (common for cloud deployments, as objects
// are stored outside of kubernetes, in object store)
func TestExtractRestoreNoObjects(t *testing.T) {
	testExtractRestore(t, false)
}

// TestExtractRestoreObjects tests extraction and restoration of objects. Note
// that since 1.8, only data in input repos is referenced by objects, so this
// tests extracting/restoring an input repo.
func TestExtractRestoreObjects(t *testing.T) {
	testExtractRestore(t, true)
}

func TestExtractRestoreHeadlessBranches(t *testing.T) {
	if testing.Short() {
		t.Skip("Skipping integration tests in short mode")
	}

	c := getPachClient(t)
	require.NoError(t, c.DeleteAll())

	dataRepo := tu.UniqueString("TestExtractRestore_data")
	require.NoError(t, c.CreateRepo(dataRepo))

	// create a headless branch
	require.NoError(t, c.CreateBranch(dataRepo, "headless", "", nil))

	ops, err := c.ExtractAll(false)
	require.NoError(t, err)
	require.NoError(t, c.DeleteAll())
	require.NoError(t, c.Restore(ops))

	bis, err := c.ListBranch(dataRepo)
	require.NoError(t, err)
	require.Equal(t, 1, len(bis))
	require.Equal(t, "headless", bis[0].Branch.Name)
}

func TestExtractVersion(t *testing.T) {
	if testing.Short() {
		t.Skip("Skipping integration tests in short mode")
	}

	c := getPachClient(t)
	require.NoError(t, c.DeleteAll())

	dataRepo := tu.UniqueString("TestExtractRestore_data")
	require.NoError(t, c.CreateRepo(dataRepo))

	r := rand.New(rand.NewSource(45))
	_, err := c.PutFile(dataRepo, "master", "file", strings.NewReader(workload.RandString(r, 40*MB)))
	require.NoError(t, err)

	pipeline := tu.UniqueString("TestExtractRestore")
	require.NoError(t, c.CreatePipeline(
		pipeline,
		"",
		[]string{"bash"},
		[]string{
			fmt.Sprintf("cp /pfs/%s/* /pfs/out/", dataRepo),
		},
		&pps.ParallelismSpec{
			Constant: 1,
		},
		client.NewPFSInput(dataRepo, "/*"),
		"",
		false,
	))

	ops, err := c.ExtractAll(false)
	require.NoError(t, err)
	require.True(t, len(ops) > 0)

	// Check that every Op looks right; the version set matches pachd's version
	for _, op := range ops {
		opV := reflect.ValueOf(op).Elem()
		var versions, nonemptyVersions int
		for i := 0; i < opV.NumField(); i++ {
			fDesc := opV.Type().Field(i)
			if !strings.HasPrefix(fDesc.Name, "Op") {
				continue
			}
			versions++
			if strings.HasSuffix(fDesc.Name,
				fmt.Sprintf("%d_%d", versionlib.MajorVersion, versionlib.MinorVersion)) {
				require.False(t, opV.Field(i).IsNil())
				nonemptyVersions++
			} else {
				require.True(t, opV.Field(i).IsNil())
			}
		}
		require.Equal(t, 1, nonemptyVersions)
		require.True(t, versions > 1)
	}
}

func TestMigrateFrom1_7(t *testing.T) {
	if testing.Short() {
		t.Skip("Skipping integration tests in short mode")
	}

	// Clear pachyderm cluster (so that next cluster starts up in a clean environment)
	c := getPachClient(t)
	require.NoError(t, c.DeleteAll())

	// Restore dumped metadata (now that objects are present)
	md, err := os.Open(path.Join(os.Getenv("GOPATH"),
		"src/github.com/pachyderm/pachyderm/etc/testing/migration/1_7/sort.metadata"))
	require.NoError(t, err)
	require.NoError(t, c.RestoreReader(snappy.NewReader(md)))
	require.NoError(t, md.Close())

	// Wait for final imported commit to be processed
	commitIter, err := c.FlushCommit([]*pfs.Commit{client.NewCommit("left", "master")}, nil)
	require.NoError(t, err)
	commitInfos := collectCommitInfos(t, commitIter)
	// filter-left and filter-right both compute a join of left and
	// right--depending on when the final commit to 'left' was added, it may have
	// been processed multiple times (should be n * 3, as there are 3 pipelines)
	require.True(t, len(commitInfos) >= 3)

	// Inspect input
	commits, err := c.ListCommit("left", "master", "", 0)
	require.NoError(t, err)
	require.Equal(t, 3, len(commits))
	commits, err = c.ListCommit("right", "master", "", 0)
	require.NoError(t, err)
	require.Equal(t, 3, len(commits))

	// Inspect output
	repos, err := c.ListRepo()
	require.NoError(t, err)
	require.ElementsEqualUnderFn(t,
		[]string{"left", "right", "copy", "sort"},
		repos, RepoInfoToName)

	// make sure all numbers 0-99 are in /nums
	var buf bytes.Buffer
	require.NoError(t, c.GetFile("sort", "master", "/nums", 0, 0, &buf))
	s := bufio.NewScanner(&buf)
	numbers := make(map[string]struct{})
	for s.Scan() {
		numbers[s.Text()] = struct{}{}
	}
	require.Equal(t, 100, len(numbers)) // job processed all inputs

	// Confirm stats commits are present
	commits, err = c.ListCommit("sort", "stats", "", 0)
	require.NoError(t, err)
	require.Equal(t, 6, len(commits))
}

func int64p(i int64) *int64 {
	return &i
}
<commit_msg>Fix comment<commit_after>package server

import (
	"bufio"
	"bytes"
	"fmt"
	"hash/fnv"
	"io"
	"math/rand"
	"os"
	"path"
	"reflect"
	"strings"
	"sync"
	"testing"

	"github.com/pachyderm/pachyderm/src/client"
	"github.com/pachyderm/pachyderm/src/client/pfs"
	"github.com/pachyderm/pachyderm/src/client/pkg/require"
	"github.com/pachyderm/pachyderm/src/client/pps"
	versionlib "github.com/pachyderm/pachyderm/src/client/version"
	"github.com/pachyderm/pachyderm/src/server/pkg/backoff"
	tu "github.com/pachyderm/pachyderm/src/server/pkg/testutil"
	"github.com/pachyderm/pachyderm/src/server/pkg/workload"

	"github.com/golang/snappy"
)

const (
	KB = 1024
	MB = 1024 * KB
)

var pachClient *client.APIClient
var getPachClientOnce sync.Once

func getPachClient(t testing.TB) *client.APIClient {
	getPachClientOnce.Do(func() {
		var err error
		if addr := os.Getenv("PACHD_PORT_650_TCP_ADDR"); addr != "" {
			pachClient, err = client.NewInCluster()
		} else {
			pachClient, err = client.NewOnUserMachine(false, false, "user")
		}
		require.NoError(t, err)
	})
	return pachClient
}

func collectCommitInfos(t testing.TB, commitInfoIter client.CommitInfoIterator) []*pfs.CommitInfo {
	var commitInfos []*pfs.CommitInfo
	for {
		commitInfo, err := commitInfoIter.Next()
		if err == io.EOF {
			return commitInfos
		}
		require.NoError(t, err)
		commitInfos = append(commitInfos, commitInfo)
	}
}

// TODO(msteffen) equivalent to funciton in src/server/auth/server/admin_test.go.
// These should be unified.
func RepoInfoToName(repoInfo interface{}) interface{} {
	return repoInfo.(*pfs.RepoInfo).Repo.Name
}

// testExtractRestored effectively implements both TestExtractRestoreObjects
// TestExtractRestoreNoObjects, as their logic is mostly the same
func testExtractRestore(t *testing.T, testObjects bool) {
	if testing.Short() {
		t.Skip("Skipping integration tests in short mode")
	}

	c := getPachClient(t)
	require.NoError(t, c.DeleteAll())

	dataRepo := tu.UniqueString("TestExtractRestoreObjects-in-")
	require.NoError(t, c.CreateRepo(dataRepo))

	// Create input data
	nCommits := 2
	r := rand.New(rand.NewSource(45))
	fileHashes := make([]string, 0, nCommits)
	for i := 0; i < nCommits; i++ {
		hash := fnv.New64a()
		fileContent := workload.RandString(r, 40*MB)
		_, err := c.PutFile(dataRepo, "master", fmt.Sprintf("file-%d", i),
			io.TeeReader(strings.NewReader(fileContent), hash))
		require.NoError(t, err)
		fileHashes = append(fileHashes, string(hash.Sum(nil)))
	}

	// Create test pipelines
	numPipelines := 3
	var input, pipeline string
	input = dataRepo
	for i := 0; i < numPipelines; i++ {
		pipeline = tu.UniqueString(fmt.Sprintf("TestExtractRestoreObjects-P%d-", i))
		require.NoError(t, c.CreatePipeline(
			pipeline,
			"",
			[]string{"bash"},
			[]string{
				fmt.Sprintf("cp /pfs/%s/* /pfs/out/", input),
			},
			&pps.ParallelismSpec{
				Constant: 1,
			},
			client.NewPFSInput(input, "/*"),
			"",
			false,
		))
		input = pipeline
	}

	// Wait for pipelines to process input data
	commitIter, err := c.FlushCommit([]*pfs.Commit{client.NewCommit(dataRepo, "master")}, nil)
	require.NoError(t, err)
	commitInfos := collectCommitInfos(t, commitIter)
	require.Equal(t, numPipelines, len(commitInfos))

	// Extract existing cluster state
	ops, err := c.ExtractAll(testObjects)
	require.NoError(t, err)

	// Delete existing metadata
	require.NoError(t, c.DeleteAll())

	if testObjects {
		// Delete existing objects
		require.NoError(t, c.GarbageCollect(10000))
	}

	// Restore metadata and possibly objects
	require.NoError(t, c.Restore(ops))

	// Wait for re-created pipelines to process recreated input data
	commitIter, err = c.FlushCommit([]*pfs.Commit{client.NewCommit(dataRepo, "master")}, nil)
	require.NoError(t, err)
	commitInfos = collectCommitInfos(t, commitIter)
	require.Equal(t, numPipelines, len(commitInfos))

	// Make sure recreated jobs all succeeded
	jis, err := c.ListJob("", nil, nil) // make sure jobs all succeeded
	require.NoError(t, err)
	for _, ji := range jis {
		require.Equal(t, pps.JobState_JOB_SUCCESS, ji.State)
	}

	// Make sure all branches were recreated
	bis, err := c.ListBranch(dataRepo)
	require.NoError(t, err)
	require.Equal(t, 1, len(bis))

	// Check input data
	// This check uses a backoff because sometimes GetFile causes pachd to OOM
	var restoredFileHashes []string
	require.NoError(t, backoff.Retry(func() error {
		restoredFileHashes = make([]string, 0, nCommits)
		for i := 0; i < nCommits; i++ {
			hash := fnv.New64a()
			err := c.GetFile(dataRepo, "master", fmt.Sprintf("file-%d", i), 0, 0, hash)
			if err != nil {
				return err
			}
			restoredFileHashes = append(restoredFileHashes, string(hash.Sum(nil)))
		}
		return nil
	}, backoff.NewTestingBackOff()))
	require.ElementsEqual(t, fileHashes, restoredFileHashes)

	// Check output data
	// This check uses a backoff because sometimes GetFile causes pachd to OOM
	require.NoError(t, backoff.Retry(func() error {
		restoredFileHashes = make([]string, 0, nCommits)
		for i := 0; i < nCommits; i++ {
			hash := fnv.New64a()
			err := c.GetFile(pipeline, "master", fmt.Sprintf("file-%d", i), 0, 0, hash)
			if err != nil {
				return err
			}
			restoredFileHashes = append(restoredFileHashes, string(hash.Sum(nil)))
		}
		return nil
	}, backoff.NewTestingBackOff()))
	require.ElementsEqual(t, fileHashes, restoredFileHashes)
}

// TestExtractRestoreNoObjects tests extraction and restoration in the case
// where existing objects are re-used (common for cloud deployments, as objects
// are stored outside of kubernetes, in object store)
func TestExtractRestoreNoObjects(t *testing.T) {
	testExtractRestore(t, false)
}

// TestExtractRestoreObjects tests extraction and restoration of objects. Note
// that since 1.8, only data in input repos is referenced by objects, so this
// tests extracting/restoring an input repo.
func TestExtractRestoreObjects(t *testing.T) {
	testExtractRestore(t, true)
}

func TestExtractRestoreHeadlessBranches(t *testing.T) {
	if testing.Short() {
		t.Skip("Skipping integration tests in short mode")
	}

	c := getPachClient(t)
	require.NoError(t, c.DeleteAll())

	dataRepo := tu.UniqueString("TestExtractRestore_data")
	require.NoError(t, c.CreateRepo(dataRepo))

	// create a headless branch
	require.NoError(t, c.CreateBranch(dataRepo, "headless", "", nil))

	ops, err := c.ExtractAll(false)
	require.NoError(t, err)
	require.NoError(t, c.DeleteAll())
	require.NoError(t, c.Restore(ops))

	bis, err := c.ListBranch(dataRepo)
	require.NoError(t, err)
	require.Equal(t, 1, len(bis))
	require.Equal(t, "headless", bis[0].Branch.Name)
}

func TestExtractVersion(t *testing.T) {
	if testing.Short() {
		t.Skip("Skipping integration tests in short mode")
	}

	c := getPachClient(t)
	require.NoError(t, c.DeleteAll())

	dataRepo := tu.UniqueString("TestExtractRestore_data")
	require.NoError(t, c.CreateRepo(dataRepo))

	r := rand.New(rand.NewSource(45))
	_, err := c.PutFile(dataRepo, "master", "file", strings.NewReader(workload.RandString(r, 40*MB)))
	require.NoError(t, err)

	pipeline := tu.UniqueString("TestExtractRestore")
	require.NoError(t, c.CreatePipeline(
		pipeline,
		"",
		[]string{"bash"},
		[]string{
			fmt.Sprintf("cp /pfs/%s/* /pfs/out/", dataRepo),
		},
		&pps.ParallelismSpec{
			Constant: 1,
		},
		client.NewPFSInput(dataRepo, "/*"),
		"",
		false,
	))

	ops, err := c.ExtractAll(false)
	require.NoError(t, err)
	require.True(t, len(ops) > 0)

	// Check that every Op looks right; the version set matches pachd's version
	for _, op := range ops {
		opV := reflect.ValueOf(op).Elem()
		var versions, nonemptyVersions int
		for i := 0; i < opV.NumField(); i++ {
			fDesc := opV.Type().Field(i)
			if !strings.HasPrefix(fDesc.Name, "Op") {
				continue
			}
			versions++
			if strings.HasSuffix(fDesc.Name,
				fmt.Sprintf("%d_%d", versionlib.MajorVersion, versionlib.MinorVersion)) {
				require.False(t, opV.Field(i).IsNil())
				nonemptyVersions++
			} else {
				require.True(t, opV.Field(i).IsNil())
			}
		}
		require.Equal(t, 1, nonemptyVersions)
		require.True(t, versions > 1)
	}
}

func TestMigrateFrom1_7(t *testing.T) {
	if testing.Short() {
		t.Skip("Skipping integration tests in short mode")
	}

	// Clear pachyderm cluster (so that next cluster starts up in a clean environment)
	c := getPachClient(t)
	require.NoError(t, c.DeleteAll())

	// Restore dumped metadata (now that objects are present)
	md, err := os.Open(path.Join(os.Getenv("GOPATH"),
		"src/github.com/pachyderm/pachyderm/etc/testing/migration/1_7/sort.metadata"))
	require.NoError(t, err)
	require.NoError(t, c.RestoreReader(snappy.NewReader(md)))
	require.NoError(t, md.Close())

	// Wait for final imported commit to be processed
	commitIter, err := c.FlushCommit([]*pfs.Commit{client.NewCommit("left", "master")}, nil)
	require.NoError(t, err)
	commitInfos := collectCommitInfos(t, commitIter)
	// filter-left and filter-right both compute a join of left and
	// right--depending on when the final commit to 'left' was added, it may have
	// been processed multiple times (should be n * 3, as there are 3 pipelines)
	require.True(t, len(commitInfos) >= 3)

	// Inspect input
	commits, err := c.ListCommit("left", "master", "", 0)
	require.NoError(t, err)
	require.Equal(t, 3, len(commits))
	commits, err = c.ListCommit("right", "master", "", 0)
	require.NoError(t, err)
	require.Equal(t, 3, len(commits))

	// Inspect output
	repos, err := c.ListRepo()
	require.NoError(t, err)
	require.ElementsEqualUnderFn(t,
		[]string{"left", "right", "copy", "sort"},
		repos, RepoInfoToName)

	// make sure all numbers 0-99 are in /nums
	var buf bytes.Buffer
	require.NoError(t, c.GetFile("sort", "master", "/nums", 0, 0, &buf))
	s := bufio.NewScanner(&buf)
	numbers := make(map[string]struct{})
	for s.Scan() {
		numbers[s.Text()] = struct{}{}
	}
	require.Equal(t, 100, len(numbers)) // job processed all inputs

	// Confirm stats commits are present
	commits, err = c.ListCommit("sort", "stats", "", 0)
	require.NoError(t, err)
	require.Equal(t, 6, len(commits))
}

func int64p(i int64) *int64 {
	return &i
}
<|endoftext|>


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

package main

import (
	"context"
	"crypto/hmac"
	"crypto/md5"
	"encoding/json"
	"errors"
	"fmt"
	"html"
	"log"
	"net/http"
	"strconv"
	"strings"
	"unicode/utf8"

	"cloud.google.com/go/datastore"
	"golang.org/x/build/app/key"
)

const (
	commitsPerPage = 30
	builderVersion = 1 // must match x/build/cmd/coordinator/dash.go's value
)

// resultHandler records a build result.
// It reads a JSON-encoded Result value from the request body,
// creates a new Result entity, and creates or updates the relevant Commit entity.
// If the Log field is not empty, resultHandler creates a new Log entity
// and updates the LogHash field before putting the Commit entity.
func resultHandler(r *http.Request) (interface{}, error) {
	if r.Method != "POST" {
		return nil, errBadMethod(r.Method)
	}

	v, _ := strconv.Atoi(r.FormValue("version"))
	if v != builderVersion {
		return nil, fmt.Errorf("rejecting POST from builder; need version %v instead of %v",
			builderVersion, v)
	}

	ctx := r.Context()
	res := new(Result)
	defer r.Body.Close()
	if err := json.NewDecoder(r.Body).Decode(res); err != nil {
		return nil, fmt.Errorf("decoding Body: %v", err)
	}
	if err := res.Valid(); err != nil {
		return nil, fmt.Errorf("validating Result: %v", err)
	}
	// store the Log text if supplied
	if len(res.Log) > 0 {
		hash, err := PutLog(ctx, res.Log)
		if err != nil {
			return nil, fmt.Errorf("putting Log: %v", err)
		}
		res.LogHash = hash
	}
	tx := func(tx *datastore.Transaction) error {
		if _, err := getOrMakePackageInTx(ctx, tx, res.PackagePath); err != nil {
			return fmt.Errorf("GetPackage: %v", err)
		}
		// put Result
		if _, err := tx.Put(res.Key(), res); err != nil {
			return fmt.Errorf("putting Result: %v", err)
		}
		// add Result to Commit
		com := &Commit{PackagePath: res.PackagePath, Hash: res.Hash}
		if err := com.AddResult(tx, res); err != nil {
			return fmt.Errorf("AddResult: %v", err)
		}
		return nil
	}
	_, err := datastoreClient.RunInTransaction(ctx, tx)
	return nil, err
}

// logHandler displays log text for a given hash.
// It handles paths like "/log/hash".
func logHandler(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-type", "text/plain; charset=utf-8")
	c := r.Context()
	hash := r.URL.Path[strings.LastIndex(r.URL.Path, "/")+1:]
	key := dsKey("Log", hash, nil)
	l := new(Log)
	if err := datastoreClient.Get(c, key, l); err != nil {
		if err == datastore.ErrNoSuchEntity {
			// Fall back to default namespace;
			// maybe this was on the old dashboard.
			key := dsKey("Log", hash, nil)
			err = datastoreClient.Get(c, key, l)
		}
		if err != nil {
			logErr(w, r, err)
			return
		}
	}
	b, err := l.Text()
	if err != nil {
		logErr(w, r, err)
		return
	}
	w.Write(b)
}

// clearResultsHandler purge a single build failure from the dashboard.
// It currently only supports the main Go repo.
func clearResultsHandler(r *http.Request) (interface{}, error) {
	if r.Method != "POST" {
		return nil, errBadMethod(r.Method)
	}
	builder := r.FormValue("builder")
	hash := r.FormValue("hash")
	if builder == "" {
		return nil, errors.New("missing 'builder'")
	}
	if hash == "" {
		return nil, errors.New("missing 'hash'")
	}

	ctx := r.Context()

	_, err := datastoreClient.RunInTransaction(ctx, func(tx *datastore.Transaction) error {
		c := &Commit{
			PackagePath: "", // TODO(adg): support clearing sub-repos
			Hash:        hash,
		}
		err := tx.Get(c.Key(), c)
		err = filterDatastoreError(err)
		if err == datastore.ErrNoSuchEntity {
			// Doesn't exist, so no build to clear.
			return nil
		}
		if err != nil {
			return err
		}

		r := c.Result(builder, "")
		if r == nil {
			// No result, so nothing to clear.
			return nil
		}
		c.RemoveResult(r)
		_, err = tx.Put(c.Key(), c)
		if err != nil {
			return err
		}
		return tx.Delete(r.Key())
	})
	return nil, err
}

type dashHandler func(*http.Request) (interface{}, error)

type dashResponse struct {
	Response interface{}
	Error    string
}

// errBadMethod is returned by a dashHandler when
// the request has an unsuitable method.
type errBadMethod string

func (e errBadMethod) Error() string {
	return "bad method: " + string(e)
}

func builderKeyRevoked(builder string) bool {
	switch builder {
	case "plan9-amd64-mischief":
		// Broken and unmaintained for months.
		// It's polluting the dashboard.
		return true
	case "linux-arm-onlinenet":
		// Requested to be revoked by Dave Cheney.
		// The machine is in a fail+report loop
		// and can't be accessed. Revoke it for now.
		return true
	}
	return false
}

// AuthHandler wraps a http.HandlerFunc with a handler that validates the
// supplied key and builder query parameters.
func AuthHandler(h dashHandler) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		c := r.Context()

		// Put the URL Query values into r.Form to avoid parsing the
		// request body when calling r.FormValue.
		r.Form = r.URL.Query()

		var err error
		var resp interface{}

		// Validate key query parameter for POST requests only.
		key := r.FormValue("key")
		builder := r.FormValue("builder")
		if r.Method == "POST" && !validKey(c, key, builder) {
			err = fmt.Errorf("invalid key %q for builder %q", key, builder)
		}

		// Call the original HandlerFunc and return the response.
		if err == nil {
			resp, err = h(r)
		}

		// Write JSON response.
		dashResp := &dashResponse{Response: resp}
		if err != nil {
			log.Printf("%v", err)
			dashResp.Error = err.Error()
		}
		w.Header().Set("Content-Type", "application/json")
		if err = json.NewEncoder(w).Encode(dashResp); err != nil {
			log.Printf("encoding response: %v", err)
		}
	}
}

// validHash reports whether hash looks like a valid git commit hash.
func validHash(hash string) bool {
	// TODO: correctly validate a hash: check that it's exactly 40
	// lowercase hex digits. But this is what we historically did:
	return hash != ""
}

func validKey(c context.Context, key, builder string) bool {
	if isMasterKey(c, key) {
		return true
	}
	if builderKeyRevoked(builder) {
		return false
	}
	return key == builderKey(c, builder)
}

func isMasterKey(c context.Context, k string) bool {
	return k == key.Secret(datastoreClient, c)
}

func builderKey(c context.Context, builder string) string {
	h := hmac.New(md5.New, []byte(key.Secret(datastoreClient, c)))
	h.Write([]byte(builder))
	return fmt.Sprintf("%x", h.Sum(nil))
}

func logErr(w http.ResponseWriter, r *http.Request, err error) {
	log.Printf("Error: %v", err)
	w.WriteHeader(http.StatusInternalServerError)
	fmt.Fprint(w, "Error: ", html.EscapeString(err.Error()))
}

// limitStringLength essentially does return s[:max],
// but it ensures that we dot not split UTF-8 rune in half.
// Otherwise appengine python scripts will break badly.
func limitStringLength(s string, max int) string {
	if len(s) <= max {
		return s
	}
	for {
		s = s[:max]
		r, size := utf8.DecodeLastRuneInString(s)
		if r != utf8.RuneError || size != 1 {
			return s
		}
		max--
	}
}
<commit_msg>app/appengine: fix log handler's fallback to the old Mercurial namespace<commit_after>// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"context"
	"crypto/hmac"
	"crypto/md5"
	"encoding/json"
	"errors"
	"fmt"
	"html"
	"log"
	"net/http"
	"strconv"
	"strings"
	"unicode/utf8"

	"cloud.google.com/go/datastore"
	"golang.org/x/build/app/key"
)

const (
	commitsPerPage = 30
	builderVersion = 1 // must match x/build/cmd/coordinator/dash.go's value
)

// resultHandler records a build result.
// It reads a JSON-encoded Result value from the request body,
// creates a new Result entity, and creates or updates the relevant Commit entity.
// If the Log field is not empty, resultHandler creates a new Log entity
// and updates the LogHash field before putting the Commit entity.
func resultHandler(r *http.Request) (interface{}, error) {
	if r.Method != "POST" {
		return nil, errBadMethod(r.Method)
	}

	v, _ := strconv.Atoi(r.FormValue("version"))
	if v != builderVersion {
		return nil, fmt.Errorf("rejecting POST from builder; need version %v instead of %v",
			builderVersion, v)
	}

	ctx := r.Context()
	res := new(Result)
	defer r.Body.Close()
	if err := json.NewDecoder(r.Body).Decode(res); err != nil {
		return nil, fmt.Errorf("decoding Body: %v", err)
	}
	if err := res.Valid(); err != nil {
		return nil, fmt.Errorf("validating Result: %v", err)
	}
	// store the Log text if supplied
	if len(res.Log) > 0 {
		hash, err := PutLog(ctx, res.Log)
		if err != nil {
			return nil, fmt.Errorf("putting Log: %v", err)
		}
		res.LogHash = hash
	}
	tx := func(tx *datastore.Transaction) error {
		if _, err := getOrMakePackageInTx(ctx, tx, res.PackagePath); err != nil {
			return fmt.Errorf("GetPackage: %v", err)
		}
		// put Result
		if _, err := tx.Put(res.Key(), res); err != nil {
			return fmt.Errorf("putting Result: %v", err)
		}
		// add Result to Commit
		com := &Commit{PackagePath: res.PackagePath, Hash: res.Hash}
		if err := com.AddResult(tx, res); err != nil {
			return fmt.Errorf("AddResult: %v", err)
		}
		return nil
	}
	_, err := datastoreClient.RunInTransaction(ctx, tx)
	return nil, err
}

// logHandler displays log text for a given hash.
// It handles paths like "/log/hash".
func logHandler(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-type", "text/plain; charset=utf-8")
	c := r.Context()
	hash := r.URL.Path[strings.LastIndex(r.URL.Path, "/")+1:]
	key := dsKey("Log", hash, nil)
	l := new(Log)
	if err := datastoreClient.Get(c, key, l); err != nil {
		if err == datastore.ErrNoSuchEntity {
			// Fall back to default namespace;
			// maybe this was on the old dashboard.
			key.Namespace = ""
			err = datastoreClient.Get(c, key, l)
		}
		if err != nil {
			logErr(w, r, err)
			return
		}
	}
	b, err := l.Text()
	if err != nil {
		logErr(w, r, err)
		return
	}
	w.Write(b)
}

// clearResultsHandler purge a single build failure from the dashboard.
// It currently only supports the main Go repo.
func clearResultsHandler(r *http.Request) (interface{}, error) {
	if r.Method != "POST" {
		return nil, errBadMethod(r.Method)
	}
	builder := r.FormValue("builder")
	hash := r.FormValue("hash")
	if builder == "" {
		return nil, errors.New("missing 'builder'")
	}
	if hash == "" {
		return nil, errors.New("missing 'hash'")
	}

	ctx := r.Context()

	_, err := datastoreClient.RunInTransaction(ctx, func(tx *datastore.Transaction) error {
		c := &Commit{
			PackagePath: "", // TODO(adg): support clearing sub-repos
			Hash:        hash,
		}
		err := tx.Get(c.Key(), c)
		err = filterDatastoreError(err)
		if err == datastore.ErrNoSuchEntity {
			// Doesn't exist, so no build to clear.
			return nil
		}
		if err != nil {
			return err
		}

		r := c.Result(builder, "")
		if r == nil {
			// No result, so nothing to clear.
			return nil
		}
		c.RemoveResult(r)
		_, err = tx.Put(c.Key(), c)
		if err != nil {
			return err
		}
		return tx.Delete(r.Key())
	})
	return nil, err
}

type dashHandler func(*http.Request) (interface{}, error)

type dashResponse struct {
	Response interface{}
	Error    string
}

// errBadMethod is returned by a dashHandler when
// the request has an unsuitable method.
type errBadMethod string

func (e errBadMethod) Error() string {
	return "bad method: " + string(e)
}

func builderKeyRevoked(builder string) bool {
	switch builder {
	case "plan9-amd64-mischief":
		// Broken and unmaintained for months.
		// It's polluting the dashboard.
		return true
	case "linux-arm-onlinenet":
		// Requested to be revoked by Dave Cheney.
		// The machine is in a fail+report loop
		// and can't be accessed. Revoke it for now.
		return true
	}
	return false
}

// AuthHandler wraps a http.HandlerFunc with a handler that validates the
// supplied key and builder query parameters.
func AuthHandler(h dashHandler) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		c := r.Context()

		// Put the URL Query values into r.Form to avoid parsing the
		// request body when calling r.FormValue.
		r.Form = r.URL.Query()

		var err error
		var resp interface{}

		// Validate key query parameter for POST requests only.
		key := r.FormValue("key")
		builder := r.FormValue("builder")
		if r.Method == "POST" && !validKey(c, key, builder) {
			err = fmt.Errorf("invalid key %q for builder %q", key, builder)
		}

		// Call the original HandlerFunc and return the response.
		if err == nil {
			resp, err = h(r)
		}

		// Write JSON response.
		dashResp := &dashResponse{Response: resp}
		if err != nil {
			log.Printf("%v", err)
			dashResp.Error = err.Error()
		}
		w.Header().Set("Content-Type", "application/json")
		if err = json.NewEncoder(w).Encode(dashResp); err != nil {
			log.Printf("encoding response: %v", err)
		}
	}
}

// validHash reports whether hash looks like a valid git commit hash.
func validHash(hash string) bool {
	// TODO: correctly validate a hash: check that it's exactly 40
	// lowercase hex digits. But this is what we historically did:
	return hash != ""
}

func validKey(c context.Context, key, builder string) bool {
	if isMasterKey(c, key) {
		return true
	}
	if builderKeyRevoked(builder) {
		return false
	}
	return key == builderKey(c, builder)
}

func isMasterKey(c context.Context, k string) bool {
	return k == key.Secret(datastoreClient, c)
}

func builderKey(c context.Context, builder string) string {
	h := hmac.New(md5.New, []byte(key.Secret(datastoreClient, c)))
	h.Write([]byte(builder))
	return fmt.Sprintf("%x", h.Sum(nil))
}

func logErr(w http.ResponseWriter, r *http.Request, err error) {
	log.Printf("Error: %v", err)
	w.WriteHeader(http.StatusInternalServerError)
	fmt.Fprint(w, "Error: ", html.EscapeString(err.Error()))
}

// limitStringLength essentially does return s[:max],
// but it ensures that we dot not split UTF-8 rune in half.
// Otherwise appengine python scripts will break badly.
func limitStringLength(s string, max int) string {
	if len(s) <= max {
		return s
	}
	for {
		s = s[:max]
		r, size := utf8.DecodeLastRuneInString(s)
		if r != utf8.RuneError || size != 1 {
			return s
		}
		max--
	}
}
<|endoftext|>


<commit_before>// Package cq provides tools for interacting with the CQ tools.
package cq

import (
	"bytes"
	"fmt"
	"sync"
	"time"

	"github.com/golang/protobuf/proto"
	"go.skia.org/infra/go/sklog"

	"go.skia.org/infra/go/buildbucket"
	"go.skia.org/infra/go/common"
	"go.skia.org/infra/go/gerrit"
	"go.skia.org/infra/go/gitiles"
	"go.skia.org/infra/go/metrics2"
	"go.skia.org/infra/go/util"
)

const (
	CQ_CFG_FILE_PATH = "infra/branch-config/cq.cfg"

	// Constants for in-flight metrics.
	INFLIGHT_METRIC_NAME     = "in_flight"
	INFLIGHT_TRYBOT_DURATION = "trybot_duration"
	INFLIGHT_TRYBOT_NUM      = "trybot_num"
	INFLIGHT_WAITING_IN_CQ   = "waiting_in_cq"

	// Constants for landed metrics.
	LANDED_METRIC_NAME     = "after_commit"
	LANDED_TRYBOT_DURATION = "trybot_duration"
	LANDED_TOTAL_DURATION  = "total_duration"

	// Thresholds after which errors are logged.
	CQ_TRYBOT_DURATION_SECS_THRESHOLD = 2700
	CQ_TRYBOTS_COUNT_THRESHOLD        = 35
)

var (
	// Slice of all known presubmit bot names.
	PRESUBMIT_BOTS = []string{"skia_presubmit-Trybot"}

	// Mutext to control access to the slice of CQ trybots.
	cqTryBotsMutex sync.RWMutex
)

// NewClient creates a new client for interacting with CQ tools.
func NewClient(gerritClient *gerrit.Gerrit, cqTryBotsFunc GetCQTryBots, metricName string) (*Client, error) {
	cqTryBots, err := cqTryBotsFunc()
	if err != nil {
		return nil, err
	}
	return &Client{gerritClient, cqTryBots, cqTryBotsFunc, metricName}, err
}

// GetCQTryBots is an interface for returing the CQ trybots of a project.
type GetCQTryBots func() ([]string, error)

type Client struct {
	gerritClient  *gerrit.Gerrit
	cqTryBots     []string
	cqTryBotsFunc GetCQTryBots
	metricName    string
}

// GetSkiaCQTryBots is a Skia implementation of GetCQTryBots.
func GetSkiaCQTryBots() ([]string, error) {
	return getCQTryBots(common.REPO_SKIA)
}

// GetSkiaInfraCQTryBots is a Skia Infra implementation of GetCQTryBots.
func GetSkiaInfraCQTryBots() ([]string, error) {
	return getCQTryBots(common.REPO_SKIA)
}

// getCQTryBots is a convenience method for the Skia and Skia Infra CQ TryBots.
func getCQTryBots(repo string) ([]string, error) {
	var buf bytes.Buffer
	if err := gitiles.NewRepo(repo).ReadFile(CQ_CFG_FILE_PATH, &buf); err != nil {
		return nil, err
	}
	var cqCfg Config
	if err := proto.UnmarshalText(buf.String(), &cqCfg); err != nil {
		return nil, err
	}
	tryJobs := []string{}
	for _, bucket := range cqCfg.Verifiers.GetTryJob().GetBuckets() {
		for _, builder := range bucket.GetBuilders() {
			if builder.GetExperimentPercentage() > 0 && builder.GetExperimentPercentage() < 100 {
				// Exclude experimental builders, unless running for all CLs.
				continue
			}
			if util.ContainsAny(builder.GetName(), PRESUBMIT_BOTS) {
				// Exclude presubmit bots because they could fail or be delayed
				// due to factors such as owners approval and other project
				// specific checks.
				continue
			}
			tryJobs = append(tryJobs, builder.GetName())
		}
	}
	sklog.Infof("The list of CQ trybots is: %s", tryJobs)
	return tryJobs, nil
}

// RefreshCQTryBots refreshes the slice of CQ trybots on the instance. Access
// to the trybots is protected by a RWMutex.
func (c *Client) RefreshCQTryBots() error {
	tryBots, err := c.cqTryBotsFunc()
	if err != nil {
		return err
	}

	cqTryBotsMutex.Lock()
	defer cqTryBotsMutex.Unlock()
	c.cqTryBots = tryBots
	return nil
}

// ReportCQStats reports all relevant stats for the specified Gerrit change.
// Note: Different stats are reported depending on whether the change has been
// merged or not.
func (c *Client) ReportCQStats(change int64) error {
	changeInfo, err := c.gerritClient.GetIssueProperties(change)
	if err != nil {
		return err
	}
	patchsetIds := changeInfo.GetPatchsetIDs()
	latestPatchsetId := patchsetIds[len(patchsetIds)-1]
	if changeInfo.Committed {
		// TODO(rmistry): The last patchset in Gerrit does not contain trybot
		// information so we have to look at the one immediately before it.
		// This will be fixed with crbug.com/634944.
		latestPatchsetId = patchsetIds[len(patchsetIds)-2]
	}

	builds, err := c.gerritClient.GetTrybotResults(change, latestPatchsetId)
	if err != nil {
		return err
	}
	// Consider only CQ bots.
	cqBuilds := []*buildbucket.Build{}
	for _, b := range builds {
		if c.isCQTryBot(b.Parameters.BuilderName) {
			cqBuilds = append(cqBuilds, b)
		}
	}
	gerritURL := fmt.Sprintf("%s/c/%d/%d", gerrit.GERRIT_SKIA_URL, change, latestPatchsetId)
	if len(cqBuilds) == 0 {
		sklog.Infof("No trybot results were found for %s", gerritURL)
		return nil
	}

	sklog.Infof("Starting processing %s. Merged status: %t", gerritURL, changeInfo.Committed)

	if changeInfo.Committed {
		c.ReportCQStatsForLandedCL(cqBuilds, gerritURL)
	} else {
		c.ReportCQStatsForInFlightCL(cqBuilds, gerritURL)
	}
	return nil
}

// ReportCQStatsForLandedCL reports the following metrics for the specified
// change and patchsetID:
// * The total time the change spent waiting for CQ trybots to complete.
// * The time each CQ trybot took to complete.
func (c *Client) ReportCQStatsForLandedCL(cqBuilds []*buildbucket.Build, gerritURL string) {
	endTimeOfCQBots := time.Time{}
	maximumTrybotDuration := int64(0)
	for _, b := range cqBuilds {
		createdTime := time.Time(b.Created).UTC()
		completedTime := time.Time(b.Completed).UTC()
		if (completedTime == time.Time{}.UTC()) {
			sklog.Warningf("Skipping %s on %s. The correct completed time has not shown up in Buildbucket yet.", b.Parameters.BuilderName, gerritURL)
			continue
		}
		if endTimeOfCQBots.Before(completedTime) {
			endTimeOfCQBots = completedTime
		}

		durationTags := map[string]string{
			"trybot": b.Parameters.BuilderName,
		}
		duration := int64(completedTime.Sub(createdTime).Seconds())
		sklog.Infof("%s was created at %s by %s and completed at %s. Total duration: %d", b.Parameters.BuilderName, createdTime, gerritURL, completedTime, duration)
		landedTrybotDurationMetric := metrics2.GetInt64Metric(fmt.Sprintf("%s_%s_%s", c.metricName, LANDED_METRIC_NAME, LANDED_TRYBOT_DURATION), durationTags)
		landedTrybotDurationMetric.Update(duration)

		if duration > maximumTrybotDuration {
			maximumTrybotDuration = duration
		}
	}

	sklog.Infof("Maximum trybot duration for %s: %d", gerritURL, maximumTrybotDuration)
	sklog.Infof("Furthest completion time for %s: %s", gerritURL, endTimeOfCQBots)
	landedTotalDurationMetric := metrics2.GetInt64Metric(fmt.Sprintf("%s_%s_%s", c.metricName, LANDED_METRIC_NAME, LANDED_TOTAL_DURATION), map[string]string{})
	landedTotalDurationMetric.Update(maximumTrybotDuration)
}

// ReportCQStatsForInFlightCL reports the following metrics for the specified
// change and patchsetID:
// * How long CQ trybots have been running for.
// * How many CQ trybots have been triggered.
func (c *Client) ReportCQStatsForInFlightCL(cqBuilds []*buildbucket.Build, gerritURL string) {
	totalTriggeredCQBots := int(0)
	currentTime := time.Now()
	for _, b := range cqBuilds {
		totalTriggeredCQBots++

		createdTime := time.Time(b.Created).UTC()
		completedTime := time.Time(b.Completed).UTC()
		if (completedTime != time.Time{}.UTC()) {
			if time.Hour*24 < time.Now().UTC().Sub(createdTime) {
				// The build has completed more than a day ago. Do not include it
				// in totalTriggeredCQBots. See skbug.com/7340.
				totalTriggeredCQBots--
			}
			// The build has completed so move on.
			continue
		}

		duration := int64(currentTime.Sub(createdTime).Seconds())
		durationTags := map[string]string{
			"trybot": b.Parameters.BuilderName,
		}
		if duration > CQ_TRYBOT_DURATION_SECS_THRESHOLD {
			sklog.Errorf("CQTrybotDurationError: %s was triggered by %s and is still running after %d seconds. Threshold is %d seconds.", b.Parameters.BuilderName, gerritURL, duration, CQ_TRYBOT_DURATION_SECS_THRESHOLD)
		}
		inflightTrybotDurationMetric := metrics2.GetInt64Metric(fmt.Sprintf("%s_%s_%s", c.metricName, INFLIGHT_METRIC_NAME, INFLIGHT_TRYBOT_DURATION), durationTags)
		inflightTrybotDurationMetric.Update(duration)

	}
	cqTryBotsMutex.RLock()
	cqTryBotsMutex.RUnlock()
	if totalTriggeredCQBots > CQ_TRYBOTS_COUNT_THRESHOLD {
		sklog.Errorf("CQCLsCountError: %d trybots have been triggered by %s. Threshold is %d trybots.", totalTriggeredCQBots, gerritURL, CQ_TRYBOTS_COUNT_THRESHOLD)
	}
	trybotNumDurationMetric := metrics2.GetInt64Metric(fmt.Sprintf("%s_%s_%s", c.metricName, INFLIGHT_METRIC_NAME, INFLIGHT_TRYBOT_NUM), map[string]string{})
	trybotNumDurationMetric.Update(int64(totalTriggeredCQBots))
}

func (c *Client) isCQTryBot(builderName string) bool {
	cqTryBotsMutex.RLock()
	isCQTrybot := util.In(builderName, c.cqTryBots)
	cqTryBotsMutex.RUnlock()
	return isCQTrybot
}
<commit_msg>Purge metrics of trybots no longer in cq.cfg<commit_after>// Package cq provides tools for interacting with the CQ tools.
package cq

import (
	"bytes"
	"fmt"
	"sync"
	"time"

	"github.com/golang/protobuf/proto"
	"go.skia.org/infra/go/sklog"

	"go.skia.org/infra/go/buildbucket"
	"go.skia.org/infra/go/common"
	"go.skia.org/infra/go/gerrit"
	"go.skia.org/infra/go/gitiles"
	"go.skia.org/infra/go/metrics2"
	"go.skia.org/infra/go/util"
)

const (
	CQ_CFG_FILE_PATH = "infra/branch-config/cq.cfg"

	// Constants for in-flight metrics.
	INFLIGHT_METRIC_NAME     = "in_flight"
	INFLIGHT_TRYBOT_DURATION = "trybot_duration"
	INFLIGHT_TRYBOT_NUM      = "trybot_num"
	INFLIGHT_WAITING_IN_CQ   = "waiting_in_cq"

	// Constants for landed metrics.
	LANDED_METRIC_NAME     = "after_commit"
	LANDED_TRYBOT_DURATION = "trybot_duration"
	LANDED_TOTAL_DURATION  = "total_duration"

	// Thresholds after which errors are logged.
	CQ_TRYBOT_DURATION_SECS_THRESHOLD = 2700
	CQ_TRYBOTS_COUNT_THRESHOLD        = 35
)

var (
	// Slice of all known presubmit bot names.
	PRESUBMIT_BOTS = []string{"skia_presubmit-Trybot"}

	// Mutext to control access to the slice of CQ trybots.
	cqTryBotsMutex sync.RWMutex
)

// NewClient creates a new client for interacting with CQ tools.
func NewClient(gerritClient *gerrit.Gerrit, cqTryBotsFunc GetCQTryBots, metricName string) (*Client, error) {
	cqTryBots, err := cqTryBotsFunc()
	if err != nil {
		return nil, err
	}
	return &Client{gerritClient, util.NewStringSet(cqTryBots), cqTryBotsFunc, metricName}, err
}

// GetCQTryBots is an interface for returing the CQ trybots of a project.
type GetCQTryBots func() ([]string, error)

type Client struct {
	gerritClient  *gerrit.Gerrit
	cqTryBots     util.StringSet
	cqTryBotsFunc GetCQTryBots
	metricName    string
}

// GetSkiaCQTryBots is a Skia implementation of GetCQTryBots.
func GetSkiaCQTryBots() ([]string, error) {
	return getCQTryBots(common.REPO_SKIA)
}

// GetSkiaInfraCQTryBots is a Skia Infra implementation of GetCQTryBots.
func GetSkiaInfraCQTryBots() ([]string, error) {
	return getCQTryBots(common.REPO_SKIA)
}

// getCQTryBots is a convenience method for the Skia and Skia Infra CQ TryBots.
func getCQTryBots(repo string) ([]string, error) {
	var buf bytes.Buffer
	if err := gitiles.NewRepo(repo).ReadFile(CQ_CFG_FILE_PATH, &buf); err != nil {
		return nil, err
	}
	var cqCfg Config
	if err := proto.UnmarshalText(buf.String(), &cqCfg); err != nil {
		return nil, err
	}
	tryJobs := []string{}
	for _, bucket := range cqCfg.Verifiers.GetTryJob().GetBuckets() {
		for _, builder := range bucket.GetBuilders() {
			if builder.GetExperimentPercentage() > 0 && builder.GetExperimentPercentage() < 100 {
				// Exclude experimental builders, unless running for all CLs.
				continue
			}
			if util.ContainsAny(builder.GetName(), PRESUBMIT_BOTS) {
				// Exclude presubmit bots because they could fail or be delayed
				// due to factors such as owners approval and other project
				// specific checks.
				continue
			}
			tryJobs = append(tryJobs, builder.GetName())
		}
	}
	sklog.Infof("The list of CQ trybots is: %s", tryJobs)
	return tryJobs, nil
}

// RefreshCQTryBots refreshes the slice of CQ trybots on the instance. Access
// to the trybots is protected by a RWMutex.
// Trybots that no longer exist in the new list will have their metrics purged.
func (c *Client) RefreshCQTryBots() error {
	tryBots, err := c.cqTryBotsFunc()
	if err != nil {
		return err
	}
	tryBotsSet := util.NewStringSet(tryBots)

	cqTryBotsMutex.Lock()
	defer cqTryBotsMutex.Unlock()
	// Gather all trybots that no longer exist and purge their metrics.
	deletedTryBots := c.cqTryBots.Complement(tryBotsSet)
	if err := c.purgeMetrics(deletedTryBots); err != nil {
		return fmt.Errorf("Could not purge metrics of %s: %s", deletedTryBots, err)
	}

	c.cqTryBots = tryBotsSet
	return nil
}

// ReportCQStats reports all relevant stats for the specified Gerrit change.
// Note: Different stats are reported depending on whether the change has been
// merged or not.
func (c *Client) ReportCQStats(change int64) error {
	changeInfo, err := c.gerritClient.GetIssueProperties(change)
	if err != nil {
		return err
	}
	patchsetIds := changeInfo.GetPatchsetIDs()
	latestPatchsetId := patchsetIds[len(patchsetIds)-1]
	if changeInfo.Committed {
		// TODO(rmistry): The last patchset in Gerrit does not contain trybot
		// information so we have to look at the one immediately before it.
		// This will be fixed with crbug.com/634944.
		latestPatchsetId = patchsetIds[len(patchsetIds)-2]
	}

	builds, err := c.gerritClient.GetTrybotResults(change, latestPatchsetId)
	if err != nil {
		return err
	}
	// Consider only CQ bots.
	cqBuilds := []*buildbucket.Build{}
	for _, b := range builds {
		if c.isCQTryBot(b.Parameters.BuilderName) {
			cqBuilds = append(cqBuilds, b)
		}
	}
	gerritURL := fmt.Sprintf("%s/c/%d/%d", gerrit.GERRIT_SKIA_URL, change, latestPatchsetId)
	if len(cqBuilds) == 0 {
		sklog.Infof("No trybot results were found for %s", gerritURL)
		return nil
	}

	sklog.Infof("Starting processing %s. Merged status: %t", gerritURL, changeInfo.Committed)

	if changeInfo.Committed {
		c.ReportCQStatsForLandedCL(cqBuilds, gerritURL)
	} else {
		c.ReportCQStatsForInFlightCL(cqBuilds, gerritURL)
	}
	return nil
}

// ReportCQStatsForLandedCL reports the following metrics for the specified
// change and patchsetID:
// * The total time the change spent waiting for CQ trybots to complete.
// * The time each CQ trybot took to complete.
func (c *Client) ReportCQStatsForLandedCL(cqBuilds []*buildbucket.Build, gerritURL string) {
	endTimeOfCQBots := time.Time{}
	maximumTrybotDuration := int64(0)
	for _, b := range cqBuilds {
		createdTime := time.Time(b.Created).UTC()
		completedTime := time.Time(b.Completed).UTC()
		if (completedTime == time.Time{}.UTC()) {
			sklog.Warningf("Skipping %s on %s. The correct completed time has not shown up in Buildbucket yet.", b.Parameters.BuilderName, gerritURL)
			continue
		}
		if endTimeOfCQBots.Before(completedTime) {
			endTimeOfCQBots = completedTime
		}

		duration := int64(completedTime.Sub(createdTime).Seconds())
		sklog.Infof("%s was created at %s by %s and completed at %s. Total duration: %d", b.Parameters.BuilderName, createdTime, gerritURL, completedTime, duration)
		landedTrybotDurationMetric := c.getLandedTrybotDurationMetric(b.Parameters.BuilderName)
		landedTrybotDurationMetric.Update(duration)

		if duration > maximumTrybotDuration {
			maximumTrybotDuration = duration
		}
	}

	sklog.Infof("Maximum trybot duration for %s: %d", gerritURL, maximumTrybotDuration)
	sklog.Infof("Furthest completion time for %s: %s", gerritURL, endTimeOfCQBots)
	landedTotalDurationMetric := metrics2.GetInt64Metric(fmt.Sprintf("%s_%s_%s", c.metricName, LANDED_METRIC_NAME, LANDED_TOTAL_DURATION), map[string]string{})
	landedTotalDurationMetric.Update(maximumTrybotDuration)
}

// ReportCQStatsForInFlightCL reports the following metrics for the specified
// change and patchsetID:
// * How long CQ trybots have been running for.
// * How many CQ trybots have been triggered.
func (c *Client) ReportCQStatsForInFlightCL(cqBuilds []*buildbucket.Build, gerritURL string) {
	totalTriggeredCQBots := int(0)
	currentTime := time.Now()
	for _, b := range cqBuilds {
		totalTriggeredCQBots++

		createdTime := time.Time(b.Created).UTC()
		completedTime := time.Time(b.Completed).UTC()
		if (completedTime != time.Time{}.UTC()) {
			if time.Hour*24 < time.Now().UTC().Sub(createdTime) {
				// The build has completed more than a day ago. Do not include it
				// in totalTriggeredCQBots. See skbug.com/7340.
				totalTriggeredCQBots--
			}
			// The build has completed so move on.
			continue
		}

		duration := int64(currentTime.Sub(createdTime).Seconds())
		if duration > CQ_TRYBOT_DURATION_SECS_THRESHOLD {
			sklog.Errorf("CQTrybotDurationError: %s was triggered by %s and is still running after %d seconds. Threshold is %d seconds.", b.Parameters.BuilderName, gerritURL, duration, CQ_TRYBOT_DURATION_SECS_THRESHOLD)
		}
		inflightTrybotDurationMetric := c.getInflightTrybotDurationMetric(b.Parameters.BuilderName)
		inflightTrybotDurationMetric.Update(duration)
	}

	cqTryBotsMutex.RLock()
	cqTryBotsMutex.RUnlock()
	if totalTriggeredCQBots > CQ_TRYBOTS_COUNT_THRESHOLD {
		sklog.Errorf("CQCLsCountError: %d trybots have been triggered by %s. Threshold is %d trybots.", totalTriggeredCQBots, gerritURL, CQ_TRYBOTS_COUNT_THRESHOLD)
	}
	trybotNumDurationMetric := metrics2.GetInt64Metric(fmt.Sprintf("%s_%s_%s", c.metricName, INFLIGHT_METRIC_NAME, INFLIGHT_TRYBOT_NUM), map[string]string{})
	trybotNumDurationMetric.Update(int64(totalTriggeredCQBots))
}

func (c *Client) purgeMetrics(tryBots util.StringSet) error {
	for _, b := range tryBots.Keys() {
		inflightTrybotDurationMetric := c.getInflightTrybotDurationMetric(b)
		if err := inflightTrybotDurationMetric.Delete(); err != nil {
			return fmt.Errorf("Could not delete inflight trybot metric: %s", err)
		}
		landedTrybotDurationMetric := c.getLandedTrybotDurationMetric(b)
		if err := landedTrybotDurationMetric.Delete(); err != nil {
			return fmt.Errorf("Could not delete landed trybot metric: %s", err)
		}
		sklog.Infof("Deleted inflight and landed metrics of %s", b)
	}
	return nil
}

func (c *Client) getInflightTrybotDurationMetric(tryBot string) metrics2.Int64Metric {
	metricName := fmt.Sprintf("%s_%s_%s", c.metricName, INFLIGHT_METRIC_NAME, INFLIGHT_TRYBOT_DURATION)
	tags := map[string]string{
		"trybot": tryBot,
	}
	return metrics2.GetInt64Metric(metricName, tags)
}

func (c *Client) getLandedTrybotDurationMetric(tryBot string) metrics2.Int64Metric {
	metricName := fmt.Sprintf("%s_%s_%s", c.metricName, LANDED_METRIC_NAME, LANDED_TRYBOT_DURATION)
	tags := map[string]string{
		"trybot": tryBot,
	}
	return metrics2.GetInt64Metric(metricName, tags)
}

func (c *Client) isCQTryBot(builderName string) bool {
	cqTryBotsMutex.RLock()
	isCQTrybot := c.cqTryBots[builderName]
	cqTryBotsMutex.RUnlock()
	return isCQTrybot
}
<|endoftext|>


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

// Package router provides interfaces that need to be satisfied in order to
// implement a new router on tsuru.
package router

import (
	"fmt"
	"github.com/globocom/tsuru/db"
	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
)

var routers = make(map[string]Router)

// Register registers a new router.
func Register(name string, r Router) {
	routers[name] = r
}

// Get gets the named router from the registry.
func Get(name string) (Router, error) {
	r, ok := routers[name]
	if !ok {
		return nil, fmt.Errorf("Unknown router: %q.", name)
	}
	return r, nil
}

// Router is the basic interface of this package. It provides methods for
// managing backends and routes. Each backend can have multiple routes.
type Router interface {
	AddBackend(name string) error
	RemoveBackend(name string) error
	AddRoute(name, address string) error
	RemoveRoute(name, address string) error
	SetCName(cname, name string) error
	UnsetCName(cname, name string) error
	Addr(name string) (string, error)

	// Swap change the router between two backends.
	Swap(string, string) error

	// Routes returns a list of routes of a backend.
	Routes(name string) ([]string, error)
}

func collection() (*mgo.Collection, error) {
	conn, err := db.Conn()
	if err != nil {
		return nil, err
	}
	return conn.Collection("routers"), nil
}

// Store stores the app name related with the
// router name.
func Store(appName, routerName string) error {
	coll, err := collection()
	if err != nil {
		return err
	}
	data := map[string]string{
		"app":    appName,
		"router": routerName,
	}
	return coll.Insert(&data)
}

func Retrieve(appName string) (string, error) {
	coll, err := collection()
	if err != nil {
		return "", err
	}
	data := map[string]string{}
	err = coll.Find(bson.M{"app": appName}).One(&data)
	if err != nil {
		return "", err
	}
	return data["router"], nil
}

func Remove(appName string) error {
	coll, err := collection()
	if err != nil {
		return err
	}
	return coll.Remove(bson.M{"app": appName})
}

func swapBackendName(backend1, backend2 string) error {
	coll, err := collection()
	if err != nil {
		return err
	}
	router1, err := Retrieve(backend1)
	if err != nil {
		return err
	}
	router2, err := Retrieve(backend2)
	if err != nil {
		return err
	}
	update := bson.M{"$set": bson.M{"router": router2}}
	err = coll.Update(bson.M{"app": backend1}, update)
	if err != nil {
		return err
	}
	update = bson.M{"$set": bson.M{"router": router1}}
	err = coll.Update(bson.M{"app": backend2}, update)
	var result []interface{}
	coll.Find(nil).All(&result)
	return err
}

func Swap(r Router, backend1, backend2 string) error {
	routes1, err := r.Routes(backend1)
	if err != nil {
		return err
	}
	routes2, err := r.Routes(backend2)
	if err != nil {
		return err
	}
	for _, route := range routes1 {
		err = r.AddRoute(backend2, route)
		if err != nil {
			return err
		}
		err = r.RemoveRoute(backend1, route)
		if err != nil {
			return err
		}
	}
	for _, route := range routes2 {
		err = r.AddRoute(backend1, route)
		if err != nil {
			return err
		}
		err = r.RemoveRoute(backend2, route)
		if err != nil {
			return err
		}
	}
	return swapBackendName(backend1, backend2)
}
<commit_msg>router: use db.Collection instead of mgo.Collection<commit_after>// Copyright 2013 tsuru authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package router provides interfaces that need to be satisfied in order to
// implement a new router on tsuru.
package router

import (
	"fmt"
	"github.com/globocom/tsuru/db"
	"labix.org/v2/mgo/bson"
)

var routers = make(map[string]Router)

// Register registers a new router.
func Register(name string, r Router) {
	routers[name] = r
}

// Get gets the named router from the registry.
func Get(name string) (Router, error) {
	r, ok := routers[name]
	if !ok {
		return nil, fmt.Errorf("Unknown router: %q.", name)
	}
	return r, nil
}

// Router is the basic interface of this package. It provides methods for
// managing backends and routes. Each backend can have multiple routes.
type Router interface {
	AddBackend(name string) error
	RemoveBackend(name string) error
	AddRoute(name, address string) error
	RemoveRoute(name, address string) error
	SetCName(cname, name string) error
	UnsetCName(cname, name string) error
	Addr(name string) (string, error)

	// Swap change the router between two backends.
	Swap(string, string) error

	// Routes returns a list of routes of a backend.
	Routes(name string) ([]string, error)
}

func collection() (*db.Collection, error) {
	conn, err := db.Conn()
	if err != nil {
		return nil, err
	}
	return conn.Collection("routers"), nil
}

// Store stores the app name related with the
// router name.
func Store(appName, routerName string) error {
	coll, err := collection()
	if err != nil {
		return err
	}
	data := map[string]string{
		"app":    appName,
		"router": routerName,
	}
	return coll.Insert(&data)
}

func Retrieve(appName string) (string, error) {
	coll, err := collection()
	if err != nil {
		return "", err
	}
	data := map[string]string{}
	err = coll.Find(bson.M{"app": appName}).One(&data)
	if err != nil {
		return "", err
	}
	return data["router"], nil
}

func Remove(appName string) error {
	coll, err := collection()
	if err != nil {
		return err
	}
	return coll.Remove(bson.M{"app": appName})
}

func swapBackendName(backend1, backend2 string) error {
	coll, err := collection()
	if err != nil {
		return err
	}
	router1, err := Retrieve(backend1)
	if err != nil {
		return err
	}
	router2, err := Retrieve(backend2)
	if err != nil {
		return err
	}
	update := bson.M{"$set": bson.M{"router": router2}}
	err = coll.Update(bson.M{"app": backend1}, update)
	if err != nil {
		return err
	}
	update = bson.M{"$set": bson.M{"router": router1}}
	err = coll.Update(bson.M{"app": backend2}, update)
	var result []interface{}
	coll.Find(nil).All(&result)
	return err
}

func Swap(r Router, backend1, backend2 string) error {
	routes1, err := r.Routes(backend1)
	if err != nil {
		return err
	}
	routes2, err := r.Routes(backend2)
	if err != nil {
		return err
	}
	for _, route := range routes1 {
		err = r.AddRoute(backend2, route)
		if err != nil {
			return err
		}
		err = r.RemoveRoute(backend1, route)
		if err != nil {
			return err
		}
	}
	for _, route := range routes2 {
		err = r.AddRoute(backend1, route)
		if err != nil {
			return err
		}
		err = r.RemoveRoute(backend2, route)
		if err != nil {
			return err
		}
	}
	return swapBackendName(backend1, backend2)
}
<|endoftext|>


<commit_before>package main

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

const (
	startPrefix = "=== RUN "
	passPrefix = "--- PASS: "
	failPrefix = "--- FAIL: "
)

var endRegexp *regexp.Regexp = regexp.MustCompile(`([^ ]+) \((\d+\.\d+)`)

type Test struct {
	Name, Time, Message string
	Failed bool
}

func parseEnd(prefix, line string) (string, string, error) {
	matches := endRegexp.FindStringSubmatch(line[len(prefix):])

	if len(matches) == 0 {
		return "", "", fmt.Errorf("can't parse %s", line)
	}

	return matches[1], matches[2], nil
}

func parseOutput(rd io.Reader) ([]*Test, error) {
	tests := []*Test{}

	reader := bufio.NewReader(rd)
	var test *Test = nil
	for {
		/* FIXME: Handle isPrefix */
		buf, _, err := reader.ReadLine()

		switch err {
		case io.EOF:
			if test != nil {
				tests = append(tests, test)
			}
			return tests, nil
		case nil:
			;
		default:
			return nil, err
		}

		line := string(buf)
		switch {
		case strings.HasPrefix(line, startPrefix):
			if test != nil {
				tests = append(tests, test)
			}
			test = &Test{Name: line[len(startPrefix):]}
		case strings.HasPrefix(line, failPrefix):
			if test == nil {
				return nil, fmt.Errorf("fail not inside test")
			}
			test.Failed = true
			name, time, err := parseEnd(failPrefix, line)
			if err != nil {
				return nil, err
			}
			if name != test.Name {
				return nil, fmt.Errorf("wrong test end (%s!=%s)", name, test.Name)
			}
			test.Time = time
		case strings.HasPrefix(line, passPrefix):
			if test == nil {
				return nil, fmt.Errorf("pass not inside test")
			}
			test.Failed = false
			name, time, err := parseEnd(passPrefix, line)
			if err != nil {
				return nil, err
			}
			if name != test.Name {
				return nil, fmt.Errorf("wrong test end (%s!=%s)", name, test.Name)
			}
			test.Time = time
		default:
			if test != nil {
				test.Message += line + "\n"
			}
		}
	}

	return tests, nil
}

func getInput(filename string) (io.Reader, error) {
	if filename == "-" || filename == "" {
		return os.Stdin, nil
	}

	return os.Open(filename)
}

func numFailures(tests []*Test) int {
	count := 0
	for _, test := range tests {
		if test.Failed {
			count ++
		}
	}

	return count
}

func writeXML(tests []*Test, out io.Writer) {
	newline := func() { fmt.Fprintln(out) }

	fmt.Fprintf(out, `<?xml version="1.0" encoding="utf-8"?>`)
	newline()
	fmt.Fprintf(out, `<testsuite name="go2xunit" tests="%d" errors="0" failures="%d" skip="0">`,
					 len(tests), numFailures(tests))
	newline()
	for _, test := range(tests) {
		fmt.Fprintf(out, `  <testcase classname="go2xunit" name="%s" time="%s"`,
		                 test.Name, test.Time)
		if !test.Failed {
			fmt.Fprintf(out, " />\n")
			continue
		}
		fmt.Fprintln(out, ">")
		fmt.Fprintf(out,  `    <failure type="go.error" message="error">`)
		newline()
		fmt.Fprintf(out, "<![CDATA[%s]]>\n", test.Message)
		fmt.Fprintln(out, "    </failure>")
		fmt.Fprintln(out, "  </testcase>")
	}
	fmt.Fprintln(out, "</testsuite>")
}

func getOutput(filename string) (io.Writer, error) {
	if filename == "-" || filename == "" {
		return os.Stdout, nil
	}

	return os.Create(filename)
}

func getIO(inputFile, outputFile string) (io.Reader, io.Writer, error) {
	input, err := getInput(inputFile)
	if err != nil {
		return nil, nil, fmt.Errorf("can't open %s for reading: %s", inputFile, err)
	}

	output, err := getOutput(outputFile)
	if err != nil {
		return nil, nil, fmt.Errorf("can't open %s for writing: %s", outputFile, err)
	}

	return input, output, nil
}

func main() {
	inputFile := flag.String("input", "", "input file (default to stdin)")
	outputFile := flag.String("output", "", "output file (default to stdout)")
	fail := flag.Bool("fail", false, "fail (non zero exit) if any test failed")
	flag.Parse()

	log.SetFlags(0)

	if flag.NArg() > 0 {
		log.Fatalf("error: %s does not take parameters (did you mean -input?)", os.Args[0])
	}

	input, output, err := getIO(*inputFile, *outputFile)
	if err != nil {
		log.Fatalf("error: %s", err)
	}

	tests, err := parseOutput(input)
	if err != nil {
		log.Fatalf("error: %s", err)
	}
	if len(tests) == 0 {
		log.Fatalf("error: no tests found")
		os.Exit(1)
	}

	writeXML(tests, output)
	if *fail && numFailures(tests) > 0 {
		os.Exit(1)
	}
}
<commit_msg>-version<commit_after>package main

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

const (
	startPrefix = "=== RUN "
	passPrefix = "--- PASS: "
	failPrefix = "--- FAIL: "

	version = "0.1.0"
)

var endRegexp *regexp.Regexp = regexp.MustCompile(`([^ ]+) \((\d+\.\d+)`)

type Test struct {
	Name, Time, Message string
	Failed bool
}

func parseEnd(prefix, line string) (string, string, error) {
	matches := endRegexp.FindStringSubmatch(line[len(prefix):])

	if len(matches) == 0 {
		return "", "", fmt.Errorf("can't parse %s", line)
	}

	return matches[1], matches[2], nil
}

func parseOutput(rd io.Reader) ([]*Test, error) {
	tests := []*Test{}

	reader := bufio.NewReader(rd)
	var test *Test = nil
	for {
		/* FIXME: Handle isPrefix */
		buf, _, err := reader.ReadLine()

		switch err {
		case io.EOF:
			if test != nil {
				tests = append(tests, test)
			}
			return tests, nil
		case nil:
			;
		default:
			return nil, err
		}

		line := string(buf)
		switch {
		case strings.HasPrefix(line, startPrefix):
			if test != nil {
				tests = append(tests, test)
			}
			test = &Test{Name: line[len(startPrefix):]}
		case strings.HasPrefix(line, failPrefix):
			if test == nil {
				return nil, fmt.Errorf("fail not inside test")
			}
			test.Failed = true
			name, time, err := parseEnd(failPrefix, line)
			if err != nil {
				return nil, err
			}
			if name != test.Name {
				return nil, fmt.Errorf("wrong test end (%s!=%s)", name, test.Name)
			}
			test.Time = time
		case strings.HasPrefix(line, passPrefix):
			if test == nil {
				return nil, fmt.Errorf("pass not inside test")
			}
			test.Failed = false
			name, time, err := parseEnd(passPrefix, line)
			if err != nil {
				return nil, err
			}
			if name != test.Name {
				return nil, fmt.Errorf("wrong test end (%s!=%s)", name, test.Name)
			}
			test.Time = time
		default:
			if test != nil {
				test.Message += line + "\n"
			}
		}
	}

	return tests, nil
}

func getInput(filename string) (io.Reader, error) {
	if filename == "-" || filename == "" {
		return os.Stdin, nil
	}

	return os.Open(filename)
}

func numFailures(tests []*Test) int {
	count := 0
	for _, test := range tests {
		if test.Failed {
			count ++
		}
	}

	return count
}

func writeXML(tests []*Test, out io.Writer) {
	newline := func() { fmt.Fprintln(out) }

	fmt.Fprintf(out, `<?xml version="1.0" encoding="utf-8"?>`)
	newline()
	fmt.Fprintf(out, `<testsuite name="go2xunit" tests="%d" errors="0" failures="%d" skip="0">`,
					 len(tests), numFailures(tests))
	newline()
	for _, test := range(tests) {
		fmt.Fprintf(out, `  <testcase classname="go2xunit" name="%s" time="%s"`,
		                 test.Name, test.Time)
		if !test.Failed {
			fmt.Fprintf(out, " />\n")
			continue
		}
		fmt.Fprintln(out, ">")
		fmt.Fprintf(out,  `    <failure type="go.error" message="error">`)
		newline()
		fmt.Fprintf(out, "<![CDATA[%s]]>\n", test.Message)
		fmt.Fprintln(out, "    </failure>")
		fmt.Fprintln(out, "  </testcase>")
	}
	fmt.Fprintln(out, "</testsuite>")
}

func getOutput(filename string) (io.Writer, error) {
	if filename == "-" || filename == "" {
		return os.Stdout, nil
	}

	return os.Create(filename)
}

func getIO(inputFile, outputFile string) (io.Reader, io.Writer, error) {
	input, err := getInput(inputFile)
	if err != nil {
		return nil, nil, fmt.Errorf("can't open %s for reading: %s", inputFile, err)
	}

	output, err := getOutput(outputFile)
	if err != nil {
		return nil, nil, fmt.Errorf("can't open %s for writing: %s", outputFile, err)
	}

	return input, output, nil
}

func main() {
	inputFile := flag.String("input", "", "input file (default to stdin)")
	outputFile := flag.String("output", "", "output file (default to stdout)")
	fail := flag.Bool("fail", false, "fail (non zero exit) if any test failed")
	showVersion := flag.Bool("version", false, "print version and exit")
	flag.Parse()

	if *showVersion {
		fmt.Println(version)
		os.Exit(0)
	}

	log.SetFlags(0)

	if flag.NArg() > 0 {
		log.Fatalf("error: %s does not take parameters (did you mean -input?)", os.Args[0])
	}

	input, output, err := getIO(*inputFile, *outputFile)
	if err != nil {
		log.Fatalf("error: %s", err)
	}

	tests, err := parseOutput(input)
	if err != nil {
		log.Fatalf("error: %s", err)
	}
	if len(tests) == 0 {
		log.Fatalf("error: no tests found")
		os.Exit(1)
	}

	writeXML(tests, output)
	if *fail && numFailures(tests) > 0 {
		os.Exit(1)
	}
}
<|endoftext|>


<commit_before>package main

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

const (
	startPrefix = "=== RUN "
	passPrefix  = "--- PASS: "
	failPrefix  = "--- FAIL: "

	version = "0.1.0"
)

// "end of test" regexp for name and time, examples:
// --- PASS: TestSub (0.00 seconds)
// --- FAIL: TestSubFail (0.00 seconds)
var endRegexp *regexp.Regexp = regexp.MustCompile(`([^ ]+) \((\d+\.\d+)`)

type Test struct {
	Name, Time, Message string
	Failed              bool
}

// parseEnd parses "end of test" line and returns (name, time, error)
func parseEnd(prefix, line string) (string, string, error) {
	matches := endRegexp.FindStringSubmatch(line[len(prefix):])

	if len(matches) == 0 {
		return "", "", fmt.Errorf("can't parse %s", line)
	}

	return matches[1], matches[2], nil
}

// parseOutput parses output of "go test -v", returns a list of tests
func parseOutput(rd io.Reader) ([]*Test, error) {
	tests := []*Test{}

	reader := bufio.NewReader(rd)
	var test *Test = nil
	for {
		buf, _, err := reader.ReadLine()

		switch err {
		case io.EOF:
			if test != nil {
				tests = append(tests, test)
			}
			return tests, nil
		case nil:

		default:
			return nil, err
		}

		line := string(buf)

		switch {
		case strings.HasPrefix(line, startPrefix):
		case strings.HasPrefix(line, failPrefix):
			// We are switching to the next test, so store the
			// current one.
			if test != nil {
				tests = append(tests, test)
			}

			// Extract the test name and the duration:
			name, time, err := parseEnd(passPrefix, line)
			if err != nil {
				return nil, err
			}

			test = &Test{
				Name: name,
				Time: time,
				Failed: true,
			}

		case strings.HasPrefix(line, passPrefix):
			// We are switching to the next test, so store the
			// current one.
			if test != nil {
				tests = append(tests, test)
			}
			// Extract the test name and the duration:
			name, time, err := parseEnd(passPrefix, line)
			if err != nil {
				return nil, err
			}

			// Create the test structure and store it.
			tests = append(tests, &Test{
				Name: name,
				Time: time,
				Failed: false,
			})
			test = nil
		default:
			if test != nil { // test != nil marks we're in the middle of a test
				test.Message += line + "\n"
			}
		}
	}

	return tests, nil
}

// numFailures count how man tests failed
func numFailures(tests []*Test) int {
	count := 0
	for _, test := range tests {
		if test.Failed {
			count++
		}
	}

	return count
}

// writeXML exits xunit XML of tests to out
func writeXML(tests []*Test, out io.Writer) {
	newline := func() { fmt.Fprintln(out) }

	fmt.Fprintf(out, `<?xml version="1.0" encoding="utf-8"?>`)
	newline()
	fmt.Fprintf(out, `<testsuite name="go2xunit" tests="%d" errors="0" failures="%d" skip="0">`,
		len(tests), numFailures(tests))
	newline()
	for _, test := range tests {
		fmt.Fprintf(out, `  <testcase classname="go2xunit" name="%s" time="%s"`,
			test.Name, test.Time)
		if !test.Failed {
			fmt.Fprintf(out, " />\n")
			continue
		}
		fmt.Fprintln(out, ">")
		fmt.Fprintf(out, `    <failure type="go.error" message="error">`)
		newline()
		fmt.Fprintf(out, "<![CDATA[%s]]>\n", test.Message)
		fmt.Fprintln(out, "    </failure>")
		fmt.Fprintln(out, "  </testcase>")
	}
	fmt.Fprintln(out, "</testsuite>")
}

// getInput return input io.Reader from file name, if file name is - it will
// return os.Stdin
func getInput(filename string) (io.Reader, error) {
	if filename == "-" || filename == "" {
		return os.Stdin, nil
	}

	return os.Open(filename)
}

// getInput return output io.Writer from file name, if file name is - it will
// return os.Stdout
func getOutput(filename string) (io.Writer, error) {
	if filename == "-" || filename == "" {
		return os.Stdout, nil
	}

	return os.Create(filename)
}

// getIO returns input and output streams from file names
func getIO(inputFile, outputFile string) (io.Reader, io.Writer, error) {
	input, err := getInput(inputFile)
	if err != nil {
		return nil, nil, fmt.Errorf("can't open %s for reading: %s", inputFile, err)
	}

	output, err := getOutput(outputFile)
	if err != nil {
		return nil, nil, fmt.Errorf("can't open %s for writing: %s", outputFile, err)
	}

	return input, output, nil
}

func main() {
	inputFile := flag.String("input", "", "input file (default to stdin)")
	outputFile := flag.String("output", "", "output file (default to stdout)")
	fail := flag.Bool("fail", false, "fail (non zero exit) if any test failed")
	showVersion := flag.Bool("version", false, "print version and exit")
	flag.Parse()

	if *showVersion {
		fmt.Println(version)
		os.Exit(0)
	}

	// No time ... prefix for error messages
	log.SetFlags(0)

	if flag.NArg() > 0 {
		log.Fatalf("error: %s does not take parameters (did you mean -input?)", os.Args[0])
	}

	input, output, err := getIO(*inputFile, *outputFile)
	if err != nil {
		log.Fatalf("error: %s", err)
	}

	tests, err := parseOutput(input)
	if err != nil {
		log.Fatalf("error: %s", err)
	}
	if len(tests) == 0 {
		log.Fatalf("error: no tests found")
		os.Exit(1)
	}

	writeXML(tests, output)
	if *fail && numFailures(tests) > 0 {
		os.Exit(1)
	}
}
<commit_msg>Handle the edge-case of the last test failing.<commit_after>package main

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

const (
	startPrefix = "=== RUN "
	passPrefix  = "--- PASS: "
	failPrefix  = "--- FAIL: "

	version = "0.1.0"
)

// "end of test" regexp for name and time, examples:
// --- PASS: TestSub (0.00 seconds)
// --- FAIL: TestSubFail (0.00 seconds)
var endRegexp *regexp.Regexp = regexp.MustCompile(`([^ ]+) \((\d+\.\d+)`)

type Test struct {
	Name, Time, Message string
	Failed              bool
}

// parseEnd parses "end of test" line and returns (name, time, error)
func parseEnd(prefix, line string) (string, string, error) {
	matches := endRegexp.FindStringSubmatch(line[len(prefix):])

	if len(matches) == 0 {
		return "", "", fmt.Errorf("can't parse %s", line)
	}

	return matches[1], matches[2], nil
}

// parseOutput parses output of "go test -v", returns a list of tests
func parseOutput(rd io.Reader) ([]*Test, error) {
	tests := []*Test{}

	reader := bufio.NewReader(rd)
	var test *Test = nil
	for {
		buf, _, err := reader.ReadLine()

		switch err {
		case io.EOF:
			if test != nil {
				tests = append(tests, test)
			}
			return tests, nil
		case nil:

		default:
			return nil, err
		}

		line := string(buf)

		switch {
		case strings.HasPrefix(line, startPrefix):
		case strings.HasPrefix(line, failPrefix):
			// We are switching to the next test, so store the
			// current one.
			if test != nil {
				tests = append(tests, test)
			}

			// Extract the test name and the duration:
			name, time, err := parseEnd(passPrefix, line)
			if err != nil {
				return nil, err
			}

			test = &Test{
				Name: name,
				Time: time,
				Failed: true,
			}

		case strings.HasPrefix(line, passPrefix):
			// We are switching to the next test, so store the
			// current one.
			if test != nil {
				tests = append(tests, test)
			}
			// Extract the test name and the duration:
			name, time, err := parseEnd(passPrefix, line)
			if err != nil {
				return nil, err
			}

			// Create the test structure and store it.
			tests = append(tests, &Test{
				Name: name,
				Time: time,
				Failed: false,
			})
			test = nil
		case line == "FAIL":
			// Handle the edge case of the last test failing: the
			// following line will be "FAIL", so we just stop there.
			tests = append(tests, test)
			test = nil
		default:
			if test != nil { // test != nil marks we're in the middle of a test
				test.Message += line + "\n"
			}
		}
	}

	return tests, nil
}

// numFailures count how man tests failed
func numFailures(tests []*Test) int {
	count := 0
	for _, test := range tests {
		if test.Failed {
			count++
		}
	}

	return count
}

// writeXML exits xunit XML of tests to out
func writeXML(tests []*Test, out io.Writer) {
	newline := func() { fmt.Fprintln(out) }

	fmt.Fprintf(out, `<?xml version="1.0" encoding="utf-8"?>`)
	newline()
	fmt.Fprintf(out, `<testsuite name="go2xunit" tests="%d" errors="0" failures="%d" skip="0">`,
		len(tests), numFailures(tests))
	newline()
	for _, test := range tests {
		fmt.Fprintf(out, `  <testcase classname="go2xunit" name="%s" time="%s"`,
			test.Name, test.Time)
		if !test.Failed {
			fmt.Fprintf(out, " />\n")
			continue
		}
		fmt.Fprintln(out, ">")
		fmt.Fprintf(out, `    <failure type="go.error" message="error">`)
		newline()
		fmt.Fprintf(out, "<![CDATA[%s]]>\n", test.Message)
		fmt.Fprintln(out, "    </failure>")
		fmt.Fprintln(out, "  </testcase>")
	}
	fmt.Fprintln(out, "</testsuite>")
}

// getInput return input io.Reader from file name, if file name is - it will
// return os.Stdin
func getInput(filename string) (io.Reader, error) {
	if filename == "-" || filename == "" {
		return os.Stdin, nil
	}

	return os.Open(filename)
}

// getInput return output io.Writer from file name, if file name is - it will
// return os.Stdout
func getOutput(filename string) (io.Writer, error) {
	if filename == "-" || filename == "" {
		return os.Stdout, nil
	}

	return os.Create(filename)
}

// getIO returns input and output streams from file names
func getIO(inputFile, outputFile string) (io.Reader, io.Writer, error) {
	input, err := getInput(inputFile)
	if err != nil {
		return nil, nil, fmt.Errorf("can't open %s for reading: %s", inputFile, err)
	}

	output, err := getOutput(outputFile)
	if err != nil {
		return nil, nil, fmt.Errorf("can't open %s for writing: %s", outputFile, err)
	}

	return input, output, nil
}

func main() {
	inputFile := flag.String("input", "", "input file (default to stdin)")
	outputFile := flag.String("output", "", "output file (default to stdout)")
	fail := flag.Bool("fail", false, "fail (non zero exit) if any test failed")
	showVersion := flag.Bool("version", false, "print version and exit")
	flag.Parse()

	if *showVersion {
		fmt.Println(version)
		os.Exit(0)
	}

	// No time ... prefix for error messages
	log.SetFlags(0)

	if flag.NArg() > 0 {
		log.Fatalf("error: %s does not take parameters (did you mean -input?)", os.Args[0])
	}

	input, output, err := getIO(*inputFile, *outputFile)
	if err != nil {
		log.Fatalf("error: %s", err)
	}

	tests, err := parseOutput(input)
	if err != nil {
		log.Fatalf("error: %s", err)
	}
	if len(tests) == 0 {
		log.Fatalf("error: no tests found")
		os.Exit(1)
	}

	writeXML(tests, output)
	if *fail && numFailures(tests) > 0 {
		os.Exit(1)
	}
}
<|endoftext|>


<commit_before>package gofigure

import (
	"github.com/droundy/goopt"
	"os"
)

// Config contains the configuration options that may be set by
// command line flags and environment variables.
type Config struct {
	Description string
	Version     string
	EnvPrefix   string
	options     map[string]*Option
	flags       map[string]*string
	values		map[string]string
}

// Returns a new Config instance.
func New() *Config {
	return &Config{
		options: make(map[string]*Option),
		flags:  make(map[string]*string),
		values: make(map[string]string),
	}
}

// Adds a configuration option, returns an Option instance for
// easily setting the corresponding environment variable, default
// value, and description.
func (c *Config) Add(name string) *Option {
	c.options[name] = &Option{
		envVar: "",
		def:    "",
		desc:   "",
	}
	return c.options[name]
}

// Returns a configuration option by flag name.
func (c *Config) Get(name string) string {
	return c.values[name]
}

// Parses the configuration options into defined flags, sets the value
// accordingly. Options are read first from command line flags, then
// from environment variables, and falls back to the default value if
// neither are set.
//
// See https://github.com/rakyll/globalconf/blob/master/globalconf.go
func (c *Config) Parse() {
	goopt.Description = func() string {
		return c.Description
	}

	goopt.Version = c.Version

	// Sets the flags from the configuration options.
	for name, o := range c.options {
		c.flags[name] = goopt.String([]string{"--"+name}, "", o.desc)
		c.values[name] = o.def
	}

	goopt.Parse(nil)

	// Gather the flags passed through command line.
	passed := make(map[string]bool)
	for name, f := range c.flags {
		if *f != "" {
			passed[name] = true
			c.values[name] = *f
		}
	}

	for name, f := range c.options {

		// Skip flags passed through the command line as the option is
		// already set and takes precedence over environment variables.
		if passed[name] {
			return
		}

		// Some options shouldn't be set via environment variables.
		if f.envVar == "" {
			return
		}

		// If the configuration option was not passed via the command line,
		// check the corresponding environment variable.
		envVar := c.EnvPrefix + f.envVar
		if val := os.Getenv(envVar); val != "" {
			c.values[name] = val
		}
	}
}

// Option contains the details of a configuration options,
// e.g. corresponding environment variable, default value,
// description.
type Option struct {
	envVar, def, desc string
}

// Sets the configuration option's default value.
func (o *Option) Default(def string) *Option {
	o.def = def
	return o
}

// Sets the configuration option's corresponding environment variable.
func (o *Option) EnvVar(envVar string) *Option {
	o.envVar = envVar
	return o
}

// Sets the configuration options long description.
func (o *Option) Description(desc string) *Option {
	o.desc = desc
	return o
}
<commit_msg>formatting fix due to space/tab confusion<commit_after>package gofigure

import (
	"github.com/droundy/goopt"
	"os"
)

// Config contains the configuration options that may be set by
// command line flags and environment variables.
type Config struct {
	Description	string
	Version		string
	EnvPrefix	string
	options		map[string]*Option
	flags		map[string]*string
	values		map[string]string
}

// Returns a new Config instance.
func New() *Config {
	return &Config{
		options: make(map[string]*Option),
		flags:  make(map[string]*string),
		values: make(map[string]string),
	}
}

// Adds a configuration option, returns an Option instance for
// easily setting the corresponding environment variable, default
// value, and description.
func (c *Config) Add(name string) *Option {
	c.options[name] = &Option{
		envVar: "",
		def:    "",
		desc:   "",
	}
	return c.options[name]
}

// Returns a configuration option by flag name.
func (c *Config) Get(name string) string {
	return c.values[name]
}

// Parses the configuration options into defined flags, sets the value
// accordingly. Options are read first from command line flags, then
// from environment variables, and falls back to the default value if
// neither are set.
//
// See https://github.com/rakyll/globalconf/blob/master/globalconf.go
func (c *Config) Parse() {
	goopt.Description = func() string {
		return c.Description
	}

	goopt.Version = c.Version

	// Sets the flags from the configuration options.
	for name, o := range c.options {
		c.flags[name] = goopt.String([]string{"--"+name}, "", o.desc)
		c.values[name] = o.def
	}

	goopt.Parse(nil)

	// Gather the flags passed through command line.
	passed := make(map[string]bool)
	for name, f := range c.flags {
		if *f != "" {
			passed[name] = true
			c.values[name] = *f
		}
	}

	for name, f := range c.options {

		// Skip flags passed through the command line as the option is
		// already set and takes precedence over environment variables.
		if passed[name] {
			return
		}

		// Some options shouldn't be set via environment variables.
		if f.envVar == "" {
			return
		}

		// If the configuration option was not passed via the command line,
		// check the corresponding environment variable.
		envVar := c.EnvPrefix + f.envVar
		if val := os.Getenv(envVar); val != "" {
			c.values[name] = val
		}
	}
}

// Option contains the details of a configuration options,
// e.g. corresponding environment variable, default value,
// description.
type Option struct {
	envVar, def, desc string
}

// Sets the configuration option's default value.
func (o *Option) Default(def string) *Option {
	o.def = def
	return o
}

// Sets the configuration option's corresponding environment variable.
func (o *Option) EnvVar(envVar string) *Option {
	o.envVar = envVar
	return o
}

// Sets the configuration options long description.
func (o *Option) Description(desc string) *Option {
	o.desc = desc
	return o
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"html"
	"io"
	"log"
	"net/http"
	"net/http/cookiejar"
	"net/url"
	"time"

	"github.com/JustinTulloss/hut"
	"github.com/PuerkitoBio/goquery"
	"github.com/facebookgo/httpcontrol"
	"github.com/temoto/robotstxt.go"
)

var ogPrefixes = []string{"og", "airbedandbreakfast", "twitter"}

var useragent = "Gogetter (https://github.com/JustinTulloss/gogetter) (like GoogleBot and facebookexternalhit)"
var service *hut.Service
var client *http.Client

type HttpError struct {
	msg        string
	StatusCode int
}

func (e *HttpError) Error() string { return e.msg }

func buildRequest(url string) (*http.Request, error) {
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return nil, err
	}
	req.Header.Set("User-Agent", useragent)
	return req, nil
}

func checkRobotsTxt(fullUrl string) (bool, error) {
	if service.Env.GetString("CHECK_ROBOTS_TXT") == "false" {
		return true, nil
	}
	parsed, err := url.Parse(fullUrl)
	if err != nil {
		return false, err
	}
	original := parsed.Path
	parsed.Path = "robots.txt"
	parsed.RawQuery = ""
	req, err := buildRequest(parsed.String())
	if err != nil {
		return false, err
	}
	resp, err := client.Do(req)
	if err != nil {
		return false, err
	}
	defer resp.Body.Close()
	robots, err := robotstxt.FromResponse(resp)
	if robots == nil {
		// Assume we can crawl if the robots.txt file doesn't work
		return true, nil
	}
	return robots.TestAgent(original, useragent), nil
}

func parseTags(r io.Reader) (map[string]string, *HttpError) {
	doc, err := goquery.NewDocumentFromReader(r)
	if err != nil {
		return nil, &HttpError{err.Error(), 500}
	}
	results := make(map[string]string)
	// First we deal with a couple special tags to get the title
	// and the favicon
	title := doc.Find("title").Text()
	if title != "" {
		results["title"] = html.UnescapeString(title)
	}
	favicon, ok := doc.Find("link[rel~=icon]").Attr("href")
	if ok {
		results["favicon"] = html.UnescapeString(favicon)
	}
	// Find all meta tags for all different og prefixes we support
	tags := doc.Find(`meta[name="description"]`)
	for _, prefix := range ogPrefixes {
		tags = tags.Add(fmt.Sprintf(`meta[property^="%s:"]`, prefix))
		tags = tags.Add(fmt.Sprintf(`meta[name^="%s:"]`, prefix))
	}
	// For all the tags, extract the content
	tags.Each(func(i int, selection *goquery.Selection) {
		key, ok := selection.Attr("name")
		if !ok {
			key, _ = selection.Attr("property")
		}
		content, _ := selection.Attr("content")
		results[key] = html.UnescapeString(content)
	})
	return results, nil
}

func getTags(url string) (map[string]string, *HttpError) {
	permitted, err := checkRobotsTxt(url)
	if err != nil {
		return nil, &HttpError{err.Error(), 500}
	}
	if !permitted {
		msg := fmt.Sprintf("Not permitted to fetch %s as a robot", url)
		log.Println(msg)
		return nil, &HttpError{msg, 403}
	}
	req, _ := buildRequest(url)
	resp, err := client.Do(req)
	if err != nil {
		return nil, &HttpError{err.Error(), 500}
	}
	log.Printf("Fetched %s\n", url)
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return nil, &HttpError{
			fmt.Sprintf("Could not fetch %s, request was: %v", url, req),
			resp.StatusCode,
		}
	}
	return parseTags(resp.Body)
}

func handler(w http.ResponseWriter, r *http.Request) {
	if err := r.ParseForm(); err != nil {
		service.ErrorReply(err, w)
		return
	}
	decodedUrl, err := url.QueryUnescape(r.Form.Get("url"))
	if err != nil {
		service.ErrorReply(err, w)
		return
	}
	tags, httpErr := getTags(decodedUrl)
	if httpErr != nil {
		service.HttpErrorReply(w, httpErr.Error(), httpErr.StatusCode)
		return
	}
	service.Reply(tags, w)
}

func main() {
	var err error
	service = hut.NewService(nil)
	jar, err := cookiejar.New(nil)
	if err != nil {
		service.Log.Error().Printf("Could not create cookie jar: %s\n", err)
	}
	service.Router.HandleFunc("/", handler)

	client = &http.Client{
		Transport: &httpcontrol.Transport{
			RequestTimeout: 10 * time.Second,
			MaxTries:       3,
		},
		Jar: jar,
	}

	flag.Parse()
	protocol := service.Env.Get("protocol")
	if protocol == "http" {
		service.Start()
	} else if len(flag.Args()) != 0 {
		tags, err := getTags(flag.Arg(0))
		if err != nil {
			fmt.Printf("Could not fetch: %s\n", err)
			return
		}
		for prop, val := range tags {
			fmt.Printf("%s -- %s\n", prop, val)
		}
	} else {
		panic("Need to use this properly and I need to print usage info!")
	}
}
<commit_msg>Handle non-html more elegantly<commit_after>package main

import (
	"flag"
	"fmt"
	"html"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/http/cookiejar"
	"net/url"
	"strings"
	"time"

	"github.com/JustinTulloss/hut"
	"github.com/PuerkitoBio/goquery"
	"github.com/facebookgo/httpcontrol"
	"github.com/temoto/robotstxt.go"
)

var ogPrefixes = []string{"og", "airbedandbreakfast", "twitter"}

var useragent = "Gogetter (https://github.com/JustinTulloss/gogetter) (like GoogleBot and facebookexternalhit)"
var service *hut.Service
var client *http.Client

type HttpError struct {
	msg        string
	StatusCode int
}

func (e *HttpError) Error() string { return e.msg }

func buildRequest(url string) (*http.Request, error) {
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return nil, err
	}
	req.Header.Set("User-Agent", useragent)
	return req, nil
}

func checkRobotsTxt(fullUrl string) (bool, error) {
	if service.Env.GetString("CHECK_ROBOTS_TXT") == "false" {
		return true, nil
	}
	parsed, err := url.Parse(fullUrl)
	if err != nil {
		return false, err
	}
	original := parsed.Path
	parsed.Path = "robots.txt"
	parsed.RawQuery = ""
	req, err := buildRequest(parsed.String())
	if err != nil {
		return false, err
	}
	resp, err := client.Do(req)
	if err != nil {
		return false, err
	}
	defer resp.Body.Close()
	robots, err := robotstxt.FromResponse(resp)
	if robots == nil {
		// Assume we can crawl if the robots.txt file doesn't work
		return true, nil
	}
	return robots.TestAgent(original, useragent), nil
}

func parseTags(r io.Reader) (map[string]string, *HttpError) {
	doc, err := goquery.NewDocumentFromReader(r)
	if err != nil {
		return nil, &HttpError{err.Error(), 500}
	}
	results := make(map[string]string)
	// First we deal with a couple special tags to get the title
	// and the favicon
	title := doc.Find("title").Text()
	if title != "" {
		results["title"] = html.UnescapeString(title)
	}
	favicon, ok := doc.Find("link[rel~=icon]").Attr("href")
	if ok {
		results["favicon"] = html.UnescapeString(favicon)
	}
	// Find all meta tags for all different og prefixes we support
	tags := doc.Find(`meta[name="description"]`)
	for _, prefix := range ogPrefixes {
		tags = tags.Add(fmt.Sprintf(`meta[property^="%s:"]`, prefix))
		tags = tags.Add(fmt.Sprintf(`meta[name^="%s:"]`, prefix))
	}
	// For all the tags, extract the content
	tags.Each(func(i int, selection *goquery.Selection) {
		key, ok := selection.Attr("name")
		if !ok {
			key, _ = selection.Attr("property")
		}
		content, _ := selection.Attr("content")
		results[key] = html.UnescapeString(content)
	})
	return results, nil
}

func getTags(url string) (map[string]string, *HttpError) {
	permitted, err := checkRobotsTxt(url)
	if err != nil {
		return nil, &HttpError{err.Error(), 500}
	}
	if !permitted {
		msg := fmt.Sprintf("Not permitted to fetch %s as a robot", url)
		log.Println(msg)
		return nil, &HttpError{msg, 403}
	}
	req, _ := buildRequest(url)
	resp, err := client.Do(req)
	if err != nil {
		return nil, &HttpError{err.Error(), 500}
	}
	log.Printf("Fetched %s\n", url)
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return nil, &HttpError{
			fmt.Sprintf("Could not fetch %s, request was: %v", url, req),
			resp.StatusCode,
		}
	}
	// We can't really trust the Content-Type header, so we take
	// a look at what actually gets returned.
	contentStart, err := ioutil.ReadAll(io.LimitReader(resp.Body, 512))
	contentType := http.DetectContentType(contentStart)
	switch {
	case strings.Contains(contentType, "text/html"):
		return parseTags(resp.Body)
	default:
		return map[string]string{
			"mimeType": contentType,
		}, nil
	}
}

func handler(w http.ResponseWriter, r *http.Request) {
	if err := r.ParseForm(); err != nil {
		service.ErrorReply(err, w)
		return
	}
	decodedUrl, err := url.QueryUnescape(r.Form.Get("url"))
	if err != nil {
		service.ErrorReply(err, w)
		return
	}
	tags, httpErr := getTags(decodedUrl)
	if httpErr != nil {
		service.HttpErrorReply(w, httpErr.Error(), httpErr.StatusCode)
		return
	}
	service.Reply(tags, w)
}

func main() {
	var err error
	service = hut.NewService(nil)
	jar, err := cookiejar.New(nil)
	if err != nil {
		service.Log.Error().Printf("Could not create cookie jar: %s\n", err)
	}
	service.Router.HandleFunc("/", handler)

	client = &http.Client{
		Transport: &httpcontrol.Transport{
			RequestTimeout: 10 * time.Second,
			MaxTries:       3,
		},
		Jar: jar,
	}

	flag.Parse()
	protocol := service.Env.Get("protocol")
	if protocol == "http" {
		service.Start()
	} else if len(flag.Args()) != 0 {
		tags, err := getTags(flag.Arg(0))
		if err != nil {
			fmt.Printf("Could not fetch: %s\n", err)
			return
		}
		for prop, val := range tags {
			fmt.Printf("%s -- %s\n", prop, val)
		}
	} else {
		panic("Need to use this properly and I need to print usage info!")
	}
}
<|endoftext|>