<commit_before>package desugar

import (
	"fmt"

	"github.com/tisp-lang/tisp/src/lib/ast"
	"github.com/tisp-lang/tisp/src/lib/debug"
	"github.com/tisp-lang/tisp/src/lib/gensym"
)

func desugarMutualRecursionStatement(s interface{}) []interface{} {
	switch s := s.(type) {
	case ast.MutualRecursion:
		return desugarMutualRecursion(s)
	default:
		return []interface{}{s}
	}
}

func desugarMutualRecursion(mr ast.MutualRecursion) []interface{} {
	fs := mr.LetFunctions()
	unrecs := make([]interface{}, 0, len(fs))

	for _, f := range fs {
		arg := gensym.GenSym("mr", "functions", "argument")
		n2i := indexLetFunctions(fs...)

		ls := make([]interface{}, 0, len(f.Lets()))

		for _, l := range f.Lets() {
			switch l := l.(type) {
			case ast.LetFunction:
				delete(n2i, l.Name())
			case ast.LetVar:
				delete(n2i, l.Name())
			default:
				panic("Unreachable")
			}

			ls = append(ls, replaceNames(arg, n2i, l, mr.DebugInfo()))
		}

		unrecs = append(
			unrecs,
			ast.NewLetFunction(
				gensym.GenSym("mr", "unrec", f.Name()),
				prependPosReqsToSig([]string{arg}, f.Signature()),
				ls,
				replaceNames(arg, deleteNamesDefinedByLets(n2i, f.Lets()), f.Body(), mr.DebugInfo()),
				f.DebugInfo()))
	}

	recsList := gensym.GenSym("ys", "mr", "functions")
	recs := make([]interface{}, 0, len(fs))

	for i, f := range fs {
		recs = append(
			recs,
			ast.NewLetVar(
				f.Name(),
				ast.NewPApp(recsList, []interface{}{fmt.Sprint(i)}, f.DebugInfo())))
	}

	return append(
		unrecs,
		append(
			[]interface{}{ast.NewLetVar(
				recsList,
				ast.NewPApp("$ys", stringsToAnys(letStatementsToNames(unrecs)), mr.DebugInfo()))},
			recs...)...)
}

func indexLetFunctions(fs ...ast.LetFunction) map[string]int {
	n2i := make(map[string]int)

	for i, f := range fs {
		n2i[f.Name()] = i
	}

	if len(n2i) != len(fs) {
		panic(fmt.Errorf("Duplicate names were found among mutually-recursive functions"))
	}

	return n2i
}

func replaceNames(funcList string, n2i map[string]int, x interface{}, di debug.Info) interface{} {
	replaceWithNameToIndex := func(n2i map[string]int) func(x interface{}) interface{} {
		return func(x interface{}) interface{} {
			return replaceNames(funcList, n2i, x, di)
		}
	}

	replace := replaceWithNameToIndex(n2i)

	switch x := x.(type) {
	case ast.LetFunction:
		n2i := copyNameToIndex(n2i)

		delete(n2i, x.Name())
		for n := range signatureToNames(x.Signature()) {
			delete(n2i, n)
		}

		return ast.NewLetFunction(
			x.Name(),
			x.Signature(),
			replaceWithNameToIndex(n2i)(x.Lets()).([]interface{}),
			replaceWithNameToIndex(deleteNamesDefinedByLets(n2i, x.Lets()))(x.Body()),
			x.DebugInfo())
	case ast.LetVar:
		n2i := copyNameToIndex(n2i)
		delete(n2i, x.Name())
		return ast.NewLetVar(x.Name(), replaceWithNameToIndex(n2i)(x.Expr()))
	case ast.App:
		return ast.NewApp(replace(x.Function()), replace(x.Arguments()).(ast.Arguments), x.DebugInfo())
	case ast.Arguments:
		ps := make([]ast.PositionalArgument, 0, len(x.Positionals()))

		for _, p := range x.Positionals() {
			ps = append(ps, ast.NewPositionalArgument(replace(p.Value()), p.Expanded()))
		}

		ks := make([]ast.KeywordArgument, 0, len(x.Keywords()))

		for _, k := range x.Keywords() {
			ks = append(ks, ast.NewKeywordArgument(k.Name(), replace(k.Value())))
		}

		ds := make([]interface{}, 0, len(x.ExpandedDicts()))

		for _, d := range x.ExpandedDicts() {
			ds = append(ds, replace(d))
		}

		return ast.NewArguments(ps, ks, ds)
	case string:
		if i, ok := n2i[x]; ok {
			return ast.NewPApp(funcList, []interface{}{fmt.Sprint(i)}, di)
		}

		return x
	}

	panic(fmt.Errorf("Invalid value: %#v", x))
}

func copyNameToIndex(n2i map[string]int) map[string]int {
	new := make(map[string]int)

	for n, i := range n2i {
		new[n] = i
	}

	return new
}

func deleteNamesDefinedByLets(n2i map[string]int, ls []interface{}) map[string]int {
	n2i = copyNameToIndex(n2i)

	for _, n := range letStatementsToNames(ls) {
		delete(n2i, n)
	}

	return n2i
}

func letStatementsToNames(ls []interface{}) []string {
	ns := make([]string, 0, len(ls))

	for _, l := range ls {
		switch l := l.(type) {
		case ast.LetFunction:
			ns = append(ns, l.Name())
		case ast.LetVar:
			ns = append(ns, l.Name())
		default:
			panic("Unreachable")
		}
	}

	return ns
}

func stringsToAnys(ss []string) []interface{} {
	xs := make([]interface{}, 0, len(ss))

	for _, s := range ss {
		xs = append(xs, s)
	}

	return xs
}
<commit_msg>Create createUnrecursiveFunction function<commit_after>package desugar

import (
	"fmt"

	"github.com/tisp-lang/tisp/src/lib/ast"
	"github.com/tisp-lang/tisp/src/lib/debug"
	"github.com/tisp-lang/tisp/src/lib/gensym"
)

func desugarMutualRecursionStatement(s interface{}) []interface{} {
	switch s := s.(type) {
	case ast.MutualRecursion:
		return desugarMutualRecursion(s)
	default:
		return []interface{}{s}
	}
}

func desugarMutualRecursion(mr ast.MutualRecursion) []interface{} {
	fs := mr.LetFunctions()
	unrecs := make([]interface{}, 0, len(fs))

	for _, f := range fs {
		unrecs = append(unrecs, createUnrecursiveFunction(indexLetFunctions(fs...), f))
	}

	recsList := gensym.GenSym("ys", "mr", "functions")
	recs := make([]interface{}, 0, len(fs))

	for i, f := range fs {
		recs = append(
			recs,
			ast.NewLetVar(
				f.Name(),
				ast.NewPApp(recsList, []interface{}{fmt.Sprint(i)}, f.DebugInfo())))
	}

	return append(
		unrecs,
		append(
			[]interface{}{ast.NewLetVar(
				recsList,
				ast.NewPApp("$ys", stringsToAnys(letStatementsToNames(unrecs)), mr.DebugInfo()))},
			recs...)...)
}

func createUnrecursiveFunction(n2i map[string]int, f ast.LetFunction) ast.LetFunction {
	arg := gensym.GenSym("mr", "functions", "argument")
	n2i = copyNameToIndex(n2i)

	ls := make([]interface{}, 0, len(f.Lets()))

	for _, l := range f.Lets() {
		switch l := l.(type) {
		case ast.LetFunction:
			delete(n2i, l.Name())
		case ast.LetVar:
			delete(n2i, l.Name())
		default:
			panic("Unreachable")
		}

		ls = append(ls, replaceNames(arg, n2i, l, f.DebugInfo()))
	}

	return ast.NewLetFunction(
		gensym.GenSym("mr", "unrec", f.Name()),
		prependPosReqsToSig([]string{arg}, f.Signature()),
		ls,
		replaceNames(arg, n2i, f.Body(), f.DebugInfo()),
		f.DebugInfo())
}

func indexLetFunctions(fs ...ast.LetFunction) map[string]int {
	n2i := make(map[string]int)

	for i, f := range fs {
		n2i[f.Name()] = i
	}

	if len(n2i) != len(fs) {
		panic(fmt.Errorf("Duplicate names were found among mutually-recursive functions"))
	}

	return n2i
}

func replaceNames(funcList string, n2i map[string]int, x interface{}, di debug.Info) interface{} {
	replaceWithNameToIndex := func(n2i map[string]int) func(x interface{}) interface{} {
		return func(x interface{}) interface{} {
			return replaceNames(funcList, n2i, x, di)
		}
	}

	replace := replaceWithNameToIndex(n2i)

	switch x := x.(type) {
	case ast.LetFunction:
		n2i := copyNameToIndex(n2i)

		delete(n2i, x.Name())
		for n := range signatureToNames(x.Signature()) {
			delete(n2i, n)
		}

		return ast.NewLetFunction(
			x.Name(),
			x.Signature(),
			replaceWithNameToIndex(n2i)(x.Lets()).([]interface{}),
			replaceWithNameToIndex(deleteNamesDefinedByLets(n2i, x.Lets()))(x.Body()),
			x.DebugInfo())
	case ast.LetVar:
		n2i := copyNameToIndex(n2i)
		delete(n2i, x.Name())
		return ast.NewLetVar(x.Name(), replaceWithNameToIndex(n2i)(x.Expr()))
	case ast.App:
		return ast.NewApp(replace(x.Function()), replace(x.Arguments()).(ast.Arguments), x.DebugInfo())
	case ast.Arguments:
		ps := make([]ast.PositionalArgument, 0, len(x.Positionals()))

		for _, p := range x.Positionals() {
			ps = append(ps, ast.NewPositionalArgument(replace(p.Value()), p.Expanded()))
		}

		ks := make([]ast.KeywordArgument, 0, len(x.Keywords()))

		for _, k := range x.Keywords() {
			ks = append(ks, ast.NewKeywordArgument(k.Name(), replace(k.Value())))
		}

		ds := make([]interface{}, 0, len(x.ExpandedDicts()))

		for _, d := range x.ExpandedDicts() {
			ds = append(ds, replace(d))
		}

		return ast.NewArguments(ps, ks, ds)
	case string:
		if i, ok := n2i[x]; ok {
			return ast.NewPApp(funcList, []interface{}{fmt.Sprint(i)}, di)
		}

		return x
	}

	panic(fmt.Errorf("Invalid value: %#v", x))
}

func copyNameToIndex(n2i map[string]int) map[string]int {
	new := make(map[string]int)

	for n, i := range n2i {
		new[n] = i
	}

	return new
}

func deleteNamesDefinedByLets(n2i map[string]int, ls []interface{}) map[string]int {
	n2i = copyNameToIndex(n2i)

	for _, n := range letStatementsToNames(ls) {
		delete(n2i, n)
	}

	return n2i
}

func letStatementsToNames(ls []interface{}) []string {
	ns := make([]string, 0, len(ls))

	for _, l := range ls {
		switch l := l.(type) {
		case ast.LetFunction:
			ns = append(ns, l.Name())
		case ast.LetVar:
			ns = append(ns, l.Name())
		default:
			panic("Unreachable")
		}
	}

	return ns
}

func stringsToAnys(ss []string) []interface{} {
	xs := make([]interface{}, 0, len(ss))

	for _, s := range ss {
		xs = append(xs, s)
	}

	return xs
}
<|endoftext|>


<commit_before>// Copyright 2017 The WPT Dashboard Project. 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 (
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"strconv"
	"sync"
	"time"

	"cloud.google.com/go/bigquery"
	"cloud.google.com/go/datastore"
	gcs "cloud.google.com/go/storage"
	"github.com/w3c/wptdashboard/metrics"
	"github.com/w3c/wptdashboard/metrics/compute"
	"github.com/w3c/wptdashboard/metrics/storage"
	base "github.com/w3c/wptdashboard/shared"
	"golang.org/x/net/context"
)

var wptDataPath *string
var projectId *string
var inputGcsBucket *string
var outputGcsBucket *string
var wptdHost *string
var outputBQMetadataDataset *string
var outputBQDataDataset *string
var outputBQPassRateTable *string
var outputBQPassRateMetadataTable *string
var outputBQFailuresTable *string
var outputBQFailuresMetadataTable *string

func init() {
	unixNow := time.Now().Unix()
	wptDataPath = flag.String("wpt_data_path", os.Getenv("HOME")+
		"/wpt-data", "Path to data directory for local data copied "+
		"from Google Cloud Storage")
	projectId = flag.String("project_id", "wptdashboard",
		"Google Cloud Platform project id")
	inputGcsBucket = flag.String("input_gcs_bucket", "wptd",
		"Google Cloud Storage bucket where test results are stored")
	outputGcsBucket = flag.String("output_gcs_bucket", "wptd-metrics",
		"Google Cloud Storage bucket where metrics are stored")
	outputBQMetadataDataset = flag.String("output_bq_metadata_dataset",
		fmt.Sprintf("wptd_metrics_%d", unixNow),
		"BigQuery dataset where metrics metadata are stored")
	outputBQDataDataset = flag.String("output_bq_data_dataset",
		fmt.Sprintf("wptd_metrics_%d", unixNow),
		"BigQuery dataset where metrics data are stored")
	outputBQPassRateTable = flag.String("output_bq_pass_rate_table",
		fmt.Sprintf("PassRates_%d", unixNow),
		"BigQuery table where pass rate metrics are stored")
	outputBQPassRateMetadataTable = flag.String("output_bq_pass_rate_metadata_table",
		fmt.Sprintf("PassRateMetadata_%d", unixNow),
		"BigQuery table where pass rate metrics are stored")
	outputBQFailuresTable = flag.String("output_bq_failures_table",
		fmt.Sprintf("Failures_%d", unixNow),
		"BigQuery table where test failure lists are stored")
	outputBQFailuresMetadataTable = flag.String("output_bq_failures_metadata_table",
		fmt.Sprintf("FailuresMetadata_%d", unixNow),
		"BigQuery table where pass rate metrics are stored")
	wptdHost = flag.String("wptd_host", "wpt.fyi",
		"Hostname of endpoint that serves WPT Dashboard data API")
}

func main() {
	log.SetFlags(log.LstdFlags | log.Lshortfile)

	logFileName := "current_metrics.log"
	logFile, err := os.OpenFile(logFileName, os.O_RDWR|os.O_CREATE|
		os.O_APPEND, 0666)
	if err != nil {
		log.Fatalf("Error opening log file: %v", err)
	}
	defer logFile.Close()
	log.Printf("Logs appended to %s\n", logFileName)
	log.SetOutput(logFile)

	ctx := context.Background()
	gcsClient, err := gcs.NewClient(ctx)
	if err != nil {
		log.Fatal(err)
	}
	inputBucket := gcsClient.Bucket(*inputGcsBucket)
	inputCtx := storage.GCSDatastoreContext{
		Context: ctx,
		Bucket: storage.Bucket{
			Name:   *inputGcsBucket,
			Handle: inputBucket,
		},
	}

	log.Println("Reading test results from Google Cloud Storage bucket: " +
		*inputGcsBucket)

	readStartTime := time.Now()
	runs := getRuns()
	allResults := storage.LoadTestRunResults(&inputCtx, runs)
	readEndTime := time.Now()

	log.Println("Read test results from Google Cloud Storage bucket: " +
		*inputGcsBucket)
	log.Println("Consolidating results")

	resultsById := compute.GatherResultsById(&allResults)

	log.Println("Consolidated results")
	log.Println("Computing metrics")

	var totals map[string]int
	var passRateMetric map[string][]int
	failuresMetrics := make(map[string][][]*metrics.TestId)
	var wg sync.WaitGroup
	wg.Add(2 + len(runs))
	go func() {
		defer wg.Done()
		totals = compute.ComputeTotals(&resultsById)
	}()
	go func() {
		defer wg.Done()
		passRateMetric = compute.ComputePassRateMetric(len(runs),
			&resultsById, compute.OkAndUnknownOrPasses)
	}()
	for _, run := range runs {
		go func(browserName string) {
			defer wg.Done()
			// TODO: Check that browser names are different
			failuresMetrics[browserName] =
				compute.ComputeBrowserFailureList(len(runs),
					browserName, &resultsById,
					compute.OkAndUnknownOrPasses)
		}(run.BrowserName)
	}
	wg.Wait()

	log.Println("Computed metrics")
	log.Println("Uploading metrics")

	outputBucket := gcsClient.Bucket(*outputGcsBucket)
	datastoreClient, err := datastore.NewClient(ctx, *projectId)
	if err != nil {
		log.Fatal(err)
	}
	bigqueryClient, err := bigquery.NewClient(ctx, *projectId)
	if err != nil {
		log.Fatal(err)
	}
	outputters := [2]storage.Outputter{
		storage.GCSDatastoreContext{
			Context: ctx,
			Bucket: storage.Bucket{
				Name:   *outputGcsBucket,
				Handle: outputBucket,
			},
			Client: datastoreClient,
		},
		storage.BQContext{
			Context: ctx,
			Client:  bigqueryClient,
		},
	}

	gcsDir := fmt.Sprintf("%d-%d", readStartTime.Unix(),
		readEndTime.Unix())
	passRatesBasename := "pass-rates"
	passRateGCSPath := fmt.Sprintf("%s/%s.json.gz", gcsDir,
		passRatesBasename)
	passRatesUrl := fmt.Sprintf(
		"https://storage.cloud.google.com/%s/%s",
		*outputGcsBucket,
		passRateGCSPath)
	failuresBasenamef := func(browserName string) string {
		return fmt.Sprintf("%s-failures", browserName)
	}
	failuresGCSPathf := func(browserName string) string {
		return fmt.Sprintf("%s/%s.json.gz", gcsDir,
			failuresBasenamef(browserName))
	}
	failuresUrlf := func(browserName string) string {
		return fmt.Sprintf(
			"https://storage.cloud.google.com/%s/%s",
			*outputGcsBucket,
			failuresGCSPathf(browserName))
	}
	passRateMetadata := metrics.PassRateMetadata{
		StartTime: readStartTime,
		EndTime:   readEndTime,
		TestRuns:  runs,
		DataUrl:   passRatesUrl,
	}

	wg.Add((1 + len(failuresMetrics)) * len(outputters))
	processUploadErrors := func(errs []error) {
		for _, err := range errs {
			log.Printf("Upload error: %v", err)
		}
		if len(errs) > 0 {
			log.Fatal(errs[len(errs)-1])
		}
	}
	for _, outputter := range outputters {
		go func(outputter storage.Outputter) {
			defer wg.Done()
			outputId := storage.OutputId{
				MetadataLocation: storage.OutputLocation{
					BQDatasetName: *outputBQMetadataDataset,
					BQTableName:   *outputBQPassRateMetadataTable,
				},
				DataLocation: storage.OutputLocation{
					GCSObjectPath: passRateGCSPath,
					BQDatasetName: *outputBQDataDataset,
					BQTableName:   *outputBQPassRateTable,
				},
			}
			_, _, errs := uploadTotalsAndPassRateMetric(
				&passRateMetadata, outputter, outputId, totals,
				passRateMetric)
			processUploadErrors(errs)
		}(outputter)
		for browserName, failuresMetric := range failuresMetrics {
			go func(browserName string, failuresMetric [][]*metrics.TestId, outputter storage.Outputter) {
				defer wg.Done()
				failuresMetadata := metrics.FailuresMetadata{
					StartTime:   readStartTime,
					EndTime:     readEndTime,
					TestRuns:    runs,
					DataUrl:     failuresUrlf(browserName),
					BrowserName: browserName,
				}
				outputId := storage.OutputId{
					MetadataLocation: storage.OutputLocation{
						BQDatasetName: *outputBQMetadataDataset,
						BQTableName:   *outputBQFailuresMetadataTable,
					},
					DataLocation: storage.OutputLocation{
						GCSObjectPath: gcsDir +
							"/" +
							failuresBasenamef(browserName) +
							".json.gz",
						BQDatasetName: *outputBQDataDataset,
						BQTableName:   *outputBQFailuresTable,
					},
				}
				_, _, errs := uploadFailureLists(&failuresMetadata,
					outputter, outputId, browserName,
					failuresMetric)
				processUploadErrors(errs)
			}(browserName, failuresMetric, outputter)
		}
	}
	wg.Wait()

	log.Printf("Uploaded metrics")
}

func getRuns() []base.TestRun {
	url := "https://" + *wptdHost + "/api/runs"
	resp, err := http.Get(url)
	if err != nil {
		log.Fatal(err)
	}
	if resp.StatusCode != 200 {
		log.Fatal(errors.New("Bad response code from " + url + ": " +
			strconv.Itoa(resp.StatusCode)))
	}
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Fatal(err)
	}
	var runs []base.TestRun
	if err := json.Unmarshal(body, &runs); err != nil {
		log.Fatal(err)
	}
	return runs
}

func failureListsToRows(browserName string, failureLists [][]*metrics.TestId) (
	rows []interface{}) {
	type FailureListsRow struct {
		BrowserName      string         `json:"browser_name"`
		NumOtherFailures int            `json:"num_other_failures"`
		Tests            metrics.TestId `json:"test"`
	}
	numRows := 0
	for _, failureList := range failureLists {
		numRows += len(failureList)
	}
	rows = make([]interface{}, 0, numRows)
	for i, failuresPtrList := range failureLists {
		for _, failure := range failuresPtrList {
			rows = append(rows, FailureListsRow{
				browserName,
				i,
				*failure,
			})
		}
	}
	return rows
}

func totalsAndPassRateMetricToRows(totals map[string]int,
	passRateMetric map[string][]int) (
	rows []interface{}) {
	type PassRateMetricRow struct {
		Dir       string `json:"dir"`
		PassRates []int  `json:"pass_rates"`
		Total     int    `json:"total"`
	}
	rows = make([]interface{}, 0, len(passRateMetric))
	for dir, passRates := range passRateMetric {
		rows = append(rows, PassRateMetricRow{dir, passRates,
			totals[dir]})
	}
	return rows
}

func uploadTotalsAndPassRateMetric(metricsRun *metrics.PassRateMetadata,
	outputter storage.Outputter, id storage.OutputId,
	totals map[string]int, passRateMetric map[string][]int) (
	interface{}, []interface{}, []error) {
	rows := totalsAndPassRateMetricToRows(totals, passRateMetric)
	return outputter.Output(id, metricsRun, rows)
}

func uploadFailureLists(metricsRun *metrics.FailuresMetadata,
	outputter storage.Outputter, id storage.OutputId,
	browserName string, failureLists [][]*metrics.TestId) (
	interface{}, []interface{}, []error) {
	rows := failureListsToRows(browserName, failureLists)
	return outputter.Output(id, metricsRun, rows)
}
<commit_msg>Switch to fast, public (unauthenticated) endpoint for pulling GCS data in metrics collection code<commit_after>// Copyright 2017 The WPT Dashboard Project. 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 (
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"strconv"
	"sync"
	"time"

	"cloud.google.com/go/bigquery"
	"cloud.google.com/go/datastore"
	gcs "cloud.google.com/go/storage"
	"github.com/w3c/wptdashboard/metrics"
	"github.com/w3c/wptdashboard/metrics/compute"
	"github.com/w3c/wptdashboard/metrics/storage"
	base "github.com/w3c/wptdashboard/shared"
	"golang.org/x/net/context"
)

var wptDataPath *string
var projectId *string
var inputGcsBucket *string
var outputGcsBucket *string
var wptdHost *string
var outputBQMetadataDataset *string
var outputBQDataDataset *string
var outputBQPassRateTable *string
var outputBQPassRateMetadataTable *string
var outputBQFailuresTable *string
var outputBQFailuresMetadataTable *string

func init() {
	unixNow := time.Now().Unix()
	wptDataPath = flag.String("wpt_data_path", os.Getenv("HOME")+
		"/wpt-data", "Path to data directory for local data copied "+
		"from Google Cloud Storage")
	projectId = flag.String("project_id", "wptdashboard",
		"Google Cloud Platform project id")
	inputGcsBucket = flag.String("input_gcs_bucket", "wptd",
		"Google Cloud Storage bucket where test results are stored")
	outputGcsBucket = flag.String("output_gcs_bucket", "wptd-metrics",
		"Google Cloud Storage bucket where metrics are stored")
	outputBQMetadataDataset = flag.String("output_bq_metadata_dataset",
		fmt.Sprintf("wptd_metrics_%d", unixNow),
		"BigQuery dataset where metrics metadata are stored")
	outputBQDataDataset = flag.String("output_bq_data_dataset",
		fmt.Sprintf("wptd_metrics_%d", unixNow),
		"BigQuery dataset where metrics data are stored")
	outputBQPassRateTable = flag.String("output_bq_pass_rate_table",
		fmt.Sprintf("PassRates_%d", unixNow),
		"BigQuery table where pass rate metrics are stored")
	outputBQPassRateMetadataTable = flag.String("output_bq_pass_rate_metadata_table",
		fmt.Sprintf("PassRateMetadata_%d", unixNow),
		"BigQuery table where pass rate metrics are stored")
	outputBQFailuresTable = flag.String("output_bq_failures_table",
		fmt.Sprintf("Failures_%d", unixNow),
		"BigQuery table where test failure lists are stored")
	outputBQFailuresMetadataTable = flag.String("output_bq_failures_metadata_table",
		fmt.Sprintf("FailuresMetadata_%d", unixNow),
		"BigQuery table where pass rate metrics are stored")
	wptdHost = flag.String("wptd_host", "wpt.fyi",
		"Hostname of endpoint that serves WPT Dashboard data API")
}

func main() {
	log.SetFlags(log.LstdFlags | log.Lshortfile)

	logFileName := "current_metrics.log"
	logFile, err := os.OpenFile(logFileName, os.O_RDWR|os.O_CREATE|
		os.O_APPEND, 0666)
	if err != nil {
		log.Fatalf("Error opening log file: %v", err)
	}
	defer logFile.Close()
	log.Printf("Logs appended to %s\n", logFileName)
	log.SetOutput(logFile)

	ctx := context.Background()
	gcsClient, err := gcs.NewClient(ctx)
	if err != nil {
		log.Fatal(err)
	}
	inputBucket := gcsClient.Bucket(*inputGcsBucket)
	inputCtx := storage.GCSDatastoreContext{
		Context: ctx,
		Bucket: storage.Bucket{
			Name:   *inputGcsBucket,
			Handle: inputBucket,
		},
	}

	log.Println("Reading test results from Google Cloud Storage bucket: " +
		*inputGcsBucket)

	readStartTime := time.Now()
	runs := getRuns()
	allResults := storage.LoadTestRunResults(&inputCtx, runs)
	readEndTime := time.Now()

	log.Println("Read test results from Google Cloud Storage bucket: " +
		*inputGcsBucket)
	log.Println("Consolidating results")

	resultsById := compute.GatherResultsById(&allResults)

	log.Println("Consolidated results")
	log.Println("Computing metrics")

	var totals map[string]int
	var passRateMetric map[string][]int
	failuresMetrics := make(map[string][][]*metrics.TestId)
	var wg sync.WaitGroup
	wg.Add(2 + len(runs))
	go func() {
		defer wg.Done()
		totals = compute.ComputeTotals(&resultsById)
	}()
	go func() {
		defer wg.Done()
		passRateMetric = compute.ComputePassRateMetric(len(runs),
			&resultsById, compute.OkAndUnknownOrPasses)
	}()
	for _, run := range runs {
		go func(browserName string) {
			defer wg.Done()
			// TODO: Check that browser names are different
			failuresMetrics[browserName] =
				compute.ComputeBrowserFailureList(len(runs),
					browserName, &resultsById,
					compute.OkAndUnknownOrPasses)
		}(run.BrowserName)
	}
	wg.Wait()

	log.Println("Computed metrics")
	log.Println("Uploading metrics")

	outputBucket := gcsClient.Bucket(*outputGcsBucket)
	datastoreClient, err := datastore.NewClient(ctx, *projectId)
	if err != nil {
		log.Fatal(err)
	}
	bigqueryClient, err := bigquery.NewClient(ctx, *projectId)
	if err != nil {
		log.Fatal(err)
	}
	outputters := [2]storage.Outputter{
		storage.GCSDatastoreContext{
			Context: ctx,
			Bucket: storage.Bucket{
				Name:   *outputGcsBucket,
				Handle: outputBucket,
			},
			Client: datastoreClient,
		},
		storage.BQContext{
			Context: ctx,
			Client:  bigqueryClient,
		},
	}

	gcsDir := fmt.Sprintf("%d-%d", readStartTime.Unix(),
		readEndTime.Unix())
	passRatesBasename := "pass-rates"
	passRateGCSPath := fmt.Sprintf("%s/%s.json.gz", gcsDir,
		passRatesBasename)
	passRatesUrl := fmt.Sprintf(
		"https://storage.googleapis.com/%s/%s",
		*outputGcsBucket,
		passRateGCSPath)
	failuresBasenamef := func(browserName string) string {
		return fmt.Sprintf("%s-failures", browserName)
	}
	failuresGCSPathf := func(browserName string) string {
		return fmt.Sprintf("%s/%s.json.gz", gcsDir,
			failuresBasenamef(browserName))
	}
	failuresUrlf := func(browserName string) string {
		return fmt.Sprintf(
			"https://storage.googleapis.com/%s/%s",
			*outputGcsBucket,
			failuresGCSPathf(browserName))
	}
	passRateMetadata := metrics.PassRateMetadata{
		StartTime: readStartTime,
		EndTime:   readEndTime,
		TestRuns:  runs,
		DataUrl:   passRatesUrl,
	}

	wg.Add((1 + len(failuresMetrics)) * len(outputters))
	processUploadErrors := func(errs []error) {
		for _, err := range errs {
			log.Printf("Upload error: %v", err)
		}
		if len(errs) > 0 {
			log.Fatal(errs[len(errs)-1])
		}
	}
	for _, outputter := range outputters {
		go func(outputter storage.Outputter) {
			defer wg.Done()
			outputId := storage.OutputId{
				MetadataLocation: storage.OutputLocation{
					BQDatasetName: *outputBQMetadataDataset,
					BQTableName:   *outputBQPassRateMetadataTable,
				},
				DataLocation: storage.OutputLocation{
					GCSObjectPath: passRateGCSPath,
					BQDatasetName: *outputBQDataDataset,
					BQTableName:   *outputBQPassRateTable,
				},
			}
			_, _, errs := uploadTotalsAndPassRateMetric(
				&passRateMetadata, outputter, outputId, totals,
				passRateMetric)
			processUploadErrors(errs)
		}(outputter)
		for browserName, failuresMetric := range failuresMetrics {
			go func(browserName string, failuresMetric [][]*metrics.TestId, outputter storage.Outputter) {
				defer wg.Done()
				failuresMetadata := metrics.FailuresMetadata{
					StartTime:   readStartTime,
					EndTime:     readEndTime,
					TestRuns:    runs,
					DataUrl:     failuresUrlf(browserName),
					BrowserName: browserName,
				}
				outputId := storage.OutputId{
					MetadataLocation: storage.OutputLocation{
						BQDatasetName: *outputBQMetadataDataset,
						BQTableName:   *outputBQFailuresMetadataTable,
					},
					DataLocation: storage.OutputLocation{
						GCSObjectPath: gcsDir +
							"/" +
							failuresBasenamef(browserName) +
							".json.gz",
						BQDatasetName: *outputBQDataDataset,
						BQTableName:   *outputBQFailuresTable,
					},
				}
				_, _, errs := uploadFailureLists(&failuresMetadata,
					outputter, outputId, browserName,
					failuresMetric)
				processUploadErrors(errs)
			}(browserName, failuresMetric, outputter)
		}
	}
	wg.Wait()

	log.Printf("Uploaded metrics")
}

func getRuns() []base.TestRun {
	url := "https://" + *wptdHost + "/api/runs"
	resp, err := http.Get(url)
	if err != nil {
		log.Fatal(err)
	}
	if resp.StatusCode != 200 {
		log.Fatal(errors.New("Bad response code from " + url + ": " +
			strconv.Itoa(resp.StatusCode)))
	}
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Fatal(err)
	}
	var runs []base.TestRun
	if err := json.Unmarshal(body, &runs); err != nil {
		log.Fatal(err)
	}
	return runs
}

func failureListsToRows(browserName string, failureLists [][]*metrics.TestId) (
	rows []interface{}) {
	type FailureListsRow struct {
		BrowserName      string         `json:"browser_name"`
		NumOtherFailures int            `json:"num_other_failures"`
		Tests            metrics.TestId `json:"test"`
	}
	numRows := 0
	for _, failureList := range failureLists {
		numRows += len(failureList)
	}
	rows = make([]interface{}, 0, numRows)
	for i, failuresPtrList := range failureLists {
		for _, failure := range failuresPtrList {
			rows = append(rows, FailureListsRow{
				browserName,
				i,
				*failure,
			})
		}
	}
	return rows
}

func totalsAndPassRateMetricToRows(totals map[string]int,
	passRateMetric map[string][]int) (
	rows []interface{}) {
	type PassRateMetricRow struct {
		Dir       string `json:"dir"`
		PassRates []int  `json:"pass_rates"`
		Total     int    `json:"total"`
	}
	rows = make([]interface{}, 0, len(passRateMetric))
	for dir, passRates := range passRateMetric {
		rows = append(rows, PassRateMetricRow{dir, passRates,
			totals[dir]})
	}
	return rows
}

func uploadTotalsAndPassRateMetric(metricsRun *metrics.PassRateMetadata,
	outputter storage.Outputter, id storage.OutputId,
	totals map[string]int, passRateMetric map[string][]int) (
	interface{}, []interface{}, []error) {
	rows := totalsAndPassRateMetricToRows(totals, passRateMetric)
	return outputter.Output(id, metricsRun, rows)
}

func uploadFailureLists(metricsRun *metrics.FailuresMetadata,
	outputter storage.Outputter, id storage.OutputId,
	browserName string, failureLists [][]*metrics.TestId) (
	interface{}, []interface{}, []error) {
	rows := failureListsToRows(browserName, failureLists)
	return outputter.Output(id, metricsRun, rows)
}
<|endoftext|>


<commit_before>// Package mux provides a simple Discord message route multiplexer that
// parses messages and then executes a matching registered handler, if found.
// mux can be used with both Disgord and the DiscordGo library.
package mux

import (
	"fmt"
	"log"
	"regexp"
	"strings"

	"github.com/bwmarrin/discordgo"
)

// Route holds information about a specific message route handler
type Route struct {
	Pattern     string      // match pattern that should trigger this route handler
	Description string      // short description of this route
	Help        string      // detailed help string for this route
	Run         HandlerFunc // route handler function to call
}

// Context holds a bit of extra data we pass along to route handlers
// This way processing some of this only needs to happen once.
type Context struct {
	Fields          []string
	Content         string
	GuildID         string
	IsDirected      bool
	IsPrivate       bool
	HasPrefix       bool
	HasMention      bool
	HasMentionFirst bool
}

// HandlerFunc is the function signature required for a message route handler.
type HandlerFunc func(*discordgo.Session, *discordgo.Message, *Context)

// Mux is the main struct for all mux methods.
type Mux struct {
	Routes  []*Route
	Default *Route
	Prefix  string
}

// New returns a new Discord message route mux
func New() *Mux {
	m := &Mux{}
	m.Prefix = "-dg "
	return m
}

// Route allows you to register a route
func (m *Mux) Route(pattern, desc string, cb HandlerFunc) (*Route, error) {

	r := Route{}
	r.Pattern = pattern
	r.Description = desc
	r.Run = cb
	m.Routes = append(m.Routes, &r)

	return &r, nil
}

// FuzzyMatch attepts to find the best route match for a givin message.
func (m *Mux) FuzzyMatch(msg string) (*Route, []string) {

	// Tokenize the msg string into a slice of words
	fields := strings.Fields(msg)

	// no point to continue if there's no fields
	if len(fields) == 0 {
		return nil, nil
	}

	// Search though the command list for a match
	var r *Route
	var rank int

	var fk int
	for fk, fv := range fields {

		for _, rv := range m.Routes {

			// If we find an exact match, return that immediately.
			if rv.Pattern == fv {
				return rv, fields[fk:]
			}

			// Some "Fuzzy" searching...
			if strings.HasPrefix(rv.Pattern, fv) {
				if len(fv) > rank {
					r = rv
					rank = len(fv)
				}
			}
		}
	}
	return r, fields[fk:]
}

// OnMessageCreate is a DiscordGo Event Handler function.  This must be
// registered using the DiscordGo.Session.AddHandler function.  This function
// will receive all Discord messages and parse them for matches to registered
// routes.
func (m *Mux) OnMessageCreate(ds *discordgo.Session, mc *discordgo.MessageCreate) {

	var err error

	// Ignore all messages created by the Bot account itself
	if mc.Author.ID == ds.State.User.ID {
		return
	}

	// Create Context struct that we can put various infos into
	ctx := &Context{
		Content: strings.TrimSpace(mc.Content),
	}

	// Fetch the channel for this Message
	var c *discordgo.Channel
	c, err = ds.State.Channel(mc.ChannelID)
	if err != nil {
		// Try fetching via REST API
		c, err = ds.Channel(mc.ChannelID)
		if err != nil {
			log.Printf("unable to fetch Channel for Message,", err)
		} else {
			// Attempt to add this channel into our State
			err = ds.State.ChannelAdd(c)
			if err != nil {
				log.Printf("error updating State with Channel,", err)
			}
			// Add Channel info into Context
			ctx.GuildID = c.GuildID
			if c.Type == discordgo.ChannelTypeDM {
				ctx.IsPrivate = true
				ctx.IsDirected = true
			}
		}
	}

	// Detect @name or @nick mentions
	if !ctx.IsDirected {

		// Detect if Bot was @mentioned
		for _, v := range mc.Mentions {

			if v.ID == ds.State.User.ID {

				ctx.IsDirected, ctx.HasMention = true, true

				reg := regexp.MustCompile(fmt.Sprintf("<@!?(%s)>", ds.State.User.ID))

				// Was the @mention the first part of the string?
				if reg.FindStringIndex(ctx.Content)[0] == 0 {
					ctx.HasMentionFirst = true
				}

				// strip bot mention tags from content string
				ctx.Content = reg.ReplaceAllString(ctx.Content, "")

				break
			}
		}
	}

	// Detect prefix mention
	if !ctx.IsDirected && len(m.Prefix) > 0 {

		// TODO : Must be changed to support a per-guild user defined prefix
		if strings.HasPrefix(ctx.Content, m.Prefix) {
			ctx.IsDirected, ctx.HasPrefix, ctx.HasMentionFirst = true, true, true
			ctx.Content = strings.TrimPrefix(ctx.Content, m.Prefix)
		}
	}

	// For now, if we're not specifically mentioned we do nothing.
	// later I might add an option for global non-mentioned command words
	if !ctx.IsDirected {
		return
	}

	// Try to find the "best match" command out of the message.
	r, fl := m.FuzzyMatch(ctx.Content)
	if r != nil {
		ctx.Fields = fl
		r.Run(ds, mc.Message, ctx)
		return
	}

	// If no command match was found, call the default.
	// Ignore if only @mentioned in the middle of a message
	if m.Default != nil && (ctx.HasMentionFirst) {
		// TODO: This could use a ratelimit
		// or should the ratelimit be inside the cmd handler?..
		// In the case of "talking" to another bot, this can create an endless
		// loop.  Probably most common in private messages.
		m.Default.Run(ds, mc.Message, ctx)
	}

}
<commit_msg>Removed GuildID from mux ctx<commit_after>// Package mux provides a simple Discord message route multiplexer that
// parses messages and then executes a matching registered handler, if found.
// mux can be used with both Disgord and the DiscordGo library.
package mux

import (
	"fmt"
	"log"
	"regexp"
	"strings"

	"github.com/bwmarrin/discordgo"
)

// Route holds information about a specific message route handler
type Route struct {
	Pattern     string      // match pattern that should trigger this route handler
	Description string      // short description of this route
	Help        string      // detailed help string for this route
	Run         HandlerFunc // route handler function to call
}

// Context holds a bit of extra data we pass along to route handlers
// This way processing some of this only needs to happen once.
type Context struct {
	Fields          []string
	Content         string
	IsDirected      bool
	IsPrivate       bool
	HasPrefix       bool
	HasMention      bool
	HasMentionFirst bool
}

// HandlerFunc is the function signature required for a message route handler.
type HandlerFunc func(*discordgo.Session, *discordgo.Message, *Context)

// Mux is the main struct for all mux methods.
type Mux struct {
	Routes  []*Route
	Default *Route
	Prefix  string
}

// New returns a new Discord message route mux
func New() *Mux {
	m := &Mux{}
	m.Prefix = "-dg "
	return m
}

// Route allows you to register a route
func (m *Mux) Route(pattern, desc string, cb HandlerFunc) (*Route, error) {

	r := Route{}
	r.Pattern = pattern
	r.Description = desc
	r.Run = cb
	m.Routes = append(m.Routes, &r)

	return &r, nil
}

// FuzzyMatch attepts to find the best route match for a givin message.
func (m *Mux) FuzzyMatch(msg string) (*Route, []string) {

	// Tokenize the msg string into a slice of words
	fields := strings.Fields(msg)

	// no point to continue if there's no fields
	if len(fields) == 0 {
		return nil, nil
	}

	// Search though the command list for a match
	var r *Route
	var rank int

	var fk int
	for fk, fv := range fields {

		for _, rv := range m.Routes {

			// If we find an exact match, return that immediately.
			if rv.Pattern == fv {
				return rv, fields[fk:]
			}

			// Some "Fuzzy" searching...
			if strings.HasPrefix(rv.Pattern, fv) {
				if len(fv) > rank {
					r = rv
					rank = len(fv)
				}
			}
		}
	}
	return r, fields[fk:]
}

// OnMessageCreate is a DiscordGo Event Handler function.  This must be
// registered using the DiscordGo.Session.AddHandler function.  This function
// will receive all Discord messages and parse them for matches to registered
// routes.
func (m *Mux) OnMessageCreate(ds *discordgo.Session, mc *discordgo.MessageCreate) {

	var err error

	// Ignore all messages created by the Bot account itself
	if mc.Author.ID == ds.State.User.ID {
		return
	}

	// Create Context struct that we can put various infos into
	ctx := &Context{
		Content: strings.TrimSpace(mc.Content),
	}

	// Fetch the channel for this Message
	var c *discordgo.Channel
	c, err = ds.State.Channel(mc.ChannelID)
	if err != nil {
		// Try fetching via REST API
		c, err = ds.Channel(mc.ChannelID)
		if err != nil {
			log.Printf("unable to fetch Channel for Message,", err)
		} else {
			// Attempt to add this channel into our State
			err = ds.State.ChannelAdd(c)
			if err != nil {
				log.Printf("error updating State with Channel,", err)
			}
			// Add Channel info into Context
			if c.Type == discordgo.ChannelTypeDM {
				ctx.IsPrivate = true
				ctx.IsDirected = true
			}
		}
	}

	// Detect @name or @nick mentions
	if !ctx.IsDirected {

		// Detect if Bot was @mentioned
		for _, v := range mc.Mentions {

			if v.ID == ds.State.User.ID {

				ctx.IsDirected, ctx.HasMention = true, true

				reg := regexp.MustCompile(fmt.Sprintf("<@!?(%s)>", ds.State.User.ID))

				// Was the @mention the first part of the string?
				if reg.FindStringIndex(ctx.Content)[0] == 0 {
					ctx.HasMentionFirst = true
				}

				// strip bot mention tags from content string
				ctx.Content = reg.ReplaceAllString(ctx.Content, "")

				break
			}
		}
	}

	// Detect prefix mention
	if !ctx.IsDirected && len(m.Prefix) > 0 {

		// TODO : Must be changed to support a per-guild user defined prefix
		if strings.HasPrefix(ctx.Content, m.Prefix) {
			ctx.IsDirected, ctx.HasPrefix, ctx.HasMentionFirst = true, true, true
			ctx.Content = strings.TrimPrefix(ctx.Content, m.Prefix)
		}
	}

	// For now, if we're not specifically mentioned we do nothing.
	// later I might add an option for global non-mentioned command words
	if !ctx.IsDirected {
		return
	}

	// Try to find the "best match" command out of the message.
	r, fl := m.FuzzyMatch(ctx.Content)
	if r != nil {
		ctx.Fields = fl
		r.Run(ds, mc.Message, ctx)
		return
	}

	// If no command match was found, call the default.
	// Ignore if only @mentioned in the middle of a message
	if m.Default != nil && (ctx.HasMentionFirst) {
		// TODO: This could use a ratelimit
		// or should the ratelimit be inside the cmd handler?..
		// In the case of "talking" to another bot, this can create an endless
		// loop.  Probably most common in private messages.
		m.Default.Run(ds, mc.Message, ctx)
	}

}
<|endoftext|>


<commit_before>package mgnstr

import (
	"testing"
	"github.com/stretchr/testify/assert"
)

func TestIncludesAny(t *testing.T) {
	targets := []string{"hage", "hoge", "hug"}
	assert.True(t,  IncludesAny("hoge", targets))
	assert.False(t, IncludesAny("ha",   targets))
}

func TestContainsAny(t *testing.T) {
	targets := []string{"__MACOSX", ".DS_Store"}
	trues := []string{"__MACOSX", "__MACOSX/hogehoge", "hoge/.DS_Store"}

	for _, sample := range trues {
		assert.True(t, ContainsAny(sample, targets))
	}

	assert.False(t, ContainsAny("hoge", targets))
}

func TestUniversalExt(t *testing.T) {
	samples := []string{"hoge.txt", "hoge.TXT"}
	for _, sample := range samples {
		assert.Equal(t,  UniversalExt(sample), ".txt")
	}

	assert.NotEqual(t,  UniversalExt("hoge.TXT"), ".TXT")
}

func TestGetUUID(t *testing.T) {
	sample := GetUUID()
	assert.Equal(t, len(sample), 36)
	assert.NotEqual(t, GetUUID(), sample)
}

func TestGet8UUID(t *testing.T) {
	sample := Get8UUID()
	assert.Equal(t, len(sample), 8)
	assert.NotEqual(t, Get8UUID(), sample)
}
<commit_msg>add TestGetSha256<commit_after>package mgnstr

import (
	"testing"
	"github.com/stretchr/testify/assert"
)

func TestIncludesAny(t *testing.T) {
	targets := []string{"hage", "hoge", "hug"}
	assert.True(t,  IncludesAny("hoge", targets))
	assert.False(t, IncludesAny("ha",   targets))
}

func TestContainsAny(t *testing.T) {
	targets := []string{"__MACOSX", ".DS_Store"}
	trues := []string{"__MACOSX", "__MACOSX/hogehoge", "hoge/.DS_Store"}

	for _, sample := range trues {
		assert.True(t, ContainsAny(sample, targets))
	}

	assert.False(t, ContainsAny("hoge", targets))
}

func TestUniversalExt(t *testing.T) {
	samples := []string{"hoge.txt", "hoge.TXT"}
	for _, sample := range samples {
		assert.Equal(t,  UniversalExt(sample), ".txt")
	}

	assert.NotEqual(t,  UniversalExt("hoge.TXT"), ".TXT")
}

func TestGetUUID(t *testing.T) {
	sample := GetUUID()
	assert.Equal(t, len(sample), 36)
	assert.NotEqual(t, GetUUID(), sample)
}

func TestGet8UUID(t *testing.T) {
	sample := Get8UUID()
	assert.Equal(t, len(sample), 8)
	assert.NotEqual(t, Get8UUID(), sample)
}

func TestGetSha256(t *testing.T) {
	sample, err := GetSha256()
	assert.Nil(t, err)
	assert.Equal(t, len(sample), 32)
	sample2, _ :=  GetSha256()
	assert.NotEqual(t, sample2, sample)
}
<|endoftext|>


<commit_before>// Copyright 2015 The etcd Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package command

import (
	"bufio"
	"context"
	"errors"
	"fmt"
	"os"
	"os/exec"
	"strings"

	"github.com/coreos/etcd/clientv3"

	"github.com/spf13/cobra"
)

var (
	errBadArgsNum              = errors.New("bad number of arguments")
	errBadArgsNumConflictEnv   = errors.New("bad number of arguments (found conflicting environment key)")
	errBadArgsNumSeparator     = errors.New("bad number of arguments (found separator --, but no commands)")
	errBadArgsInteractiveWatch = errors.New("args[0] must be 'watch' for interactive calls")
)

var (
	watchRev         int64
	watchPrefix      bool
	watchInteractive bool
	watchPrevKey     bool
)

// NewWatchCommand returns the cobra command for "watch".
func NewWatchCommand() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "watch [options] [key or prefix] [range_end] [--] [exec-command arg1 arg2 ...]",
		Short: "Watches events stream on keys or prefixes",
		Run:   watchCommandFunc,
	}

	cmd.Flags().BoolVarP(&watchInteractive, "interactive", "i", false, "Interactive mode")
	cmd.Flags().BoolVar(&watchPrefix, "prefix", false, "Watch on a prefix if prefix is set")
	cmd.Flags().Int64Var(&watchRev, "rev", 0, "Revision to start watching")
	cmd.Flags().BoolVar(&watchPrevKey, "prev-kv", false, "get the previous key-value pair before the event happens")

	return cmd
}

// watchCommandFunc executes the "watch" command.
func watchCommandFunc(cmd *cobra.Command, args []string) {
	envKey, envRange := os.Getenv("ETCDCTL_WATCH_KEY"), os.Getenv("ETCDCTL_WATCH_RANGE_END")
	if envKey == "" && envRange != "" {
		ExitWithError(ExitBadArgs, fmt.Errorf("ETCDCTL_WATCH_KEY is empty but got ETCDCTL_WATCH_RANGE_END=%q", envRange))
	}

	if watchInteractive {
		watchInteractiveFunc(cmd, os.Args, envKey, envRange)
		return
	}

	watchArgs, execArgs, err := parseWatchArgs(os.Args, args, envKey, envRange, false)
	if err != nil {
		ExitWithError(ExitBadArgs, err)
	}

	c := mustClientFromCmd(cmd)
	wc, err := getWatchChan(c, watchArgs)
	if err != nil {
		ExitWithError(ExitBadArgs, err)
	}

	printWatchCh(c, wc, execArgs)
	if err = c.Close(); err != nil {
		ExitWithError(ExitBadConnection, err)
	}
	ExitWithError(ExitInterrupted, fmt.Errorf("watch is canceled by the server"))
}

func watchInteractiveFunc(cmd *cobra.Command, osArgs []string, envKey, envRange string) {
	c := mustClientFromCmd(cmd)

	reader := bufio.NewReader(os.Stdin)

	for {
		l, err := reader.ReadString('\n')
		if err != nil {
			ExitWithError(ExitInvalidInput, fmt.Errorf("Error reading watch request line: %v", err))
		}
		l = strings.TrimSuffix(l, "\n")

		args := argify(l)
		if len(args) < 2 && envKey == "" {
			fmt.Fprintf(os.Stderr, "Invalid command %s (command type or key is not provided)\n", l)
			continue
		}

		if args[0] != "watch" {
			fmt.Fprintf(os.Stderr, "Invalid command %s (only support watch)\n", l)
			continue
		}

		watchArgs, execArgs, perr := parseWatchArgs(osArgs, args, envKey, envRange, true)
		if perr != nil {
			ExitWithError(ExitBadArgs, perr)
		}

		ch, err := getWatchChan(c, watchArgs)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Invalid command %s (%v)\n", l, err)
			continue
		}
		go printWatchCh(c, ch, execArgs)
	}
}

func getWatchChan(c *clientv3.Client, args []string) (clientv3.WatchChan, error) {
	if len(args) < 1 {
		return nil, errBadArgsNum
	}

	key := args[0]
	opts := []clientv3.OpOption{clientv3.WithRev(watchRev)}
	if len(args) == 2 {
		if watchPrefix {
			return nil, fmt.Errorf("`range_end` and `--prefix` are mutually exclusive")
		}
		opts = append(opts, clientv3.WithRange(args[1]))
	}
	if watchPrefix {
		opts = append(opts, clientv3.WithPrefix())
	}
	if watchPrevKey {
		opts = append(opts, clientv3.WithPrevKV())
	}
	return c.Watch(clientv3.WithRequireLeader(context.Background()), key, opts...), nil
}

func printWatchCh(c *clientv3.Client, ch clientv3.WatchChan, execArgs []string) {
	for resp := range ch {
		if resp.Canceled {
			fmt.Fprintf(os.Stderr, "watch was canceled (%v)\n", resp.Err())
		}
		display.Watch(resp)

		if len(execArgs) > 0 {
			for _, ev := range resp.Events {
				cmd := exec.CommandContext(c.Ctx(), execArgs[0], execArgs[1:]...)
				cmd.Env = os.Environ()
				cmd.Env = append(cmd.Env, fmt.Sprintf("ETCD_WATCH_KEY=%q", ev.Kv.Key))
				cmd.Env = append(cmd.Env, fmt.Sprintf("ETCD_WATCH_VALUE=%q", ev.Kv.Value))
				cmd.Stdout, cmd.Stderr = os.Stdout, os.Stderr
				if err := cmd.Run(); err != nil {
					fmt.Fprintf(os.Stderr, "command %q error (%v)\n", execArgs, err)
				}
			}
		}
	}
}

// "commandArgs" is the command arguments after "spf13/cobra" parses
// all "watch" command flags, strips out special characters (e.g. "--").
// "orArgs" is the raw arguments passed to "watch" command
// (e.g. ./bin/etcdctl watch foo --rev 1 bar).
// "--" characters are invalid arguments for "spf13/cobra" library,
// so no need to handle such cases.
func parseWatchArgs(osArgs, commandArgs []string, envKey, envRange string, interactive bool) (watchArgs []string, execArgs []string, err error) {
	watchArgs = commandArgs

	// remove preceding commands (e.g. "watch foo bar" in interactive mode)
	idx := 0
	for idx = range watchArgs {
		if watchArgs[idx] == "watch" {
			break
		}
	}
	if idx < len(watchArgs)-1 || envKey != "" {
		if idx < len(watchArgs)-1 {
			watchArgs = watchArgs[idx+1:]
		}

		execIdx, execExist := 0, false
		for execIdx = range osArgs {
			v := osArgs[execIdx]
			if v == "--" && execIdx != len(osArgs)-1 {
				execExist = true
				break
			}
		}

		if idx == len(watchArgs)-1 && envKey != "" {
			if len(watchArgs) > 0 && !interactive {
				// "watch --rev 1 -- echo Hello World" has no conflict
				if !execExist {
					// "watch foo" with ETCDCTL_WATCH_KEY=foo
					// (watchArgs==["foo"])
					return nil, nil, errBadArgsNumConflictEnv
				}
			}
			// otherwise, watch with no argument and environment key is set
			// if interactive, first "watch" command string should be removed
			if interactive {
				watchArgs = []string{}
			}
		}

		// "watch foo -- echo hello" with ETCDCTL_WATCH_KEY=foo
		// (watchArgs==["foo","echo","hello"])
		if envKey != "" && execExist {
			widx, oidx := 0, len(osArgs)-1
			for widx = len(watchArgs) - 1; widx >= 0; widx-- {
				if watchArgs[widx] == osArgs[oidx] {
					oidx--
					continue
				}
				if oidx == execIdx { // watchArgs has extra
					return nil, nil, errBadArgsNumConflictEnv
				}
			}
		}
	} else if interactive { // "watch" not found
		return nil, nil, errBadArgsInteractiveWatch
	}
	if len(watchArgs) < 1 && envKey == "" {
		return nil, nil, errBadArgsNum
	}

	// remove preceding commands (e.g. ./bin/etcdctl watch)
	for idx = range osArgs {
		if osArgs[idx] == "watch" {
			break
		}
	}
	if idx < len(osArgs)-1 {
		osArgs = osArgs[idx+1:]
	} else if envKey == "" {
		return nil, nil, errBadArgsNum
	}

	argsWithSep := osArgs
	if interactive { // interactive mode pass "--" to the command args
		argsWithSep = watchArgs
	}
	foundSep := false
	for idx = range argsWithSep {
		if argsWithSep[idx] == "--" {
			foundSep = true
			break
		}
	}
	if interactive {
		flagset := NewWatchCommand().Flags()
		if err := flagset.Parse(argsWithSep); err != nil {
			return nil, nil, err
		}
		watchArgs = flagset.Args()
	}

	// "watch -- echo hello" with ETCDCTL_WATCH_KEY=foo
	// should be translated to "watch foo -- echo hello"
	// (watchArgs=["echo","hello"] should be ["foo","echo","hello"])
	if envKey != "" {
		tmp := []string{envKey}
		if envRange != "" {
			tmp = append(tmp, envRange)
		}
		watchArgs = append(tmp, watchArgs...)
	}

	if !foundSep {
		return watchArgs, nil, nil
	}

	if idx == len(argsWithSep)-1 {
		// "watch foo bar --" should error
		return nil, nil, errBadArgsNumSeparator
	}
	execArgs = argsWithSep[idx+1:]

	// "watch foo bar --rev 1 -- echo hello" or "watch foo --rev 1 bar -- echo hello",
	// then "watchArgs" is "foo bar echo hello"
	// so need ignore args after "argsWithSep[idx]", which is "--"
	endIdx := 0
	for endIdx = len(watchArgs) - 1; endIdx >= 0; endIdx-- {
		if watchArgs[endIdx] == argsWithSep[idx+1] {
			break
		}
	}
	watchArgs = watchArgs[:endIdx]

	return watchArgs, execArgs, nil
}
<commit_msg>ctlv3: exit on exec watch error<commit_after>// Copyright 2015 The etcd Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package command

import (
	"bufio"
	"context"
	"errors"
	"fmt"
	"os"
	"os/exec"
	"strings"

	"github.com/coreos/etcd/clientv3"

	"github.com/spf13/cobra"
)

var (
	errBadArgsNum              = errors.New("bad number of arguments")
	errBadArgsNumConflictEnv   = errors.New("bad number of arguments (found conflicting environment key)")
	errBadArgsNumSeparator     = errors.New("bad number of arguments (found separator --, but no commands)")
	errBadArgsInteractiveWatch = errors.New("args[0] must be 'watch' for interactive calls")
)

var (
	watchRev         int64
	watchPrefix      bool
	watchInteractive bool
	watchPrevKey     bool
)

// NewWatchCommand returns the cobra command for "watch".
func NewWatchCommand() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "watch [options] [key or prefix] [range_end] [--] [exec-command arg1 arg2 ...]",
		Short: "Watches events stream on keys or prefixes",
		Run:   watchCommandFunc,
	}

	cmd.Flags().BoolVarP(&watchInteractive, "interactive", "i", false, "Interactive mode")
	cmd.Flags().BoolVar(&watchPrefix, "prefix", false, "Watch on a prefix if prefix is set")
	cmd.Flags().Int64Var(&watchRev, "rev", 0, "Revision to start watching")
	cmd.Flags().BoolVar(&watchPrevKey, "prev-kv", false, "get the previous key-value pair before the event happens")

	return cmd
}

// watchCommandFunc executes the "watch" command.
func watchCommandFunc(cmd *cobra.Command, args []string) {
	envKey, envRange := os.Getenv("ETCDCTL_WATCH_KEY"), os.Getenv("ETCDCTL_WATCH_RANGE_END")
	if envKey == "" && envRange != "" {
		ExitWithError(ExitBadArgs, fmt.Errorf("ETCDCTL_WATCH_KEY is empty but got ETCDCTL_WATCH_RANGE_END=%q", envRange))
	}

	if watchInteractive {
		watchInteractiveFunc(cmd, os.Args, envKey, envRange)
		return
	}

	watchArgs, execArgs, err := parseWatchArgs(os.Args, args, envKey, envRange, false)
	if err != nil {
		ExitWithError(ExitBadArgs, err)
	}

	c := mustClientFromCmd(cmd)
	wc, err := getWatchChan(c, watchArgs)
	if err != nil {
		ExitWithError(ExitBadArgs, err)
	}

	printWatchCh(c, wc, execArgs)
	if err = c.Close(); err != nil {
		ExitWithError(ExitBadConnection, err)
	}
	ExitWithError(ExitInterrupted, fmt.Errorf("watch is canceled by the server"))
}

func watchInteractiveFunc(cmd *cobra.Command, osArgs []string, envKey, envRange string) {
	c := mustClientFromCmd(cmd)

	reader := bufio.NewReader(os.Stdin)

	for {
		l, err := reader.ReadString('\n')
		if err != nil {
			ExitWithError(ExitInvalidInput, fmt.Errorf("Error reading watch request line: %v", err))
		}
		l = strings.TrimSuffix(l, "\n")

		args := argify(l)
		if len(args) < 2 && envKey == "" {
			fmt.Fprintf(os.Stderr, "Invalid command %s (command type or key is not provided)\n", l)
			continue
		}

		if args[0] != "watch" {
			fmt.Fprintf(os.Stderr, "Invalid command %s (only support watch)\n", l)
			continue
		}

		watchArgs, execArgs, perr := parseWatchArgs(osArgs, args, envKey, envRange, true)
		if perr != nil {
			ExitWithError(ExitBadArgs, perr)
		}

		ch, err := getWatchChan(c, watchArgs)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Invalid command %s (%v)\n", l, err)
			continue
		}
		go printWatchCh(c, ch, execArgs)
	}
}

func getWatchChan(c *clientv3.Client, args []string) (clientv3.WatchChan, error) {
	if len(args) < 1 {
		return nil, errBadArgsNum
	}

	key := args[0]
	opts := []clientv3.OpOption{clientv3.WithRev(watchRev)}
	if len(args) == 2 {
		if watchPrefix {
			return nil, fmt.Errorf("`range_end` and `--prefix` are mutually exclusive")
		}
		opts = append(opts, clientv3.WithRange(args[1]))
	}
	if watchPrefix {
		opts = append(opts, clientv3.WithPrefix())
	}
	if watchPrevKey {
		opts = append(opts, clientv3.WithPrevKV())
	}
	return c.Watch(clientv3.WithRequireLeader(context.Background()), key, opts...), nil
}

func printWatchCh(c *clientv3.Client, ch clientv3.WatchChan, execArgs []string) {
	for resp := range ch {
		if resp.Canceled {
			fmt.Fprintf(os.Stderr, "watch was canceled (%v)\n", resp.Err())
		}
		display.Watch(resp)

		if len(execArgs) > 0 {
			for _, ev := range resp.Events {
				cmd := exec.CommandContext(c.Ctx(), execArgs[0], execArgs[1:]...)
				cmd.Env = os.Environ()
				cmd.Env = append(cmd.Env, fmt.Sprintf("ETCD_WATCH_KEY=%q", ev.Kv.Key))
				cmd.Env = append(cmd.Env, fmt.Sprintf("ETCD_WATCH_VALUE=%q", ev.Kv.Value))
				cmd.Stdout, cmd.Stderr = os.Stdout, os.Stderr
				if err := cmd.Run(); err != nil {
					fmt.Fprintf(os.Stderr, "command %q error (%v)\n", execArgs, err)
					os.Exit(1)
				}
			}
		}
	}
}

// "commandArgs" is the command arguments after "spf13/cobra" parses
// all "watch" command flags, strips out special characters (e.g. "--").
// "orArgs" is the raw arguments passed to "watch" command
// (e.g. ./bin/etcdctl watch foo --rev 1 bar).
// "--" characters are invalid arguments for "spf13/cobra" library,
// so no need to handle such cases.
func parseWatchArgs(osArgs, commandArgs []string, envKey, envRange string, interactive bool) (watchArgs []string, execArgs []string, err error) {
	watchArgs = commandArgs

	// remove preceding commands (e.g. "watch foo bar" in interactive mode)
	idx := 0
	for idx = range watchArgs {
		if watchArgs[idx] == "watch" {
			break
		}
	}
	if idx < len(watchArgs)-1 || envKey != "" {
		if idx < len(watchArgs)-1 {
			watchArgs = watchArgs[idx+1:]
		}

		execIdx, execExist := 0, false
		for execIdx = range osArgs {
			v := osArgs[execIdx]
			if v == "--" && execIdx != len(osArgs)-1 {
				execExist = true
				break
			}
		}

		if idx == len(watchArgs)-1 && envKey != "" {
			if len(watchArgs) > 0 && !interactive {
				// "watch --rev 1 -- echo Hello World" has no conflict
				if !execExist {
					// "watch foo" with ETCDCTL_WATCH_KEY=foo
					// (watchArgs==["foo"])
					return nil, nil, errBadArgsNumConflictEnv
				}
			}
			// otherwise, watch with no argument and environment key is set
			// if interactive, first "watch" command string should be removed
			if interactive {
				watchArgs = []string{}
			}
		}

		// "watch foo -- echo hello" with ETCDCTL_WATCH_KEY=foo
		// (watchArgs==["foo","echo","hello"])
		if envKey != "" && execExist {
			widx, oidx := 0, len(osArgs)-1
			for widx = len(watchArgs) - 1; widx >= 0; widx-- {
				if watchArgs[widx] == osArgs[oidx] {
					oidx--
					continue
				}
				if oidx == execIdx { // watchArgs has extra
					return nil, nil, errBadArgsNumConflictEnv
				}
			}
		}
	} else if interactive { // "watch" not found
		return nil, nil, errBadArgsInteractiveWatch
	}
	if len(watchArgs) < 1 && envKey == "" {
		return nil, nil, errBadArgsNum
	}

	// remove preceding commands (e.g. ./bin/etcdctl watch)
	for idx = range osArgs {
		if osArgs[idx] == "watch" {
			break
		}
	}
	if idx < len(osArgs)-1 {
		osArgs = osArgs[idx+1:]
	} else if envKey == "" {
		return nil, nil, errBadArgsNum
	}

	argsWithSep := osArgs
	if interactive { // interactive mode pass "--" to the command args
		argsWithSep = watchArgs
	}
	foundSep := false
	for idx = range argsWithSep {
		if argsWithSep[idx] == "--" {
			foundSep = true
			break
		}
	}
	if interactive {
		flagset := NewWatchCommand().Flags()
		if err := flagset.Parse(argsWithSep); err != nil {
			return nil, nil, err
		}
		watchArgs = flagset.Args()
	}

	// "watch -- echo hello" with ETCDCTL_WATCH_KEY=foo
	// should be translated to "watch foo -- echo hello"
	// (watchArgs=["echo","hello"] should be ["foo","echo","hello"])
	if envKey != "" {
		tmp := []string{envKey}
		if envRange != "" {
			tmp = append(tmp, envRange)
		}
		watchArgs = append(tmp, watchArgs...)
	}

	if !foundSep {
		return watchArgs, nil, nil
	}

	if idx == len(argsWithSep)-1 {
		// "watch foo bar --" should error
		return nil, nil, errBadArgsNumSeparator
	}
	execArgs = argsWithSep[idx+1:]

	// "watch foo bar --rev 1 -- echo hello" or "watch foo --rev 1 bar -- echo hello",
	// then "watchArgs" is "foo bar echo hello"
	// so need ignore args after "argsWithSep[idx]", which is "--"
	endIdx := 0
	for endIdx = len(watchArgs) - 1; endIdx >= 0; endIdx-- {
		if watchArgs[endIdx] == argsWithSep[idx+1] {
			break
		}
	}
	watchArgs = watchArgs[:endIdx]

	return watchArgs, execArgs, nil
}
<|endoftext|>


<commit_before>// Copyright 2015 The etcd Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package command

import (
	"bufio"
	"context"
	"fmt"
	"os"
	"strings"

	"github.com/coreos/etcd/clientv3"

	"github.com/spf13/cobra"
)

var (
	watchRev         int64
	watchPrefix      bool
	watchInteractive bool
	watchPrevKey     bool
)

// NewWatchCommand returns the cobra command for "watch".
func NewWatchCommand() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "watch [options] [key or prefix] [range_end]",
		Short: "Watches events stream on keys or prefixes",
		Run:   watchCommandFunc,
	}

	cmd.Flags().BoolVarP(&watchInteractive, "interactive", "i", false, "Interactive mode")
	cmd.Flags().BoolVar(&watchPrefix, "prefix", false, "Watch on a prefix if prefix is set")
	cmd.Flags().Int64Var(&watchRev, "rev", 0, "Revision to start watching")
	cmd.Flags().BoolVar(&watchPrevKey, "prev-kv", false, "get the previous key-value pair before the event happens")

	return cmd
}

// watchCommandFunc executes the "watch" command.
func watchCommandFunc(cmd *cobra.Command, args []string) {
	if watchInteractive {
		watchInteractiveFunc(cmd, args)
		return
	}

	c := mustClientFromCmd(cmd)
	wc, err := getWatchChan(c, args)
	if err != nil {
		ExitWithError(ExitBadArgs, err)
	}

	printWatchCh(wc)
	if err = c.Close(); err != nil {
		ExitWithError(ExitBadConnection, err)
	}
	ExitWithError(ExitInterrupted, fmt.Errorf("watch is canceled by the server"))
}

func watchInteractiveFunc(cmd *cobra.Command, args []string) {
	c := mustClientFromCmd(cmd)

	reader := bufio.NewReader(os.Stdin)

	for {
		l, err := reader.ReadString('\n')
		if err != nil {
			ExitWithError(ExitInvalidInput, fmt.Errorf("Error reading watch request line: %v", err))
		}
		l = strings.TrimSuffix(l, "\n")

		args := argify(l)
		if len(args) < 2 {
			fmt.Fprintf(os.Stderr, "Invalid command %s (command type or key is not provided)\n", l)
			continue
		}

		if args[0] != "watch" {
			fmt.Fprintf(os.Stderr, "Invalid command %s (only support watch)\n", l)
			continue
		}

		flagset := NewWatchCommand().Flags()
		err = flagset.Parse(args[1:])
		if err != nil {
			fmt.Fprintf(os.Stderr, "Invalid command %s (%v)\n", l, err)
			continue
		}
		ch, err := getWatchChan(c, flagset.Args())
		if err != nil {
			fmt.Fprintf(os.Stderr, "Invalid command %s (%v)\n", l, err)
			continue
		}
		go printWatchCh(ch)
	}
}

func getWatchChan(c *clientv3.Client, args []string) (clientv3.WatchChan, error) {
	if len(args) < 1 || len(args) > 2 {
		return nil, fmt.Errorf("bad number of arguments")
	}
	key := args[0]
	opts := []clientv3.OpOption{clientv3.WithRev(watchRev)}
	if len(args) == 2 {
		if watchPrefix {
			return nil, fmt.Errorf("`range_end` and `--prefix` are mutually exclusive")
		}
		opts = append(opts, clientv3.WithRange(args[1]))
	}
	if watchPrefix {
		opts = append(opts, clientv3.WithPrefix())
	}
	if watchPrevKey {
		opts = append(opts, clientv3.WithPrevKV())
	}
	return c.Watch(context.TODO(), key, opts...), nil
}

func printWatchCh(ch clientv3.WatchChan) {
	for resp := range ch {
		if resp.Canceled {
			fmt.Fprintf(os.Stderr, "watch was canceled (%v)\n", resp.Err())
		}
		display.Watch(resp)
	}
}
<commit_msg>etcdctl/ctlv3: enable 'require-leader' for 'watch' command<commit_after>// Copyright 2015 The etcd Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package command

import (
	"bufio"
	"context"
	"fmt"
	"os"
	"strings"

	"github.com/coreos/etcd/clientv3"

	"github.com/spf13/cobra"
)

var (
	watchRev         int64
	watchPrefix      bool
	watchInteractive bool
	watchPrevKey     bool
)

// NewWatchCommand returns the cobra command for "watch".
func NewWatchCommand() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "watch [options] [key or prefix] [range_end]",
		Short: "Watches events stream on keys or prefixes",
		Run:   watchCommandFunc,
	}

	cmd.Flags().BoolVarP(&watchInteractive, "interactive", "i", false, "Interactive mode")
	cmd.Flags().BoolVar(&watchPrefix, "prefix", false, "Watch on a prefix if prefix is set")
	cmd.Flags().Int64Var(&watchRev, "rev", 0, "Revision to start watching")
	cmd.Flags().BoolVar(&watchPrevKey, "prev-kv", false, "get the previous key-value pair before the event happens")

	return cmd
}

// watchCommandFunc executes the "watch" command.
func watchCommandFunc(cmd *cobra.Command, args []string) {
	if watchInteractive {
		watchInteractiveFunc(cmd, args)
		return
	}

	c := mustClientFromCmd(cmd)
	wc, err := getWatchChan(c, args)
	if err != nil {
		ExitWithError(ExitBadArgs, err)
	}

	printWatchCh(wc)
	if err = c.Close(); err != nil {
		ExitWithError(ExitBadConnection, err)
	}
	ExitWithError(ExitInterrupted, fmt.Errorf("watch is canceled by the server"))
}

func watchInteractiveFunc(cmd *cobra.Command, args []string) {
	c := mustClientFromCmd(cmd)

	reader := bufio.NewReader(os.Stdin)

	for {
		l, err := reader.ReadString('\n')
		if err != nil {
			ExitWithError(ExitInvalidInput, fmt.Errorf("Error reading watch request line: %v", err))
		}
		l = strings.TrimSuffix(l, "\n")

		args := argify(l)
		if len(args) < 2 {
			fmt.Fprintf(os.Stderr, "Invalid command %s (command type or key is not provided)\n", l)
			continue
		}

		if args[0] != "watch" {
			fmt.Fprintf(os.Stderr, "Invalid command %s (only support watch)\n", l)
			continue
		}

		flagset := NewWatchCommand().Flags()
		err = flagset.Parse(args[1:])
		if err != nil {
			fmt.Fprintf(os.Stderr, "Invalid command %s (%v)\n", l, err)
			continue
		}
		ch, err := getWatchChan(c, flagset.Args())
		if err != nil {
			fmt.Fprintf(os.Stderr, "Invalid command %s (%v)\n", l, err)
			continue
		}
		go printWatchCh(ch)
	}
}

func getWatchChan(c *clientv3.Client, args []string) (clientv3.WatchChan, error) {
	if len(args) < 1 || len(args) > 2 {
		return nil, fmt.Errorf("bad number of arguments")
	}
	key := args[0]
	opts := []clientv3.OpOption{clientv3.WithRev(watchRev)}
	if len(args) == 2 {
		if watchPrefix {
			return nil, fmt.Errorf("`range_end` and `--prefix` are mutually exclusive")
		}
		opts = append(opts, clientv3.WithRange(args[1]))
	}
	if watchPrefix {
		opts = append(opts, clientv3.WithPrefix())
	}
	if watchPrevKey {
		opts = append(opts, clientv3.WithPrevKV())
	}
	return c.Watch(clientv3.WithRequireLeader(context.Background()), key, opts...), nil
}

func printWatchCh(ch clientv3.WatchChan) {
	for resp := range ch {
		if resp.Canceled {
			fmt.Fprintf(os.Stderr, "watch was canceled (%v)\n", resp.Err())
		}
		display.Watch(resp)
	}
}
<|endoftext|>


<commit_before>package sisito

import (
	. "."
	"github.com/bouk/monkey"
	"github.com/stretchr/testify/assert"
	"gopkg.in/gorp.v1"
	"reflect"
	"testing"
)

func TestRecentlyBounced(t *testing.T) {
	assert := assert.New(t)
	driver := &Driver{DbMap: &gorp.DbMap{}}

	var guard *monkey.PatchGuard
	guard = monkey.PatchInstanceMethod(
		reflect.TypeOf(driver.DbMap), "Select",
		func(_ *gorp.DbMap, i interface{}, query string, args ...interface{}) ([]interface{}, error) {
			defer guard.Unpatch()
			guard.Restore()

			assert.Equal(`
    SELECT bm.*, IF(wm.id IS NULL, 0, 1) AS whitelisted
      FROM bounce_mails bm LEFT JOIN whitelist_mails wm
        ON bm.recipient = wm.recipient AND bm.senderdomain = wm.senderdomain
     WHERE bm.recipient = ?
       AND bm.senderdomain = ?
  ORDER BY bm.id DESC
     LIMIT 1`, query)

			assert.Equal([]interface{}{"foo@example.com", "example.net"}, args)

			rows := i.(*[]BounceMail)
			*rows = append(*rows, BounceMail{Id: 1})

			return nil, nil
		})

	rows, _ := driver.RecentlyBounced("recipient", "foo@example.com", "example.net")

	assert.Equal([]BounceMail{BounceMail{Id: 1}}, rows)
}

func TestRecentlyBouncedWithoutSenderdomain(t *testing.T) {
	assert := assert.New(t)
	driver := &Driver{DbMap: &gorp.DbMap{}}

	var guard *monkey.PatchGuard
	guard = monkey.PatchInstanceMethod(
		reflect.TypeOf(driver.DbMap), "Select",
		func(_ *gorp.DbMap, i interface{}, query string, args ...interface{}) ([]interface{}, error) {
			defer guard.Unpatch()
			guard.Restore()

			assert.Equal(`
    SELECT bm.*, IF(wm.id IS NULL, 0, 1) AS whitelisted
      FROM bounce_mails bm LEFT JOIN whitelist_mails wm
        ON bm.recipient = wm.recipient AND bm.senderdomain = wm.senderdomain
     WHERE bm.recipient = ?
  ORDER BY bm.id DESC
     LIMIT 1`, query)

			assert.Equal([]interface{}{"foo@example.com"}, args)

			rows := i.(*[]BounceMail)
			*rows = append(*rows, BounceMail{Id: 1})

			return nil, nil
		})

	rows, _ := driver.RecentlyBounced("recipient", "foo@example.com", "")

	assert.Equal([]BounceMail{BounceMail{Id: 1}}, rows)
}

func TestIsBounced(t *testing.T) {
	assert := assert.New(t)
	driver := &Driver{DbMap: &gorp.DbMap{}}

	var guard *monkey.PatchGuard
	guard = monkey.PatchInstanceMethod(
		reflect.TypeOf(driver.DbMap), "SelectInt",
		func(_ *gorp.DbMap, query string, args ...interface{}) (int64, error) {
			defer guard.Unpatch()
			guard.Restore()

			assert.Equal(`
    SELECT 1
      FROM bounce_mails bm LEFT JOIN whitelist_mails wm
        ON bm.recipient = wm.recipient AND bm.senderdomain = wm.senderdomain
     WHERE bm.recipient = ?
       AND bm.senderdomain = ?
       AND wm.id IS NULL
     LIMIT 1`, query)

			assert.Equal([]interface{}{"foo@example.com", "example.net"}, args)

			return 1, nil
		})

	count, _ := driver.IsBounced("recipient", "foo@example.com", "example.net")

	assert.Equal(count, true)
}

func TestIsBouncedWithoutSenderdomain(t *testing.T) {
	assert := assert.New(t)
	driver := &Driver{DbMap: &gorp.DbMap{}}

	var guard *monkey.PatchGuard
	guard = monkey.PatchInstanceMethod(
		reflect.TypeOf(driver.DbMap), "SelectInt",
		func(_ *gorp.DbMap, query string, args ...interface{}) (int64, error) {
			defer guard.Unpatch()
			guard.Restore()

			assert.Equal(`
    SELECT 1
      FROM bounce_mails bm LEFT JOIN whitelist_mails wm
        ON bm.recipient = wm.recipient AND bm.senderdomain = wm.senderdomain
     WHERE bm.recipient = ?
       AND wm.id IS NULL
     LIMIT 1`, query)

			assert.Equal([]interface{}{"foo@example.com"}, args)

			return 1, nil
		})

	count, _ := driver.IsBounced("recipient", "foo@example.com", "")

	assert.Equal(count, true)
}

func TestBlacklistRecipients(t *testing.T) {
	assert := assert.New(t)
	driver := &Driver{DbMap: &gorp.DbMap{}}

	var guard *monkey.PatchGuard
	guard = monkey.PatchInstanceMethod(
		reflect.TypeOf(driver.DbMap), "Select",
		func(_ *gorp.DbMap, i interface{}, query string, args ...interface{}) ([]interface{}, error) {
			defer guard.Unpatch()
			guard.Restore()

			assert.Equal(`
    SELECT bm.recipient
      FROM bounce_mails bm LEFT JOIN whitelist_mails wm
        ON bm.recipient = wm.recipient AND bm.senderdomain = wm.senderdomain
     WHERE wm.id IS NULL
       AND bm.senderdomain = ?
       AND bm.reason IN (?,?)
       AND bm.softbounce = ?
  GROUP BY recipient
  ORDER BY recipient
     LIMIT ?
    OFFSET ?`, query)

			assert.Equal([]interface{}{
				"example.net", "userunknown", "filtered", false, uint64(100), uint64(100)}, args)

			rows := i.(*[]string)
			*rows = append(*rows, "foo@example.com")

			return nil, nil
		})

	recipients, _ := driver.BlacklistRecipients(
		"example.net", []string{"userunknown", "filtered"}, new(bool), 100, 100)

	assert.Equal(recipients, []string{"foo@example.com"})
}

func TestBlacklistRecipientsWithoutOptions(t *testing.T) {
	assert := assert.New(t)
	driver := &Driver{DbMap: &gorp.DbMap{}}

	var guard *monkey.PatchGuard
	guard = monkey.PatchInstanceMethod(
		reflect.TypeOf(driver.DbMap), "Select",
		func(_ *gorp.DbMap, i interface{}, query string, args ...interface{}) ([]interface{}, error) {
			defer guard.Unpatch()
			guard.Restore()

			assert.Equal(`
    SELECT bm.recipient
      FROM bounce_mails bm LEFT JOIN whitelist_mails wm
        ON bm.recipient = wm.recipient AND bm.senderdomain = wm.senderdomain
     WHERE wm.id IS NULL
  GROUP BY recipient
  ORDER BY recipient`, query)

			assert.Equal([]interface{}{}, args)

			rows := i.(*[]string)
			*rows = append(*rows, "foo@example.com")

			return nil, nil
		})

	recipients, _ := driver.BlacklistRecipients(nil, nil, 0, 0)

	assert.Equal(recipients, []string{"foo@example.com"})
}
<commit_msg>Add driver test<commit_after>package sisito

import (
	. "."
	"github.com/bouk/monkey"
	"github.com/stretchr/testify/assert"
	"gopkg.in/gorp.v1"
	"reflect"
	"testing"
)

func TestRecentlyBounced(t *testing.T) {
	assert := assert.New(t)
	driver := &Driver{DbMap: &gorp.DbMap{}}

	var guard *monkey.PatchGuard
	guard = monkey.PatchInstanceMethod(
		reflect.TypeOf(driver.DbMap), "Select",
		func(_ *gorp.DbMap, i interface{}, query string, args ...interface{}) ([]interface{}, error) {
			defer guard.Unpatch()
			guard.Restore()

			assert.Equal(`
    SELECT bm.*, IF(wm.id IS NULL, 0, 1) AS whitelisted
      FROM bounce_mails bm LEFT JOIN whitelist_mails wm
        ON bm.recipient = wm.recipient AND bm.senderdomain = wm.senderdomain
     WHERE bm.recipient = ?
       AND bm.senderdomain = ?
  ORDER BY bm.id DESC
     LIMIT 1`, query)

			assert.Equal([]interface{}{"foo@example.com", "example.net"}, args)

			rows := i.(*[]BounceMail)
			*rows = append(*rows, BounceMail{Id: 1})

			return nil, nil
		})

	rows, _ := driver.RecentlyBounced("recipient", "foo@example.com", "example.net")

	assert.Equal([]BounceMail{BounceMail{Id: 1}}, rows)
}

func TestRecentlyBouncedWithoutSenderdomain(t *testing.T) {
	assert := assert.New(t)
	driver := &Driver{DbMap: &gorp.DbMap{}}

	var guard *monkey.PatchGuard
	guard = monkey.PatchInstanceMethod(
		reflect.TypeOf(driver.DbMap), "Select",
		func(_ *gorp.DbMap, i interface{}, query string, args ...interface{}) ([]interface{}, error) {
			defer guard.Unpatch()
			guard.Restore()

			assert.Equal(`
    SELECT bm.*, IF(wm.id IS NULL, 0, 1) AS whitelisted
      FROM bounce_mails bm LEFT JOIN whitelist_mails wm
        ON bm.recipient = wm.recipient AND bm.senderdomain = wm.senderdomain
     WHERE bm.recipient = ?
  ORDER BY bm.id DESC
     LIMIT 1`, query)

			assert.Equal([]interface{}{"foo@example.com"}, args)

			rows := i.(*[]BounceMail)
			*rows = append(*rows, BounceMail{Id: 1})

			return nil, nil
		})

	rows, _ := driver.RecentlyBounced("recipient", "foo@example.com", "")

	assert.Equal([]BounceMail{BounceMail{Id: 1}}, rows)
}

func TestIsBounced(t *testing.T) {
	assert := assert.New(t)
	driver := &Driver{DbMap: &gorp.DbMap{}}

	var guard *monkey.PatchGuard
	guard = monkey.PatchInstanceMethod(
		reflect.TypeOf(driver.DbMap), "SelectInt",
		func(_ *gorp.DbMap, query string, args ...interface{}) (int64, error) {
			defer guard.Unpatch()
			guard.Restore()

			assert.Equal(`
    SELECT 1
      FROM bounce_mails bm LEFT JOIN whitelist_mails wm
        ON bm.recipient = wm.recipient AND bm.senderdomain = wm.senderdomain
     WHERE bm.recipient = ?
       AND bm.senderdomain = ?
       AND wm.id IS NULL
     LIMIT 1`, query)

			assert.Equal([]interface{}{"foo@example.com", "example.net"}, args)

			return 1, nil
		})

	count, _ := driver.IsBounced("recipient", "foo@example.com", "example.net")

	assert.Equal(count, true)
}

func TestIsBouncedWithoutSenderdomain(t *testing.T) {
	assert := assert.New(t)
	driver := &Driver{DbMap: &gorp.DbMap{}}

	var guard *monkey.PatchGuard
	guard = monkey.PatchInstanceMethod(
		reflect.TypeOf(driver.DbMap), "SelectInt",
		func(_ *gorp.DbMap, query string, args ...interface{}) (int64, error) {
			defer guard.Unpatch()
			guard.Restore()

			assert.Equal(`
    SELECT 1
      FROM bounce_mails bm LEFT JOIN whitelist_mails wm
        ON bm.recipient = wm.recipient AND bm.senderdomain = wm.senderdomain
     WHERE bm.recipient = ?
       AND wm.id IS NULL
     LIMIT 1`, query)

			assert.Equal([]interface{}{"foo@example.com"}, args)

			return 1, nil
		})

	count, _ := driver.IsBounced("recipient", "foo@example.com", "")

	assert.Equal(count, true)
}

func TestBlacklistRecipients(t *testing.T) {
	assert := assert.New(t)
	driver := &Driver{DbMap: &gorp.DbMap{}}

	var guard *monkey.PatchGuard
	guard = monkey.PatchInstanceMethod(
		reflect.TypeOf(driver.DbMap), "Select",
		func(_ *gorp.DbMap, i interface{}, query string, args ...interface{}) ([]interface{}, error) {
			defer guard.Unpatch()
			guard.Restore()

			assert.Equal(`
    SELECT bm.recipient
      FROM bounce_mails bm LEFT JOIN whitelist_mails wm
        ON bm.recipient = wm.recipient AND bm.senderdomain = wm.senderdomain
     WHERE wm.id IS NULL
       AND bm.senderdomain = ?
       AND bm.reason IN (?,?)
       AND bm.softbounce = ?
  GROUP BY recipient
  ORDER BY recipient
     LIMIT ?
    OFFSET ?`, query)

			assert.Equal([]interface{}{
				"example.net", "userunknown", "filtered", false, uint64(100), uint64(100)}, args)

			rows := i.(*[]string)
			*rows = append(*rows, "foo@example.com")

			return nil, nil
		})

	recipients, _ := driver.BlacklistRecipients(
		"example.net", []string{"userunknown", "filtered"}, new(bool), 100, 100)

	assert.Equal(recipients, []string{"foo@example.com"})
}

func TestBlacklistRecipientsWithoutOptions(t *testing.T) {
	assert := assert.New(t)
	driver := &Driver{DbMap: &gorp.DbMap{}}

	var guard *monkey.PatchGuard
	guard = monkey.PatchInstanceMethod(
		reflect.TypeOf(driver.DbMap), "Select",
		func(_ *gorp.DbMap, i interface{}, query string, args ...interface{}) ([]interface{}, error) {
			defer guard.Unpatch()
			guard.Restore()

			assert.Equal(`
    SELECT bm.recipient
      FROM bounce_mails bm LEFT JOIN whitelist_mails wm
        ON bm.recipient = wm.recipient AND bm.senderdomain = wm.senderdomain
     WHERE wm.id IS NULL
  GROUP BY recipient
  ORDER BY recipient`, query)

			assert.Equal([]interface{}{}, args)

			rows := i.(*[]string)
			*rows = append(*rows, "foo@example.com")

			return nil, nil
		})

	recipients, _ := driver.BlacklistRecipients("", nil, nil, 0, 0)

	assert.Equal(recipients, []string{"foo@example.com"})
}
<|endoftext|>


<commit_before>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/

package sync

import (
	"bytes"
	"code.google.com/p/go.crypto/hkdf"
	"crypto/aes"
	"crypto/cipher"
	"crypto/hmac"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"errors"
	"github.com/st3fan/gofxa/fxa"
	"io"
	"io/ioutil"
	"net/http"
)

type EncryptedPayload struct {
	CipherText string `json:"ciphertext"`
	IV         string `json:"IV"`
	HMAC       string `json:"HMAC"`
}

type Record struct {
	Id       string  `json:"id"`
	Modified float64 `json:"modified"`
	Payload  string  `json:"payload"`
}

func (r *Record) Decrypt(keyBundle KeyBundle) error {
	encryptedPayload := &EncryptedPayload{}
	if err := json.Unmarshal([]byte(r.Payload), encryptedPayload); err != nil {
		return errors.New("Record does not have an encrypted payload")
	}

	HMAC, err := hex.DecodeString(encryptedPayload.HMAC)
	if err != nil {
		return errors.New("Malformed HMAC in record")
	}

	// Verify the payload

	mac := hmac.New(sha256.New, keyBundle.ValidationKey)
	mac.Write([]byte(encryptedPayload.CipherText))
	if !bytes.Equal(mac.Sum(nil), HMAC) {
		return errors.New("Record validation failed")
	}

	// Decrypt the payload

	ciphertext, err := base64.StdEncoding.DecodeString(encryptedPayload.CipherText)
	if err != nil {
		return errors.New("Malformed ciphertext in record")
	}

	iv, err := base64.StdEncoding.DecodeString(encryptedPayload.IV)
	if err != nil {
		return errors.New("Malformed IV in record")
	}

	block, err := aes.NewCipher(keyBundle.EncryptionKey)
	if err != nil {
		return err
	}

	mode := cipher.NewCBCDecrypter(block, iv)
	mode.CryptBlocks(ciphertext, ciphertext)

	// TODO: Is there a stdlib something to deal with PKCS7 padding?
	length := len(ciphertext)
	unpadding := int(ciphertext[length-1])
	ciphertext = ciphertext[:(length - unpadding)]

	r.Payload = string(ciphertext)

	return nil
}

//

type KeyBundle struct {
	EncryptionKey []byte
	ValidationKey []byte
}

func NewKeyBundle(key []byte) (KeyBundle, error) {
	secret := make([]byte, 2*32)
	if _, err := io.ReadFull(hkdf.New(sha256.New, key, nil, []byte("identity.mozilla.com/picl/v1/oldsync")), secret); err != nil {
		return KeyBundle{}, err
	}
	return KeyBundle{
		EncryptionKey: secret[0:32],
		ValidationKey: secret[32:64],
	}, nil
}

//

type StorageClient struct {
	endpoint  string
	hawkKeyId string
	hawkKey   string
	secret    []byte
}

func NewStorageClient(endpoint, hawkKeyId, hawkKey string, secret []byte) (*StorageClient, error) {
	return &StorageClient{
		endpoint:  endpoint,
		hawkKeyId: hawkKeyId,
		hawkKey:   hawkKey,
		secret:    secret,
	}, nil
}

type KeysPayload struct {
	Default []string `json:"default"`
}

func (sc *StorageClient) FetchKeys() (KeyBundle, error) {
	globalKeyBundle, err := NewKeyBundle(sc.secret)
	if err != nil {
		return KeyBundle{}, err
	}

	record, err := sc.GetEncryptedRecord("crypto", "keys", globalKeyBundle)
	if err != nil {
		return KeyBundle{}, err
	}

	keysPayload := KeysPayload{}
	if err = json.Unmarshal([]byte(record.Payload), &keysPayload); err != nil {
		return KeyBundle{}, err
	}

	encryptionKey, err := base64.StdEncoding.DecodeString(keysPayload.Default[0])
	if err != nil {
		return KeyBundle{}, err
	}

	validationKey, err := base64.StdEncoding.DecodeString(keysPayload.Default[1])
	if err != nil {
		return KeyBundle{}, err
	}

	return KeyBundle{
		EncryptionKey: encryptionKey,
		ValidationKey: validationKey,
	}, nil
}

func (sc *StorageClient) GetRecord(collectionName, recordId string) (Record, error) {
	req, err := http.NewRequest("GET", sc.endpoint+"/storage/"+collectionName+"/"+recordId, nil)
	if err != nil {
		return Record{}, err
	}

	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Accept", "application/json")

	hawkCredentials := fxa.NewHawkCredentials(sc.hawkKeyId, []byte(sc.hawkKey))
	if err := hawkCredentials.AuthorizeRequest(req, nil, ""); err != nil {
		return Record{}, err
	}

	client := &http.Client{}
	res, err := client.Do(req)
	if err != nil {
		return Record{}, err
	}
	defer res.Body.Close()

	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return Record{}, err
	}

	if res.StatusCode != http.StatusOK {
		return Record{}, errors.New(res.Status) // TODO: Proper errors based on what the server returns
	}

	record := Record{}
	if err = json.Unmarshal(body, &record); err != nil {
		return Record{}, err
	}

	return record, nil
}

func (sc *StorageClient) GetEncryptedRecord(collectionName, recordId string, keyBundle KeyBundle) (Record, error) {
	record, err := sc.GetRecord(collectionName, recordId)
	if err != nil {
		return Record{}, err
	}

	if err := record.Decrypt(keyBundle); err != nil {
		return Record{}, err
	}

	return record, nil
}

func (sc *StorageClient) GetEncryptedRecords(collectionName string, keyBundle KeyBundle) ([]Record, error) {
	req, err := http.NewRequest("GET", sc.endpoint+"/storage/"+collectionName+"?full=1", nil)
	if err != nil {
		return nil, err
	}

	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Accept", "application/json")

	hawkCredentials := fxa.NewHawkCredentials(sc.hawkKeyId, []byte(sc.hawkKey))
	if err := hawkCredentials.AuthorizeRequest(req, nil, ""); err != nil {
		return nil, err
	}

	client := &http.Client{}
	res, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	defer res.Body.Close()

	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return nil, err
	}

	if res.StatusCode != http.StatusOK {
		return nil, errors.New(res.Status) // TODO: Proper errors based on what the server returns
	}

	records := []Record{}
	if err = json.Unmarshal(body, &records); err != nil {
		return nil, err
	}

	for i := 0; i < len(records); i++ {
		if err := records[i].Decrypt(keyBundle); err != nil {
			return nil, err
		}
	}

	return records, nil
}
<commit_msg>Refactored storageClient code to have a Login func<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/

package sync

import (
	"bytes"
	"code.google.com/p/go.crypto/hkdf"
	"crypto/aes"
	"crypto/cipher"
	"crypto/hmac"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"errors"
	"github.com/st3fan/gofxa/fxa"
	"io"
	"io/ioutil"
	"net/http"
	"strconv"
)

type EncryptedPayload struct {
	CipherText string `json:"ciphertext"`
	IV         string `json:"IV"`
	HMAC       string `json:"HMAC"`
}

type Record struct {
	Id       string  `json:"id"`
	Modified float64 `json:"modified"`
	Payload  string  `json:"payload"`
}

func (r *Record) Decrypt(keyBundle *KeyBundle) error {
	encryptedPayload := &EncryptedPayload{}
	if err := json.Unmarshal([]byte(r.Payload), encryptedPayload); err != nil {
		return errors.New("Record does not have an encrypted payload")
	}

	HMAC, err := hex.DecodeString(encryptedPayload.HMAC)
	if err != nil {
		return errors.New("Malformed HMAC in record")
	}

	// Verify the payload

	mac := hmac.New(sha256.New, keyBundle.ValidationKey)
	mac.Write([]byte(encryptedPayload.CipherText))
	if !bytes.Equal(mac.Sum(nil), HMAC) {
		return errors.New("Record validation failed")
	}

	// Decrypt the payload

	ciphertext, err := base64.StdEncoding.DecodeString(encryptedPayload.CipherText)
	if err != nil {
		return errors.New("Malformed ciphertext in record")
	}

	iv, err := base64.StdEncoding.DecodeString(encryptedPayload.IV)
	if err != nil {
		return errors.New("Malformed IV in record")
	}

	block, err := aes.NewCipher(keyBundle.EncryptionKey)
	if err != nil {
		return err
	}

	mode := cipher.NewCBCDecrypter(block, iv)
	mode.CryptBlocks(ciphertext, ciphertext)

	// TODO: Is there a stdlib something to deal with PKCS7 padding?
	length := len(ciphertext)
	unpadding := int(ciphertext[length-1])
	ciphertext = ciphertext[:(length - unpadding)]

	r.Payload = string(ciphertext)

	return nil
}

//

type KeyBundle struct {
	EncryptionKey []byte
	ValidationKey []byte
}

func NewKeyBundle(key []byte) (*KeyBundle, error) {
	secret := make([]byte, 2*32)
	if _, err := io.ReadFull(hkdf.New(sha256.New, key, nil, []byte("identity.mozilla.com/picl/v1/oldsync")), secret); err != nil {
		return nil, err
	}
	return &KeyBundle{
		EncryptionKey: secret[0:32],
		ValidationKey: secret[32:64],
	}, nil
}

//

type StorageClient struct {
	endpoint  string
	hawkKeyId string
	hawkKey   string
	secret    []byte
	keyBundle KeyBundle
}

func NewStorageClient(endpoint, hawkKeyId, hawkKey string, secret []byte) (*StorageClient, error) {
	return &StorageClient{
		endpoint:  endpoint,
		hawkKeyId: hawkKeyId,
		hawkKey:   hawkKey,
		secret:    secret,
	}, nil
}

type KeysPayload struct {
	Default []string `json:"default"`
}

func (sc *StorageClient) FetchKeys() (KeyBundle, error) {
	globalKeyBundle, err := NewKeyBundle(sc.secret)
	if err != nil {
		return KeyBundle{}, err
	}

	record, err := sc.GetEncryptedRecord("crypto", "keys", globalKeyBundle)
	if err != nil {
		return KeyBundle{}, err
	}

	keysPayload := KeysPayload{}
	if err = json.Unmarshal([]byte(record.Payload), &keysPayload); err != nil {
		return KeyBundle{}, err
	}

	encryptionKey, err := base64.StdEncoding.DecodeString(keysPayload.Default[0])
	if err != nil {
		return KeyBundle{}, err
	}

	validationKey, err := base64.StdEncoding.DecodeString(keysPayload.Default[1])
	if err != nil {
		return KeyBundle{}, err
	}

	return KeyBundle{
		EncryptionKey: encryptionKey,
		ValidationKey: validationKey,
	}, nil
}

func (sc *StorageClient) GetRecord(collectionName, recordId string) (Record, error) {
	req, err := http.NewRequest("GET", sc.endpoint+"/storage/"+collectionName+"/"+recordId, nil)
	if err != nil {
		return Record{}, err
	}

	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Accept", "application/json")

	hawkCredentials := fxa.NewHawkCredentials(sc.hawkKeyId, []byte(sc.hawkKey))
	if err := hawkCredentials.AuthorizeRequest(req, nil, ""); err != nil {
		return Record{}, err
	}

	client := &http.Client{}
	res, err := client.Do(req)
	if err != nil {
		return Record{}, err
	}
	defer res.Body.Close()

	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return Record{}, err
	}

	if res.StatusCode != http.StatusOK {
		return Record{}, errors.New(res.Status) // TODO: Proper errors based on what the server returns
	}

	record := Record{}
	if err = json.Unmarshal(body, &record); err != nil {
		return Record{}, err
	}

	return record, nil
}

func (sc *StorageClient) GetEncryptedRecord(collectionName, recordId string, keyBundle *KeyBundle) (Record, error) {
	record, err := sc.GetRecord(collectionName, recordId)
	if err != nil {
		return Record{}, err
	}

	if keyBundle == nil {
		keyBundle = &sc.keyBundle
	}

	if err := record.Decrypt(keyBundle); err != nil {
		return Record{}, err
	}

	return record, nil
}

type GetRecordsOptions struct {
	Limit int
	Sort  string
}

func (sc *StorageClient) GetEncryptedRecords(collectionName string, keyBundle *KeyBundle, options *GetRecordsOptions) ([]Record, error) {
	url := sc.endpoint + "/storage/" + collectionName + "?full=1"
	if options != nil {
		if options.Limit != 0 {
			url += "&limit=" + strconv.Itoa(options.Limit)
		}
		if options.Sort != "" {
			url += "&sort=" + options.Sort
		}
	}

	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return nil, err
	}

	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Accept", "application/json")

	hawkCredentials := fxa.NewHawkCredentials(sc.hawkKeyId, []byte(sc.hawkKey))
	if err := hawkCredentials.AuthorizeRequest(req, nil, ""); err != nil {
		return nil, err
	}

	client := &http.Client{}
	res, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	defer res.Body.Close()

	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return nil, err
	}

	if res.StatusCode != http.StatusOK {
		return nil, errors.New(res.Status) // TODO: Proper errors based on what the server returns
	}

	records := []Record{}
	if err = json.Unmarshal(body, &records); err != nil {
		return nil, err
	}

	if keyBundle == nil {
		keyBundle = &sc.keyBundle
	}

	for i := 0; i < len(records); i++ {
		if err := records[i].Decrypt(keyBundle); err != nil {
			return nil, err
		}
	}

	return records, nil
}

func (sc *StorageClient) Login() error {
	keyBundle, err := sc.FetchKeys()
	if err != nil {
		return err
	}
	sc.keyBundle = keyBundle
	return nil
}
<|endoftext|>