<commit_before>package api

import (
	"net/http"
	"time"

	"nimona.io/internal/http/router"
	"nimona.io/internal/store/kv"
	"nimona.io/pkg/context"
	"nimona.io/pkg/crypto"
	"nimona.io/pkg/errors"
	"nimona.io/pkg/object"
	"nimona.io/pkg/stream"
)

func (api *API) HandleGetObjects(c *router.Context) {
	// TODO this will be replaced by manager.Subscribe()
	ms, err := api.objectStore.Heads()
	if err != nil {
		c.AbortWithError(500, err) // nolint: errcheck
		return
	}
	// ms := []interface{}{}
	// for _, objectHash := range objectHashs {
	// 	b, err := api.objectStore.Get(objectHash)
	// 	if err != nil {
	// 		c.AbortWithError(500, err) // nolint: errcheck
	// 		return
	// 	}
	// 	m, err := object.Unmarshal(b)
	// 	if err != nil {
	// 		c.AbortWithError(500, err) // nolint: errcheck
	// 		return
	// 	}
	// 	ms = append(ms, api.mapObject(m))
	// }
	c.JSON(http.StatusOK, api.mapObjects(ms))
	// c.JSON(http.StatusNotImplemented, nil)
}

func (api *API) HandleGetObject(c *router.Context) {
	objectHash := c.Param("objectHash")
	if objectHash == "" {
		c.AbortWithError(400, errors.New("missing object hash")) // nolint: errcheck
	}
	o, err := api.objectStore.Get(objectHash)
	if err != nil && err != kv.ErrNotFound {
		c.AbortWithError(500, err) // nolint: errcheck
		return
	} else if err == nil {
		ms := api.mapObject(o)
		c.JSON(http.StatusOK, ms)
		return
	}

	ctx := context.New(context.WithTimeout(time.Second * 5))
	defer ctx.Cancel()

	h, _ := object.HashFromCompact(objectHash)
	ps, err := api.discovery.FindByContent(ctx, h)
	if err != nil {
		c.AbortWithError(500, err) // nolint: errcheck
		return
	}
	addrs := []string{}
	for _, p := range ps {
		addrs = append(addrs, p.Address())
	}
	hs := []*object.Hash{h}
	os, err := api.orchestrator.Sync(ctx, hs, addrs)
	if err != nil {
		if err == kv.ErrNotFound {
			c.AbortWithError(404, err) // nolint: errcheck
			return
		}
		c.AbortWithError(500, err) // nolint: errcheck
		return
	}
	if len(os.Objects) == 0 {
		c.AbortWithError(404, err) // nolint: errcheck
		return
	}
	ms := api.mapObject(os.Objects[0])
	c.JSON(http.StatusOK, ms)
}

func (api *API) HandlePostObject(c *router.Context) {
	req := map[string]interface{}{}
	if err := c.BindBody(&req); err != nil {
		c.AbortWithError(400, err) // nolint: errcheck
		return
	}

	o := object.FromMap(req)
	op := stream.Policies(o)
	if len(op) == 0 {
		c.AbortWithError(400, errors.New("missing policy")) // nolint: errcheck
		return
	}

	if err := crypto.Sign(o, api.local.GetPeerKey()); err != nil {
		c.AbortWithError(500, errors.New("could not sign object")) // nolint: errcheck
		return
	}

	for _, p := range op {
		for _, s := range p.Subjects {
			ctx := context.Background()
			addr := "peer:" + s
			if err := api.exchange.Send(ctx, o, addr); err != nil {
				c.AbortWithError(500, err) // nolint: errcheck
				return
			}
		}
	}

	c.JSON(http.StatusOK, nil)
}
<commit_msg>chore(api): make sending object to subjects async<commit_after>package api

import (
	"net/http"
	"time"

	"nimona.io/internal/http/router"
	"nimona.io/internal/store/kv"
	"nimona.io/pkg/context"
	"nimona.io/pkg/crypto"
	"nimona.io/pkg/errors"
	"nimona.io/pkg/object"
	"nimona.io/pkg/stream"
)

func (api *API) HandleGetObjects(c *router.Context) {
	// TODO this will be replaced by manager.Subscribe()
	ms, err := api.objectStore.Heads()
	if err != nil {
		c.AbortWithError(500, err) // nolint: errcheck
		return
	}
	// ms := []interface{}{}
	// for _, objectHash := range objectHashs {
	// 	b, err := api.objectStore.Get(objectHash)
	// 	if err != nil {
	// 		c.AbortWithError(500, err) // nolint: errcheck
	// 		return
	// 	}
	// 	m, err := object.Unmarshal(b)
	// 	if err != nil {
	// 		c.AbortWithError(500, err) // nolint: errcheck
	// 		return
	// 	}
	// 	ms = append(ms, api.mapObject(m))
	// }
	c.JSON(http.StatusOK, api.mapObjects(ms))
	// c.JSON(http.StatusNotImplemented, nil)
}

func (api *API) HandleGetObject(c *router.Context) {
	objectHash := c.Param("objectHash")
	if objectHash == "" {
		c.AbortWithError(400, errors.New("missing object hash")) // nolint: errcheck
	}
	o, err := api.objectStore.Get(objectHash)
	if err != nil && err != kv.ErrNotFound {
		c.AbortWithError(500, err) // nolint: errcheck
		return
	} else if err == nil {
		ms := api.mapObject(o)
		c.JSON(http.StatusOK, ms)
		return
	}

	ctx := context.New(context.WithTimeout(time.Second * 5))
	defer ctx.Cancel()

	h, _ := object.HashFromCompact(objectHash)
	ps, err := api.discovery.FindByContent(ctx, h)
	if err != nil {
		c.AbortWithError(500, err) // nolint: errcheck
		return
	}
	addrs := []string{}
	for _, p := range ps {
		addrs = append(addrs, p.Address())
	}
	hs := []*object.Hash{h}
	os, err := api.orchestrator.Sync(ctx, hs, addrs)
	if err != nil {
		if err == kv.ErrNotFound {
			c.AbortWithError(404, err) // nolint: errcheck
			return
		}
		c.AbortWithError(500, err) // nolint: errcheck
		return
	}
	if len(os.Objects) == 0 {
		c.AbortWithError(404, err) // nolint: errcheck
		return
	}
	ms := api.mapObject(os.Objects[0])
	c.JSON(http.StatusOK, ms)
}

func (api *API) HandlePostObject(c *router.Context) {
	req := map[string]interface{}{}
	if err := c.BindBody(&req); err != nil {
		c.AbortWithError(400, err) // nolint: errcheck
		return
	}

	o := object.FromMap(req)
	op := stream.Policies(o)
	if len(op) == 0 {
		c.AbortWithError(400, errors.New("missing policy")) // nolint: errcheck
		return
	}

	if err := crypto.Sign(o, api.local.GetPeerKey()); err != nil {
		c.AbortWithError(500, errors.New("could not sign object")) // nolint: errcheck
		return
	}

	for _, p := range op {
		for _, s := range p.Subjects {
			ctx := context.New(
				context.WithTimeout(time.Second * 5),
			)
			go api.exchange.Send(ctx, o, "peer:"+s)
		}
	}

	c.JSON(http.StatusOK, nil)
}
<|endoftext|>


<commit_before>package monitor

import (
	"fmt"
	"reflect"
	"strconv"
	"testing"
	"time"

	"github.com/turnage/graw/internal/operator"
	"github.com/turnage/redditproto"
)

type handler struct {
	calls        int
	postCalls    int
	commentCalls int
	messageCalls int
}

func (h *handler) post(post *redditproto.Link) {
	h.postCalls++
	h.calls++
}

func (h *handler) comment(comment *redditproto.Comment) {
	h.commentCalls++
	h.calls++
}

func (h *handler) message(message *redditproto.Message) {
	h.messageCalls++
	h.calls++
}

func float64Pointer(val float64) *float64 {
	return &val
}

func stringPointer(val string) *string {
	return &val
}

func TestBaseFromPath(t *testing.T) {
	han := &handler{}
	mon, err := baseFromPath(
		&operator.MockOperator{
			ScrapeLinksReturn: []*redditproto.Link{
				&redditproto.Link{
					Name: stringPointer("name"),
				},
			},
		},
		"/r/self",
		nil,
		han.comment,
		nil,
		Forward,
	)
	if err != nil {
		t.Fatal(err)
	}

	b := mon.(*base)
	if b.dir != Forward {
		t.Errorf("got %d; wanted %d (Forward)", b.dir, Forward)
	}
	if b.handleComment == nil {
		t.Errorf("wanted comment handler set")
	}
	if b.path != "/r/self" {
		t.Errorf("got %s; wanted /r/self", b.path)
	}
	mon, err = baseFromPath(
		&operator.MockOperator{
			ScrapeLinksReturn: []*redditproto.Link{
				&redditproto.Link{
					Name: stringPointer("name"),
				},
			},
		},
		"/r/self",
		nil,
		nil,
		nil,
		Forward,
	)
	if err == nil {
		t.Errorf("wanted error if no handlers are provided")
	}
}

func TestMerge(t *testing.T) {
	things := merge(
		[]*redditproto.Link{
			&redditproto.Link{
				CreatedUtc: float64Pointer(1),
				Name:       stringPointer("one"),
			},
			&redditproto.Link{
				CreatedUtc: float64Pointer(2),
				Name:       stringPointer("two"),
			},
		},
		[]*redditproto.Comment{
			&redditproto.Comment{
				CreatedUtc: float64Pointer(0),
				Name:       stringPointer("zero"),
			},
			&redditproto.Comment{
				CreatedUtc: float64Pointer(3),
				Name:       stringPointer("three"),
			},
		},
		[]*redditproto.Message{
			&redditproto.Message{
				CreatedUtc: float64Pointer(4),
				Name:       stringPointer("four"),
			},
			&redditproto.Message{
				CreatedUtc: float64Pointer(5),
				Name:       stringPointer("five"),
			},
		},
		Backward,
	)

	if len(things) != 6 {
		t.Fatalf("got %d things; wanted 6", len(things))
	}

	if things[0].GetName() != "zero" {
		t.Errorf("got %s; wanted zero", things[0].GetName())
	}

	if things[1].GetName() != "one" {
		t.Errorf("got %s; wanted one", things[1].GetName())
	}

	if things[2].GetName() != "two" {
		t.Errorf("got %s; wanted two", things[2].GetName())
	}

	if things[3].GetName() != "three" {
		t.Errorf("got %s; wanted three", things[3].GetName())
	}

	if things[4].GetName() != "four" {
		t.Errorf("got %s; wanted four", things[4].GetName())
	}

	if things[5].GetName() != "five" {
		t.Errorf("got %s; wanted five", things[5].GetName())
	}

	things = merge(
		[]*redditproto.Link{
			&redditproto.Link{
				CreatedUtc: float64Pointer(2),
				Name:       stringPointer("two"),
			},
			&redditproto.Link{
				CreatedUtc: float64Pointer(1),
				Name:       stringPointer("one"),
			},
		},
		[]*redditproto.Comment{
			&redditproto.Comment{
				CreatedUtc: float64Pointer(3),
				Name:       stringPointer("three"),
			},
			&redditproto.Comment{
				CreatedUtc: float64Pointer(0),
				Name:       stringPointer("zero"),
			},
		},
		[]*redditproto.Message{
			&redditproto.Message{
				CreatedUtc: float64Pointer(5),
				Name:       stringPointer("five"),
			},
			&redditproto.Message{
				CreatedUtc: float64Pointer(4),
				Name:       stringPointer("four"),
			},
		},
		Forward,
	)

	if len(things) != 6 {
		t.Fatalf("got %d things; wanted 6", len(things))
	}

	if things[0].GetName() != "five" {
		t.Errorf("got %s; wanted five", things[0].GetName())
	}

	if things[1].GetName() != "four" {
		t.Errorf("got %s; wanted four", things[1].GetName())
	}

	if things[2].GetName() != "three" {
		t.Errorf("got %s; wanted three", things[2].GetName())
	}

	if things[3].GetName() != "two" {
		t.Errorf("got %s; wanted two", things[3].GetName())
	}

	if things[4].GetName() != "one" {
		t.Errorf("got %s; wanted one", things[4].GetName())
	}

	if things[5].GetName() != "zero" {
		t.Errorf("got %s; wanted zero", things[5].GetName())
	}
}

func TestShaveTip(t *testing.T) {
	b := &base{
		tip: []string{"1", "2"},
	}
	b.shaveTip()
	if !reflect.DeepEqual(b.tip, []string{"2"}) {
		t.Errorf("got %v\n; wanted %v", b.tip, []string{"2"})
	}

	b = &base{
		tip: nil,
	}
	b.shaveTip()
	if !reflect.DeepEqual(b.tip, []string{""}) {
		t.Errorf("got %v\n; wanted %v", b.tip, []string{""})
	}
}

func TestFixTip(t *testing.T) {
	b := &base{
		tip: []string{"1"},
	}

	broken, err := b.fixTip(
		&operator.MockOperator{
			IsThereThingErr: fmt.Errorf("an error"),
		},
	)
	if err == nil {
		t.Errorf("wanted error propagated from operator error")
	}

	broken, err = b.fixTip(
		&operator.MockOperator{
			IsThereThingReturn: false,
		},
	)
	if !broken {
		t.Errorf("got false; wanted true")
	}
}

func TestUpdateTip(t *testing.T) {
	b := &base{
		dir: Forward,
	}
	for i := 0; i < maxTipSize; i++ {
		b.tip = append(b.tip, strconv.Itoa(i))
	}

	err := b.updateTip(
		[]*redditproto.Link{
			&redditproto.Link{
				CreatedUtc: float64Pointer(2),
				Name:       stringPointer("two"),
			},
			&redditproto.Link{
				CreatedUtc: float64Pointer(1),
				Name:       stringPointer("one"),
			},
		},
		nil,
		nil,
		nil,
	)
	if err != nil {
		t.Fatal(err)
	}

	if len(b.tip) != maxTipSize {
		t.Fatalf("got size %d tip log; wanted %d", len(b.tip), maxTipSize)
	}

	if b.tip[0] != "one" {
		t.Errorf("got %s; wanted one", b.tip[0])
	}

	if b.tip[1] != "two" {
		t.Errorf("got %s; wanted two", b.tip[1])
	}

	err = b.updateTip(
		nil,
		nil,
		nil,
		&operator.MockOperator{
			IsThereThingErr: fmt.Errorf("an error"),
		},
	)
	if err == nil {
		t.Errorf("wanted error propagated from healthCheck")
	}
}

func TestHealthCheck(t *testing.T) {
	b := &base{
		blankThreshold: blankThreshold,
		tip:            []string{""},
	}
	err := b.healthCheck(&operator.MockOperator{})
	if err != nil {
		t.Fatal(err)
	}
	if b.blanks != 1 {
		t.Errorf("got %d blanks; wanted 1", b.blanks)
	}

	b.blanks = b.blankThreshold
	err = b.healthCheck(
		&operator.MockOperator{
			IsThereThingReturn: true,
		},
	)
	if err != nil {
		t.Fatal(err)
	}
	if b.blanks != 0 {
		t.Errorf("got %d blanks; wanted 0", b.blanks)
	}
	if b.blankThreshold <= blankThreshold {
		t.Errorf("got %d; wanted > %d", b.blankThreshold, blankThreshold)
	}

	b.blanks = b.blankThreshold
	err = b.healthCheck(
		&operator.MockOperator{
			IsThereThingErr: fmt.Errorf("an error"),
		},
	)
	if err == nil {
		t.Fatalf("wanted error propagated from operator")
	}
}

func TestDispatch(t *testing.T) {
	hand := &handler{}
	b := &base{
		handlePost:    hand.post,
		handleComment: hand.comment,
		handleMessage: hand.message,
	}
	b.dispatch(
		[]*redditproto.Link{
			&redditproto.Link{},
		},
		[]*redditproto.Comment{
			&redditproto.Comment{},
		},
		[]*redditproto.Message{
			&redditproto.Message{},
		},
	)

	for i := 0; i < 100 && hand.calls < 3; i++ {
		time.Sleep(10 * time.Millisecond)
	}

	if hand.postCalls != 1 {
		t.Errorf("got %d post calls; wanted 1", hand.postCalls)
	}

	if hand.commentCalls != 1 {
		t.Errorf("got %d comment calls; wanted 1", hand.commentCalls)
	}

	if hand.messageCalls != 1 {
		t.Errorf("got %d message calls; wanted 1", hand.messageCalls)
	}
}
<commit_msg>Fix post order expectation in monitor test.<commit_after>package monitor

import (
	"fmt"
	"reflect"
	"strconv"
	"testing"
	"time"

	"github.com/turnage/graw/internal/operator"
	"github.com/turnage/redditproto"
)

type handler struct {
	calls        int
	postCalls    int
	commentCalls int
	messageCalls int
}

func (h *handler) post(post *redditproto.Link) {
	h.postCalls++
	h.calls++
}

func (h *handler) comment(comment *redditproto.Comment) {
	h.commentCalls++
	h.calls++
}

func (h *handler) message(message *redditproto.Message) {
	h.messageCalls++
	h.calls++
}

func float64Pointer(val float64) *float64 {
	return &val
}

func stringPointer(val string) *string {
	return &val
}

func TestBaseFromPath(t *testing.T) {
	han := &handler{}
	mon, err := baseFromPath(
		&operator.MockOperator{
			ScrapeLinksReturn: []*redditproto.Link{
				&redditproto.Link{
					Name: stringPointer("name"),
				},
			},
		},
		"/r/self",
		nil,
		han.comment,
		nil,
		Forward,
	)
	if err != nil {
		t.Fatal(err)
	}

	b := mon.(*base)
	if b.dir != Forward {
		t.Errorf("got %d; wanted %d (Forward)", b.dir, Forward)
	}
	if b.handleComment == nil {
		t.Errorf("wanted comment handler set")
	}
	if b.path != "/r/self" {
		t.Errorf("got %s; wanted /r/self", b.path)
	}
	mon, err = baseFromPath(
		&operator.MockOperator{
			ScrapeLinksReturn: []*redditproto.Link{
				&redditproto.Link{
					Name: stringPointer("name"),
				},
			},
		},
		"/r/self",
		nil,
		nil,
		nil,
		Forward,
	)
	if err == nil {
		t.Errorf("wanted error if no handlers are provided")
	}
}

func TestMerge(t *testing.T) {
	things := merge(
		[]*redditproto.Link{
			&redditproto.Link{
				CreatedUtc: float64Pointer(1),
				Name:       stringPointer("one"),
			},
			&redditproto.Link{
				CreatedUtc: float64Pointer(2),
				Name:       stringPointer("two"),
			},
		},
		[]*redditproto.Comment{
			&redditproto.Comment{
				CreatedUtc: float64Pointer(0),
				Name:       stringPointer("zero"),
			},
			&redditproto.Comment{
				CreatedUtc: float64Pointer(3),
				Name:       stringPointer("three"),
			},
		},
		[]*redditproto.Message{
			&redditproto.Message{
				CreatedUtc: float64Pointer(4),
				Name:       stringPointer("four"),
			},
			&redditproto.Message{
				CreatedUtc: float64Pointer(5),
				Name:       stringPointer("five"),
			},
		},
		Backward,
	)

	if len(things) != 6 {
		t.Fatalf("got %d things; wanted 6", len(things))
	}

	if things[0].GetName() != "zero" {
		t.Errorf("got %s; wanted zero", things[0].GetName())
	}

	if things[1].GetName() != "one" {
		t.Errorf("got %s; wanted one", things[1].GetName())
	}

	if things[2].GetName() != "two" {
		t.Errorf("got %s; wanted two", things[2].GetName())
	}

	if things[3].GetName() != "three" {
		t.Errorf("got %s; wanted three", things[3].GetName())
	}

	if things[4].GetName() != "four" {
		t.Errorf("got %s; wanted four", things[4].GetName())
	}

	if things[5].GetName() != "five" {
		t.Errorf("got %s; wanted five", things[5].GetName())
	}

	things = merge(
		[]*redditproto.Link{
			&redditproto.Link{
				CreatedUtc: float64Pointer(2),
				Name:       stringPointer("two"),
			},
			&redditproto.Link{
				CreatedUtc: float64Pointer(1),
				Name:       stringPointer("one"),
			},
		},
		[]*redditproto.Comment{
			&redditproto.Comment{
				CreatedUtc: float64Pointer(3),
				Name:       stringPointer("three"),
			},
			&redditproto.Comment{
				CreatedUtc: float64Pointer(0),
				Name:       stringPointer("zero"),
			},
		},
		[]*redditproto.Message{
			&redditproto.Message{
				CreatedUtc: float64Pointer(5),
				Name:       stringPointer("five"),
			},
			&redditproto.Message{
				CreatedUtc: float64Pointer(4),
				Name:       stringPointer("four"),
			},
		},
		Forward,
	)

	if len(things) != 6 {
		t.Fatalf("got %d things; wanted 6", len(things))
	}

	if things[0].GetName() != "five" {
		t.Errorf("got %s; wanted five", things[0].GetName())
	}

	if things[1].GetName() != "four" {
		t.Errorf("got %s; wanted four", things[1].GetName())
	}

	if things[2].GetName() != "three" {
		t.Errorf("got %s; wanted three", things[2].GetName())
	}

	if things[3].GetName() != "two" {
		t.Errorf("got %s; wanted two", things[3].GetName())
	}

	if things[4].GetName() != "one" {
		t.Errorf("got %s; wanted one", things[4].GetName())
	}

	if things[5].GetName() != "zero" {
		t.Errorf("got %s; wanted zero", things[5].GetName())
	}
}

func TestShaveTip(t *testing.T) {
	b := &base{
		tip: []string{"1", "2"},
	}
	b.shaveTip()
	if !reflect.DeepEqual(b.tip, []string{"2"}) {
		t.Errorf("got %v\n; wanted %v", b.tip, []string{"2"})
	}

	b = &base{
		tip: nil,
	}
	b.shaveTip()
	if !reflect.DeepEqual(b.tip, []string{""}) {
		t.Errorf("got %v\n; wanted %v", b.tip, []string{""})
	}
}

func TestFixTip(t *testing.T) {
	b := &base{
		tip: []string{"1"},
	}

	broken, err := b.fixTip(
		&operator.MockOperator{
			IsThereThingErr: fmt.Errorf("an error"),
		},
	)
	if err == nil {
		t.Errorf("wanted error propagated from operator error")
	}

	broken, err = b.fixTip(
		&operator.MockOperator{
			IsThereThingReturn: false,
		},
	)
	if !broken {
		t.Errorf("got false; wanted true")
	}
}

func TestUpdateTip(t *testing.T) {
	b := &base{
		dir: Forward,
	}
	for i := 0; i < maxTipSize; i++ {
		b.tip = append(b.tip, strconv.Itoa(i))
	}

	err := b.updateTip(
		[]*redditproto.Link{
			&redditproto.Link{
				CreatedUtc: float64Pointer(2),
				Name:       stringPointer("two"),
			},
			&redditproto.Link{
				CreatedUtc: float64Pointer(1),
				Name:       stringPointer("one"),
			},
		},
		nil,
		nil,
		nil,
	)
	if err != nil {
		t.Fatal(err)
	}

	if len(b.tip) != maxTipSize {
		t.Fatalf("got size %d tip log; wanted %d", len(b.tip), maxTipSize)
	}

	if b.tip[0] != "two" {
		t.Errorf("got %s; wanted two", b.tip[0])
	}

	if b.tip[1] != "one" {
		t.Errorf("got %s; wanted one", b.tip[1])
	}

	err = b.updateTip(
		nil,
		nil,
		nil,
		&operator.MockOperator{
			IsThereThingErr: fmt.Errorf("an error"),
		},
	)
	if err == nil {
		t.Errorf("wanted error propagated from healthCheck")
	}
}

func TestHealthCheck(t *testing.T) {
	b := &base{
		blankThreshold: blankThreshold,
		tip:            []string{""},
	}
	err := b.healthCheck(&operator.MockOperator{})
	if err != nil {
		t.Fatal(err)
	}
	if b.blanks != 1 {
		t.Errorf("got %d blanks; wanted 1", b.blanks)
	}

	b.blanks = b.blankThreshold
	err = b.healthCheck(
		&operator.MockOperator{
			IsThereThingReturn: true,
		},
	)
	if err != nil {
		t.Fatal(err)
	}
	if b.blanks != 0 {
		t.Errorf("got %d blanks; wanted 0", b.blanks)
	}
	if b.blankThreshold <= blankThreshold {
		t.Errorf("got %d; wanted > %d", b.blankThreshold, blankThreshold)
	}

	b.blanks = b.blankThreshold
	err = b.healthCheck(
		&operator.MockOperator{
			IsThereThingErr: fmt.Errorf("an error"),
		},
	)
	if err == nil {
		t.Fatalf("wanted error propagated from operator")
	}
}

func TestDispatch(t *testing.T) {
	hand := &handler{}
	b := &base{
		handlePost:    hand.post,
		handleComment: hand.comment,
		handleMessage: hand.message,
	}
	b.dispatch(
		[]*redditproto.Link{
			&redditproto.Link{},
		},
		[]*redditproto.Comment{
			&redditproto.Comment{},
		},
		[]*redditproto.Message{
			&redditproto.Message{},
		},
	)

	for i := 0; i < 100 && hand.calls < 3; i++ {
		time.Sleep(10 * time.Millisecond)
	}

	if hand.postCalls != 1 {
		t.Errorf("got %d post calls; wanted 1", hand.postCalls)
	}

	if hand.commentCalls != 1 {
		t.Errorf("got %d comment calls; wanted 1", hand.commentCalls)
	}

	if hand.messageCalls != 1 {
		t.Errorf("got %d message calls; wanted 1", hand.messageCalls)
	}
}
<|endoftext|>


<commit_before>package telemetry_test

import (
	"context"
	"log"
	"strings"
	"testing"

	"golang.org/x/tools/internal/telemetry/event"
	"golang.org/x/tools/internal/telemetry/export"
)

type Hooks struct {
	A func(ctx context.Context, a *int) (context.Context, func())
	B func(ctx context.Context, b *string) (context.Context, func())
}

var (
	aValue  = event.NewInt64Key("a", "")
	bValue  = event.NewStringKey("b", "")
	aCount  = event.NewInt64Key("aCount", "Count of time A is called.")
	aStat   = event.NewInt64Key("aValue", "A value.")
	bCount  = event.NewInt64Key("B", "Count of time B is called.")
	bLength = event.NewInt64Key("BLen", "B length.")

	Baseline = Hooks{
		A: func(ctx context.Context, a *int) (context.Context, func()) {
			return ctx, func() {}
		},
		B: func(ctx context.Context, b *string) (context.Context, func()) {
			return ctx, func() {}
		},
	}

	StdLog = Hooks{
		A: func(ctx context.Context, a *int) (context.Context, func()) {
			log.Printf("start A where a=%d", *a)
			return ctx, func() {
				log.Printf("end A where a=%d", *a)
			}
		},
		B: func(ctx context.Context, b *string) (context.Context, func()) {
			log.Printf("start B where b=%q", *b)
			return ctx, func() {
				log.Printf("end B where b=%q", *b)
			}
		},
	}

	Log = Hooks{
		A: func(ctx context.Context, a *int) (context.Context, func()) {
			event.Print(ctx, "start A", aValue.Of(int64(*a)))
			return ctx, func() {
				event.Print(ctx, "end A", aValue.Of(int64(*a)))
			}
		},
		B: func(ctx context.Context, b *string) (context.Context, func()) {
			event.Print(ctx, "start B", bValue.Of(*b))
			return ctx, func() {
				event.Print(ctx, "end B", bValue.Of(*b))
			}
		},
	}

	Trace = Hooks{
		A: func(ctx context.Context, a *int) (context.Context, func()) {
			return event.StartSpan(ctx, "A")
		},
		B: func(ctx context.Context, b *string) (context.Context, func()) {
			return event.StartSpan(ctx, "B")
		},
	}

	Stats = Hooks{
		A: func(ctx context.Context, a *int) (context.Context, func()) {
			event.Record(ctx, aCount.Of(1))
			return ctx, func() {
				event.Record(ctx, aStat.Of(int64(*a)))
			}
		},
		B: func(ctx context.Context, b *string) (context.Context, func()) {
			event.Record(ctx, bCount.Of(1))
			return ctx, func() {
				event.Record(ctx, bLength.Of(int64(len(*b))))
			}
		},
	}
)

func Benchmark(b *testing.B) {
	b.Run("Baseline", Baseline.runBenchmark)
	b.Run("StdLog", StdLog.runBenchmark)
	event.SetExporter(nil)
	b.Run("LogNoExporter", Log.runBenchmark)
	b.Run("TraceNoExporter", Trace.runBenchmark)
	b.Run("StatsNoExporter", Stats.runBenchmark)

	event.SetExporter(noopExporter)
	b.Run("Log", Log.runBenchmark)
	b.Run("Trace", Trace.runBenchmark)
	b.Run("Stats", Stats.runBenchmark)
}

func A(ctx context.Context, hooks Hooks, a int) int {
	ctx, done := hooks.A(ctx, &a)
	defer done()
	if a > 0 {
		a = a * 10
	}
	return B(ctx, hooks, a, "Called from A")
}

func B(ctx context.Context, hooks Hooks, a int, b string) int {
	_, done := hooks.B(ctx, &b)
	defer done()
	b = strings.ToUpper(b)
	if len(b) > 1024 {
		b = strings.ToLower(b)
	}
	return a + len(b)
}

func (hooks Hooks) runBenchmark(b *testing.B) {
	ctx := context.Background()
	b.ReportAllocs()
	b.ResetTimer()
	var acc int
	for i := 0; i < b.N; i++ {
		for _, value := range []int{0, 10, 20, 100, 1000} {
			acc += A(ctx, hooks, value)
		}
	}
}

func init() {
	log.SetOutput(new(noopWriter))
}

type noopWriter int

func (nw *noopWriter) Write(b []byte) (int, error) {
	return len(b), nil
}

var noopExporter = export.Spans(export.LogWriter(new(noopWriter), false))
<commit_msg>internal/telemetry: improve the telemetry benchmark<commit_after>package telemetry_test

import (
	"context"
	"log"
	"testing"

	"golang.org/x/tools/internal/telemetry/event"
	"golang.org/x/tools/internal/telemetry/export"
)

type Hooks struct {
	A func(ctx context.Context, a int) (context.Context, func())
	B func(ctx context.Context, b string) (context.Context, func())
}

var (
	aValue  = event.NewIntKey("a", "")
	bValue  = event.NewStringKey("b", "")
	aCount  = event.NewInt64Key("aCount", "Count of time A is called.")
	aStat   = event.NewIntKey("aValue", "A value.")
	bCount  = event.NewInt64Key("B", "Count of time B is called.")
	bLength = event.NewIntKey("BLen", "B length.")

	Baseline = Hooks{
		A: func(ctx context.Context, a int) (context.Context, func()) {
			return ctx, func() {}
		},
		B: func(ctx context.Context, b string) (context.Context, func()) {
			return ctx, func() {}
		},
	}

	StdLog = Hooks{
		A: func(ctx context.Context, a int) (context.Context, func()) {
			log.Printf("A where a=%d", a)
			return ctx, func() {}
		},
		B: func(ctx context.Context, b string) (context.Context, func()) {
			log.Printf("B where b=%q", b)
			return ctx, func() {}
		},
	}

	Log = Hooks{
		A: func(ctx context.Context, a int) (context.Context, func()) {
			event.Print(ctx, "A", aValue.Of(a))
			return ctx, func() {}
		},
		B: func(ctx context.Context, b string) (context.Context, func()) {
			event.Print(ctx, "B", bValue.Of(b))
			return ctx, func() {}
		},
	}

	Trace = Hooks{
		A: func(ctx context.Context, a int) (context.Context, func()) {
			return event.StartSpan(ctx, "A", aValue.Of(a))
		},
		B: func(ctx context.Context, b string) (context.Context, func()) {
			return event.StartSpan(ctx, "B", bValue.Of(b))
		},
	}

	Stats = Hooks{
		A: func(ctx context.Context, a int) (context.Context, func()) {
			event.Record(ctx, aStat.Of(a))
			event.Record(ctx, aCount.Of(1))
			return ctx, func() {}
		},
		B: func(ctx context.Context, b string) (context.Context, func()) {
			event.Record(ctx, bLength.Of(len(b)))
			event.Record(ctx, bCount.Of(1))
			return ctx, func() {}
		},
	}

	initialList = []int{0, 1, 22, 333, 4444, 55555, 666666, 7777777}
	stringList  = []string{
		"A value",
		"Some other value",
		"A nice longer value but not too long",
		"V",
		"",
		"ı",
		"prime count of values",
	}
)

type namedBenchmark struct {
	name string
	test func(*testing.B)
}

func Benchmark(b *testing.B) {
	b.Run("Baseline", Baseline.runBenchmark)
	b.Run("StdLog", StdLog.runBenchmark)
	benchmarks := []namedBenchmark{
		{"Log", Log.runBenchmark},
		{"Trace", Trace.runBenchmark},
		{"Stats", Stats.runBenchmark},
	}

	event.SetExporter(nil)
	for _, t := range benchmarks {
		b.Run(t.name+"NoExporter", t.test)
	}

	event.SetExporter(noopExporter)
	for _, t := range benchmarks {
		b.Run(t.name, t.test)
	}
}

func A(ctx context.Context, hooks Hooks, a int) int {
	ctx, done := hooks.A(ctx, a)
	defer done()
	return B(ctx, hooks, a, stringList[a%len(stringList)])
}

func B(ctx context.Context, hooks Hooks, a int, b string) int {
	_, done := hooks.B(ctx, b)
	defer done()
	return a + len(b)
}

func (hooks Hooks) runBenchmark(b *testing.B) {
	ctx := context.Background()
	b.ReportAllocs()
	b.ResetTimer()
	var acc int
	for i := 0; i < b.N; i++ {
		for _, value := range initialList {
			acc += A(ctx, hooks, value)
		}
	}
}

func init() {
	log.SetOutput(new(noopWriter))
}

type noopWriter int

func (nw *noopWriter) Write(b []byte) (int, error) {
	return len(b), nil
}

var noopExporter = export.Spans(export.LogWriter(new(noopWriter), false))
<|endoftext|>


<commit_before>package main

import (
	"archive/tar"
	"bytes"
	"fmt"
	"github.com/dynport/dgtk/dockerclient"
	"github.com/dynport/dgtk/git"
	"io"
	"io/ioutil"
	"launchpad.net/goyaml"
	"os"
	"path/filepath"
	"strings"
)

type Build struct {
	GitRepository   string `yaml:"repository"`
	Tag             string `yaml:"tag"`
	Proxy           string `yaml:"proxy"`
	Root            string
	DockerHost      string
	dockerfileAdded bool
}

func (b *Build) Build() (string, error) {
	buf := &bytes.Buffer{}
	e := b.buildArchive(buf)
	if e != nil {
		return "", e
	}
	client := dockerclient.New(b.DockerHost, 4243)
	return client.Build(buf, &dockerclient.BuildImageOptions{Tag: b.Tag, Callback: callback})
}

func (build *Build) LoadConfig() error {
	b, e := ioutil.ReadFile(build.Root + "/dockerbuild.yml")
	if e != nil {
		return e
	}
	return goyaml.Unmarshal(b, build)
}

func (b *Build) buildArchive(w io.Writer) error {
	t := tar.NewWriter(w)
	defer t.Flush()
	defer t.Close()
	if b.GitRepository != "" {
		repo := &git.Repository{Origin: b.GitRepository}
		e := repo.Init()
		if e != nil {
			return e
		}
		if e := repo.Tar(t); e != nil {
			return e
		}
	}
	if e := b.addFilesToArchive(b.Root, t); e != nil {
		return e
	}
	if !b.dockerfileAdded {
		return fmt.Errorf("archive must contain a Dockerfile")
	}
	return nil
}

func (b *Build) addFilesToArchive(root string, t *tar.Writer) error {
	return filepath.Walk(root, func(p string, info os.FileInfo, e error) error {
		if e == nil && p != root {
			var e error
			name := strings.TrimPrefix(p, root+"/")
			header := &tar.Header{Name: name, ModTime: info.ModTime().UTC()}
			if info.IsDir() {
				header.Typeflag = tar.TypeDir
				header.Mode = 0755
				e = t.WriteHeader(header)
			} else {
				header.Mode = 0644
				header.Size = info.Size()
				e = t.WriteHeader(header)
				if e != nil {
					return e
				}
				f, e := os.Open(p)
				if e != nil {
					return e
				}
				defer f.Close()
				if _, e = io.Copy(t, f); e != nil {
					return e
				}
				if name == "Dockerfile" {
					b.dockerfileAdded = true
				}
			}
			if e != nil {
				return e
			}
		}
		return nil
	})
}
<commit_msg>refactoring<commit_after>package main

import (
	"archive/tar"
	"bytes"
	"fmt"
	"github.com/dynport/dgtk/dockerclient"
	"github.com/dynport/dgtk/git"
	"io"
	"io/ioutil"
	"launchpad.net/goyaml"
	"os"
	"path/filepath"
	"strings"
)

type Build struct {
	GitRepository   string `yaml:"repository"`
	Tag             string `yaml:"tag"`
	Proxy           string `yaml:"proxy"`
	Root            string
	DockerHost      string
	dockerfileAdded bool
}

func (b *Build) Build() (string, error) {
	buf := &bytes.Buffer{}
	e := b.buildArchive(buf)
	if e != nil {
		return "", e
	}
	client := dockerclient.New(b.DockerHost, 4243)
	return client.Build(buf, &dockerclient.BuildImageOptions{Tag: b.Tag, Callback: callback})
}

func (build *Build) LoadConfig() error {
	b, e := ioutil.ReadFile(build.Root + "/dockerbuild.yml")
	if e != nil {
		return e
	}
	return goyaml.Unmarshal(b, build)
}

func (b *Build) buildArchive(w io.Writer) error {
	t := tar.NewWriter(w)
	defer t.Flush()
	defer t.Close()
	if b.GitRepository != "" {
		repo := &git.Repository{Origin: b.GitRepository}
		e := repo.Init()
		if e != nil {
			return e
		}
		if e := repo.Tar(t); e != nil {
			return e
		}
	}
	if e := b.addFilesToArchive(b.Root, t); e != nil {
		return e
	}
	if !b.dockerfileAdded {
		return fmt.Errorf("archive must contain a Dockerfile")
	}
	return nil
}

func (build *Build) addFilesToArchive(root string, t *tar.Writer) error {
	return filepath.Walk(root, func(p string, info os.FileInfo, e error) error {
		if e == nil && p != root {
			var e error
			name := strings.TrimPrefix(p, root+"/")
			header := &tar.Header{Name: name, ModTime: info.ModTime().UTC()}
			if info.IsDir() {
				header.Typeflag = tar.TypeDir
				header.Mode = 0755
				e = t.WriteHeader(header)
			} else {
				header.Mode = 0644
				b, e := ioutil.ReadFile(p)
				if e != nil {
					return e
				}
				if name == "Dockerfile" {
					build.dockerfileAdded = true
				}
				header.Size = int64(len(b))
				e = t.WriteHeader(header)
				if e != nil {
					return e
				}
				_, e = t.Write(b)
				if e != nil {
					return e
				}
			}
			if e != nil {
				return e
			}
		}
		return nil
	})
}
<|endoftext|>


<commit_before>/*
   DNS-over-HTTPS
   Copyright (C) 2017-2018 Star Brilliant <m13253@hotmail.com>

   Permission is hereby granted, free of charge, to any person obtaining a
   copy of this software and associated documentation files (the "Software"),
   to deal in the Software without restriction, including without limitation
   the rights to use, copy, modify, merge, publish, distribute, sublicense,
   and/or sell copies of the Software, and to permit persons to whom the
   Software is furnished to do so, subject to the following conditions:

   The above copyright notice and this permission notice shall be included in
   all copies or substantial portions of the Software.

   THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
   IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
   FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
   AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
   LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
   FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
   DEALINGS IN THE SOFTWARE.
*/

package main

import (
	"context"
	"fmt"
	"log"
	"math/rand"
	"net"
	"net/http"
	"os"
	"strings"
	"time"

	"github.com/gorilla/handlers"
	"github.com/m13253/dns-over-https/json-dns"
	"github.com/miekg/dns"
)

type Server struct {
	conf      *config
	udpClient *dns.Client
	tcpClient *dns.Client
	servemux  *http.ServeMux
}

type DNSRequest struct {
	request         *dns.Msg
	response        *dns.Msg
	transactionID   uint16
	currentUpstream string
	isTailored      bool
	errcode         int
	errtext         string
}

func NewServer(conf *config) (*Server, error) {
	timeout := time.Duration(conf.Timeout) * time.Second
	s := &Server{
		conf: conf,
		udpClient: &dns.Client{
			Net:     "udp",
			UDPSize: dns.DefaultMsgSize,
			Timeout: timeout,
		},
		tcpClient: &dns.Client{
			Net:     "tcp",
			Timeout: timeout,
		},
		servemux: http.NewServeMux(),
	}
	if conf.LocalAddr != "" {
		udpLocalAddr, err := net.ResolveUDPAddr("udp", conf.LocalAddr)
		if err != nil {
			return nil, err
		}
		tcpLocalAddr, err := net.ResolveTCPAddr("tcp", conf.LocalAddr)
		if err != nil {
			return nil, err
		}
		s.udpClient.Dialer = &net.Dialer{
			Timeout:   timeout,
			LocalAddr: udpLocalAddr,
		}
		s.tcpClient.Dialer = &net.Dialer{
			Timeout:   timeout,
			LocalAddr: tcpLocalAddr,
		}
	}
	s.servemux.HandleFunc(conf.Path, s.handlerFunc)
	return s, nil
}

func (s *Server) Start() error {
	servemux := http.Handler(s.servemux)
	if s.conf.Verbose {
		servemux = handlers.CombinedLoggingHandler(os.Stdout, servemux)
	}
	results := make(chan error, len(s.conf.Listen))
	for _, addr := range s.conf.Listen {
		go func(addr string) {
			var err error
			if s.conf.Cert != "" || s.conf.Key != "" {
				err = http.ListenAndServeTLS(addr, s.conf.Cert, s.conf.Key, servemux)
			} else {
				err = http.ListenAndServe(addr, servemux)
			}
			if err != nil {
				log.Println(err)
			}
			results <- err
		}(addr)
	}
	// wait for all handlers
	for i := 0; i < cap(results); i++ {
		err := <-results
		if err != nil {
			return err
		}
	}
	close(results)
	return nil
}

func (s *Server) handlerFunc(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()

	w.Header().Set("Access-Control-Allow-Headers", "Content-Type")
	w.Header().Set("Access-Control-Allow-Methods", "GET, HEAD, OPTIONS, POST")
	w.Header().Set("Access-Control-Allow-Origin", "*")
	w.Header().Set("Access-Control-Max-Age", "3600")
	w.Header().Set("Server", USER_AGENT)
	w.Header().Set("X-Powered-By", USER_AGENT)

	if r.Method == "OPTIONS" {
		w.Header().Set("Content-Length", "0")
		return
	}

	if r.Form == nil {
		const maxMemory = 32 << 20 // 32 MB
		r.ParseMultipartForm(maxMemory)
	}

	for _, header := range s.conf.DebugHTTPHeaders {
		if value := r.Header.Get(header); value != "" {
			log.Printf("%s: %s\n", header, value)
		}
	}

	contentType := r.Header.Get("Content-Type")
	if ct := r.FormValue("ct"); ct != "" {
		contentType = ct
	}
	if contentType == "" {
		// Guess request Content-Type based on other parameters
		if r.FormValue("name") != "" {
			contentType = "application/dns-json"
		} else if r.FormValue("dns") != "" {
			contentType = "application/dns-message"
		}
	}
	var responseType string
	for _, responseCandidate := range strings.Split(r.Header.Get("Accept"), ",") {
		responseCandidate = strings.SplitN(responseCandidate, ";", 2)[0]
		if responseCandidate == "application/json" {
			responseType = "application/json"
			break
		} else if responseCandidate == "application/dns-udpwireformat" {
			responseType = "application/dns-message"
			break
		} else if responseCandidate == "application/dns-message" {
			responseType = "application/dns-message"
			break
		}
	}
	if responseType == "" {
		// Guess response Content-Type based on request Content-Type
		if contentType == "application/dns-json" {
			responseType = "application/json"
		} else if contentType == "application/dns-message" {
			responseType = "application/dns-message"
		} else if contentType == "application/dns-udpwireformat" {
			responseType = "application/dns-message"
		}
	}

	var req *DNSRequest
	if contentType == "application/dns-json" {
		req = s.parseRequestGoogle(ctx, w, r)
	} else if contentType == "application/dns-message" {
		req = s.parseRequestIETF(ctx, w, r)
	} else if contentType == "application/dns-udpwireformat" {
		req = s.parseRequestIETF(ctx, w, r)
	} else {
		jsonDNS.FormatError(w, fmt.Sprintf("Invalid argument value: \"ct\" = %q", contentType), 415)
		return
	}
	if req.errcode == 444 {
		return
	}
	if req.errcode != 0 {
		jsonDNS.FormatError(w, req.errtext, req.errcode)
		return
	}

	req = s.patchRootRD(req)

	var err error
	req, err = s.doDNSQuery(ctx, req)
	if err != nil {
		jsonDNS.FormatError(w, fmt.Sprintf("DNS query failure (%s)", err.Error()), 503)
		return
	}

	if responseType == "application/json" {
		s.generateResponseGoogle(ctx, w, r, req)
	} else if responseType == "application/dns-message" {
		s.generateResponseIETF(ctx, w, r, req)
	} else {
		panic("Unknown response Content-Type")
	}
}

func (s *Server) findClientIP(r *http.Request) net.IP {
	XForwardedFor := r.Header.Get("X-Forwarded-For")
	if XForwardedFor != "" {
		for _, addr := range strings.Split(XForwardedFor, ",") {
			addr = strings.TrimSpace(addr)
			ip := net.ParseIP(addr)
			if jsonDNS.IsGlobalIP(ip) {
				return ip
			}
		}
	}
	XRealIP := r.Header.Get("X-Real-IP")
	if XRealIP != "" {
		addr := strings.TrimSpace(XRealIP)
		ip := net.ParseIP(addr)
		if jsonDNS.IsGlobalIP(ip) {
			return ip
		}
	}
	remoteAddr, err := net.ResolveTCPAddr("tcp", r.RemoteAddr)
	if err != nil {
		return nil
	}
	if ip := remoteAddr.IP; jsonDNS.IsGlobalIP(ip) {
		return ip
	}
	return nil
}

// Workaround a bug causing Unbound to refuse returning anything about the root
func (s *Server) patchRootRD(req *DNSRequest) *DNSRequest {
	for _, question := range req.request.Question {
		if question.Name == "." {
			req.request.RecursionDesired = true
		}
	}
	return req
}

func (s *Server) doDNSQuery(ctx context.Context, req *DNSRequest) (resp *DNSRequest, err error) {
	// TODO(m13253): Make ctx work. Waiting for a patch for ExchangeContext from miekg/dns.
	numServers := len(s.conf.Upstream)
	for i := uint(0); i < s.conf.Tries; i++ {
		req.currentUpstream = s.conf.Upstream[rand.Intn(numServers)]
		if !s.conf.TCPOnly {
			req.response, _, err = s.udpClient.Exchange(req.request, req.currentUpstream)
			if err == nil && req.response != nil && req.response.Truncated {
				log.Println(err)
				req.response, _, err = s.tcpClient.Exchange(req.request, req.currentUpstream)
			}
		} else {
			req.response, _, err = s.tcpClient.Exchange(req.request, req.currentUpstream)
		}
		if err == nil {
			return req, nil
		}
		log.Printf("DNS error from upstream %s: %s\n", req.currentUpstream, err.Error())
	}
	return req, err
}
<commit_msg>Use TCP when appropriate for the given query type/response<commit_after>/*
   DNS-over-HTTPS
   Copyright (C) 2017-2018 Star Brilliant <m13253@hotmail.com>

   Permission is hereby granted, free of charge, to any person obtaining a
   copy of this software and associated documentation files (the "Software"),
   to deal in the Software without restriction, including without limitation
   the rights to use, copy, modify, merge, publish, distribute, sublicense,
   and/or sell copies of the Software, and to permit persons to whom the
   Software is furnished to do so, subject to the following conditions:

   The above copyright notice and this permission notice shall be included in
   all copies or substantial portions of the Software.

   THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
   IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
   FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
   AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
   LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
   FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
   DEALINGS IN THE SOFTWARE.
*/

package main

import (
	"context"
	"fmt"
	"log"
	"math/rand"
	"net"
	"net/http"
	"os"
	"strings"
	"time"

	"github.com/gorilla/handlers"
	"github.com/m13253/dns-over-https/json-dns"
	"github.com/miekg/dns"
)

type Server struct {
	conf      *config
	udpClient *dns.Client
	tcpClient *dns.Client
	servemux  *http.ServeMux
}

type DNSRequest struct {
	request         *dns.Msg
	response        *dns.Msg
	transactionID   uint16
	currentUpstream string
	isTailored      bool
	errcode         int
	errtext         string
}

func NewServer(conf *config) (*Server, error) {
	timeout := time.Duration(conf.Timeout) * time.Second
	s := &Server{
		conf: conf,
		udpClient: &dns.Client{
			Net:     "udp",
			UDPSize: dns.DefaultMsgSize,
			Timeout: timeout,
		},
		tcpClient: &dns.Client{
			Net:     "tcp",
			Timeout: timeout,
		},
		servemux: http.NewServeMux(),
	}
	if conf.LocalAddr != "" {
		udpLocalAddr, err := net.ResolveUDPAddr("udp", conf.LocalAddr)
		if err != nil {
			return nil, err
		}
		tcpLocalAddr, err := net.ResolveTCPAddr("tcp", conf.LocalAddr)
		if err != nil {
			return nil, err
		}
		s.udpClient.Dialer = &net.Dialer{
			Timeout:   timeout,
			LocalAddr: udpLocalAddr,
		}
		s.tcpClient.Dialer = &net.Dialer{
			Timeout:   timeout,
			LocalAddr: tcpLocalAddr,
		}
	}
	s.servemux.HandleFunc(conf.Path, s.handlerFunc)
	return s, nil
}

func (s *Server) Start() error {
	servemux := http.Handler(s.servemux)
	if s.conf.Verbose {
		servemux = handlers.CombinedLoggingHandler(os.Stdout, servemux)
	}
	results := make(chan error, len(s.conf.Listen))
	for _, addr := range s.conf.Listen {
		go func(addr string) {
			var err error
			if s.conf.Cert != "" || s.conf.Key != "" {
				err = http.ListenAndServeTLS(addr, s.conf.Cert, s.conf.Key, servemux)
			} else {
				err = http.ListenAndServe(addr, servemux)
			}
			if err != nil {
				log.Println(err)
			}
			results <- err
		}(addr)
	}
	// wait for all handlers
	for i := 0; i < cap(results); i++ {
		err := <-results
		if err != nil {
			return err
		}
	}
	close(results)
	return nil
}

func (s *Server) handlerFunc(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()

	w.Header().Set("Access-Control-Allow-Headers", "Content-Type")
	w.Header().Set("Access-Control-Allow-Methods", "GET, HEAD, OPTIONS, POST")
	w.Header().Set("Access-Control-Allow-Origin", "*")
	w.Header().Set("Access-Control-Max-Age", "3600")
	w.Header().Set("Server", USER_AGENT)
	w.Header().Set("X-Powered-By", USER_AGENT)

	if r.Method == "OPTIONS" {
		w.Header().Set("Content-Length", "0")
		return
	}

	if r.Form == nil {
		const maxMemory = 32 << 20 // 32 MB
		r.ParseMultipartForm(maxMemory)
	}

	for _, header := range s.conf.DebugHTTPHeaders {
		if value := r.Header.Get(header); value != "" {
			log.Printf("%s: %s\n", header, value)
		}
	}

	contentType := r.Header.Get("Content-Type")
	if ct := r.FormValue("ct"); ct != "" {
		contentType = ct
	}
	if contentType == "" {
		// Guess request Content-Type based on other parameters
		if r.FormValue("name") != "" {
			contentType = "application/dns-json"
		} else if r.FormValue("dns") != "" {
			contentType = "application/dns-message"
		}
	}
	var responseType string
	for _, responseCandidate := range strings.Split(r.Header.Get("Accept"), ",") {
		responseCandidate = strings.SplitN(responseCandidate, ";", 2)[0]
		if responseCandidate == "application/json" {
			responseType = "application/json"
			break
		} else if responseCandidate == "application/dns-udpwireformat" {
			responseType = "application/dns-message"
			break
		} else if responseCandidate == "application/dns-message" {
			responseType = "application/dns-message"
			break
		}
	}
	if responseType == "" {
		// Guess response Content-Type based on request Content-Type
		if contentType == "application/dns-json" {
			responseType = "application/json"
		} else if contentType == "application/dns-message" {
			responseType = "application/dns-message"
		} else if contentType == "application/dns-udpwireformat" {
			responseType = "application/dns-message"
		}
	}

	var req *DNSRequest
	if contentType == "application/dns-json" {
		req = s.parseRequestGoogle(ctx, w, r)
	} else if contentType == "application/dns-message" {
		req = s.parseRequestIETF(ctx, w, r)
	} else if contentType == "application/dns-udpwireformat" {
		req = s.parseRequestIETF(ctx, w, r)
	} else {
		jsonDNS.FormatError(w, fmt.Sprintf("Invalid argument value: \"ct\" = %q", contentType), 415)
		return
	}
	if req.errcode == 444 {
		return
	}
	if req.errcode != 0 {
		jsonDNS.FormatError(w, req.errtext, req.errcode)
		return
	}

	req = s.patchRootRD(req)

	var err error
	req, err = s.doDNSQuery(ctx, req)
	if err != nil {
		jsonDNS.FormatError(w, fmt.Sprintf("DNS query failure (%s)", err.Error()), 503)
		return
	}

	if responseType == "application/json" {
		s.generateResponseGoogle(ctx, w, r, req)
	} else if responseType == "application/dns-message" {
		s.generateResponseIETF(ctx, w, r, req)
	} else {
		panic("Unknown response Content-Type")
	}
}

func (s *Server) findClientIP(r *http.Request) net.IP {
	XForwardedFor := r.Header.Get("X-Forwarded-For")
	if XForwardedFor != "" {
		for _, addr := range strings.Split(XForwardedFor, ",") {
			addr = strings.TrimSpace(addr)
			ip := net.ParseIP(addr)
			if jsonDNS.IsGlobalIP(ip) {
				return ip
			}
		}
	}
	XRealIP := r.Header.Get("X-Real-IP")
	if XRealIP != "" {
		addr := strings.TrimSpace(XRealIP)
		ip := net.ParseIP(addr)
		if jsonDNS.IsGlobalIP(ip) {
			return ip
		}
	}
	remoteAddr, err := net.ResolveTCPAddr("tcp", r.RemoteAddr)
	if err != nil {
		return nil
	}
	if ip := remoteAddr.IP; jsonDNS.IsGlobalIP(ip) {
		return ip
	}
	return nil
}

// Workaround a bug causing Unbound to refuse returning anything about the root
func (s *Server) patchRootRD(req *DNSRequest) *DNSRequest {
	for _, question := range req.request.Question {
		if question.Name == "." {
			req.request.RecursionDesired = true
		}
	}
	return req
}

// Return the position index for the question of qtype from a DNS msg, otherwise return -1
func (s *Server) indexQuestionType(msg *dns.Msg, qtype uint16) int {
	for i, question := range msg.Question {
		if question.Qtype == qtype {
			return i
		}
	}
	return -1
}

func (s *Server) doDNSQuery(ctx context.Context, req *DNSRequest) (resp *DNSRequest, err error) {
	// TODO(m13253): Make ctx work. Waiting for a patch for ExchangeContext from miekg/dns.
	numServers := len(s.conf.Upstream)
	for i := uint(0); i < s.conf.Tries; i++ {
		req.currentUpstream = s.conf.Upstream[rand.Intn(numServers)]

		// Use TCP if always configured to or if the Query type dictates it (AXFR)
		if s.conf.TCPOnly || (s.indexQuestionType(req.request, dns.TypeAXFR) > -1) {
			req.response, _, err = s.tcpClient.Exchange(req.request, req.currentUpstream)
		} else {
			req.response, _, err = s.udpClient.Exchange(req.request, req.currentUpstream)
			if err == nil && req.response != nil && req.response.Truncated {
				log.Println(err)
				req.response, _, err = s.tcpClient.Exchange(req.request, req.currentUpstream)
			}

			// Retry with TCP if this was an IXFR request and we only received an SOA
			if (s.indexQuestionType(req.request, dns.TypeIXFR) > -1) &&
				(len(req.response.Answer) == 1) &&
				(req.response.Answer[0].Header().Rrtype == dns.TypeSOA) {
				req.response, _, err = s.tcpClient.Exchange(req.request, req.currentUpstream)
			}
		}
		if err == nil {
			return req, nil
		}
		log.Printf("DNS error from upstream %s: %s\n", req.currentUpstream, err.Error())
	}
	return req, err
}
<|endoftext|>


<commit_before>package herd

import (
	"bufio"
	"fmt"
	"github.com/Symantec/Dominator/lib/url"
	"net/http"
)

func (herd *Herd) listReachableSubsHandler(w http.ResponseWriter,
	req *http.Request) {
	writer := bufio.NewWriter(w)
	defer writer.Flush()
	selector, err := herd.getReachableSelector(url.ParseQuery(req.URL))
	if err != nil {
		fmt.Fprintln(writer, err)
		return
	}
	for _, sub := range herd.getSelectedSubs(selector) {
		fmt.Fprintln(writer, sub.mdb.Hostname)
	}
}

func (herd *Herd) listSubsHandler(w http.ResponseWriter, req *http.Request) {
	writer := bufio.NewWriter(w)
	defer writer.Flush()
	herd.RLock()
	subs := make([]string, 0, len(herd.subsByIndex))
	for _, sub := range herd.subsByIndex {
		subs = append(subs, sub.mdb.Hostname)
	}
	herd.RUnlock()
	for _, name := range subs {
		fmt.Fprintln(writer, name)
	}
}
<commit_msg>Add output=json support to /list{,Reachable}Subs handlers in dominator.<commit_after>package herd

import (
	"bufio"
	"fmt"
	"github.com/Symantec/Dominator/lib/json"
	"github.com/Symantec/Dominator/lib/url"
	"net/http"
)

func (herd *Herd) listReachableSubsHandler(w http.ResponseWriter,
	req *http.Request) {
	writer := bufio.NewWriter(w)
	defer writer.Flush()
	parsedQuery := url.ParseQuery(req.URL)
	selector, err := herd.getReachableSelector(parsedQuery)
	if err != nil {
		fmt.Fprintln(writer, err)
		return
	}
	subs := herd.getSelectedSubs(selector)
	switch parsedQuery.OutputType() {
	case url.OutputTypeText:
	case url.OutputTypeHtml:
		for _, sub := range subs {
			fmt.Fprintln(writer, sub.mdb.Hostname)
		}
	case url.OutputTypeJson:
		subNames := make([]string, 0, len(subs))
		for _, sub := range subs {
			subNames = append(subNames, sub.mdb.Hostname)
		}
		json.WriteWithIndent(writer, "  ", subNames)
		fmt.Fprintln(writer)
	}
}

func (herd *Herd) listSubsHandler(w http.ResponseWriter, req *http.Request) {
	writer := bufio.NewWriter(w)
	defer writer.Flush()
	herd.RLock()
	subNames := make([]string, 0, len(herd.subsByIndex))
	for _, sub := range herd.subsByIndex {
		subNames = append(subNames, sub.mdb.Hostname)
	}
	herd.RUnlock()
	parsedQuery := url.ParseQuery(req.URL)
	switch parsedQuery.OutputType() {
	case url.OutputTypeText:
	case url.OutputTypeHtml:
		for _, name := range subNames {
			fmt.Fprintln(writer, name)
		}
	case url.OutputTypeJson:
		json.WriteWithIndent(writer, "  ", subNames)
		fmt.Fprintln(writer)
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Knative Authors

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// util has constants and helper methods useful for zipkin tracing support.

package zipkin

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

	"github.com/openzipkin/zipkin-go/model"
	"go.opencensus.io/trace"
	"k8s.io/client-go/kubernetes"
	"knative.dev/pkg/test/logging"
	"knative.dev/pkg/test/monitoring"
)

const (
	// ZipkinTraceIDHeader HTTP response header key to be used to store Zipkin Trace ID.
	ZipkinTraceIDHeader = "ZIPKIN_TRACE_ID"

	// ZipkinPort is port exposed by the Zipkin Pod
	// https://github.com/knative/serving/blob/master/config/monitoring/200-common/100-zipkin.yaml#L25 configures the Zipkin Port on the cluster.
	ZipkinPort = 9411

	// ZipkinTraceEndpoint port-forwarded zipkin endpoint
	ZipkinTraceEndpoint = "http://localhost:9411/api/v2/trace/"

	// App is the name of this component.
	// This will be used as a label selector.
	app = "zipkin"

	// istioNS is the namespace we are using for istio components.
	istioNS = "istio-system"
)

var (
	zipkinPortForwardPID int

	// ZipkinTracingEnabled variable indicating if zipkin tracing is enabled.
	ZipkinTracingEnabled = false

	// sync.Once variable to ensure we execute zipkin setup only once.
	setupOnce sync.Once

	// sync.Once variable to ensure we execute zipkin cleanup only if zipkin is setup and it is executed only once.
	teardownOnce sync.Once
)

// SetupZipkinTracing sets up zipkin tracing which involves:
// 1. Setting up port-forwarding from localhost to zipkin pod on the cluster
//    (pid of the process doing Port-Forward is stored in a global variable).
// 2. Enable AlwaysSample config for tracing for the SpoofingClient.
func SetupZipkinTracing(kubeClientset *kubernetes.Clientset, logf logging.FormatLogger) bool {
	setupOnce.Do(func() {
		if err := monitoring.CheckPortAvailability(ZipkinPort); err != nil {
			logf("Zipkin port not available on the machine: %v", err)
			return
		}

		zipkinPods, err := monitoring.GetPods(kubeClientset, app, istioNS)
		if err != nil {
			logf("Error retrieving Zipkin pod details: %v", err)
			return
		}

		zipkinPortForwardPID, err = monitoring.PortForward(logf, zipkinPods, ZipkinPort, ZipkinPort, istioNS)
		if err != nil {
			logf("Error starting kubectl port-forward command: %v", err)
			return
		}

		logf("Zipkin port-forward process started with PID: %d", zipkinPortForwardPID)

		// Applying AlwaysSample config to ensure we propagate zipkin header for every request made by this client.
		trace.ApplyConfig(trace.Config{DefaultSampler: trace.AlwaysSample()})
		logf("Successfully setup SpoofingClient for Zipkin Tracing")
		ZipkinTracingEnabled = true
	})
	return ZipkinTracingEnabled
}

// CleanupZipkinTracingSetup cleans up the Zipkin tracing setup on the machine. This involves killing the process performing port-forward.
// This should be called exactly once in TestMain. Likely in the form:
//
// func TestMain(m *testing.M) {
//     os.Exit(func() int {
//       // Any setup required for the tests.
//       defer zipkin.CleanupZipkinTracingSetup(logger)
//       return m.Run()
//     }())
// }
func CleanupZipkinTracingSetup(logf logging.FormatLogger) {
	teardownOnce.Do(func() {
		// Because CleanupZipkinTracingSetup only runs once, make sure that now that it has been
		// run, SetupZipkinTracing will no longer setup any port forwarding.
		setupOnce.Do(func() {})

		if !ZipkinTracingEnabled {
			return
		}

		if err := monitoring.Cleanup(zipkinPortForwardPID); err != nil {
			logf("Encountered error killing port-forward process in CleanupZipkinTracingSetup() : %v", err)
			return
		}

		ZipkinTracingEnabled = false
	})
}

// JSONTrace returns a trace for the given traceID. It will continually try to get the trace. If the
// trace it gets has the expected number of spans, then it will be returned. If not, it will try
// again. If it reaches timeout, then it returns everything it has so far with an error.
func JSONTrace(traceID string, expected int, timeout time.Duration) ([]model.SpanModel, error) {
	return JSONTracePred(traceID, timeout, func(trace []model.SpanModel) bool { return len(trace) == expected })
}

// JSONTracePred returns a trace for the given traceID. It will
// continually try to get the trace until the trace spans satisfy the
// predicate. If the timeout is reached then the last fetched trace
// tree if available is returned along with an error.
func JSONTracePred(traceID string, timeout time.Duration, pred func([]model.SpanModel) bool) (trace []model.SpanModel, err error) {
	t := time.After(timeout)
	for !pred(trace) {
		select {
		case <-t:
			return trace, &TimeoutError{
				lastErr: err,
			}
		default:
			trace, err = jsonTrace(traceID)
		}
	}
	return trace, err
}

// TimeoutError is an error returned by JSONTrace if it times out before getting the expected number
// of traces.
type TimeoutError struct {
	lastErr error
}

func (t *TimeoutError) Error() string {
	return fmt.Sprintf("timeout getting JSONTrace, most recent error: %v", t.lastErr)
}

// jsonTrace gets a trace from Zipkin and returns it. Errors returned from this function should be
// retried, as they are likely caused by random problems communicating with Zipkin, or Zipkin
// communicating with its data store.
func jsonTrace(traceID string) ([]model.SpanModel, error) {
	var empty []model.SpanModel

	resp, err := http.Get(ZipkinTraceEndpoint + traceID)
	if err != nil {
		return empty, err
	}
	defer resp.Body.Close()

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

	var models []model.SpanModel
	err = json.Unmarshal(body, &models)
	if err != nil {
		return empty, fmt.Errorf("got an error in unmarshalling JSON %q: %w", body, err)
	}
	return models, nil
}
<commit_msg>Removed zipkin hardcoded values (#1448)<commit_after>/*
Copyright 2019 The Knative Authors

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// util has constants and helper methods useful for zipkin tracing support.

package zipkin

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"sync"
	"testing"
	"time"

	tracingconfig "knative.dev/pkg/tracing/config"

	"github.com/openzipkin/zipkin-go/model"
	"go.opencensus.io/trace"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/kubernetes"
	"knative.dev/pkg/test/logging"
	"knative.dev/pkg/test/monitoring"
)

const (
	// ZipkinTraceIDHeader HTTP response header key to be used to store Zipkin Trace ID.
	ZipkinTraceIDHeader = "ZIPKIN_TRACE_ID"

	// ZipkinPort is port exposed by the Zipkin Pod
	// https://github.com/knative/serving/blob/master/config/monitoring/200-common/100-zipkin.yaml#L25 configures the Zipkin Port on the cluster.
	ZipkinPort = 9411

	// ZipkinTraceEndpoint port-forwarded zipkin endpoint
	ZipkinTraceEndpoint = "http://localhost:9411/api/v2/trace/"

	// App is the name of this component.
	// This will be used as a label selector.
	appLabel = "zipkin"
)

var (
	zipkinPortForwardPID int

	// ZipkinTracingEnabled variable indicating if zipkin tracing is enabled.
	ZipkinTracingEnabled = false

	// sync.Once variable to ensure we execute zipkin setup only once.
	setupOnce sync.Once

	// sync.Once variable to ensure we execute zipkin cleanup only if zipkin is setup and it is executed only once.
	teardownOnce sync.Once
)

// SetupZipkinTracingFromConfigTracing setups zipkin tracing like SetupZipkinTracing but retrieving the zipkin configuration
// from config-tracing config map
func SetupZipkinTracingFromConfigTracing(kubeClientset *kubernetes.Clientset, logf logging.FormatLogger, configMapNamespace string) error {
	cm, err := kubeClientset.CoreV1().ConfigMaps(configMapNamespace).Get("config-tracing", metav1.GetOptions{})
	if err != nil {
		return fmt.Errorf("error while retrieving config-tracing config map: %w", err)
	}
	c, err := tracingconfig.NewTracingConfigFromConfigMap(cm)
	if err != nil {
		return fmt.Errorf("error while parsing config-tracing config map: %w", err)
	}
	zipkinEndpointURL, err := url.Parse(c.ZipkinEndpoint)
	if err != nil {
		return fmt.Errorf("error while parsing the zipkin endpoint in config-tracing config map: %w", err)
	}
	unparsedPort := zipkinEndpointURL.Port()
	port := uint64(80)
	if unparsedPort != "" {
		port, err = strconv.ParseUint(unparsedPort, 10, 16)
		if err != nil {
			return fmt.Errorf("error while parsing the zipkin endpoint port in config-tracing config map: %w", err)
		}
	}

	namespace, err := parseNamespaceFromHostname(zipkinEndpointURL.Host)
	if err != nil {
		return fmt.Errorf("error while parsing the Zipkin endpoint in config-tracing config map: %w", err)
	}

	return SetupZipkinTracing(kubeClientset, logf, int(port), namespace)
}

// SetupZipkinTracingFromConfigTracingOrFail is same as SetupZipkinTracingFromConfigTracing, but fails the test if an error happens
func SetupZipkinTracingFromConfigTracingOrFail(t testing.TB, kubeClientset *kubernetes.Clientset, configMapNamespace string) {
	if err := SetupZipkinTracingFromConfigTracing(kubeClientset, t.Logf, configMapNamespace); err != nil {
		t.Fatalf("Error while setup Zipkin tracing: %v", err)
	}
}

// SetupZipkinTracing sets up zipkin tracing which involves:
// 1. Setting up port-forwarding from localhost to zipkin pod on the cluster
//    (pid of the process doing Port-Forward is stored in a global variable).
// 2. Enable AlwaysSample config for tracing for the SpoofingClient.
// The zipkin deployment must have the label app=zipkin
func SetupZipkinTracing(kubeClientset *kubernetes.Clientset, logf logging.FormatLogger, zipkinRemotePort int, zipkinNamespace string) (err error) {
	setupOnce.Do(func() {
		if e := monitoring.CheckPortAvailability(zipkinRemotePort); e != nil {
			err = fmt.Errorf("Zipkin port not available on the machine: %w", err)
			return
		}

		zipkinPods, e := monitoring.GetPods(kubeClientset, appLabel, zipkinNamespace)
		if e != nil {
			err = fmt.Errorf("error retrieving Zipkin pod details: %w", err)
			return
		}

		zipkinPortForwardPID, e = monitoring.PortForward(logf, zipkinPods, ZipkinPort, zipkinRemotePort, zipkinNamespace)
		if e != nil {
			err = fmt.Errorf("error starting kubectl port-forward command: %w", err)
			return
		}

		logf("Zipkin port-forward process started with PID: %d", zipkinPortForwardPID)

		// Applying AlwaysSample config to ensure we propagate zipkin header for every request made by this client.
		trace.ApplyConfig(trace.Config{DefaultSampler: trace.AlwaysSample()})
		logf("Successfully setup SpoofingClient for Zipkin Tracing")
	})
	return
}

// SetupZipkinTracingOrFail is same as SetupZipkinTracing, but fails the test if an error happens
func SetupZipkinTracingOrFail(t testing.TB, kubeClientset *kubernetes.Clientset, zipkinRemotePort int, zipkinNamespace string) {
	if err := SetupZipkinTracing(kubeClientset, t.Logf, zipkinRemotePort, zipkinNamespace); err != nil {
		t.Fatalf("Error while setup zipkin tracing: %v", err)
	}
}

// CleanupZipkinTracingSetup cleans up the Zipkin tracing setup on the machine. This involves killing the process performing port-forward.
// This should be called exactly once in TestMain. Likely in the form:
//
// func TestMain(m *testing.M) {
//     os.Exit(func() int {
//       // Any setup required for the tests.
//       defer zipkin.CleanupZipkinTracingSetup(logger)
//       return m.Run()
//     }())
// }
func CleanupZipkinTracingSetup(logf logging.FormatLogger) {
	teardownOnce.Do(func() {
		// Because CleanupZipkinTracingSetup only runs once, make sure that now that it has been
		// run, SetupZipkinTracing will no longer setup any port forwarding.
		setupOnce.Do(func() {})

		if !ZipkinTracingEnabled {
			return
		}

		if err := monitoring.Cleanup(zipkinPortForwardPID); err != nil {
			logf("Encountered error killing port-forward process in CleanupZipkinTracingSetup() : %v", err)
			return
		}

		ZipkinTracingEnabled = false
	})
}

// JSONTrace returns a trace for the given traceID. It will continually try to get the trace. If the
// trace it gets has the expected number of spans, then it will be returned. If not, it will try
// again. If it reaches timeout, then it returns everything it has so far with an error.
func JSONTrace(traceID string, expected int, timeout time.Duration) ([]model.SpanModel, error) {
	return JSONTracePred(traceID, timeout, func(trace []model.SpanModel) bool { return len(trace) == expected })
}

// JSONTracePred returns a trace for the given traceID. It will
// continually try to get the trace until the trace spans satisfy the
// predicate. If the timeout is reached then the last fetched trace
// tree if available is returned along with an error.
func JSONTracePred(traceID string, timeout time.Duration, pred func([]model.SpanModel) bool) (trace []model.SpanModel, err error) {
	t := time.After(timeout)
	for !pred(trace) {
		select {
		case <-t:
			return trace, &TimeoutError{
				lastErr: err,
			}
		default:
			trace, err = jsonTrace(traceID)
		}
	}
	return trace, err
}

// TimeoutError is an error returned by JSONTrace if it times out before getting the expected number
// of traces.
type TimeoutError struct {
	lastErr error
}

func (t *TimeoutError) Error() string {
	return fmt.Sprintf("timeout getting JSONTrace, most recent error: %v", t.lastErr)
}

// jsonTrace gets a trace from Zipkin and returns it. Errors returned from this function should be
// retried, as they are likely caused by random problems communicating with Zipkin, or Zipkin
// communicating with its data store.
func jsonTrace(traceID string) ([]model.SpanModel, error) {
	var empty []model.SpanModel

	resp, err := http.Get(ZipkinTraceEndpoint + traceID)
	if err != nil {
		return empty, err
	}
	defer resp.Body.Close()

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

	var models []model.SpanModel
	err = json.Unmarshal(body, &models)
	if err != nil {
		return empty, fmt.Errorf("got an error in unmarshalling JSON %q: %w", body, err)
	}
	return models, nil
}

func parseNamespaceFromHostname(hostname string) (string, error) {
	parts := strings.Split(hostname, ".")
	if len(parts) < 3 || parts[2] != "svc" {
		return "", fmt.Errorf("could not extract namespace/name from %s", hostname)
	}
	return parts[1], nil
}
<|endoftext|>


<commit_before>package tests

import (
	"path/filepath"
	"testing"

	"github.com/ethereum/go-ethereum/accounts"
	"github.com/ethereum/go-ethereum/common"
	"github.com/ethereum/go-ethereum/crypto"
	"github.com/ethereum/go-ethereum/eth"
	"github.com/ethereum/go-ethereum/ethdb"
)

// TODO: refactor test setup & execution to better align with vm and tx tests
func TestBcValidBlockTests(t *testing.T) {
	runBlockTestsInFile("files/BlockTests/bcValidBlockTest.json", []string{}, t)
}

func TestBcUncleTests(t *testing.T) {
	runBlockTestsInFile("files/BlockTests/bcUncleTest.json", []string{}, t)
}

func TestBcUncleHeaderValidityTests(t *testing.T) {
	runBlockTestsInFile("files/BlockTests/bcUncleHeaderValiditiy.json", []string{}, t)
}

func TestBcInvalidHeaderTests(t *testing.T) {
	runBlockTestsInFile("files/BlockTests/bcInvalidHeaderTest.json", []string{}, t)
}

func TestBcInvalidRLPTests(t *testing.T) {
	runBlockTestsInFile("files/BlockTests/bcInvalidRLPTest.json", []string{}, t)
}

func TestBcJSAPITests(t *testing.T) {
	runBlockTestsInFile("files/BlockTests/bcJS_API_Test.json", []string{}, t)
}

func TestBcRPCAPITests(t *testing.T) {
	runBlockTestsInFile("files/BlockTests/bcRPC_API_Test.json", []string{}, t)
}

func TestBcForkBlockTests(t *testing.T) {
	runBlockTestsInFile("files/BlockTests/bcForkBlockTest.json", []string{}, t)
}

func TestBcTotalDifficulty(t *testing.T) {
	runBlockTestsInFile("files/BlockTests/bcTotalDifficultyTest.json", []string{}, t)
}

func TestBcWallet(t *testing.T) {
	runBlockTestsInFile("files/BlockTests/bcWalletTest.json", []string{}, t)
}

func runBlockTestsInFile(filepath string, snafus []string, t *testing.T) {
	bt, err := LoadBlockTests(filepath)
	if err != nil {
		t.Fatal(err)
	}

	notWorking := make(map[string]bool, 100)
	for _, name := range snafus {
		notWorking[name] = true
	}

	for name, test := range bt {
		if !notWorking[name] {
			runBlockTest(name, test, t)
		}
	}
}

func runBlockTest(name string, test *BlockTest, t *testing.T) {
	cfg := testEthConfig()
	ethereum, err := eth.New(cfg)
	if err != nil {
		t.Fatalf("%v", err)
	}

	err = ethereum.Start()
	if err != nil {
		t.Fatalf("%v", err)
	}

	// import the genesis block
	ethereum.ResetWithGenesisBlock(test.Genesis)

	// import pre accounts
	statedb, err := test.InsertPreState(ethereum)
	if err != nil {
		t.Fatalf("InsertPreState: %v", err)
	}

	err = test.TryBlocksInsert(ethereum.ChainManager())
	if err != nil {
		t.Fatal(err)
	}

	if err = test.ValidatePostState(statedb); err != nil {
		t.Fatal("post state validation failed: %v", err)
	}
	t.Log("Test passed: ", name)
}

func testEthConfig() *eth.Config {
	ks := crypto.NewKeyStorePassphrase(filepath.Join(common.DefaultDataDir(), "keystore"))

	return &eth.Config{
		DataDir:        common.DefaultDataDir(),
		Verbosity:      5,
		Etherbase:      "primary",
		AccountManager: accounts.NewManager(ks),
		NewDB:          func(path string) (common.Database, error) { return ethdb.NewMemDatabase() },
	}
}
<commit_msg>tests: removed missing block test<commit_after>package tests

import (
	"path/filepath"
	"testing"

	"github.com/ethereum/go-ethereum/accounts"
	"github.com/ethereum/go-ethereum/common"
	"github.com/ethereum/go-ethereum/crypto"
	"github.com/ethereum/go-ethereum/eth"
	"github.com/ethereum/go-ethereum/ethdb"
)

// TODO: refactor test setup & execution to better align with vm and tx tests
func TestBcValidBlockTests(t *testing.T) {
	runBlockTestsInFile("files/BlockTests/bcValidBlockTest.json", []string{}, t)
}

func TestBcUncleTests(t *testing.T) {
	runBlockTestsInFile("files/BlockTests/bcUncleTest.json", []string{}, t)
}

func TestBcUncleHeaderValidityTests(t *testing.T) {
	runBlockTestsInFile("files/BlockTests/bcUncleHeaderValiditiy.json", []string{}, t)
}

func TestBcInvalidHeaderTests(t *testing.T) {
	runBlockTestsInFile("files/BlockTests/bcInvalidHeaderTest.json", []string{}, t)
}

func TestBcInvalidRLPTests(t *testing.T) {
	runBlockTestsInFile("files/BlockTests/bcInvalidRLPTest.json", []string{}, t)
}

func TestBcRPCAPITests(t *testing.T) {
	runBlockTestsInFile("files/BlockTests/bcRPC_API_Test.json", []string{}, t)
}

func TestBcForkBlockTests(t *testing.T) {
	runBlockTestsInFile("files/BlockTests/bcForkBlockTest.json", []string{}, t)
}

func TestBcTotalDifficulty(t *testing.T) {
	runBlockTestsInFile("files/BlockTests/bcTotalDifficultyTest.json", []string{}, t)
}

func TestBcWallet(t *testing.T) {
	runBlockTestsInFile("files/BlockTests/bcWalletTest.json", []string{}, t)
}

func runBlockTestsInFile(filepath string, snafus []string, t *testing.T) {
	bt, err := LoadBlockTests(filepath)
	if err != nil {
		t.Fatal(err)
	}

	notWorking := make(map[string]bool, 100)
	for _, name := range snafus {
		notWorking[name] = true
	}

	for name, test := range bt {
		if !notWorking[name] {
			runBlockTest(name, test, t)
		}
	}
}

func runBlockTest(name string, test *BlockTest, t *testing.T) {
	cfg := testEthConfig()
	ethereum, err := eth.New(cfg)
	if err != nil {
		t.Fatalf("%v", err)
	}

	err = ethereum.Start()
	if err != nil {
		t.Fatalf("%v", err)
	}

	// import the genesis block
	ethereum.ResetWithGenesisBlock(test.Genesis)

	// import pre accounts
	statedb, err := test.InsertPreState(ethereum)
	if err != nil {
		t.Fatalf("InsertPreState: %v", err)
	}

	err = test.TryBlocksInsert(ethereum.ChainManager())
	if err != nil {
		t.Fatal(err)
	}

	if err = test.ValidatePostState(statedb); err != nil {
		t.Fatal("post state validation failed: %v", err)
	}
	t.Log("Test passed: ", name)
}

func testEthConfig() *eth.Config {
	ks := crypto.NewKeyStorePassphrase(filepath.Join(common.DefaultDataDir(), "keystore"))

	return &eth.Config{
		DataDir:        common.DefaultDataDir(),
		Verbosity:      5,
		Etherbase:      "primary",
		AccountManager: accounts.NewManager(ks),
		NewDB:          func(path string) (common.Database, error) { return ethdb.NewMemDatabase() },
	}
}
<|endoftext|>


<commit_before>package websocket_test

import (
	"bytes"
	"context"
	"testing"
	"time"

	"v2ray.com/core/common"
	"v2ray.com/core/common/net"
	"v2ray.com/core/testing/assert"
	tlsgen "v2ray.com/core/testing/tls"
	"v2ray.com/core/transport/internet"
	v2tls "v2ray.com/core/transport/internet/tls"
	. "v2ray.com/core/transport/internet/websocket"
)

func Test_listenWSAndDial(t *testing.T) {
	assert := assert.On(t)
	listen, err := ListenWS(internet.ContextWithTransportSettings(context.Background(), &Config{
		Path: "ws",
	}), net.DomainAddress("localhost"), 13146, func(ctx context.Context, conn internet.Connection) bool {
		go func(c internet.Connection) {
			defer c.Close()

			var b [1024]byte
			n, err := c.Read(b[:])
			//assert.Error(err).IsNil()
			if err != nil {
				return
			}
			assert.Bool(bytes.HasPrefix(b[:n], []byte("Test connection"))).IsTrue()

			_, err = c.Write([]byte("Response"))
			assert.Error(err).IsNil()
		}(conn)
		return true
	})
	assert.Error(err).IsNil()

	ctx := internet.ContextWithTransportSettings(context.Background(), &Config{Path: "ws"})
	conn, err := Dial(ctx, net.TCPDestination(net.DomainAddress("localhost"), 13146))

	assert.Error(err).IsNil()
	_, err = conn.Write([]byte("Test connection 1"))
	assert.Error(err).IsNil()

	var b [1024]byte
	n, err := conn.Read(b[:])
	assert.Error(err).IsNil()
	assert.String(string(b[:n])).Equals("Response")

	assert.Error(conn.Close()).IsNil()
	<-time.After(time.Second * 5)
	conn, err = Dial(ctx, net.TCPDestination(net.DomainAddress("localhost"), 13146))
	assert.Error(err).IsNil()
	_, err = conn.Write([]byte("Test connection 2"))
	assert.Error(err).IsNil()
	n, err = conn.Read(b[:])
	assert.Error(err).IsNil()
	assert.String(string(b[:n])).Equals("Response")
	assert.Error(conn.Close()).IsNil()
	<-time.After(time.Second * 15)
	conn, err = Dial(ctx, net.TCPDestination(net.DomainAddress("localhost"), 13146))
	assert.Error(err).IsNil()
	_, err = conn.Write([]byte("Test connection 3"))
	assert.Error(err).IsNil()
	n, err = conn.Read(b[:])
	assert.Error(err).IsNil()
	assert.String(string(b[:n])).Equals("Response")
	assert.Error(conn.Close()).IsNil()

	assert.Error(listen.Close()).IsNil()
}

func Test_listenWSAndDial_TLS(t *testing.T) {
	assert := assert.On(t)
	go func() {
		<-time.After(time.Second * 5)
		assert.Fail("Too slow")
	}()

	ctx := internet.ContextWithTransportSettings(context.Background(), &Config{
		Path: "wss",
	})
	ctx = internet.ContextWithSecuritySettings(ctx, &v2tls.Config{
		AllowInsecure: true,
		Certificate:   []*v2tls.Certificate{tlsgen.GenerateCertificateForTest()},
	})
	listen, err := ListenWS(ctx, net.DomainAddress("localhost"), 13143, func(ctx context.Context, conn internet.Connection) bool {
		go func() {
			common.Must(conn.Close())
		}()
		return true
	})
	assert.Error(err).IsNil()
	defer listen.Close()

	conn, err := Dial(ctx, net.TCPDestination(net.DomainAddress("localhost"), 13143))
	assert.Error(err).IsNil()
	common.Must(conn.Close())
}
<commit_msg>fix websocket test<commit_after>package websocket_test

import (
	"bytes"
	"context"
	"testing"
	"time"

	"v2ray.com/core/common/net"
	"v2ray.com/core/testing/assert"
	tlsgen "v2ray.com/core/testing/tls"
	"v2ray.com/core/transport/internet"
	v2tls "v2ray.com/core/transport/internet/tls"
	. "v2ray.com/core/transport/internet/websocket"
)

func Test_listenWSAndDial(t *testing.T) {
	assert := assert.On(t)
	listen, err := ListenWS(internet.ContextWithTransportSettings(context.Background(), &Config{
		Path: "ws",
	}), net.DomainAddress("localhost"), 13146, func(ctx context.Context, conn internet.Connection) bool {
		go func(c internet.Connection) {
			defer c.Close()

			var b [1024]byte
			n, err := c.Read(b[:])
			//assert.Error(err).IsNil()
			if err != nil {
				return
			}
			assert.Bool(bytes.HasPrefix(b[:n], []byte("Test connection"))).IsTrue()

			_, err = c.Write([]byte("Response"))
			assert.Error(err).IsNil()
		}(conn)
		return true
	})
	assert.Error(err).IsNil()

	ctx := internet.ContextWithTransportSettings(context.Background(), &Config{Path: "ws"})
	conn, err := Dial(ctx, net.TCPDestination(net.DomainAddress("localhost"), 13146))

	assert.Error(err).IsNil()
	_, err = conn.Write([]byte("Test connection 1"))
	assert.Error(err).IsNil()

	var b [1024]byte
	n, err := conn.Read(b[:])
	assert.Error(err).IsNil()
	assert.String(string(b[:n])).Equals("Response")

	assert.Error(conn.Close()).IsNil()
	<-time.After(time.Second * 5)
	conn, err = Dial(ctx, net.TCPDestination(net.DomainAddress("localhost"), 13146))
	assert.Error(err).IsNil()
	_, err = conn.Write([]byte("Test connection 2"))
	assert.Error(err).IsNil()
	n, err = conn.Read(b[:])
	assert.Error(err).IsNil()
	assert.String(string(b[:n])).Equals("Response")
	assert.Error(conn.Close()).IsNil()
	<-time.After(time.Second * 15)
	conn, err = Dial(ctx, net.TCPDestination(net.DomainAddress("localhost"), 13146))
	assert.Error(err).IsNil()
	_, err = conn.Write([]byte("Test connection 3"))
	assert.Error(err).IsNil()
	n, err = conn.Read(b[:])
	assert.Error(err).IsNil()
	assert.String(string(b[:n])).Equals("Response")
	assert.Error(conn.Close()).IsNil()

	assert.Error(listen.Close()).IsNil()
}

func Test_listenWSAndDial_TLS(t *testing.T) {
	assert := assert.On(t)
	go func() {
		<-time.After(time.Second * 5)
		assert.Fail("Too slow")
	}()

	ctx := internet.ContextWithTransportSettings(context.Background(), &Config{
		Path: "wss",
	})
	ctx = internet.ContextWithSecuritySettings(ctx, &v2tls.Config{
		AllowInsecure: true,
		Certificate:   []*v2tls.Certificate{tlsgen.GenerateCertificateForTest()},
	})
	listen, err := ListenWS(ctx, net.DomainAddress("localhost"), 13143, func(ctx context.Context, conn internet.Connection) bool {
		go func() {
			_ = conn.Close()
		}()
		return true
	})
	assert.Error(err).IsNil()
	defer listen.Close()

	conn, err := Dial(ctx, net.TCPDestination(net.DomainAddress("localhost"), 13143))
	assert.Error(err).IsNil()
	_ = conn.Close()
}
<|endoftext|>


<commit_before>package golang

import (
	"bytes"
	"text/template"

	"sync"

	"github.com/sourcegraph/go-vcsurl"
	"sourcegraph.com/sourcegraph/srcgraph/dep2"
	"sourcegraph.com/sourcegraph/srcgraph/toolchain"
)

func init() {
	// Register the Go toolchain.
	toolchain.Register("golang", defaultGoVersion)
}

// goVersion represents a Go release: where to download it, how to create graph
// references to it, etc.
type goVersion struct {
	// VersionString is the version string for this Go version, as listed at
	// https://code.google.com/p/go/downloads/list. (E.g., "go1.2.1" or
	// "go1.2rc5".)
	VersionString string

	RepositoryCloneURL string
	RepositoryVCS      vcsurl.VCS
	VCSRevision        string
	SourceUnitPrefix   string

	resolveCache   map[string]*dep2.ResolvedTarget
	resolveCacheMu sync.Mutex
}

var goVersions = map[string]*goVersion{
	"1.2.1": &goVersion{
		VersionString:      "go1.2.1",
		RepositoryCloneURL: "https://code.google.com/p/go",
		RepositoryVCS:      vcsurl.Mercurial,
		VCSRevision:        "go1.2.1",
		SourceUnitPrefix:   "src/pkg",
	},
}

var defaultGoVersion = goVersions["1.2.1"]

func (v *goVersion) baseDockerfile() ([]byte, error) {
	var buf bytes.Buffer
	err := template.Must(template.New("").Parse(baseDockerfile)).Execute(&buf, struct {
		GoVersion *goVersion
		GOPATH    string
	}{
		GoVersion: v,
		GOPATH:    containerGOPATH,
	})
	if err != nil {
		return nil, err
	}

	return buf.Bytes(), nil
}

const containerGOPATH = "/tmp/sg/gopath"

const baseDockerfile = `FROM ubuntu:14.04
RUN apt-get update
RUN apt-get install -qy curl

# Install Go {{.GoVersion.VersionString}}.
RUN curl -o /tmp/golang.tgz https://go.googlecode.com/files/{{.GoVersion.VersionString}}.linux-amd64.tar.gz
RUN tar -xzf /tmp/golang.tgz -C /usr/local
ENV GOROOT /usr/local/go

# Add "go" to the PATH.
ENV PATH /usr/local/go/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin

ENV GOPATH {{.GOPATH}}
`

type baseBuild struct {
	Stdlib *goVersion
	GOPATH string
}
<commit_msg>quieter apt-get<commit_after>package golang

import (
	"bytes"
	"text/template"

	"sync"

	"github.com/sourcegraph/go-vcsurl"
	"sourcegraph.com/sourcegraph/srcgraph/dep2"
	"sourcegraph.com/sourcegraph/srcgraph/toolchain"
)

func init() {
	// Register the Go toolchain.
	toolchain.Register("golang", defaultGoVersion)
}

// goVersion represents a Go release: where to download it, how to create graph
// references to it, etc.
type goVersion struct {
	// VersionString is the version string for this Go version, as listed at
	// https://code.google.com/p/go/downloads/list. (E.g., "go1.2.1" or
	// "go1.2rc5".)
	VersionString string

	RepositoryCloneURL string
	RepositoryVCS      vcsurl.VCS
	VCSRevision        string
	SourceUnitPrefix   string

	resolveCache   map[string]*dep2.ResolvedTarget
	resolveCacheMu sync.Mutex
}

var goVersions = map[string]*goVersion{
	"1.2.1": &goVersion{
		VersionString:      "go1.2.1",
		RepositoryCloneURL: "https://code.google.com/p/go",
		RepositoryVCS:      vcsurl.Mercurial,
		VCSRevision:        "go1.2.1",
		SourceUnitPrefix:   "src/pkg",
	},
}

var defaultGoVersion = goVersions["1.2.1"]

func (v *goVersion) baseDockerfile() ([]byte, error) {
	var buf bytes.Buffer
	err := template.Must(template.New("").Parse(baseDockerfile)).Execute(&buf, struct {
		GoVersion *goVersion
		GOPATH    string
	}{
		GoVersion: v,
		GOPATH:    containerGOPATH,
	})
	if err != nil {
		return nil, err
	}

	return buf.Bytes(), nil
}

const containerGOPATH = "/tmp/sg/gopath"

const baseDockerfile = `FROM ubuntu:14.04
RUN apt-get update -qq
RUN apt-get install -qqy curl

# Install Go {{.GoVersion.VersionString}}.
RUN curl -o /tmp/golang.tgz https://go.googlecode.com/files/{{.GoVersion.VersionString}}.linux-amd64.tar.gz
RUN tar -xzf /tmp/golang.tgz -C /usr/local
ENV GOROOT /usr/local/go

# Add "go" to the PATH.
ENV PATH /usr/local/go/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin

ENV GOPATH {{.GOPATH}}
`

type baseBuild struct {
	Stdlib *goVersion
	GOPATH string
}
<|endoftext|>


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

package ledcli

import (
	"net/http"
	"net/url"
	"strconv"
	"strings"

	"golang.org/x/net/context"

	"github.com/andygrunwald/go-gerrit"
	"github.com/maruel/subcommands"

	bbpb "go.chromium.org/luci/buildbucket/proto"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/led/job"
)

func editCrCLCmd(opts cmdBaseOptions) *subcommands.Command {
	return &subcommands.Command{
		UsageLine: "edit-cr-cl [-remove|-no-implicit-clear] URL_TO_CHANGELIST",
		ShortDesc: "sets Chromium CL-related properties on this JobDefinition (for experimenting with tryjob recipes)",
		LongDesc: `This allows you to edit a JobDefinition for some tryjob recipe
(e.g. chromium_tryjob), and associate a changelist with it, as if the recipe
was triggered via Gerrit.

Recognized URLs:
	https://<gerrit_host>/c/<path/to/project>/+/<issue>/<patchset>
	https://<gerrit_host>/c/<path/to/project>/+/<issue>/<patchset>

If you provide a CL missing <patchset> and <gerrit_host> has public read access,
this will fill in the patchset from the latest version of the issue.

By default, when adding a CL, this will clear all existing CLs on the job, unless
you pass -no-implicit-clear. Most jobs (as of 2020Q2) only expect one CL, so we
did this implicit clearing behavior for CLI ergonomic reasons.
`,

		CommandRun: func() subcommands.CommandRun {
			ret := &cmdEditCl{}
			ret.initFlags(opts)
			return ret
		},
	}
}

type cmdEditCl struct {
	cmdBase

	gerritChange    *bbpb.GerritChange
	remove          bool
	noImplicitClear bool
}

func (c *cmdEditCl) initFlags(opts cmdBaseOptions) {
	c.Flags.BoolVar(&c.remove, "remove", false, "If provided, will remove the given CL instead of adding it.")
	c.Flags.BoolVar(&c.noImplicitClear, "no-implicit-clear", false,
		"If provided, will not clear existing CLs when adding a new one.")
	c.cmdBase.initFlags(opts)
}

func (c *cmdEditCl) jobInput() bool                  { return true }
func (c *cmdEditCl) positionalRange() (min, max int) { return 1, 1 }

func parseCrChangeListURL(clURL string) (*bbpb.GerritChange, error) {
	p, err := url.Parse(clURL)
	if err != nil {
		return nil, errors.Annotate(err, "URL_TO_CHANGELIST is invalid").Err()
	}
	if !strings.HasSuffix(p.Hostname(), "-review.googlesource.com") {
		return nil, errors.Annotate(err, "Only *-review.googlesource.com URLs are supported.").Err()
	}

	var toks []string
	if trimPath := strings.Trim(p.Path, "/"); len(trimPath) > 0 {
		toks = strings.Split(trimPath, "/")
	}

	if len(toks) == 0 {
		// https://<gerrit_host>/#/c/<issue>
		// https://<gerrit_host>/#/c/<issue>/<patchset>
		return nil, errors.Reason("old gerrit URL: %q", clURL).Err()
	} else if toks[0] != "c" {
		return nil, errors.Reason("Unknown changelist URL format: %q", clURL).Err()
	}
	toks = toks[1:] // remove "c"

	// toks ==                 v --------------------------------v
	// https://<gerrit_host>/c/<issue>
	// https://<gerrit_host>/c/<issue>/<patchset>
	// https://<gerrit_host>/c/<project/path>/+/<issue>
	// https://<gerrit_host>/c/<project/path>/+/<issue>/<patchset>

	var projectToks []string
	var issuePatchsetToks []string
	for i, tok := range toks {
		if tok == "+" {
			projectToks, issuePatchsetToks = toks[:i], toks[i+1:]
			break
		}
	}

	if len(projectToks) == 0 {
		return nil, errors.Reason("gerrit URL missing project: %q", clURL).Err()
	}
	if len(issuePatchsetToks) == 0 {
		return nil, errors.Reason("gerrit URL missing issue/patchset: %q", clURL).Err()
	}

	ret := &bbpb.GerritChange{
		Host:    p.Hostname(),
		Project: strings.Join(projectToks, "/"),
	}
	ret.Change, err = strconv.ParseInt(issuePatchsetToks[0], 10, 64)
	if err != nil {
		return nil, errors.Reason("gerrit URL parsing issue %q from %q", issuePatchsetToks[0], clURL).Err()
	}
	if len(issuePatchsetToks) > 1 {
		ret.Patchset, err = strconv.ParseInt(issuePatchsetToks[1], 10, 64)
		if err != nil {
			return nil, errors.Reason("gerrit URL parsing patchset %q from %q", issuePatchsetToks[1], clURL).Err()
		}
	} else {
		gc, err := gerrit.NewClient("https://"+ret.Host, nil)
		if err != nil {
			return nil, errors.Annotate(err, "creating new gerrit client").Err()
		}

		ci, rsp, err := gc.Changes.GetChangeDetail(strconv.FormatInt(ret.Change, 10), &gerrit.ChangeOptions{
			AdditionalFields: []string{"CURRENT_REVISION"}})
		if rsp != nil && rsp.StatusCode == http.StatusUnauthorized {
			return nil, errors.Annotate(err,
				"Gerrit host %q requires authentication and no patchset was provided in CL URL %q. "+
					"Please include the patchset you want in your URL (or `0` to ignore this).",
				ret.Host, clURL,
			).Err()
		}
		if err != nil {
			return nil, errors.Annotate(err, "GetChangeDetail").Err()
		}

		// There's only one.
		for _, rd := range ci.Revisions {
			ret.Patchset = int64(rd.Number)
			break
		}
	}
	return ret, nil
}

func (c *cmdEditCl) validateFlags(ctx context.Context, positionals []string, _ subcommands.Env) (err error) {
	if c.remove && c.noImplicitClear {
		return errors.New("cannot specify both -remove and -no-implicit-clear")
	}

	c.gerritChange, err = parseCrChangeListURL(positionals[0])
	return errors.Annotate(err, "invalid URL_TO_CHANGESET").Err()
}

func (c *cmdEditCl) execute(ctx context.Context, _ *http.Client, inJob *job.Definition) (out interface{}, err error) {
	return inJob, inJob.HighLevelEdit(func(je job.HighLevelEditor) {
		if c.remove {
			je.RemoveGerritChange(c.gerritChange)
		} else {
			if !c.noImplicitClear {
				je.ClearGerritChanges()
			}
			je.AddGerritChange(c.gerritChange)
		}

		// wipe out all the old properties
		je.Properties(map[string]string{
			"blamelist":            "",
			"buildbucket":          "",
			"issue":                "",
			"patch_gerrit_url":     "",
			"patch_issue":          "",
			"patch_project":        "",
			"patch_ref":            "",
			"patch_repository_url": "",
			"patch_set":            "",
			"patch_storage":        "",
			"patchset":             "",
			"repository":           "",
			"rietveld":             "",
		}, true)
	})
}

func (c *cmdEditCl) Run(a subcommands.Application, args []string, env subcommands.Env) int {
	return c.doContextExecute(a, c, args, env)
}
<commit_msg>[led] Fix edit-cr-cl docs.<commit_after>// Copyright 2020 The LUCI Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package ledcli

import (
	"net/http"
	"net/url"
	"strconv"
	"strings"

	"golang.org/x/net/context"

	"github.com/andygrunwald/go-gerrit"
	"github.com/maruel/subcommands"

	bbpb "go.chromium.org/luci/buildbucket/proto"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/led/job"
)

func editCrCLCmd(opts cmdBaseOptions) *subcommands.Command {
	return &subcommands.Command{
		UsageLine: "edit-cr-cl [-remove|-no-implicit-clear] URL_TO_CHANGELIST",
		ShortDesc: "sets Chromium CL-related properties on this JobDefinition (for experimenting with tryjob recipes)",
		LongDesc: `This allows you to edit a JobDefinition for some tryjob recipe
(e.g. chromium_tryjob), and associate a changelist with it, as if the recipe
was triggered via Gerrit.

Recognized URLs:
	https://<gerrit_host>/#/c/<issue>
	https://<gerrit_host>/#/c/<issue>/<patchset>
	https://<gerrit_host>/c/<issue>
	https://<gerrit_host>/c/<issue>/<patchset>
	https://<gerrit_host>/c/<path/to/project>/+/<issue>
	https://<gerrit_host>/c/<path/to/project>/+/<issue>/<patchset>

If you provide a CL missing <patchset> AND <gerrit_host> has public read access,
this will fill in the patchset from the latest version of the issue. Otherwise
this will fail and ask you to provide the full CL/patchset url.

By default, when adding a CL, this will clear all existing CLs on the job, unless
you pass -no-implicit-clear. Most jobs (as of 2020Q2) only expect one CL, so we
did this implicit clearing behavior for CLI ergonomic reasons.
`,

		CommandRun: func() subcommands.CommandRun {
			ret := &cmdEditCl{}
			ret.initFlags(opts)
			return ret
		},
	}
}

type cmdEditCl struct {
	cmdBase

	gerritChange    *bbpb.GerritChange
	remove          bool
	noImplicitClear bool
}

func (c *cmdEditCl) initFlags(opts cmdBaseOptions) {
	c.Flags.BoolVar(&c.remove, "remove", false, "If provided, will remove the given CL instead of adding it.")
	c.Flags.BoolVar(&c.noImplicitClear, "no-implicit-clear", false,
		"If provided, will not clear existing CLs when adding a new one.")
	c.cmdBase.initFlags(opts)
}

func (c *cmdEditCl) jobInput() bool                  { return true }
func (c *cmdEditCl) positionalRange() (min, max int) { return 1, 1 }

func parseCrChangeListURL(clURL string) (*bbpb.GerritChange, error) {
	p, err := url.Parse(clURL)
	if err != nil {
		return nil, errors.Annotate(err, "URL_TO_CHANGELIST is invalid").Err()
	}
	if !strings.HasSuffix(p.Hostname(), "-review.googlesource.com") {
		return nil, errors.Annotate(err, "Only *-review.googlesource.com URLs are supported.").Err()
	}

	var toks []string
	if trimPath := strings.Trim(p.Path, "/"); len(trimPath) > 0 {
		toks = strings.Split(trimPath, "/")
	}

	if len(toks) == 0 {
		// https://<gerrit_host>/#/c/<issue>
		// https://<gerrit_host>/#/c/<issue>/<patchset>
		return nil, errors.Reason("old gerrit URL: %q", clURL).Err()
	} else if toks[0] != "c" {
		return nil, errors.Reason("Unknown changelist URL format: %q", clURL).Err()
	}
	toks = toks[1:] // remove "c"

	// toks ==                 v --------------------------------v
	// https://<gerrit_host>/c/<issue>
	// https://<gerrit_host>/c/<issue>/<patchset>
	// https://<gerrit_host>/c/<project/path>/+/<issue>
	// https://<gerrit_host>/c/<project/path>/+/<issue>/<patchset>

	var projectToks []string
	var issuePatchsetToks []string
	for i, tok := range toks {
		if tok == "+" {
			projectToks, issuePatchsetToks = toks[:i], toks[i+1:]
			break
		}
	}

	if len(projectToks) == 0 {
		return nil, errors.Reason("gerrit URL missing project: %q", clURL).Err()
	}
	if len(issuePatchsetToks) == 0 {
		return nil, errors.Reason("gerrit URL missing issue/patchset: %q", clURL).Err()
	}

	ret := &bbpb.GerritChange{
		Host:    p.Hostname(),
		Project: strings.Join(projectToks, "/"),
	}
	ret.Change, err = strconv.ParseInt(issuePatchsetToks[0], 10, 64)
	if err != nil {
		return nil, errors.Reason("gerrit URL parsing issue %q from %q", issuePatchsetToks[0], clURL).Err()
	}
	if len(issuePatchsetToks) > 1 {
		ret.Patchset, err = strconv.ParseInt(issuePatchsetToks[1], 10, 64)
		if err != nil {
			return nil, errors.Reason("gerrit URL parsing patchset %q from %q", issuePatchsetToks[1], clURL).Err()
		}
	} else {
		gc, err := gerrit.NewClient("https://"+ret.Host, nil)
		if err != nil {
			return nil, errors.Annotate(err, "creating new gerrit client").Err()
		}

		ci, rsp, err := gc.Changes.GetChangeDetail(strconv.FormatInt(ret.Change, 10), &gerrit.ChangeOptions{
			AdditionalFields: []string{"CURRENT_REVISION"}})
		if rsp != nil && rsp.StatusCode == http.StatusUnauthorized {
			return nil, errors.Annotate(err,
				"Gerrit host %q requires authentication and no patchset was provided in CL URL %q. "+
					"Please include the patchset you want in your URL (or `0` to ignore this).",
				ret.Host, clURL,
			).Err()
		}
		if err != nil {
			return nil, errors.Annotate(err, "GetChangeDetail").Err()
		}

		// There's only one.
		for _, rd := range ci.Revisions {
			ret.Patchset = int64(rd.Number)
			break
		}
	}
	return ret, nil
}

func (c *cmdEditCl) validateFlags(ctx context.Context, positionals []string, _ subcommands.Env) (err error) {
	if c.remove && c.noImplicitClear {
		return errors.New("cannot specify both -remove and -no-implicit-clear")
	}

	c.gerritChange, err = parseCrChangeListURL(positionals[0])
	return errors.Annotate(err, "invalid URL_TO_CHANGESET").Err()
}

func (c *cmdEditCl) execute(ctx context.Context, _ *http.Client, inJob *job.Definition) (out interface{}, err error) {
	return inJob, inJob.HighLevelEdit(func(je job.HighLevelEditor) {
		if c.remove {
			je.RemoveGerritChange(c.gerritChange)
		} else {
			if !c.noImplicitClear {
				je.ClearGerritChanges()
			}
			je.AddGerritChange(c.gerritChange)
		}

		// wipe out all the old properties
		je.Properties(map[string]string{
			"blamelist":            "",
			"buildbucket":          "",
			"issue":                "",
			"patch_gerrit_url":     "",
			"patch_issue":          "",
			"patch_project":        "",
			"patch_ref":            "",
			"patch_repository_url": "",
			"patch_set":            "",
			"patch_storage":        "",
			"patchset":             "",
			"repository":           "",
			"rietveld":             "",
		}, true)
	})
}

func (c *cmdEditCl) Run(a subcommands.Application, args []string, env subcommands.Env) int {
	return c.doContextExecute(a, c, args, env)
}
<|endoftext|>