<commit_before>package kdtree

import (
	"encoding/json"
	"fmt"
	"math"
	"math/big"
	"reflect"
	"testing"
)

type intAndFloat struct {
	i int
	f float64
}

func (inf intAndFloat) MarshalJSON() ([]byte, error) {
	return json.Marshal(map[string]interface{}{
		"i": inf.i,
		"f": inf.f,
	})
}

type embedTest struct {
	s string
	intAndFloat
}

func (e embedTest) MarshalJSON() ([]byte, error) {
	return json.Marshal(map[string]interface{}{
		"s":           e.s,
		"intAndFloat": e.intAndFloat,
	})
}

var (
	diffFptrs = func(f, g []*float64) bool {
		if len(f) != len(g) {
			return false
		}
		for i := range f {
			if f[i] == g[i] {
				return false
			}
		}
		return true
	}

	getF64ptrs = func(f []float64) (f64ptrs []*float64) {
		for i := range f {
			f64ptrs = append(f64ptrs, &f[i])
		}
		return
	}

	embedded = embedTest{"hello world", intAndFloat{0, math.E}}

	rational = big.NewRat(5, 4)

	str = "cassandra"

	mappy = map[string]interface{}{
		"a": int('a'),
		"b": string("banana"),
		"c": intAndFloat{0, math.Pi},
	}

	interfaces = []interface{}{&str, str, &rational, rational, &embedded, embedded, &mappy, mappy}

	constructorInputs = []struct {
		linked interface{}
		values []float64
	}{
		{
			&str,
			[]float64{6.0000125, 6.10000125, -1.3173, 1373},
		},
		{
			str,
			[]float64{6.0000125, 6.10000125, -1.3173, 1373},
		},
		{
			&rational,
			[]float64{1, 2, 3, 4, 5},
		},
		{
			rational,
			[]float64{1, 2, 3, 4, 5},
		},
		{
			&embedded,
			[]float64{100.3},
		},
		{
			embedded,
			[]float64{100.3},
		},
		{
			&mappy,
			[]float64{0.8050908121798804, 0.53238545404102},
		},
		{
			mappy,
			[]float64{0.8050908121798804, 0.53238545404102},
		},
	}
)

func TestDatapointConstructor(t *testing.T) {
	constructorTests := []struct {
		dp         *Datapoint
		wantData   interface{}
		wantValues []float64
	}{
		{
			dp:         NewDatapoint(constructorInputs[0].linked, constructorInputs[0].values),
			wantData:   constructorInputs[0].linked,
			wantValues: constructorInputs[0].values,
		},
		{
			dp:         NewDatapoint(constructorInputs[1].linked, constructorInputs[1].values),
			wantData:   constructorInputs[1].linked,
			wantValues: constructorInputs[1].values,
		},
		{
			dp:         NewDatapoint(constructorInputs[2].linked, constructorInputs[2].values),
			wantData:   constructorInputs[2].linked,
			wantValues: constructorInputs[2].values,
		},
		{
			dp:         NewDatapoint(constructorInputs[3].linked, constructorInputs[3].values),
			wantData:   constructorInputs[3].linked,
			wantValues: constructorInputs[3].values,
		},
	}

	for _, ct := range constructorTests {
		if reflect.DeepEqual(ct.wantData, ct.dp.Data()) == false {
			t.Fail()
		}
		if reflect.DeepEqual(ct.wantValues, ct.dp.Set()) == false {
			t.Fail()
		}
	}
}

func TestDatapointSetCopy(t *testing.T) {

	var randomPoints Datapoints
	for i := uint(1); i <= 10; i++ {
		p := RandomDatapoint(i)
		randomPoints = append(randomPoints, p)
	}

	type dpst struct {
		dp             *Datapoint
		wantValues     []float64
		duplicateAddrs []*float64
	}

	var dpSetTests []dpst

	for i := range randomPoints {
		var values []float64
		copy(randomPoints[i].set, values)
		var f64ptrs []*float64
		for k := range randomPoints[i].set {
			values = append(values, randomPoints[i].set[k])
			f64ptrs = append(f64ptrs, &randomPoints[i].set[k])
		}
		dpSetTests = append(
			dpSetTests,
			dpst{
				randomPoints[i],
				values,
				f64ptrs,
			})
	}

	for i := range dpSetTests {
		gotSet := dpSetTests[i].dp.Set()
		// verify that the values were copied correctly
		if reflect.DeepEqual(gotSet, dpSetTests[i].wantValues) == false {
			t.Errorf(`wrong values returned by Set()
            got:       %v
            want:      %v`, gotSet, dpSetTests[i].wantValues)
		}
		// but we also need to make sure it's a true copy - i.e. not referring to the same underlying array/slice
		gotAddrs := getF64ptrs(gotSet)
		if diffFptrs(gotAddrs, dpSetTests[i].duplicateAddrs) == false {
			t.Errorf(`duplicate addresses returned by Set()
		    got:       %v
		    matching:  %v`, gotAddrs, dpSetTests[i].duplicateAddrs)
		}
	}
}

func TestDatapointLinkedDataIdentical(t *testing.T) {
	for k, ct := range constructorInputs {
		newDP := NewDatapoint(ct.linked, ct.values)
		if reflect.DeepEqual(newDP.Data(), interfaces[k]) == false {
			t.Error(k, `: Datapoint.Data() does not refer to the same object
			got:      `, newDP.Data(), `
			want:     `, interfaces[k])
		}

		srcRv := reflect.ValueOf(interfaces[k])
		dpRv := reflect.ValueOf(newDP.Data())

		if srcRv.Kind() != dpRv.Kind() { //	if they aren't to the same kind, then we have already failed.
			t.Error(k, `: Datapoint.Data() does not refer to an identical reflect.Kind
			got:      `, dpRv, `
			want:     `, srcRv)
		}

		errStringGot := fmt.Sprint(`got:      `, reflect.TypeOf(newDP.Data()), ` `, dpRv.Interface())
		errStringWant := fmt.Sprint(`want:     `, reflect.TypeOf(interfaces[k]), ` `, srcRv.Interface())

		if reflect.DeepEqual(reflect.TypeOf(newDP.Data()), reflect.TypeOf(interfaces[k])) == false {
			t.Error(k, `: Datapoint.Data() does not reflect an identical reflect.Type
			`, errStringGot, `
			`, errStringWant)
		}
	}
}

func TestDatapointDimensionality(t *testing.T) {
	var dimTests = []struct {
		dp   *Datapoint
		want int
	}{
		{RandomDatapoint(3), 3},
		{RandomDatapoint(0), 0},
	}

	for _, dt := range dimTests {
		got := dt.dp.Dimensionality()
		if got != dt.want {
			t.Error(`want: `, dt.want, `
			got: `, got)
		}
	}
}

func TestDatapointStringer(t *testing.T) {
	dpStringStr := `{data: cassandra}, {set: [0:{6.0000125}, 1:{6.10000125}, 2:{-1.3173}, 3:{1373}]}`
	dpRationalStr := `{data: &{{false [5]} {false [4]}}}, {set: [0:{1}, 1:{2}, 2:{3}, 3:{4}, 4:{5}]}`
	var stringerTests = []struct {
		dp   *Datapoint
		want string
	}{
		{
			&Datapoint{
				&str,
				[]float64{6.0000125, 6.10000125, -1.3173, 1373},
			},
			dpStringStr,
		},
		{
			&Datapoint{
				&rational,
				[]float64{1, 2, 3, 4, 5},
			},
			dpRationalStr,
		},
	}

	for _, s := range stringerTests {
		got := fmt.Sprint(s.dp)
		if got != s.want {
			t.Error(`want: `, s.want, `
			got: `, got)
		}
	}

}

func TestDatapointMarshalJSON(t *testing.T) {
	for k, ct := range constructorInputs {
		newDP := NewDatapoint(ct.linked, ct.values)
		dpJSON, _ := json.Marshal(newDP)
		var want []byte
		switch k {
		case 0:
			want = []byte(`{"data":"cassandra","set":[6.0000125,6.10000125,-1.3173,1373]}`)
		case 1:
			want = []byte(`{"data":"cassandra","set":[6.0000125,6.10000125,-1.3173,1373]}`)
		case 2:
			want = []byte(`{"data":"5/4","set":[1,2,3,4,5]}`)
		case 3:
			want = []byte(`{"data":"5/4","set":[1,2,3,4,5]}`)
		case 4:
			want = []byte(`{"data":{"intAndFloat":{"f":2.718281828459045,"i":0},"s":"hello world"},"set":[100.3]}`)
		case 5:
			want = []byte(`{"data":{"intAndFloat":{"f":2.718281828459045,"i":0},"s":"hello world"},"set":[100.3]}`)
		case 6:
			want = []byte(`{"data":{"a":97,"b":"banana","c":{"f":3.141592653589793,"i":0}},"set":[0.8050908121798804,0.53238545404102]}`)
		case 7:
			want = []byte(`{"data":{"a":97,"b":"banana","c":{"f":3.141592653589793,"i":0}},"set":[0.8050908121798804,0.53238545404102]}`)
		}
		if reflect.DeepEqual(dpJSON, want) == false {
			t.Error(`want: `, string(want), `
			got: `, string(dpJSON))
		}
	}
}
<commit_msg>Test Import using Importable interface<commit_after>package kdtree

import (
	"encoding/json"
	"fmt"
	"math"
	"math/big"
	"reflect"
	"testing"
)

type intAndFloat struct {
	i int
	f float64
}

func (inf intAndFloat) MarshalJSON() ([]byte, error) {
	return json.Marshal(map[string]interface{}{
		"i": inf.i,
		"f": inf.f,
	})
}

type embedTest struct {
	s string
	intAndFloat
}

func (e embedTest) MarshalJSON() ([]byte, error) {
	return json.Marshal(map[string]interface{}{
		"s":           e.s,
		"intAndFloat": e.intAndFloat,
	})
}

var (
	diffFptrs = func(f, g []*float64) bool {
		if len(f) != len(g) {
			return false
		}
		for i := range f {
			if f[i] == g[i] {
				return false
			}
		}
		return true
	}

	getF64ptrs = func(f []float64) (f64ptrs []*float64) {
		for i := range f {
			f64ptrs = append(f64ptrs, &f[i])
		}
		return
	}

	embedded = embedTest{"hello world", intAndFloat{0, math.E}}

	rational = big.NewRat(5, 4)

	str = "cassandra"

	mappy = map[string]interface{}{
		"a": int('a'),
		"b": string("banana"),
		"c": intAndFloat{0, math.Pi},
	}

	interfaces = []interface{}{&str, str, &rational, rational, &embedded, embedded, &mappy, mappy}

	constructorInputs = []struct {
		linked interface{}
		values []float64
	}{
		{
			&str,
			[]float64{6.0000125, 6.10000125, -1.3173, 1373},
		},
		{
			str,
			[]float64{6.0000125, 6.10000125, -1.3173, 1373},
		},
		{
			&rational,
			[]float64{1, 2, 3, 4, 5},
		},
		{
			rational,
			[]float64{1, 2, 3, 4, 5},
		},
		{
			&embedded,
			[]float64{100.3},
		},
		{
			embedded,
			[]float64{100.3},
		},
		{
			&mappy,
			[]float64{0.8050908121798804, 0.53238545404102},
		},
		{
			mappy,
			[]float64{0.8050908121798804, 0.53238545404102},
		},
	}
)

func TestDatapointConstructor(t *testing.T) {
	constructorTests := []struct {
		dp         *Datapoint
		wantData   interface{}
		wantValues []float64
	}{
		{
			dp:         NewDatapoint(constructorInputs[0].linked, constructorInputs[0].values),
			wantData:   constructorInputs[0].linked,
			wantValues: constructorInputs[0].values,
		},
		{
			dp:         NewDatapoint(constructorInputs[1].linked, constructorInputs[1].values),
			wantData:   constructorInputs[1].linked,
			wantValues: constructorInputs[1].values,
		},
		{
			dp:         NewDatapoint(constructorInputs[2].linked, constructorInputs[2].values),
			wantData:   constructorInputs[2].linked,
			wantValues: constructorInputs[2].values,
		},
		{
			dp:         NewDatapoint(constructorInputs[3].linked, constructorInputs[3].values),
			wantData:   constructorInputs[3].linked,
			wantValues: constructorInputs[3].values,
		},
	}

	for _, ct := range constructorTests {
		if reflect.DeepEqual(ct.wantData, ct.dp.Data()) == false {
			t.Fail()
		}
		if reflect.DeepEqual(ct.wantValues, ct.dp.Set()) == false {
			t.Fail()
		}
	}
}

func TestDatapointSetCopy(t *testing.T) {

	var randomPoints Datapoints
	for i := uint(1); i <= 10; i++ {
		p := RandomDatapoint(i)
		randomPoints = append(randomPoints, p)
	}

	type dpst struct {
		dp             *Datapoint
		wantValues     []float64
		duplicateAddrs []*float64
	}

	var dpSetTests []dpst

	for i := range randomPoints {
		var values []float64
		copy(randomPoints[i].set, values)
		var f64ptrs []*float64
		for k := range randomPoints[i].set {
			values = append(values, randomPoints[i].set[k])
			f64ptrs = append(f64ptrs, &randomPoints[i].set[k])
		}
		dpSetTests = append(
			dpSetTests,
			dpst{
				randomPoints[i],
				values,
				f64ptrs,
			})
	}

	for i := range dpSetTests {
		gotSet := dpSetTests[i].dp.Set()
		// verify that the values were copied correctly
		if reflect.DeepEqual(gotSet, dpSetTests[i].wantValues) == false {
			t.Errorf(`wrong values returned by Set()
            got:       %v
            want:      %v`, gotSet, dpSetTests[i].wantValues)
		}
		// but we also need to make sure it's a true copy - i.e. not referring to the same underlying array/slice
		gotAddrs := getF64ptrs(gotSet)
		if diffFptrs(gotAddrs, dpSetTests[i].duplicateAddrs) == false {
			t.Errorf(`duplicate addresses returned by Set()
		    got:       %v
		    matching:  %v`, gotAddrs, dpSetTests[i].duplicateAddrs)
		}
	}
}

func TestDatapointLinkedDataIdentical(t *testing.T) {
	for k, ct := range constructorInputs {
		newDP := NewDatapoint(ct.linked, ct.values)
		if reflect.DeepEqual(newDP.Data(), interfaces[k]) == false {
			t.Error(k, `: Datapoint.Data() does not refer to the same object
			got:      `, newDP.Data(), `
			want:     `, interfaces[k])
		}

		srcRv := reflect.ValueOf(interfaces[k])
		dpRv := reflect.ValueOf(newDP.Data())

		if srcRv.Kind() != dpRv.Kind() { //	if they aren't to the same kind, then we have already failed.
			t.Error(k, `: Datapoint.Data() does not refer to an identical reflect.Kind
			got:      `, dpRv, `
			want:     `, srcRv)
		}

		errStringGot := fmt.Sprint(`got:      `, reflect.TypeOf(newDP.Data()), ` `, dpRv.Interface())
		errStringWant := fmt.Sprint(`want:     `, reflect.TypeOf(interfaces[k]), ` `, srcRv.Interface())

		if reflect.DeepEqual(reflect.TypeOf(newDP.Data()), reflect.TypeOf(interfaces[k])) == false {
			t.Error(k, `: Datapoint.Data() does not reflect an identical reflect.Type
			`, errStringGot, `
			`, errStringWant)
		}
	}
}

func TestDatapointDimensionality(t *testing.T) {
	var dimTests = []struct {
		dp   *Datapoint
		want int
	}{
		{RandomDatapoint(3), 3},
		{RandomDatapoint(0), 0},
	}

	for _, dt := range dimTests {
		got := dt.dp.Dimensionality()
		if got != dt.want {
			t.Error(`want: `, dt.want, `
			got: `, got)
		}
	}
}

func TestDatapointStringer(t *testing.T) {
	dpStringStr := `{data: cassandra}, {set: [0:{6.0000125}, 1:{6.10000125}, 2:{-1.3173}, 3:{1373}]}`
	dpRationalStr := `{data: &{{false [5]} {false [4]}}}, {set: [0:{1}, 1:{2}, 2:{3}, 3:{4}, 4:{5}]}`
	var stringerTests = []struct {
		dp   *Datapoint
		want string
	}{
		{
			&Datapoint{
				&str,
				[]float64{6.0000125, 6.10000125, -1.3173, 1373},
			},
			dpStringStr,
		},
		{
			&Datapoint{
				&rational,
				[]float64{1, 2, 3, 4, 5},
			},
			dpRationalStr,
		},
	}

	for _, s := range stringerTests {
		got := fmt.Sprint(s.dp)
		if got != s.want {
			t.Error(`want: `, s.want, `
			got: `, got)
		}
	}

}

func TestDatapointMarshalJSON(t *testing.T) {
	for k, ct := range constructorInputs {
		newDP := NewDatapoint(ct.linked, ct.values)
		dpJSON, _ := json.Marshal(newDP)
		var want []byte
		switch k {
		case 0:
			want = []byte(`{"data":"cassandra","set":[6.0000125,6.10000125,-1.3173,1373]}`)
		case 1:
			want = []byte(`{"data":"cassandra","set":[6.0000125,6.10000125,-1.3173,1373]}`)
		case 2:
			want = []byte(`{"data":"5/4","set":[1,2,3,4,5]}`)
		case 3:
			want = []byte(`{"data":"5/4","set":[1,2,3,4,5]}`)
		case 4:
			want = []byte(`{"data":{"intAndFloat":{"f":2.718281828459045,"i":0},"s":"hello world"},"set":[100.3]}`)
		case 5:
			want = []byte(`{"data":{"intAndFloat":{"f":2.718281828459045,"i":0},"s":"hello world"},"set":[100.3]}`)
		case 6:
			want = []byte(`{"data":{"a":97,"b":"banana","c":{"f":3.141592653589793,"i":0}},"set":[0.8050908121798804,0.53238545404102]}`)
		case 7:
			want = []byte(`{"data":{"a":97,"b":"banana","c":{"f":3.141592653589793,"i":0}},"set":[0.8050908121798804,0.53238545404102]}`)
		}
		if reflect.DeepEqual(dpJSON, want) == false {
			t.Error(`want: `, string(want), `
			got: `, string(dpJSON))
		}
	}
}

type char rune

func (c *char) ToDatapoint() *Datapoint {
	return &Datapoint{c, []float64{float64(*c) / 10.0, float64(*c) / 100.0}}
}

type myString string

func (str *myString) ToDatapoint() *Datapoint {
	var d Datapoint
	d.data = str
	var f []float64
	for _, s := range *str {
		f = append(f, math.Pi/float64(s))
	}
	d.set = f
	return &d
}

func TestImportableInterface1(t *testing.T) {
	A := char('a')
	B := char('b')
	S := myString("Hello!")
	importTests := []struct {
		imp  Importable
		want *Datapoint
	}{
		{
			imp:  &A,
			want: &Datapoint{&A, []float64{9.7, 0.97}},
		},
		{
			imp:  &B,
			want: &Datapoint{&B, []float64{9.8, 0.98}},
		},
		{
			imp:  &S,
			want: &Datapoint{&S, []float64{0.04363323129985824, 0.031104877758314782, 0.02908882086657216, 0.02908882086657216, 0.028302636518826967, 0.09519977738150888}},
		},
	}

	var dps Datapoints

	for k, it := range importTests {
		dps.Import(it.imp)
		if len(dps) != k+1 {
			t.Fail()
		}
		if reflect.DeepEqual(dps[k], it.want) == false {
			t.Error(`got: `, dps[k], `
		want: `, it.want)
		}
		if reflect.DeepEqual(reflect.TypeOf(dps[k].Data()), reflect.TypeOf(it.want.Data())) == false {
			t.Fail()
		}

		dpsRv := reflect.ValueOf(dps[k].Data())
		itwRv := reflect.ValueOf(it.want.Data())
		if reflect.DeepEqual(dpsRv, itwRv) == false {
			t.Errorf("\ngot:\t%v\nwant:\t%v\n", dpsRv, itwRv)
		}
		if reflect.DeepEqual(reflect.Indirect(dpsRv), reflect.Indirect(itwRv)) == false {
			t.Errorf("\ngot:\t%v\nwant:\t%v\n", reflect.Indirect(dpsRv), reflect.Indirect(itwRv))
		}
		if reflect.Indirect(dpsRv).CanSet() == false {
			t.Fail()
		}
	}
}
<|endoftext|>


<commit_before>package runner

import (
	"fmt"
	"log"
	"os"
	"os/exec"

	"github.com/Azure/acs-engine/test/e2e/config"
	"github.com/Azure/acs-engine/test/e2e/metrics"
	"github.com/kelseyhightower/envconfig"
)

// Ginkgo contains all of the information needed to run the ginkgo suite of tests
type Ginkgo struct {
	GinkgoNodes string `envconfig:"GINKGO_NODES" default:"6"`
	Config      *config.Config
	Point       *metrics.Point
}

// BuildGinkgoRunner creates a new Ginkgo object
func BuildGinkgoRunner(cfg *config.Config, pt *metrics.Point) (*Ginkgo, error) {
	g := new(Ginkgo)
	if err := envconfig.Process("ginkgo", g); err != nil {
		return nil, err
	}
	g.Config = cfg
	g.Point = pt
	return g, nil
}

// Run will execute an orchestrator suite of tests
func (g *Ginkgo) Run() error {
	g.Point.SetTestStart()
	testDir := fmt.Sprintf("test/e2e/%s", g.Config.Orchestrator)
	cmd := exec.Command("ginkgo", "-slowSpecThreshold", "180", "-r", testDir)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	err := cmd.Start()
	if err != nil {
		g.Point.RecordTestError()
		log.Printf("Error while trying to start ginkgo:%s\n", err)
		return err
	}

	err = cmd.Wait()
	if err != nil {
		g.Point.RecordTestError()
		if g.Config.IsKubernetes() {
			out, _ := exec.Command("kubectl", "get", "all", "-o", "wide").CombinedOutput()
			log.Printf("Running kubectl get all:\n%s\n", out)
			out, _ = exec.Command("kubectl", "get", "nodes", "-o", "wide").CombinedOutput()
			log.Printf("Running kubectl get nodes:\n%s\n", out)
		}
		return err
	}
	g.Point.RecordTestSuccess()
	return nil
}
<commit_msg>fix(runner/ginkgo.go): add --all-namespaces to debug output (#1747)<commit_after>package runner

import (
	"fmt"
	"log"
	"os"
	"os/exec"

	"github.com/Azure/acs-engine/test/e2e/config"
	"github.com/Azure/acs-engine/test/e2e/metrics"
	"github.com/kelseyhightower/envconfig"
)

// Ginkgo contains all of the information needed to run the ginkgo suite of tests
type Ginkgo struct {
	GinkgoNodes string `envconfig:"GINKGO_NODES" default:"6"`
	Config      *config.Config
	Point       *metrics.Point
}

// BuildGinkgoRunner creates a new Ginkgo object
func BuildGinkgoRunner(cfg *config.Config, pt *metrics.Point) (*Ginkgo, error) {
	g := new(Ginkgo)
	if err := envconfig.Process("ginkgo", g); err != nil {
		return nil, err
	}
	g.Config = cfg
	g.Point = pt
	return g, nil
}

// Run will execute an orchestrator suite of tests
func (g *Ginkgo) Run() error {
	g.Point.SetTestStart()
	testDir := fmt.Sprintf("test/e2e/%s", g.Config.Orchestrator)
	cmd := exec.Command("ginkgo", "-slowSpecThreshold", "180", "-r", testDir)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	err := cmd.Start()
	if err != nil {
		g.Point.RecordTestError()
		log.Printf("Error while trying to start ginkgo:%s\n", err)
		return err
	}

	err = cmd.Wait()
	if err != nil {
		g.Point.RecordTestError()
		if g.Config.IsKubernetes() {
			out, _ := exec.Command("kubectl", "get", "all", "--all-namespaces", "-o", "wide").CombinedOutput()
			log.Printf("Running kubectl get all:\n%s\n", out)
			out, _ = exec.Command("kubectl", "get", "nodes", "-o", "wide").CombinedOutput()
			log.Printf("Running kubectl get nodes:\n%s\n", out)
		}
		return err
	}
	g.Point.RecordTestSuccess()
	return nil
}
<|endoftext|>


<commit_before>/* https://leetcode.com/problems/sudoku-solver/#/description
Write a program to solve a Sudoku puzzle by filling the empty cells.

Empty cells are indicated by the character '.'.

You may assume that there will be only one unique solution.

http://upload.wikimedia.org/wikipedia/commons/thumb/f/ff/Sudoku-by-L2G-20050714.svg/250px-Sudoku-by-L2G-20050714.svg.png

A sudoku puzzle...

http://upload.wikimedia.org/wikipedia/commons/thumb/3/31/Sudoku-by-L2G-20050714_solution.svg/250px-Sudoku-by-L2G-20050714_solution.svg.png

..and its solution numbers marked in red.
*/

package leetcode

func solveSudoku(board [][]byte) {
	isValidSudoku := func(board [][]byte) bool {
		var rowMask, colMask, areaMask [9][9]bool
		for r := range board {
			for c := range board[r] {
				if board[r][c] != '.' {
					digit := board[r][c] - '0' - 1
					area := 3*(r/3) + c/3
					if rowMask[r][digit] || colMask[c][digit] || areaMask[area][digit] {
						return false
					}
					rowMask[r][digit], colMask[c][digit], areaMask[area][digit] = true, true, true
				}
			}
		}
		return true
	}

	var helper func(board [][]byte) bool
	helper = func(board [][]byte) bool {
		for row := range board {
			for col := range board[row] {
				if board[row][col] == '.' {
					for _, num := range []byte("123456789") {
						board[row][col] = num
						if !isValidSudoku(board) {
							board[row][col] = '.'
						} else {
							if helper(board) {
								return true
							} else {
								board[row][col] = '.'
							}
						}
					}
					return false
				}
			}
		}
		return true
	}

	helper(board)
}
<commit_msg>enhance isValid func, 113ms --> 12ms, exist 0ms solution<commit_after>/* https://leetcode.com/problems/sudoku-solver/#/description
Write a program to solve a Sudoku puzzle by filling the empty cells.

Empty cells are indicated by the character '.'.

You may assume that there will be only one unique solution.

http://upload.wikimedia.org/wikipedia/commons/thumb/f/ff/Sudoku-by-L2G-20050714.svg/250px-Sudoku-by-L2G-20050714.svg.png

A sudoku puzzle...

http://upload.wikimedia.org/wikipedia/commons/thumb/3/31/Sudoku-by-L2G-20050714_solution.svg/250px-Sudoku-by-L2G-20050714_solution.svg.png

..and its solution numbers marked in red.
*/

package leetcode

type Board [][]byte

func (board Board) isValid(row, col int, c byte) bool {
	i, j := row/3*3, col/3*3
	for k := 0; k < 9; k++ {
		if board[row][k] == c {
			return false
		}
		if board[k][col] == c {
			return false
		}
		if board[i+k%3][j+k/3] == c {
			return false
		}
	}
	return true
}

func (board Board) Solve() bool {
	for row := range board {
		for col := range board[row] {
			if board[row][col] != '.' {
				continue
			}
			for _, num := range []byte("123456789") {
				if !board.isValid(row, col, num) {
					continue
				}
				board[row][col] = num
				if board.Solve() {
					return true
				} else {
					board[row][col] = '.'
				}
			}
			return false
		}
	}
	return true
}

func solveSudoku(board [][]byte) {
	newBoard := Board(board)
	newBoard.Solve()
}
<|endoftext|>


<commit_before>package acceptance_test

import (
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"os"
	"os/exec"
	"vxlan-policy-agent/config"

	"code.cloudfoundry.org/garden"
	"code.cloudfoundry.org/garden/gardenfakes"
	"code.cloudfoundry.org/garden/server"
	"code.cloudfoundry.org/lager/lagertest"

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

var mockPolicyServer *httptest.Server

var _ = Describe("VXLAN Policy Agent", func() {
	var (
		session         *gexec.Session
		gardenBackend   *gardenfakes.FakeBackend
		gardenContainer *gardenfakes.FakeContainer
		gardenServer    *server.GardenServer
		logger          *lagertest.TestLogger
		subnetFile      *os.File
		configFilePath  string
	)
	BeforeEach(func() {
		mockPolicyServer = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			if r.URL.Path == "/networking/v0/internal/policies" {
				w.WriteHeader(http.StatusOK)
				w.Write([]byte(`{"policies": [
				{"source": {"id":"some-app-guid", "tag":"A"},
				"destination": {"id": "some-other-app-guid", "tag":"B", "protocol":"tcp", "port":3333}},
				{"source": {"id":"another-app-guid", "tag":"C"},
				"destination": {"id": "some-app-guid", "tag":"A", "protocol":"tcp", "port":9999}}
					]}`))
				return
			}

			w.WriteHeader(http.StatusNotFound)
			return
		}))

		logger = lagertest.NewTestLogger("fake garden server")
		gardenBackend = &gardenfakes.FakeBackend{}
		gardenContainer = &gardenfakes.FakeContainer{}
		gardenContainer.InfoReturns(garden.ContainerInfo{
			ContainerIP: "10.255.100.21",
			Properties:  garden.Properties{"network.app_id": "some-app-guid"},
		}, nil)
		gardenContainer.HandleReturns("some-handle")

		gardenBackend.CreateReturns(gardenContainer, nil)
		gardenBackend.LookupReturns(gardenContainer, nil)
		gardenBackend.ContainersReturns([]garden.Container{gardenContainer}, nil)

		gardenServer = server.New("tcp", ":60123", 0, gardenBackend, logger)
		Expect(gardenServer.Start()).To(Succeed())

		var err error
		subnetFile, err = ioutil.TempFile("", "")
		Expect(err).NotTo(HaveOccurred())
		Expect(ioutil.WriteFile(subnetFile.Name(), []byte("FLANNEL_NETWORK=10.255.0.0/16\nFLANNEL_SUBNET=10.255.100.1/24"), os.ModePerm))

		conf := config.VxlanPolicyAgent{
			PollInterval:      1,
			PolicyServerURL:   mockPolicyServer.URL,
			GardenAddress:     ":60123",
			GardenProtocol:    "tcp",
			VNI:               42,
			FlannelSubnetFile: subnetFile.Name(),
		}
		configFilePath = WriteConfigFile(conf)
	})

	AfterEach(func() {
		session.Interrupt()
		Eventually(session, DEFAULT_TIMEOUT).Should(gexec.Exit())

		gardenServer.Stop()

		_ = RunIptablesCommand("filter", "F")
		_ = RunIptablesCommand("filter", "X")
		_ = RunIptablesCommand("nat", "F")
		_ = RunIptablesCommand("nat", "X")
	})

	Describe("boring daemon behavior", func() {
		It("should boot and gracefully terminate", func() {
			session = StartAgent(vxlanPolicyAgentPath, configFilePath)
			Consistently(session).ShouldNot(gexec.Exit())

			session.Interrupt()
			Eventually(session, DEFAULT_TIMEOUT).Should(gexec.Exit())
		})
	})

	Describe("Default rules", func() {
		BeforeEach(func() {
			session = StartAgent(vxlanPolicyAgentPath, configFilePath)
		})
		It("writes the masquerade rule", func() {
			Consistently(session).ShouldNot(gexec.Exit())

			ipTablesRules := RunIptablesCommand("nat", "S")

			Expect(ipTablesRules).To(ContainSubstring("-s 10.255.100.0/24 ! -d 10.255.0.0/16 -j MASQUERADE"))
		})

		It("writes the default remote rules", func() {
			Consistently(session).ShouldNot(gexec.Exit())

			ipTablesRules := RunIptablesCommand("filter", "S")

			Expect(ipTablesRules).To(ContainSubstring("-i flannel.42 -m state --state RELATED,ESTABLISHED -j ACCEPT"))
			Expect(ipTablesRules).To(ContainSubstring("-i flannel.42 -j REJECT --reject-with icmp-port-unreachable"))
		})

		It("writes the default local rules", func() {
			Consistently(session).ShouldNot(gexec.Exit())

			ipTablesRules := RunIptablesCommand("filter", "S")

			Expect(ipTablesRules).To(ContainSubstring("-i cni-flannel0 -m state --state RELATED,ESTABLISHED -j ACCEPT"))
			Expect(ipTablesRules).To(ContainSubstring("-s 10.255.100.0/24 -d 10.255.100.0/24 -i cni-flannel0 -j REJECT --reject-with icmp-port-unreachable"))
		})
	})

	Describe("policy enforcement", func() {
		BeforeEach(func() {
			session = StartAgent(vxlanPolicyAgentPath, configFilePath)
		})
		It("writes the mark rule", func() {
			Consistently(session, DEFAULT_TIMEOUT).ShouldNot(gexec.Exit())

			ipTablesRules := RunIptablesCommand("filter", "S")

			Expect(ipTablesRules).To(ContainSubstring(`-s 10.255.100.21/32 -m comment --comment "src:some-app-guid" -j MARK --set-xmark 0xa/0xffffffff`))
		})
		It("enforces policies", func() {
			Consistently(session, DEFAULT_TIMEOUT).ShouldNot(gexec.Exit())

			ipTablesRules := RunIptablesCommand("filter", "S")

			Expect(ipTablesRules).To(ContainSubstring(`-d 10.255.100.21/32 -p tcp -m tcp --dport 9999 -m mark --mark 0xc -m comment --comment "src:another-app-guid dst:some-app-guid" -j ACCEPT`))
		})
	})

	Context("when the policy server is unavailable", func() {
		BeforeEach(func() {
			conf := config.VxlanPolicyAgent{
				PollInterval:      1,
				PolicyServerURL:   "",
				GardenAddress:     ":60123",
				GardenProtocol:    "tcp",
				VNI:               42,
				FlannelSubnetFile: subnetFile.Name(),
			}
			configFilePath = WriteConfigFile(conf)
			session = StartAgent(vxlanPolicyAgentPath, configFilePath)
		})

		It("still writes the default rules", func() {
			Consistently(session, DEFAULT_TIMEOUT).ShouldNot(gexec.Exit())

			ipTablesRules := RunIptablesCommand("filter", "S")
			Expect(ipTablesRules).To(ContainSubstring("-i flannel.42 -m state --state RELATED,ESTABLISHED -j ACCEPT"))
			Expect(ipTablesRules).To(ContainSubstring("-i flannel.42 -j REJECT --reject-with icmp-port-unreachable"))
			Expect(ipTablesRules).To(ContainSubstring("-i cni-flannel0 -m state --state RELATED,ESTABLISHED -j ACCEPT"))
			Expect(ipTablesRules).To(ContainSubstring("-s 10.255.100.0/24 -d 10.255.100.0/24 -i cni-flannel0 -j REJECT --reject-with icmp-port-unreachable"))

			ipTablesRules = RunIptablesCommand("nat", "S")
			Expect(ipTablesRules).To(ContainSubstring("-s 10.255.100.0/24 ! -d 10.255.0.0/16 -j MASQUERADE"))
		})
	})
})

func RunIptablesCommand(table, flag string) string {
	iptCmd := exec.Command("iptables", "-t", table, "-"+flag)
	iptablesSession, err := gexec.Start(iptCmd, GinkgoWriter, GinkgoWriter)
	Expect(err).NotTo(HaveOccurred())
	Eventually(iptablesSession).Should(gexec.Exit(0))
	return string(iptablesSession.Out.Contents())
}

func StartAgent(binaryPath, configPath string) *gexec.Session {
	cmd := exec.Command(binaryPath, "-config-file", configPath)
	session, err := gexec.Start(cmd, GinkgoWriter, GinkgoWriter)
	Expect(err).NotTo(HaveOccurred())
	return session
}
<commit_msg>Fix flakes in vxlan-policy-agent acceptance tests<commit_after>package acceptance_test

import (
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"os"
	"os/exec"
	"vxlan-policy-agent/config"

	"code.cloudfoundry.org/garden"
	"code.cloudfoundry.org/garden/gardenfakes"
	"code.cloudfoundry.org/garden/server"
	"code.cloudfoundry.org/lager/lagertest"

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

var mockPolicyServer *httptest.Server

var _ = Describe("VXLAN Policy Agent", func() {
	var (
		session         *gexec.Session
		gardenBackend   *gardenfakes.FakeBackend
		gardenContainer *gardenfakes.FakeContainer
		gardenServer    *server.GardenServer
		logger          *lagertest.TestLogger
		subnetFile      *os.File
		configFilePath  string
	)
	BeforeEach(func() {
		mockPolicyServer = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			if r.URL.Path == "/networking/v0/internal/policies" {
				w.WriteHeader(http.StatusOK)
				w.Write([]byte(`{"policies": [
				{"source": {"id":"some-app-guid", "tag":"A"},
				"destination": {"id": "some-other-app-guid", "tag":"B", "protocol":"tcp", "port":3333}},
				{"source": {"id":"another-app-guid", "tag":"C"},
				"destination": {"id": "some-app-guid", "tag":"A", "protocol":"tcp", "port":9999}}
					]}`))
				return
			}

			w.WriteHeader(http.StatusNotFound)
			return
		}))

		logger = lagertest.NewTestLogger("fake garden server")
		gardenBackend = &gardenfakes.FakeBackend{}
		gardenContainer = &gardenfakes.FakeContainer{}
		gardenContainer.InfoReturns(garden.ContainerInfo{
			ContainerIP: "10.255.100.21",
			Properties:  garden.Properties{"network.app_id": "some-app-guid"},
		}, nil)
		gardenContainer.HandleReturns("some-handle")

		gardenBackend.CreateReturns(gardenContainer, nil)
		gardenBackend.LookupReturns(gardenContainer, nil)
		gardenBackend.ContainersReturns([]garden.Container{gardenContainer}, nil)

		gardenServer = server.New("tcp", ":60123", 0, gardenBackend, logger)
		Expect(gardenServer.Start()).To(Succeed())

		var err error
		subnetFile, err = ioutil.TempFile("", "")
		Expect(err).NotTo(HaveOccurred())
		Expect(ioutil.WriteFile(subnetFile.Name(), []byte("FLANNEL_NETWORK=10.255.0.0/16\nFLANNEL_SUBNET=10.255.100.1/24"), os.ModePerm))

		conf := config.VxlanPolicyAgent{
			PollInterval:      1,
			PolicyServerURL:   mockPolicyServer.URL,
			GardenAddress:     ":60123",
			GardenProtocol:    "tcp",
			VNI:               42,
			FlannelSubnetFile: subnetFile.Name(),
		}
		configFilePath = WriteConfigFile(conf)
	})

	AfterEach(func() {
		session.Interrupt()
		Eventually(session, DEFAULT_TIMEOUT).Should(gexec.Exit())

		gardenServer.Stop()

		_ = RunIptablesCommand("filter", "F")
		_ = RunIptablesCommand("filter", "X")
		_ = RunIptablesCommand("nat", "F")
		_ = RunIptablesCommand("nat", "X")
	})

	Describe("boring daemon behavior", func() {
		It("should boot and gracefully terminate", func() {
			session = StartAgent(vxlanPolicyAgentPath, configFilePath)
			Consistently(session).ShouldNot(gexec.Exit())

			session.Interrupt()
			Eventually(session, DEFAULT_TIMEOUT).Should(gexec.Exit())
		})
	})

	var waitUntilPollLoop = func(numComplete int) {
		Eventually(gardenBackend.ContainersCallCount, DEFAULT_TIMEOUT).Should(BeNumerically(">=", numComplete+1))
	}

	Describe("Default rules", func() {
		BeforeEach(func() {
			session = StartAgent(vxlanPolicyAgentPath, configFilePath)
		})

		It("writes the masquerade rule", func() {
			waitUntilPollLoop(1)

			ipTablesRules := RunIptablesCommand("nat", "S")

			Expect(ipTablesRules).To(ContainSubstring("-s 10.255.100.0/24 ! -d 10.255.0.0/16 -j MASQUERADE"))
		})

		It("writes the default remote rules", func() {
			waitUntilPollLoop(1)

			ipTablesRules := RunIptablesCommand("filter", "S")

			Expect(ipTablesRules).To(ContainSubstring("-i flannel.42 -m state --state RELATED,ESTABLISHED -j ACCEPT"))
			Expect(ipTablesRules).To(ContainSubstring("-i flannel.42 -j REJECT --reject-with icmp-port-unreachable"))
		})

		It("writes the default local rules", func() {
			waitUntilPollLoop(1)

			ipTablesRules := RunIptablesCommand("filter", "S")

			Expect(ipTablesRules).To(ContainSubstring("-i cni-flannel0 -m state --state RELATED,ESTABLISHED -j ACCEPT"))
			Expect(ipTablesRules).To(ContainSubstring("-s 10.255.100.0/24 -d 10.255.100.0/24 -i cni-flannel0 -j REJECT --reject-with icmp-port-unreachable"))
		})
	})

	Describe("policy enforcement", func() {
		BeforeEach(func() {
			session = StartAgent(vxlanPolicyAgentPath, configFilePath)
		})
		It("writes the mark rule", func() {
			waitUntilPollLoop(2) // wait for a second one so we know the first enforcement completed

			ipTablesRules := RunIptablesCommand("filter", "S")

			Expect(ipTablesRules).To(ContainSubstring(`-s 10.255.100.21/32 -m comment --comment "src:some-app-guid" -j MARK --set-xmark 0xa/0xffffffff`))
		})
		It("enforces policies", func() {
			waitUntilPollLoop(2) // wait for a second one so we know the first enforcement completed

			ipTablesRules := RunIptablesCommand("filter", "S")

			Expect(ipTablesRules).To(ContainSubstring(`-d 10.255.100.21/32 -p tcp -m tcp --dport 9999 -m mark --mark 0xc -m comment --comment "src:another-app-guid dst:some-app-guid" -j ACCEPT`))
		})
	})

	Context("when the policy server is unavailable", func() {
		BeforeEach(func() {
			conf := config.VxlanPolicyAgent{
				PollInterval:      1,
				PolicyServerURL:   "",
				GardenAddress:     ":60123",
				GardenProtocol:    "tcp",
				VNI:               42,
				FlannelSubnetFile: subnetFile.Name(),
			}
			configFilePath = WriteConfigFile(conf)
			session = StartAgent(vxlanPolicyAgentPath, configFilePath)
		})

		It("still writes the default rules", func() {
			waitUntilPollLoop(2) // wait for a second one so we know the first enforcement completed

			ipTablesRules := RunIptablesCommand("filter", "S")
			Expect(ipTablesRules).To(ContainSubstring("-i flannel.42 -m state --state RELATED,ESTABLISHED -j ACCEPT"))
			Expect(ipTablesRules).To(ContainSubstring("-i flannel.42 -j REJECT --reject-with icmp-port-unreachable"))
			Expect(ipTablesRules).To(ContainSubstring("-i cni-flannel0 -m state --state RELATED,ESTABLISHED -j ACCEPT"))
			Expect(ipTablesRules).To(ContainSubstring("-s 10.255.100.0/24 -d 10.255.100.0/24 -i cni-flannel0 -j REJECT --reject-with icmp-port-unreachable"))

			ipTablesRules = RunIptablesCommand("nat", "S")
			Expect(ipTablesRules).To(ContainSubstring("-s 10.255.100.0/24 ! -d 10.255.0.0/16 -j MASQUERADE"))
		})
	})
})

func RunIptablesCommand(table, flag string) string {
	iptCmd := exec.Command("iptables", "-w", "-t", table, "-"+flag)
	iptablesSession, err := gexec.Start(iptCmd, GinkgoWriter, GinkgoWriter)
	Expect(err).NotTo(HaveOccurred())
	Eventually(iptablesSession, DEFAULT_TIMEOUT).Should(gexec.Exit(0))
	return string(iptablesSession.Out.Contents())
}

func StartAgent(binaryPath, configPath string) *gexec.Session {
	cmd := exec.Command(binaryPath, "-config-file", configPath)
	session, err := gexec.Start(cmd, GinkgoWriter, GinkgoWriter)
	Expect(err).NotTo(HaveOccurred())
	return session
}
<|endoftext|>


<commit_before>package gotest

import (
	"fmt"
	"strings"
	"time"

	"github.com/jstemmer/go-junit-report/v2/gtr"
	"github.com/jstemmer/go-junit-report/v2/parser/gotest/internal/collector"
)

const (
	globalID = 0
)

// reportBuilder helps build a test Report from a collection of events.
//
// The reportBuilder keeps track of the active context whenever a test or build
// error is created. This is necessary because the test parser do not contain
// any state themselves and simply just emit an event for every line that is
// read. By tracking the active context, any output that is appended to the
// reportBuilder gets attributed to the correct test or build error.
type reportBuilder struct {
	packages    []gtr.Package
	tests       map[int]gtr.Test
	buildErrors map[int]gtr.Error

	// state
	nextID    int               // next free unused id
	lastID    int               // most recently created id
	output    *collector.Output // output collected for each id
	coverage  float64           // coverage percentage
	parentIDs map[int]struct{}  // set of test id's that contain subtests

	// options
	packageName   string
	subtestMode   SubtestMode
	timestampFunc func() time.Time
}

// newReportBuilder creates a new reportBuilder.
func newReportBuilder() *reportBuilder {
	return &reportBuilder{
		tests:         make(map[int]gtr.Test),
		buildErrors:   make(map[int]gtr.Error),
		nextID:        1,
		output:        collector.New(),
		parentIDs:     make(map[int]struct{}),
		timestampFunc: time.Now,
	}
}

// ProcessEvent gives an event to this reportBuilder to be processed for this
// report.
func (b *reportBuilder) ProcessEvent(ev Event) {
	switch ev.Type {
	case "run_test":
		b.CreateTest(ev.Name)
	case "pause_test":
		b.PauseTest(ev.Name)
	case "cont_test":
		b.ContinueTest(ev.Name)
	case "end_test":
		b.EndTest(ev.Name, ev.Result, ev.Duration, ev.Indent)
	case "run_benchmark":
		b.CreateTest(ev.Name)
	case "benchmark":
		b.BenchmarkResult(ev.Name, ev.Iterations, ev.NsPerOp, ev.MBPerSec, ev.BytesPerOp, ev.AllocsPerOp)
	case "end_benchmark":
		b.EndTest(ev.Name, ev.Result, 0, 0)
	case "status":
		b.End()
	case "summary":
		b.CreatePackage(ev.Name, ev.Result, ev.Duration, ev.Data)
	case "coverage":
		b.Coverage(ev.CovPct, ev.CovPackages)
	case "build_output":
		b.CreateBuildError(ev.Name)
	case "output":
		b.AppendOutput(ev.Data)
	default:
		fmt.Printf("reportBuilder: unhandled event type: %v\n", ev.Type)
	}
}

// newID returns a new unique id and sets the active context this id.
func (b *reportBuilder) newID() int {
	id := b.nextID
	b.lastID = id
	b.nextID++
	return id
}

// flush creates a new package in this report containing any tests we've
// collected so far. This is necessary when a test did not end with a summary.
func (b *reportBuilder) flush() {
	if len(b.tests) > 0 {
		b.CreatePackage(b.packageName, "", 0, "")
	}
}

// Build returns the new Report containing all the tests and output created so
// far.
func (b *reportBuilder) Build() gtr.Report {
	b.flush()
	return gtr.Report{Packages: b.packages}
}

// CreateTest adds a test with the given name to the report, and marks it as
// active.
func (b *reportBuilder) CreateTest(name string) int {
	if parentID, ok := b.findTestParentID(name); ok {
		b.parentIDs[parentID] = struct{}{}
	}
	id := b.newID()
	b.tests[id] = gtr.NewTest(id, name)
	return id
}

// PauseTest marks the active context as no longer active. Any results or
// output added to the report after calling PauseTest will no longer be assumed
// to belong to this test.
func (b *reportBuilder) PauseTest(name string) {
	b.lastID = 0
}

// ContinueTest finds the test with the given name and marks it as active. If
// more than one test exist with this name, the most recently created test will
// be used.
func (b *reportBuilder) ContinueTest(name string) {
	b.lastID, _ = b.findTest(name)
}

// EndTest finds the test with the given name, sets the result, duration and
// level. If more than one test exists with this name, the most recently
// created test will be used. If no test exists with this name, a new test is
// created.
func (b *reportBuilder) EndTest(name, result string, duration time.Duration, level int) {
	id, ok := b.findTest(name)
	if !ok {
		// test did not exist, create one
		// TODO: Likely reason is that the user ran go test without the -v
		// flag, should we report this somewhere?
		id = b.CreateTest(name)
	}

	t := b.tests[id]
	t.Result = parseResult(result)
	t.Duration = duration
	t.Level = level
	b.tests[id] = t
	b.lastID = 0
}

// End marks the active context as no longer active.
func (b *reportBuilder) End() {
	b.lastID = 0
}

// BenchmarkResult updates an existing or adds a new test with the given
// results and marks it as active. If an existing test with this name exists
// but without result, then that one is updated. Otherwise a new one is added
// to the report.
func (b *reportBuilder) BenchmarkResult(name string, iterations int64, nsPerOp, mbPerSec float64, bytesPerOp, allocsPerOp int64) {
	id, ok := b.findTest(name)
	if !ok || b.tests[id].Result != gtr.Unknown {
		id = b.CreateTest(name)
	}

	benchmark := Benchmark{iterations, nsPerOp, mbPerSec, bytesPerOp, allocsPerOp}
	test := gtr.NewTest(id, name)
	test.Result = gtr.Pass
	test.Duration = benchmark.ApproximateDuration()
	SetBenchmarkData(&test, benchmark)
	b.tests[id] = test
}

// CreateBuildError creates a new build error and marks it as active.
func (b *reportBuilder) CreateBuildError(packageName string) {
	id := b.newID()
	b.buildErrors[id] = gtr.Error{ID: id, Name: packageName}
}

// CreatePackage adds a new package with the given name to the Report. This
// package contains all the build errors, output, tests and benchmarks created
// so far. Afterwards all state is reset.
func (b *reportBuilder) CreatePackage(name, result string, duration time.Duration, data string) {
	pkg := gtr.Package{
		Name:     name,
		Duration: duration,
	}

	if b.timestampFunc != nil {
		pkg.Timestamp = b.timestampFunc()
	}

	// Build errors are treated somewhat differently. Rather than having a
	// single package with all build errors collected so far, we only care
	// about the build errors for this particular package.
	for id, buildErr := range b.buildErrors {
		if buildErr.Name == name {
			if len(b.tests) > 0 {
				panic("unexpected tests found in build error package")
			}
			buildErr.ID = id
			buildErr.Duration = duration
			buildErr.Cause = data
			buildErr.Output = b.output.Get(id)

			pkg.BuildError = buildErr
			b.packages = append(b.packages, pkg)

			delete(b.buildErrors, id)
			// TODO: reset state
			// TODO: buildErrors shouldn't reset/use nextID/lastID, they're more like a global cache
			return
		}
	}

	// If we've collected output, but there were no tests then either there
	// actually were no tests, or there was some other non-build error.
	if b.output.Contains(globalID) && len(b.tests) == 0 {
		if parseResult(result) == gtr.Fail {
			pkg.RunError = gtr.Error{
				Name:   name,
				Output: b.output.Get(globalID),
			}
		} else if b.output.Contains(globalID) {
			pkg.Output = b.output.Get(globalID)
		}
		b.packages = append(b.packages, pkg)
		b.output.Clear(globalID)
		return
	}

	// If the summary result says we failed, but there were no failing tests
	// then something else must have failed.
	if parseResult(result) == gtr.Fail && len(b.tests) > 0 && !b.containsFailures() {
		pkg.RunError = gtr.Error{
			Name:   name,
			Output: b.output.Get(globalID),
		}
		b.output.Clear(globalID)
	}

	// Collect tests for this package, maintaining insertion order.
	var tests []gtr.Test
	for id := 1; id < b.nextID; id++ {
		if t, ok := b.tests[id]; ok {
			if b.isParent(id) {
				if b.subtestMode == IgnoreParentResults {
					t.Result = gtr.Pass
				} else if b.subtestMode == ExcludeParents {
					b.output.Merge(id, globalID)
					continue
				}
			}
			t.Output = b.output.Get(id)
			tests = append(tests, t)
			continue
		}
	}
	tests = b.groupBenchmarksByName(tests)

	pkg.Coverage = b.coverage
	pkg.Output = b.output.Get(globalID)
	pkg.Tests = tests
	b.packages = append(b.packages, pkg)

	// reset state, except for nextID to ensure all id's are unique.
	b.lastID = 0
	b.output.Clear(globalID)
	b.coverage = 0
	b.tests = make(map[int]gtr.Test)
	b.parentIDs = make(map[int]struct{})
}

// Coverage sets the code coverage percentage.
func (b *reportBuilder) Coverage(pct float64, packages []string) {
	b.coverage = pct
}

// AppendOutput appends the given text to the currently active context. If no
// active context exists, the output is assumed to belong to the package.
func (b *reportBuilder) AppendOutput(text string) {
	b.output.Append(b.lastID, text)
}

// findTest returns the id of the most recently created test with the given
// name if it exists.
func (b *reportBuilder) findTest(name string) (int, bool) {
	// check if this test was lastID
	if t, ok := b.tests[b.lastID]; ok && t.Name == name {
		return b.lastID, true
	}
	for i := b.nextID; i > 0; i-- {
		if test, ok := b.tests[i]; ok && test.Name == name {
			return i, true
		}
	}
	return 0, false
}

func (b *reportBuilder) findTestParentID(name string) (int, bool) {
	parent := dropLastSegment(name)
	for parent != "" {
		if id, ok := b.findTest(parent); ok {
			return id, true
		}
		parent = dropLastSegment(parent)
	}
	return 0, false
}

func (b *reportBuilder) isParent(id int) bool {
	_, ok := b.parentIDs[id]
	return ok
}

func dropLastSegment(name string) string {
	if idx := strings.LastIndexByte(name, '/'); idx >= 0 {
		return name[:idx]
	}
	return ""
}

// containsFailures return true if the current list of tests contains at least
// one failing test or an unknown result.
func (b *reportBuilder) containsFailures() bool {
	for _, test := range b.tests {
		if test.Result == gtr.Fail || test.Result == gtr.Unknown {
			return true
		}
	}
	return false
}

// parseResult returns a Result for the given string r.
func parseResult(r string) gtr.Result {
	switch r {
	case "PASS":
		return gtr.Pass
	case "FAIL":
		return gtr.Fail
	case "SKIP":
		return gtr.Skip
	case "BENCH":
		return gtr.Pass
	default:
		return gtr.Unknown
	}
}

func (b *reportBuilder) groupBenchmarksByName(tests []gtr.Test) []gtr.Test {
	if len(tests) == 0 {
		return nil
	}

	var grouped []gtr.Test
	byName := make(map[string][]gtr.Test)
	for _, test := range tests {
		if !strings.HasPrefix(test.Name, "Benchmark") {
			// If this test is not a benchmark, we won't group it by name but
			// just add it to the final result.
			grouped = append(grouped, test)
			continue
		}
		if _, ok := byName[test.Name]; !ok {
			grouped = append(grouped, gtr.NewTest(test.ID, test.Name))
		}
		byName[test.Name] = append(byName[test.Name], test)
	}

	for i, group := range grouped {
		if !strings.HasPrefix(group.Name, "Benchmark") {
			continue
		}
		var (
			ids   []int
			total Benchmark
			count int
		)
		for _, test := range byName[group.Name] {
			ids = append(ids, test.ID)
			if test.Result != gtr.Pass {
				continue
			}

			if bench, ok := GetBenchmarkData(test); ok {
				total.Iterations += bench.Iterations
				total.NsPerOp += bench.NsPerOp
				total.MBPerSec += bench.MBPerSec
				total.BytesPerOp += bench.BytesPerOp
				total.AllocsPerOp += bench.AllocsPerOp
				count++
			}
		}

		group.Duration = combinedDuration(byName[group.Name])
		group.Result = groupResults(byName[group.Name])
		group.Output = b.output.GetAll(ids...)
		if count > 0 {
			total.Iterations /= int64(count)
			total.NsPerOp /= float64(count)
			total.MBPerSec /= float64(count)
			total.BytesPerOp /= int64(count)
			total.AllocsPerOp /= int64(count)
			SetBenchmarkData(&group, total)
		}
		grouped[i] = group
	}
	return grouped
}

func combinedDuration(tests []gtr.Test) time.Duration {
	var total time.Duration
	for _, test := range tests {
		total += test.Duration
	}
	return total
}

func groupResults(tests []gtr.Test) gtr.Result {
	var result gtr.Result
	for _, test := range tests {
		if test.Result == gtr.Fail {
			return gtr.Fail
		}
		if result != gtr.Pass {
			result = test.Result
		}
	}
	return result
}
<commit_msg>parser/gotest: Switch to output collector for active id tracking<commit_after>package gotest

import (
	"fmt"
	"strings"
	"time"

	"github.com/jstemmer/go-junit-report/v2/gtr"
	"github.com/jstemmer/go-junit-report/v2/parser/gotest/internal/collector"
)

const (
	globalID = 0
)

// reportBuilder helps build a test Report from a collection of events.
//
// The reportBuilder keeps track of the active context whenever a test or build
// error is created. This is necessary because the test parser do not contain
// any state themselves and simply just emit an event for every line that is
// read. By tracking the active context, any output that is appended to the
// reportBuilder gets attributed to the correct test or build error.
type reportBuilder struct {
	packages    []gtr.Package
	tests       map[int]gtr.Test
	buildErrors map[int]gtr.Error

	// state
	nextID    int               // next free unused id
	output    *collector.Output // output collected for each id
	coverage  float64           // coverage percentage
	parentIDs map[int]struct{}  // set of test id's that contain subtests

	// options
	packageName   string
	subtestMode   SubtestMode
	timestampFunc func() time.Time
}

// newReportBuilder creates a new reportBuilder.
func newReportBuilder() *reportBuilder {
	return &reportBuilder{
		tests:         make(map[int]gtr.Test),
		buildErrors:   make(map[int]gtr.Error),
		nextID:        1,
		output:        collector.New(),
		parentIDs:     make(map[int]struct{}),
		timestampFunc: time.Now,
	}
}

// ProcessEvent gives an event to this reportBuilder to be processed for this
// report.
func (b *reportBuilder) ProcessEvent(ev Event) {
	switch ev.Type {
	case "run_test":
		b.CreateTest(ev.Name)
	case "pause_test":
		b.PauseTest(ev.Name)
	case "cont_test":
		b.ContinueTest(ev.Name)
	case "end_test":
		b.EndTest(ev.Name, ev.Result, ev.Duration, ev.Indent)
	case "run_benchmark":
		b.CreateTest(ev.Name)
	case "benchmark":
		b.BenchmarkResult(ev.Name, ev.Iterations, ev.NsPerOp, ev.MBPerSec, ev.BytesPerOp, ev.AllocsPerOp)
	case "end_benchmark":
		b.EndTest(ev.Name, ev.Result, 0, 0)
	case "status":
		b.End()
	case "summary":
		b.CreatePackage(ev.Name, ev.Result, ev.Duration, ev.Data)
	case "coverage":
		b.Coverage(ev.CovPct, ev.CovPackages)
	case "build_output":
		b.CreateBuildError(ev.Name)
	case "output":
		b.AppendOutput(ev.Data)
	default:
		fmt.Printf("reportBuilder: unhandled event type: %v\n", ev.Type)
	}
}

// newID returns a new unique id and sets the active context this id.
func (b *reportBuilder) newID() int {
	id := b.nextID
	b.nextID++
	return id
}

// flush creates a new package in this report containing any tests we've
// collected so far. This is necessary when a test did not end with a summary.
func (b *reportBuilder) flush() {
	if len(b.tests) > 0 {
		b.CreatePackage(b.packageName, "", 0, "")
	}
}

// Build returns the new Report containing all the tests and output created so
// far.
func (b *reportBuilder) Build() gtr.Report {
	b.flush()
	return gtr.Report{Packages: b.packages}
}

// CreateTest adds a test with the given name to the report, and marks it as
// active.
func (b *reportBuilder) CreateTest(name string) int {
	if parentID, ok := b.findTestParentID(name); ok {
		b.parentIDs[parentID] = struct{}{}
	}
	id := b.newID()
	b.output.SetActiveID(id)
	b.tests[id] = gtr.NewTest(id, name)
	return id
}

// PauseTest marks the active context as no longer active. Any results or
// output added to the report after calling PauseTest will no longer be assumed
// to belong to this test.
func (b *reportBuilder) PauseTest(name string) {
	b.output.SetActiveID(0)
}

// ContinueTest finds the test with the given name and marks it as active. If
// more than one test exist with this name, the most recently created test will
// be used.
func (b *reportBuilder) ContinueTest(name string) {
	id, _ := b.findTest(name)
	b.output.SetActiveID(id)
}

// EndTest finds the test with the given name, sets the result, duration and
// level. If more than one test exists with this name, the most recently
// created test will be used. If no test exists with this name, a new test is
// created.
func (b *reportBuilder) EndTest(name, result string, duration time.Duration, level int) {
	id, ok := b.findTest(name)
	if !ok {
		// test did not exist, create one
		// TODO: Likely reason is that the user ran go test without the -v
		// flag, should we report this somewhere?
		id = b.CreateTest(name)
	}

	t := b.tests[id]
	t.Result = parseResult(result)
	t.Duration = duration
	t.Level = level
	b.tests[id] = t
	b.output.SetActiveID(0)
}

// End marks the active context as no longer active.
func (b *reportBuilder) End() {
	b.output.SetActiveID(0)
}

// BenchmarkResult updates an existing or adds a new test with the given
// results and marks it as active. If an existing test with this name exists
// but without result, then that one is updated. Otherwise a new one is added
// to the report.
func (b *reportBuilder) BenchmarkResult(name string, iterations int64, nsPerOp, mbPerSec float64, bytesPerOp, allocsPerOp int64) {
	id, ok := b.findTest(name)
	if !ok || b.tests[id].Result != gtr.Unknown {
		id = b.CreateTest(name)
	}

	benchmark := Benchmark{iterations, nsPerOp, mbPerSec, bytesPerOp, allocsPerOp}
	test := gtr.NewTest(id, name)
	test.Result = gtr.Pass
	test.Duration = benchmark.ApproximateDuration()
	SetBenchmarkData(&test, benchmark)
	b.tests[id] = test
}

// CreateBuildError creates a new build error and marks it as active.
func (b *reportBuilder) CreateBuildError(packageName string) {
	id := b.newID()
	b.output.SetActiveID(id)
	b.buildErrors[id] = gtr.Error{ID: id, Name: packageName}
}

// CreatePackage adds a new package with the given name to the Report. This
// package contains all the build errors, output, tests and benchmarks created
// so far. Afterwards all state is reset.
func (b *reportBuilder) CreatePackage(name, result string, duration time.Duration, data string) {
	pkg := gtr.Package{
		Name:     name,
		Duration: duration,
	}

	if b.timestampFunc != nil {
		pkg.Timestamp = b.timestampFunc()
	}

	// Build errors are treated somewhat differently. Rather than having a
	// single package with all build errors collected so far, we only care
	// about the build errors for this particular package.
	for id, buildErr := range b.buildErrors {
		if buildErr.Name == name {
			if len(b.tests) > 0 {
				panic("unexpected tests found in build error package")
			}
			buildErr.ID = id
			buildErr.Duration = duration
			buildErr.Cause = data
			buildErr.Output = b.output.Get(id)

			pkg.BuildError = buildErr
			b.packages = append(b.packages, pkg)

			delete(b.buildErrors, id)
			// TODO: reset state
			// TODO: buildErrors shouldn't reset/use nextID/lastID, they're more like a global cache
			return
		}
	}

	// If we've collected output, but there were no tests then either there
	// actually were no tests, or there was some other non-build error.
	if b.output.Contains(globalID) && len(b.tests) == 0 {
		if parseResult(result) == gtr.Fail {
			pkg.RunError = gtr.Error{
				Name:   name,
				Output: b.output.Get(globalID),
			}
		} else if b.output.Contains(globalID) {
			pkg.Output = b.output.Get(globalID)
		}
		b.packages = append(b.packages, pkg)
		b.output.Clear(globalID)
		return
	}

	// If the summary result says we failed, but there were no failing tests
	// then something else must have failed.
	if parseResult(result) == gtr.Fail && len(b.tests) > 0 && !b.containsFailures() {
		pkg.RunError = gtr.Error{
			Name:   name,
			Output: b.output.Get(globalID),
		}
		b.output.Clear(globalID)
	}

	// Collect tests for this package, maintaining insertion order.
	var tests []gtr.Test
	for id := 1; id < b.nextID; id++ {
		if t, ok := b.tests[id]; ok {
			if b.isParent(id) {
				if b.subtestMode == IgnoreParentResults {
					t.Result = gtr.Pass
				} else if b.subtestMode == ExcludeParents {
					b.output.Merge(id, globalID)
					continue
				}
			}
			t.Output = b.output.Get(id)
			tests = append(tests, t)
			continue
		}
	}
	tests = b.groupBenchmarksByName(tests)

	pkg.Coverage = b.coverage
	pkg.Output = b.output.Get(globalID)
	pkg.Tests = tests
	b.packages = append(b.packages, pkg)

	// reset state, except for nextID to ensure all id's are unique.
	b.output.SetActiveID(0)
	b.output.Clear(globalID)
	b.coverage = 0
	b.tests = make(map[int]gtr.Test)
	b.parentIDs = make(map[int]struct{})
}

// Coverage sets the code coverage percentage.
func (b *reportBuilder) Coverage(pct float64, packages []string) {
	b.coverage = pct
}

// AppendOutput appends the given text to the currently active context. If no
// active context exists, the output is assumed to belong to the package.
func (b *reportBuilder) AppendOutput(text string) {
	b.output.Append(text)
}

// findTest returns the id of the most recently created test with the given
// name if it exists.
func (b *reportBuilder) findTest(name string) (int, bool) {
	for i := b.nextID; i > 0; i-- {
		if test, ok := b.tests[i]; ok && test.Name == name {
			return i, true
		}
	}
	return 0, false
}

func (b *reportBuilder) findTestParentID(name string) (int, bool) {
	parent := dropLastSegment(name)
	for parent != "" {
		if id, ok := b.findTest(parent); ok {
			return id, true
		}
		parent = dropLastSegment(parent)
	}
	return 0, false
}

func (b *reportBuilder) isParent(id int) bool {
	_, ok := b.parentIDs[id]
	return ok
}

func dropLastSegment(name string) string {
	if idx := strings.LastIndexByte(name, '/'); idx >= 0 {
		return name[:idx]
	}
	return ""
}

// containsFailures return true if the current list of tests contains at least
// one failing test or an unknown result.
func (b *reportBuilder) containsFailures() bool {
	for _, test := range b.tests {
		if test.Result == gtr.Fail || test.Result == gtr.Unknown {
			return true
		}
	}
	return false
}

// parseResult returns a Result for the given string r.
func parseResult(r string) gtr.Result {
	switch r {
	case "PASS":
		return gtr.Pass
	case "FAIL":
		return gtr.Fail
	case "SKIP":
		return gtr.Skip
	case "BENCH":
		return gtr.Pass
	default:
		return gtr.Unknown
	}
}

func (b *reportBuilder) groupBenchmarksByName(tests []gtr.Test) []gtr.Test {
	if len(tests) == 0 {
		return nil
	}

	var grouped []gtr.Test
	byName := make(map[string][]gtr.Test)
	for _, test := range tests {
		if !strings.HasPrefix(test.Name, "Benchmark") {
			// If this test is not a benchmark, we won't group it by name but
			// just add it to the final result.
			grouped = append(grouped, test)
			continue
		}
		if _, ok := byName[test.Name]; !ok {
			grouped = append(grouped, gtr.NewTest(test.ID, test.Name))
		}
		byName[test.Name] = append(byName[test.Name], test)
	}

	for i, group := range grouped {
		if !strings.HasPrefix(group.Name, "Benchmark") {
			continue
		}
		var (
			ids   []int
			total Benchmark
			count int
		)
		for _, test := range byName[group.Name] {
			ids = append(ids, test.ID)
			if test.Result != gtr.Pass {
				continue
			}

			if bench, ok := GetBenchmarkData(test); ok {
				total.Iterations += bench.Iterations
				total.NsPerOp += bench.NsPerOp
				total.MBPerSec += bench.MBPerSec
				total.BytesPerOp += bench.BytesPerOp
				total.AllocsPerOp += bench.AllocsPerOp
				count++
			}
		}

		group.Duration = combinedDuration(byName[group.Name])
		group.Result = groupResults(byName[group.Name])
		group.Output = b.output.GetAll(ids...)
		if count > 0 {
			total.Iterations /= int64(count)
			total.NsPerOp /= float64(count)
			total.MBPerSec /= float64(count)
			total.BytesPerOp /= int64(count)
			total.AllocsPerOp /= int64(count)
			SetBenchmarkData(&group, total)
		}
		grouped[i] = group
	}
	return grouped
}

func combinedDuration(tests []gtr.Test) time.Duration {
	var total time.Duration
	for _, test := range tests {
		total += test.Duration
	}
	return total
}

func groupResults(tests []gtr.Test) gtr.Result {
	var result gtr.Result
	for _, test := range tests {
		if test.Result == gtr.Fail {
			return gtr.Fail
		}
		if result != gtr.Pass {
			result = test.Result
		}
	}
	return result
}
<|endoftext|>


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

package main

import (
	"fmt"
	"os"
	"regexp"
	"strings"
)

func mail(args []string) {
	var (
		diff   = flags.Bool("diff", false, "show change commit diff and don't upload or mail")
		force  = flags.Bool("f", false, "mail even if there are staged changes")
		rList  = new(stringList) // installed below
		ccList = new(stringList) // installed below
	)
	flags.Var(rList, "r", "comma-separated list of reviewers")
	flags.Var(ccList, "cc", "comma-separated list of people to CC:")

	flags.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %s mail %s [-r reviewer,...] [-cc mail,...]\n", os.Args[0], globalFlags)
	}
	flags.Parse(args)
	if len(flags.Args()) != 0 {
		flags.Usage()
		os.Exit(2)
	}

	b := CurrentBranch()
	if b.ChangeID() == "" {
		dief("no pending change; can't mail.")
	}

	if *diff {
		run("git", "diff", "HEAD^..HEAD")
		return
	}

	if !*force && HasStagedChanges() {
		dief("there are staged changes; aborting.\n" +
			"Use 'review change' to include them or 'review mail -f' to force it.")
	}

	refSpec := "HEAD:refs/for/master"
	start := "%"
	if *rList != "" {
		refSpec += mailList(start, "r", string(*rList))
		start = ","
	}
	if *ccList != "" {
		refSpec += mailList(start, "cc", string(*ccList))
	}
	run("git", "push", "-q", "origin", refSpec)
}

// mailAddressRE matches the mail addresses we admit. It's restrictive but admits
// all the addresses in the Go CONTRIBUTORS file at time of writing (tested separately).
var mailAddressRE = regexp.MustCompile(`^[a-zA-Z0-9][-_.a-zA-Z0-9]*@[-_.a-zA-Z0-9]+$`)

// mailList turns the list of mail addresses from the flag value into the format
// expected by gerrit. The start argument is a % or , depending on where we
// are in the processing sequence.
func mailList(start, tag string, flagList string) string {
	spec := start
	for i, addr := range strings.Split(flagList, ",") {
		if !mailAddressRE.MatchString(addr) {
			dief("%q is not a valid reviewer mail address", addr)
		}
		if i > 0 {
			spec += ","
		}
		spec += tag + "=" + addr
	}
	return spec
}

// stringList is a flag.Value that is like flag.String, but if repeated
// keeps appending to the old value, inserting commas as separators.
// This allows people to write -r rsc,adg (like the old hg command)
// but also -r rsc -r adg (like standard git commands).
// This does change the meaning of -r rsc -r adg (it used to mean just adg).
type stringList string

func (x *stringList) String() string {
	return string(*x)
}

func (x *stringList) Set(s string) error {
	if *x != "" && s != "" {
		*x += ","
	}
	*x += stringList(s)
	return nil
}
<commit_msg>git-review: save reference to mailed revision using tag<commit_after>// Copyright 2014 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"fmt"
	"os"
	"regexp"
	"strings"
)

func mail(args []string) {
	var (
		diff   = flags.Bool("diff", false, "show change commit diff and don't upload or mail")
		force  = flags.Bool("f", false, "mail even if there are staged changes")
		rList  = new(stringList) // installed below
		ccList = new(stringList) // installed below
	)
	flags.Var(rList, "r", "comma-separated list of reviewers")
	flags.Var(ccList, "cc", "comma-separated list of people to CC:")

	flags.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %s mail %s [-r reviewer,...] [-cc mail,...]\n", os.Args[0], globalFlags)
	}
	flags.Parse(args)
	if len(flags.Args()) != 0 {
		flags.Usage()
		os.Exit(2)
	}

	b := CurrentBranch()
	if b.ChangeID() == "" {
		dief("no pending change; can't mail.")
	}

	if *diff {
		run("git", "diff", "HEAD^..HEAD")
		return
	}

	if !*force && HasStagedChanges() {
		dief("there are staged changes; aborting.\n" +
			"Use 'review change' to include them or 'review mail -f' to force it.")
	}

	refSpec := "HEAD:refs/for/master"
	start := "%"
	if *rList != "" {
		refSpec += mailList(start, "r", string(*rList))
		start = ","
	}
	if *ccList != "" {
		refSpec += mailList(start, "cc", string(*ccList))
	}
	run("git", "push", "-q", "origin", refSpec)

	// Create local tag for mailed change.
	// If in the 'work' branch, this creates or updates work.mailed.
	// Older mailings are in the reflog, so work.mailed is newest,
	// work.mailed@{1} is the one before that, work.mailed@{2} before that,
	// and so on.
	// Git doesn't actually have a concept of a local tag,
	// but Gerrit won't let people push tags to it, so the tag
	// can't propagate out of the local client into the official repo.
	// There is no conflict with the branch names people are using
	// for work, because git change rejects any name containing a dot.
	// The space of names with dots is ours (the Go team's) to define.
	run("git", "tag", "-f", b.Name+".mailed")
}

// mailAddressRE matches the mail addresses we admit. It's restrictive but admits
// all the addresses in the Go CONTRIBUTORS file at time of writing (tested separately).
var mailAddressRE = regexp.MustCompile(`^[a-zA-Z0-9][-_.a-zA-Z0-9]*@[-_.a-zA-Z0-9]+$`)

// mailList turns the list of mail addresses from the flag value into the format
// expected by gerrit. The start argument is a % or , depending on where we
// are in the processing sequence.
func mailList(start, tag string, flagList string) string {
	spec := start
	for i, addr := range strings.Split(flagList, ",") {
		if !mailAddressRE.MatchString(addr) {
			dief("%q is not a valid reviewer mail address", addr)
		}
		if i > 0 {
			spec += ","
		}
		spec += tag + "=" + addr
	}
	return spec
}

// stringList is a flag.Value that is like flag.String, but if repeated
// keeps appending to the old value, inserting commas as separators.
// This allows people to write -r rsc,adg (like the old hg command)
// but also -r rsc -r adg (like standard git commands).
// This does change the meaning of -r rsc -r adg (it used to mean just adg).
type stringList string

func (x *stringList) String() string {
	return string(*x)
}

func (x *stringList) Set(s string) error {
	if *x != "" && s != "" {
		*x += ","
	}
	*x += stringList(s)
	return nil
}
<|endoftext|>


<commit_before>package goprocessctx

import (
	"context"

	goprocess "github.com/jbenet/goprocess"
)

// WithContext constructs and returns a Process that respects
// given context. It is the equivalent of:
//
//   func ProcessWithContext(ctx context.Context) goprocess.Process {
//     p := goprocess.WithParent(goprocess.Background())
//     CloseAfterContext(p, ctx)
//     return p
//   }
//
func WithContext(ctx context.Context) goprocess.Process {
	p := goprocess.WithParent(goprocess.Background())
	CloseAfterContext(p, ctx)
	return p
}

// WithContextAndTeardown is a helper function to set teardown at initiation
// of WithContext
func WithContextAndTeardown(ctx context.Context, tf goprocess.TeardownFunc) goprocess.Process {
	p := goprocess.WithTeardown(tf)
	CloseAfterContext(p, ctx)
	return p
}

// WaitForContext makes p WaitFor ctx. When Closing, p waits for
// ctx.Done(), before being Closed(). It is simply:
//
//   p.WaitFor(goprocess.WithContext(ctx))
//
func WaitForContext(ctx context.Context, p goprocess.Process) {
	p.WaitFor(WithContext(ctx))
}

// CloseAfterContext schedules the process to close after the given
// context is done. It is the equivalent of:
//
//   func CloseAfterContext(p goprocess.Process, ctx context.Context) {
//     go func() {
//       <-ctx.Done()
//       p.Close()
//     }()
//   }
//
func CloseAfterContext(p goprocess.Process, ctx context.Context) {
	if p == nil {
		panic("nil Process")
	}
	if ctx == nil {
		panic("nil Context")
	}

	// context.Background(). if ctx.Done() is nil, it will never be done.
	// we check for this to avoid wasting a goroutine forever.
	if ctx.Done() == nil {
		return
	}

	go func() {
		<-ctx.Done()
		p.Close()
	}()
}

// WithProcessClosing returns a context.Context derived from ctx that
// is cancelled as p is Closing (after: <-p.Closing()). It is simply:
//
//   func WithProcessClosing(ctx context.Context, p goprocess.Process) context.Context {
//     ctx, cancel := context.WithCancel(ctx)
//     go func() {
//       <-p.Closing()
//       cancel()
//     }()
//     return ctx
//   }
//
func WithProcessClosing(ctx context.Context, p goprocess.Process) context.Context {
	ctx, cancel := context.WithCancel(ctx)
	goprocess.WithTeardown(func() error {
		cancel()
		return nil
	})
	return ctx
}

// WithProcessClosed returns a context.Context that is cancelled
// after Process p is Closed. It is the equivalent of:
//
//   func WithProcessClosed(ctx context.Context, p goprocess.Process) context.Context {
//     ctx, cancel := context.WithCancel(ctx)
//     go func() {
//       <-p.Closed()
//       cancel()
//     }()
//     return ctx
//   }
//
func WithProcessClosed(ctx context.Context, p goprocess.Process) context.Context {
	ctx, cancel := context.WithCancel(ctx)
	p.AddChildNoWait(goprocess.WithTeardown(func() error {
		<-p.Closed()
		cancel()
		return nil
	}))
	return ctx
}
<commit_msg>goprocessctx: actually cancel the context<commit_after>package goprocessctx

import (
	"context"

	goprocess "github.com/jbenet/goprocess"
)

// WithContext constructs and returns a Process that respects
// given context. It is the equivalent of:
//
//   func ProcessWithContext(ctx context.Context) goprocess.Process {
//     p := goprocess.WithParent(goprocess.Background())
//     CloseAfterContext(p, ctx)
//     return p
//   }
//
func WithContext(ctx context.Context) goprocess.Process {
	p := goprocess.WithParent(goprocess.Background())
	CloseAfterContext(p, ctx)
	return p
}

// WithContextAndTeardown is a helper function to set teardown at initiation
// of WithContext
func WithContextAndTeardown(ctx context.Context, tf goprocess.TeardownFunc) goprocess.Process {
	p := goprocess.WithTeardown(tf)
	CloseAfterContext(p, ctx)
	return p
}

// WaitForContext makes p WaitFor ctx. When Closing, p waits for
// ctx.Done(), before being Closed(). It is simply:
//
//   p.WaitFor(goprocess.WithContext(ctx))
//
func WaitForContext(ctx context.Context, p goprocess.Process) {
	p.WaitFor(WithContext(ctx))
}

// CloseAfterContext schedules the process to close after the given
// context is done. It is the equivalent of:
//
//   func CloseAfterContext(p goprocess.Process, ctx context.Context) {
//     go func() {
//       <-ctx.Done()
//       p.Close()
//     }()
//   }
//
func CloseAfterContext(p goprocess.Process, ctx context.Context) {
	if p == nil {
		panic("nil Process")
	}
	if ctx == nil {
		panic("nil Context")
	}

	// context.Background(). if ctx.Done() is nil, it will never be done.
	// we check for this to avoid wasting a goroutine forever.
	if ctx.Done() == nil {
		return
	}

	go func() {
		<-ctx.Done()
		p.Close()
	}()
}

// WithProcessClosing returns a context.Context derived from ctx that
// is cancelled as p is Closing (after: <-p.Closing()). It is simply:
//
//   func WithProcessClosing(ctx context.Context, p goprocess.Process) context.Context {
//     ctx, cancel := context.WithCancel(ctx)
//     go func() {
//       <-p.Closing()
//       cancel()
//     }()
//     return ctx
//   }
//
func WithProcessClosing(ctx context.Context, p goprocess.Process) context.Context {
	ctx, cancel := context.WithCancel(ctx)
	p.AddChildNoWait(goprocess.WithTeardown(func() error {
		cancel()
		return nil
	}))
	return ctx
}

// WithProcessClosed returns a context.Context that is cancelled
// after Process p is Closed. It is the equivalent of:
//
//   func WithProcessClosed(ctx context.Context, p goprocess.Process) context.Context {
//     ctx, cancel := context.WithCancel(ctx)
//     go func() {
//       <-p.Closed()
//       cancel()
//     }()
//     return ctx
//   }
//
func WithProcessClosed(ctx context.Context, p goprocess.Process) context.Context {
	ctx, cancel := context.WithCancel(ctx)
	p.AddChildNoWait(goprocess.WithTeardown(func() error {
		<-p.Closed()
		cancel()
		return nil
	}))
	return ctx
}
<|endoftext|>


<commit_before>package typhon

import (
	"bytes"
	"context"
	"encoding/json"
	"google.golang.org/protobuf/proto"
	"io/ioutil"
	"strings"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"github.com/monzo/typhon/prototest"
)

// TestRequestDecodeCloses verifies that a request body is closed after calling Decode()
func TestRequestDecodeCloses(t *testing.T) {
	t.Parallel()
	req := NewRequest(nil, "GET", "/", nil)
	b := []byte("{\"a\":\"b\"}\n")
	r := newDoneReader(ioutil.NopCloser(bytes.NewReader(b)), -1)
	req.Body = r

	bout := map[string]string{}
	req.Decode(&bout)
	select {
	case <-r.closed:
	default:
		assert.Fail(t, "response body was not closed after Decode()")
	}
}

// TestRequestEncodeReader verifies that passing an io.Reader to request.Encode() uses it properly as the body, and
// does not attempt to encode it as JSON
func TestRequestEncodeReader(t *testing.T) {
	t.Parallel()

	// io.ReadCloser: this should be used with no modification
	rc := ioutil.NopCloser(strings.NewReader("hello world"))
	req := NewRequest(nil, "GET", "/", nil)
	req.Encode(rc)
	assert.Equal(t, req.Body, rc)
	assert.EqualValues(t, -1, req.ContentLength)
	assert.Empty(t, req.Header.Get("Content-Type"))

	// io.Reader: this should be wrapped in an ioutil.NopCloser
	r := strings.NewReader("hello world, again")
	req = NewRequest(nil, "GET", "/", nil)
	req.Encode(r)
	assert.EqualValues(t, -1, req.ContentLength)
	assert.Empty(t, req.Header.Get("Content-Type"))
	body, err := ioutil.ReadAll(req.Body)
	require.NoError(t, err)
	assert.Equal(t, []byte("hello world, again"), body)

	// an io.ReadCloser that happens to implement json.Marshaler should not be used directly and should be marshaled
	jm := jsonMarshalerReader{
		ReadCloser: ioutil.NopCloser(strings.NewReader("this should never see the light of day"))}
	req = NewRequest(nil, "GET", "/", nil)
	req.Encode(jm)
	assert.EqualValues(t, 3, req.ContentLength)
	assert.Equal(t, "application/json", req.Header.Get("Content-Type"))
	body, err = ioutil.ReadAll(req.Body)
	require.NoError(t, err)
	assert.Equal(t, []byte("{}\n"), body)
}

func TestRequestEncodeProtobuf(t *testing.T) {
	g := &prototest.Greeting{
		Message:  "Hello world!",
		Priority: 1}

	protoContentForComparison, err := proto.Marshal(g)
	require.NoError(t, err)

	req := NewRequest(nil, "GET", "/", nil)
	req.EncodeAsProtobuf(g)

	bodyBytes, err := req.BodyBytes(false)
	require.NoError(t, err)

	assert.Equal(t, "application/protobuf", req.Header.Get("Content-Type"))
	assert.EqualValues(t, bodyBytes, protoContentForComparison)

}

func TestRequestEncodeJSON(t *testing.T) {
	message := map[string]interface{} {
		"foo": "bar",
		"bar": 3,
	}

	jsonContentForComparison, err := json.Marshal(message)
	require.NoError(t, err)

	req := NewRequest(nil, "GET", "/", nil)
	req.EncodeAsJSON(message)

	bodyBytes, err := req.BodyBytes(false)
	require.NoError(t, err)

	assert.Equal(t, "application/json", req.Header.Get("Content-Type"))
	assert.EqualValues(t, bodyBytes, jsonContentForComparison)
}

func TestRequestSetMetadata(t *testing.T) {
	t.Parallel()

	ctx := AppendMetadataToContext(context.Background(), NewMetadata(map[string]string{
		"meta": "data",
	}))

	req := NewRequest(ctx, "GET", "/", nil)

	assert.Equal(t, []string{"data"}, req.Request.Header["meta"])
}
<commit_msg>Consume body with ioutil<commit_after>package typhon

import (
	"bytes"
	"context"
	"encoding/json"
	"google.golang.org/protobuf/proto"
	"io/ioutil"
	"strings"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"github.com/monzo/typhon/prototest"
)

// TestRequestDecodeCloses verifies that a request body is closed after calling Decode()
func TestRequestDecodeCloses(t *testing.T) {
	t.Parallel()
	req := NewRequest(nil, "GET", "/", nil)
	b := []byte("{\"a\":\"b\"}\n")
	r := newDoneReader(ioutil.NopCloser(bytes.NewReader(b)), -1)
	req.Body = r

	bout := map[string]string{}
	req.Decode(&bout)
	select {
	case <-r.closed:
	default:
		assert.Fail(t, "response body was not closed after Decode()")
	}
}

// TestRequestEncodeReader verifies that passing an io.Reader to request.Encode() uses it properly as the body, and
// does not attempt to encode it as JSON
func TestRequestEncodeReader(t *testing.T) {
	t.Parallel()

	// io.ReadCloser: this should be used with no modification
	rc := ioutil.NopCloser(strings.NewReader("hello world"))
	req := NewRequest(nil, "GET", "/", nil)
	req.Encode(rc)
	assert.Equal(t, req.Body, rc)
	assert.EqualValues(t, -1, req.ContentLength)
	assert.Empty(t, req.Header.Get("Content-Type"))

	// io.Reader: this should be wrapped in an ioutil.NopCloser
	r := strings.NewReader("hello world, again")
	req = NewRequest(nil, "GET", "/", nil)
	req.Encode(r)
	assert.EqualValues(t, -1, req.ContentLength)
	assert.Empty(t, req.Header.Get("Content-Type"))
	body, err := ioutil.ReadAll(req.Body)
	require.NoError(t, err)
	assert.Equal(t, []byte("hello world, again"), body)

	// an io.ReadCloser that happens to implement json.Marshaler should not be used directly and should be marshaled
	jm := jsonMarshalerReader{
		ReadCloser: ioutil.NopCloser(strings.NewReader("this should never see the light of day"))}
	req = NewRequest(nil, "GET", "/", nil)
	req.Encode(jm)
	assert.EqualValues(t, 3, req.ContentLength)
	assert.Equal(t, "application/json", req.Header.Get("Content-Type"))
	body, err = ioutil.ReadAll(req.Body)
	require.NoError(t, err)
	assert.Equal(t, []byte("{}\n"), body)
}

func TestRequestEncodeProtobuf(t *testing.T) {
	g := &prototest.Greeting{
		Message:  "Hello world!",
		Priority: 1}

	protoContentForComparison, err := proto.Marshal(g)
	require.NoError(t, err)

	req := NewRequest(nil, "GET", "/", nil)
	req.EncodeAsProtobuf(g)

	bodyBytes, err := ioutil.ReadAll(req.Body)
	require.NoError(t, err)

	assert.Equal(t, "application/protobuf", req.Header.Get("Content-Type"))
	assert.EqualValues(t, bodyBytes, protoContentForComparison)

}

func TestRequestEncodeJSON(t *testing.T) {
	message := map[string]interface{} {
		"foo": "bar",
		"bar": 3,
	}

	jsonContentForComparison, err := json.Marshal(message)
	require.NoError(t, err)

	req := NewRequest(nil, "GET", "/", nil)
	req.EncodeAsJSON(message)

	bodyBytes, err := ioutil.ReadAll(req.Body)
	require.NoError(t, err)

	assert.Equal(t, "application/json", req.Header.Get("Content-Type"))
	assert.EqualValues(t, bodyBytes, jsonContentForComparison)
}

func TestRequestSetMetadata(t *testing.T) {
	t.Parallel()

	ctx := AppendMetadataToContext(context.Background(), NewMetadata(map[string]string{
		"meta": "data",
	}))

	req := NewRequest(ctx, "GET", "/", nil)

	assert.Equal(t, []string{"data"}, req.Request.Header["meta"])
}
<|endoftext|>


<commit_before>// +build linux,386

package system

import (
	"syscall"
)

// Setuid sets the uid of the calling thread to the specified uid.
func Setuid(uid int) (err error) {
	_, _, e1 := syscall.RawSyscall(syscall.SYS_SETUID, uintptr(uid), 0, 0)
	if e1 != 0 {
		err = e1
	}
	return
}

// Setgid sets the gid of the calling thread to the specified gid.
func Setgid(gid int) (err error) {
	_, _, e1 := syscall.RawSyscall(syscall.SYS_SETGID32, uintptr(gid), 0, 0)
	if e1 != 0 {
		err = e1
	}
	return
}
<commit_msg>Support 32 bit UID on i386<commit_after>// +build linux,386

package system

import (
	"syscall"
)

// Setuid sets the uid of the calling thread to the specified uid.
func Setuid(uid int) (err error) {
	_, _, e1 := syscall.RawSyscall(syscall.SYS_SETUID32, uintptr(uid), 0, 0)
	if e1 != 0 {
		err = e1
	}
	return
}

// Setgid sets the gid of the calling thread to the specified gid.
func Setgid(gid int) (err error) {
	_, _, e1 := syscall.RawSyscall(syscall.SYS_SETGID32, uintptr(gid), 0, 0)
	if e1 != 0 {
		err = e1
	}
	return
}
<|endoftext|>


<commit_before>package pxlocal

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/http/httputil"
	"strconv"
	"sync"

	"github.com/gorilla/websocket"
	"github.com/qiniu/log"
)

// :thinking
// start server HTTP service
// start agent
//  - agent connect server with websocket
//  - agent convert http request to conn
// need ref: revproxy

const (
	TCP_MIN_PORT = 13000
	TCP_MAX_PORT = 14000

	TYPE_NEWCONN = iota + 1
	TYPE_MESSAGE
)

var (
	upgrader = websocket.Upgrader{
		ReadBufferSize:  1024,
		WriteBufferSize: 1024,
	}
	namedConnection = make(map[string]chan net.Conn, 10)
	proxyStats      = &ProxyStats{}
)

type Msg struct {
	Type int
	Name string
	Body string
}

type Tunnel struct {
	wsconn *websocket.Conn
	sync.Mutex
	index    int64
	freeport *FreePort
}

func (t *Tunnel) uniqName() string {
	t.Lock()
	defer t.Unlock()
	t.index += 1
	return fmt.Sprintf("%d", t.index)
}

func (t *Tunnel) RequestNewConn(remoteAddr string) (net.Conn, error) {
	connCh := make(chan net.Conn)
	namedConnection[remoteAddr] = connCh
	defer delete(namedConnection, remoteAddr)

	// request a reverse connection
	var msg = Msg{Type: TYPE_NEWCONN, Name: remoteAddr}
	t.wsconn.WriteJSON(msg)
	lconn := <-connCh
	if lconn == nil {
		return nil, errors.New("maybe hijack not supported, failed")
	}
	return lconn, nil
}

// used for httputil reverse proxy
func (t *Tunnel) generateTransportDial() func(network, addr string) (net.Conn, error) {
	return func(network, addr string) (net.Conn, error) {
		log.Println("transport", network, addr)
		return t.RequestNewConn(t.uniqName())
	}
}

/*
func listenTcpInRangePort(port, minPort, maxPort int) (finnalPort int, lis *net.TCPListener, err error) {
	if port != 0 {
		laddr, err := net.ResolveTCPAddr("tcp", fmt.Sprintf(":%d", port))
		if err != nil {
			return 0, nil, err
		}
		lis, err = net.ListenTCP("tcp", laddr)
		return port, lis, err
	}
	for port = minPort; port < maxPort; port++ {
		laddr, err := net.ResolveTCPAddr("tcp", fmt.Sprintf(":%d", port))
		if err != nil {
			return 0, nil, err
		}
		lis, err := net.ListenTCP("tcp", laddr)
		if err == nil {
			return port, lis, nil
		}
	}
	return 0, nil, errors.New("No port avaliable")
}
*/

// Listen and forward connections
func NewTcpProxyListener(tunnel *Tunnel, port int) (listener *net.TCPListener, err error) {
	var laddr *net.TCPAddr
	if port != 0 {
		laddr, _ = net.ResolveTCPAddr("tcp", ":"+strconv.Itoa(port))
		listener, err = net.ListenTCP("tcp", laddr)
	} else {
		laddr, listener, err = tunnel.freeport.ListenTCP()
	}
	if err != nil {
		return nil, err
	}
	port = laddr.Port
	//port, listener, err = listenTcpInRangePort(port, TCP_MIN_PORT, TCP_MAX_PORT)
	//if err != nil {
	//return nil, err
	//}
	// hook here
	err = hook(HOOK_TCP_POST_CONNECT, []string{
		"PORT=" + strconv.Itoa(port),
		"CLIENT_ADDRESS=" + tunnel.wsconn.RemoteAddr().String(),
	})
	if err != nil {
		listener.Close()
		return
	}

	go func() {
		for {
			rconn, err := listener.AcceptTCP()
			if err != nil {
				log.Println(err)
				break
			}
			// find proxy to where
			log.Println("Receive new connections from", rconn.RemoteAddr())
			lconn, err := tunnel.RequestNewConn(rconn.RemoteAddr().String())
			if err != nil {
				log.Println("request new conn err:", err)
				rconn.Close()
				continue
			}

			log.Println("request new conn:", lconn, err)
			pc := &ProxyConn{
				lconn: lconn,
				rconn: rconn,
				stats: proxyStats,
			}
			go pc.start()
		}
	}()
	return listener, nil
}

func parseConnectRequest(r *http.Request) (protocal, subdomain string, port int) {
	protocal = r.FormValue("protocal")
	if protocal == "" {
		protocal = "http"
	}
	reqPort := r.FormValue("port")
	if reqPort == "" {
		port = 0
	} else {
		fmt.Sscanf(reqPort, "%d", &port)
	}
	subdomain = r.FormValue("subdomain")
	return
}

type HijactRW struct {
	*net.TCPConn
	bufrw *bufio.ReadWriter
}

func (this *HijactRW) Write(data []byte) (int, error) {
	nn, err := this.bufrw.Write(data)
	this.bufrw.Flush()
	return nn, err
}

func (this *HijactRW) Read(p []byte) (int, error) {
	return this.bufrw.Read(p)
}

func NewHijackReadWriteCloser(conn *net.TCPConn, bufrw *bufio.ReadWriter) net.Conn {
	return &HijactRW{
		bufrw:   bufrw,
		TCPConn: conn,
	}
}

func proxyHandler(w http.ResponseWriter, r *http.Request) {
	hj, ok := w.(http.Hijacker)
	if !ok {
		http.Error(w, "webserver don't support hijacking", http.StatusInternalServerError)
		return
	}

	var proxyFor = r.Header.Get("X-Proxy-For")
	log.Println("proxy name:", proxyFor)

	connCh, ok := namedConnection[proxyFor]
	if !ok {
		http.Error(w, "inside error: proxy not ready to receive conn", http.StatusInternalServerError)
		return
	}
	hjconn, bufrw, err := hj.Hijack()
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		connCh <- nil
		return
	}
	if _, ok := hjconn.(*net.TCPConn); ok {
		log.Println("Hijack is tcp conn")
	}

	conn := NewHijackReadWriteCloser(hjconn.(*net.TCPConn), bufrw)
	connCh <- conn
}

type ProxyServer struct {
	domain string
	*http.ServeMux
	revProxies map[string]*httputil.ReverseProxy
	sync.RWMutex
}

func wsSendMessage(conn *websocket.Conn, message string) error {
	return conn.WriteJSON(&Msg{Type: TYPE_MESSAGE, Body: message})
}

func (ps *ProxyServer) newHomepageHandler() func(w http.ResponseWriter, r *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		w.Header().Add("Content-Type", "text/html")
		io.WriteString(w, fmt.Sprintf("TCP: recvBytes: %d, sendBytes: %d
",
			proxyStats.receivedBytes, proxyStats.sentBytes))
		io.WriteString(w, fmt.Sprintf("HTTP: ...
"))
		io.WriteString(w, "
")
		for pname, _ := range ps.revProxies {
			io.WriteString(w, fmt.Sprintf("http proxy: %s
", pname))
		}
	}
}

func (ps *ProxyServer) newControlHandler() func(w http.ResponseWriter, r *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		// read listen port from request
		protocal, subdomain, port := parseConnectRequest(r)
		log.Println("proxy listen addr:", protocal, subdomain, port)

		// create websocket connection
		conn, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			http.Error(w, err.Error(), 502)
			return
		}
		defer conn.Close()
		log.Println(conn.RemoteAddr())

		tunnel := &Tunnel{
			wsconn:   conn,
			freeport: NewFreePort(TCP_MIN_PORT, TCP_MAX_PORT),
		}
		// TCP: create new port to listen
		switch protocal {
		case "tcp":
			// proxyAddr := fmt.Sprintf("0.0.0.0:%d", port)
			listener, err := NewTcpProxyListener(tunnel, port)
			if err != nil {
				log.Warnf("new tcp proxy err: %v", err)
				http.Error(w, err.Error(), 501)
				return
			}
			defer listener.Close()
			host, _, _ := net.SplitHostPort(ps.domain)
			_, port, _ := net.SplitHostPort(listener.Addr().String())
			wsSendMessage(conn, fmt.Sprintf(
				"Local tcp conn is now publicly available via:\n%v:%v\n", host, port))
		case "http", "https":
			log.Println("start http proxy")
			tr := &http.Transport{
				Dial: tunnel.generateTransportDial(),
			}
			revProxy := &httputil.ReverseProxy{
				Director: func(req *http.Request) {
					log.Println("director:", req.RequestURI)
				},
				Transport: tr,
			}
			// should hook here
			// hook(HOOK_CREATE_HTTP_SUBDOMAIN, subdomain)
			// generate a uniq domain
			if subdomain == "" {
				subdomain = uniqName(5) + ".t"
			}
			pxDomain := subdomain + "." + ps.domain
			log.Println("http px use domain:", pxDomain)
			if _, exists := ps.revProxies[pxDomain]; exists {
				wsSendMessage(conn, fmt.Sprintf("subdomain [%s] has already been taken", pxDomain))
				return
			}
			ps.Lock()
			ps.revProxies[pxDomain] = revProxy
			ps.Unlock()
			wsSendMessage(conn, fmt.Sprintf(
				"Local server is now publicly available via:\nhttp://%s\n", pxDomain))

			defer func() {
				ps.Lock()
				delete(ps.revProxies, pxDomain)
				ps.Unlock()
			}()
		default:
			log.Println("unknown protocal:", protocal)
			return
		}
		// HTTP: use httputil.ReverseProxy
		for {
			var msg Msg
			if err := conn.ReadJSON(&msg); err != nil {
				log.Println(err)
				break
			}
			log.Println("recv json:", msg)
		}
	}
}

func (p *ProxyServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	log.Println("request info:", r.Method, r.Host, r.RequestURI)
	//host, _, _ := net.SplitHostPort(r.Host)
	// http://stackoverflow.com/questions/6899069/why-are-request-url-host-and-scheme-blank-in-the-development-server
	r.URL.Scheme = "http" // ??
	r.URL.Host = r.Host   // ??
	log.Println("URL path:", r.URL.Path)
	log.Printf("pxies: %v", p.revProxies)
	if rpx, ok := p.revProxies[r.Host]; ok {
		log.Println("server http rev proxy")
		rpx.ServeHTTP(w, r)
		return
	}

	//if p.domain != host {
	//	http.Error(w, fmt.Sprintf("%s not ready", host), 504)
	//	return
	//}

	h, _ := p.Handler(r)
	h.ServeHTTP(w, r)
}

// domain, ex shengxiang.me
// dns should set *.shengxiang.me
func NewProxyServer(domain string) *ProxyServer {
	if domain == "" {
		domain = "localhost"
	}
	p := &ProxyServer{
		domain:     domain,
		ServeMux:   http.NewServeMux(),
		revProxies: make(map[string]*httputil.ReverseProxy),
	}
	p.HandleFunc("/", p.newHomepageHandler())
	p.HandleFunc("/ws", p.newControlHandler())
	p.HandleFunc("/proxyhijack", proxyHandler)

	return p
}
<commit_msg>global freeport<commit_after>package pxlocal

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/http/httputil"
	"strconv"
	"sync"

	"github.com/gorilla/websocket"
	"github.com/qiniu/log"
)

// :thinking
// start server HTTP service
// start agent
//  - agent connect server with websocket
//  - agent convert http request to conn
// need ref: revproxy

const (
	TCP_MIN_PORT = 13000
	TCP_MAX_PORT = 14000

	TYPE_NEWCONN = iota + 1
	TYPE_MESSAGE
)

var (
	upgrader = websocket.Upgrader{
		ReadBufferSize:  1024,
		WriteBufferSize: 1024,
	}
	namedConnection = make(map[string]chan net.Conn, 10)
	proxyStats      = &ProxyStats{}
)

type Msg struct {
	Type int
	Name string
	Body string
}

type Tunnel struct {
	wsconn *websocket.Conn
	sync.Mutex
	index int64
}

var freeport = NewFreePort(TCP_MIN_PORT, TCP_MAX_PORT)

func (t *Tunnel) uniqName() string {
	t.Lock()
	defer t.Unlock()
	t.index += 1
	return fmt.Sprintf("%d", t.index)
}

func (t *Tunnel) RequestNewConn(remoteAddr string) (net.Conn, error) {
	connCh := make(chan net.Conn)
	namedConnection[remoteAddr] = connCh
	defer delete(namedConnection, remoteAddr)

	// request a reverse connection
	var msg = Msg{Type: TYPE_NEWCONN, Name: remoteAddr}
	t.wsconn.WriteJSON(msg)
	lconn := <-connCh
	if lconn == nil {
		return nil, errors.New("maybe hijack not supported, failed")
	}
	return lconn, nil
}

// used for httputil reverse proxy
func (t *Tunnel) generateTransportDial() func(network, addr string) (net.Conn, error) {
	return func(network, addr string) (net.Conn, error) {
		log.Println("transport", network, addr)
		return t.RequestNewConn(t.uniqName())
	}
}

// Listen and forward connections
func NewTcpProxyListener(tunnel *Tunnel, port int) (listener *net.TCPListener, err error) {
	var laddr *net.TCPAddr
	if port != 0 {
		laddr, _ = net.ResolveTCPAddr("tcp", ":"+strconv.Itoa(port))
		listener, err = net.ListenTCP("tcp", laddr)
	} else {
		laddr, listener, err = freeport.ListenTCP()
	}
	if err != nil {
		return nil, err
	}
	port = laddr.Port
	//port, listener, err = listenTcpInRangePort(port, TCP_MIN_PORT, TCP_MAX_PORT)
	//if err != nil {
	//return nil, err
	//}
	// hook here
	err = hook(HOOK_TCP_POST_CONNECT, []string{
		"PORT=" + strconv.Itoa(port),
		"CLIENT_ADDRESS=" + tunnel.wsconn.RemoteAddr().String(),
	})
	if err != nil {
		listener.Close()
		return
	}

	go func() {
		for {
			rconn, err := listener.AcceptTCP()
			if err != nil {
				log.Println(err)
				break
			}
			// find proxy to where
			log.Println("Receive new connections from", rconn.RemoteAddr())
			lconn, err := tunnel.RequestNewConn(rconn.RemoteAddr().String())
			if err != nil {
				log.Println("request new conn err:", err)
				rconn.Close()
				continue
			}

			log.Println("request new conn:", lconn, err)
			pc := &ProxyConn{
				lconn: lconn,
				rconn: rconn,
				stats: proxyStats,
			}
			go pc.start()
		}
	}()
	return listener, nil
}

func parseConnectRequest(r *http.Request) (protocal, subdomain string, port int) {
	protocal = r.FormValue("protocal")
	if protocal == "" {
		protocal = "http"
	}
	reqPort := r.FormValue("port")
	if reqPort == "" {
		port = 0
	} else {
		fmt.Sscanf(reqPort, "%d", &port)
	}
	subdomain = r.FormValue("subdomain")
	return
}

type HijactRW struct {
	*net.TCPConn
	bufrw *bufio.ReadWriter
}

func (this *HijactRW) Write(data []byte) (int, error) {
	nn, err := this.bufrw.Write(data)
	this.bufrw.Flush()
	return nn, err
}

func (this *HijactRW) Read(p []byte) (int, error) {
	return this.bufrw.Read(p)
}

func NewHijackReadWriteCloser(conn *net.TCPConn, bufrw *bufio.ReadWriter) net.Conn {
	return &HijactRW{
		bufrw:   bufrw,
		TCPConn: conn,
	}
}

func proxyHandler(w http.ResponseWriter, r *http.Request) {
	hj, ok := w.(http.Hijacker)
	if !ok {
		http.Error(w, "webserver don't support hijacking", http.StatusInternalServerError)
		return
	}

	var proxyFor = r.Header.Get("X-Proxy-For")
	log.Println("proxy name:", proxyFor)

	connCh, ok := namedConnection[proxyFor]
	if !ok {
		http.Error(w, "inside error: proxy not ready to receive conn", http.StatusInternalServerError)
		return
	}
	hjconn, bufrw, err := hj.Hijack()
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		connCh <- nil
		return
	}
	if _, ok := hjconn.(*net.TCPConn); ok {
		log.Println("Hijack is tcp conn")
	}

	conn := NewHijackReadWriteCloser(hjconn.(*net.TCPConn), bufrw)
	connCh <- conn
}

type ProxyServer struct {
	domain string
	*http.ServeMux
	revProxies map[string]*httputil.ReverseProxy
	sync.RWMutex
}

func wsSendMessage(conn *websocket.Conn, message string) error {
	return conn.WriteJSON(&Msg{Type: TYPE_MESSAGE, Body: message})
}

func (ps *ProxyServer) newHomepageHandler() func(w http.ResponseWriter, r *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		w.Header().Add("Content-Type", "text/html")
		io.WriteString(w, fmt.Sprintf("TCP: recvBytes: %d, sendBytes: %d
",
			proxyStats.receivedBytes, proxyStats.sentBytes))
		io.WriteString(w, fmt.Sprintf("HTTP: ...
"))
		io.WriteString(w, "
")
		for pname, _ := range ps.revProxies {
			io.WriteString(w, fmt.Sprintf("http proxy: %s
", pname))
		}
	}
}

func (ps *ProxyServer) newControlHandler() func(w http.ResponseWriter, r *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		// read listen port from request
		protocal, subdomain, port := parseConnectRequest(r)
		log.Println("proxy listen addr:", protocal, subdomain, port)

		// create websocket connection
		conn, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			http.Error(w, err.Error(), 502)
			return
		}
		defer conn.Close()
		log.Println(conn.RemoteAddr())

		tunnel := &Tunnel{
			wsconn: conn,
		}
		// TCP: create new port to listen
		switch protocal {
		case "tcp":
			// proxyAddr := fmt.Sprintf("0.0.0.0:%d", port)
			listener, err := NewTcpProxyListener(tunnel, port)
			if err != nil {
				log.Warnf("new tcp proxy err: %v", err)
				http.Error(w, err.Error(), 501)
				return
			}
			defer listener.Close()
			host, _, _ := net.SplitHostPort(ps.domain)
			_, port, _ := net.SplitHostPort(listener.Addr().String())
			wsSendMessage(conn, fmt.Sprintf(
				"Local tcp conn is now publicly available via:\n%v:%v\n", host, port))
		case "http", "https":
			log.Println("start http proxy")
			tr := &http.Transport{
				Dial: tunnel.generateTransportDial(),
			}
			revProxy := &httputil.ReverseProxy{
				Director: func(req *http.Request) {
					log.Println("director:", req.RequestURI)
				},
				Transport: tr,
			}
			// should hook here
			// hook(HOOK_CREATE_HTTP_SUBDOMAIN, subdomain)
			// generate a uniq domain
			if subdomain == "" {
				subdomain = uniqName(5) + ".t"
			}
			pxDomain := subdomain + "." + ps.domain
			log.Println("http px use domain:", pxDomain)
			if _, exists := ps.revProxies[pxDomain]; exists {
				wsSendMessage(conn, fmt.Sprintf("subdomain [%s] has already been taken", pxDomain))
				return
			}
			ps.Lock()
			ps.revProxies[pxDomain] = revProxy
			ps.Unlock()
			wsSendMessage(conn, fmt.Sprintf(
				"Local server is now publicly available via:\nhttp://%s\n", pxDomain))

			defer func() {
				ps.Lock()
				delete(ps.revProxies, pxDomain)
				ps.Unlock()
			}()
		default:
			log.Println("unknown protocal:", protocal)
			return
		}
		// HTTP: use httputil.ReverseProxy
		for {
			var msg Msg
			if err := conn.ReadJSON(&msg); err != nil {
				log.Println(err)
				break
			}
			log.Println("recv json:", msg)
		}
	}
}

func (p *ProxyServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	log.Println("request info:", r.Method, r.Host, r.RequestURI)
	//host, _, _ := net.SplitHostPort(r.Host)
	// http://stackoverflow.com/questions/6899069/why-are-request-url-host-and-scheme-blank-in-the-development-server
	r.URL.Scheme = "http" // ??
	r.URL.Host = r.Host   // ??
	log.Println("URL path:", r.URL.Path)
	log.Printf("pxies: %v", p.revProxies)
	if rpx, ok := p.revProxies[r.Host]; ok {
		log.Println("server http rev proxy")
		rpx.ServeHTTP(w, r)
		return
	}

	//if p.domain != host {
	//	http.Error(w, fmt.Sprintf("%s not ready", host), 504)
	//	return
	//}

	h, _ := p.Handler(r)
	h.ServeHTTP(w, r)
}

// domain, ex shengxiang.me
// dns should set *.shengxiang.me
func NewProxyServer(domain string) *ProxyServer {
	if domain == "" {
		domain = "localhost"
	}
	p := &ProxyServer{
		domain:     domain,
		ServeMux:   http.NewServeMux(),
		revProxies: make(map[string]*httputil.ReverseProxy),
	}
	p.HandleFunc("/", p.newHomepageHandler())
	p.HandleFunc("/ws", p.newControlHandler())
	p.HandleFunc("/proxyhijack", proxyHandler)

	return p
}
<|endoftext|>