<commit_before>package goahocorasick

import (
	"fmt"
)

import (
	"gitlab.baidu.com/hanshinan/godarts"
)

const FAIL_STATE = -1
const ROOT_STATE = 1

type Machine struct {
	trie    *godarts.DoubleArrayTrie
	failure []int
	output  map[int]([][]rune)
}

type Term struct {
	Pos  int
	Word []rune
}

func (m *Machine) Build(keywords [][]rune) (err error) {
	if len(keywords) == 0 {
		return fmt.Errorf("empty keywords")
	}

	d := new(godarts.Darts)

	trie := new(godarts.LinkedListTrie)
	m.trie, trie, err = d.Build(keywords)
	if err != nil {
		return err
	}

	m.output = make(map[int]([][]rune), 0)
	for idx, val := range d.Output {
		m.output[idx] = append(m.output[idx], val)
	}

	queue := make([](*godarts.LinkedListTrieNode), 0)
	m.failure = make([]int, len(m.trie.Base))
	for _, c := range trie.Root.Children {
		m.failure[c.Base] = godarts.ROOT_NODE_BASE
	}
	queue = append(queue, trie.Root.Children...)

	for {
		if len(queue) == 0 {
			break
		}

		node := queue[0]
		for _, n := range node.Children {
			if n.Base == godarts.END_NODE_BASE {
				continue
			}
			inState := m.f(node.Base)
		set_state:
			outState := m.g(inState, n.Code-godarts.ROOT_NODE_BASE)
			if outState == FAIL_STATE {
				inState = m.f(inState)
				goto set_state
			}
			if _, ok := m.output[outState]; ok != false {
				m.output[n.Base] = append(m.output[outState], m.output[n.Base]...)
			}
			m.setF(n.Base, outState)
		}
		queue = append(queue, node.Children...)
		queue = queue[1:]
	}

	return nil
}

func (m *Machine) PrintFailure() {
	fmt.Printf("+-----+-----+\n")
	fmt.Printf("|%5s|%5s|\n", "index", "value")
	fmt.Printf("+-----+-----+\n")
	for i, v := range m.failure {
		fmt.Printf("|%5d|%5d|\n", i, v)
	}
	fmt.Printf("+-----+-----+\n")
}

func (m *Machine) PrintOutput() {
	fmt.Printf("+-----+----------+\n")
	fmt.Printf("|%5s|%10s|\n", "index", "value")
	fmt.Printf("+-----+----------+\n")
	for i, v := range m.output {
		var val string
		for _, o := range v {
			val = val + " " + string(o)
		}
		fmt.Printf("|%5d|%10s|\n", i, val)
	}
	fmt.Printf("+-----+----------+\n")
}

func (m *Machine) g(inState int, input rune) (outState int) {
	if inState == FAIL_STATE {
		return ROOT_STATE
	}

	t := inState + int(input) + godarts.ROOT_NODE_BASE
	if t >= len(m.trie.Base) {
		if inState == ROOT_STATE {
			return ROOT_STATE
		}
		return FAIL_STATE
	}
	if inState == m.trie.Check[t] {
		return m.trie.Base[t]
	}

	if inState == ROOT_STATE {
		return ROOT_STATE
	}

	return FAIL_STATE
}

func (m *Machine) f(index int) (state int) {
	return m.failure[index]
}

func (m *Machine) setF(inState, outState int) {
	m.failure[inState] = outState
}

func (m *Machine) MultiPatternSearch(content []rune, returnImmediately bool) [](*Term) {
	terms := make([](*Term), 0)

	state := ROOT_STATE
	for pos, c := range content {
	start:
		if m.g(state, c) == FAIL_STATE {
			state = m.f(state)
			goto start
		} else {
			state = m.g(state, c)
			if val, ok := m.output[state]; ok != false {
				for _, word := range val {
					term := new(Term)
					term.Pos = pos - len(word) + 1
					term.Word = word
					terms = append(terms, term)
					if returnImmediately {
						return terms
					}
				}
			}
		}
	}

	return terms
}

func (m *Machine) ExactSearch(content []rune) [](*Term) {
	if m.trie.ExactMatchSearch(content, 0) {
		t := new(Term)
		t.Word = content
		t.Pos = 0
		return [](*Term){t}
	}

	return nil
}
<commit_msg>bug fix<commit_after>package goahocorasick

import (
	"fmt"
)

import (
	"github.com/anknown/darts"
)

const FAIL_STATE = -1
const ROOT_STATE = 1

type Machine struct {
	trie    *godarts.DoubleArrayTrie
	failure []int
	output  map[int]([][]rune)
}

type Term struct {
	Pos  int
	Word []rune
}

func (m *Machine) Build(keywords [][]rune) (err error) {
	if len(keywords) == 0 {
		return fmt.Errorf("empty keywords")
	}

	d := new(godarts.Darts)

	trie := new(godarts.LinkedListTrie)
	m.trie, trie, err = d.Build(keywords)
	if err != nil {
		return err
	}

	m.output = make(map[int]([][]rune), 0)
	for idx, val := range d.Output {
		m.output[idx] = append(m.output[idx], val)
	}

	queue := make([](*godarts.LinkedListTrieNode), 0)
	m.failure = make([]int, len(m.trie.Base))
	for _, c := range trie.Root.Children {
		m.failure[c.Base] = godarts.ROOT_NODE_BASE
	}
	queue = append(queue, trie.Root.Children...)

	for {
		if len(queue) == 0 {
			break
		}

		node := queue[0]
		for _, n := range node.Children {
			if n.Base == godarts.END_NODE_BASE {
				continue
			}
			inState := m.f(node.Base)
		set_state:
			outState := m.g(inState, n.Code-godarts.ROOT_NODE_BASE)
			if outState == FAIL_STATE {
				inState = m.f(inState)
				goto set_state
			}
			if _, ok := m.output[outState]; ok != false {
				m.output[n.Base] = append(m.output[outState], m.output[n.Base]...)
			}
			m.setF(n.Base, outState)
		}
		queue = append(queue, node.Children...)
		queue = queue[1:]
	}

	return nil
}

func (m *Machine) PrintFailure() {
	fmt.Printf("+-----+-----+\n")
	fmt.Printf("|%5s|%5s|\n", "index", "value")
	fmt.Printf("+-----+-----+\n")
	for i, v := range m.failure {
		fmt.Printf("|%5d|%5d|\n", i, v)
	}
	fmt.Printf("+-----+-----+\n")
}

func (m *Machine) PrintOutput() {
	fmt.Printf("+-----+----------+\n")
	fmt.Printf("|%5s|%10s|\n", "index", "value")
	fmt.Printf("+-----+----------+\n")
	for i, v := range m.output {
		var val string
		for _, o := range v {
			val = val + " " + string(o)
		}
		fmt.Printf("|%5d|%10s|\n", i, val)
	}
	fmt.Printf("+-----+----------+\n")
}

func (m *Machine) g(inState int, input rune) (outState int) {
	if inState == FAIL_STATE {
		return ROOT_STATE
	}

	t := inState + int(input) + godarts.ROOT_NODE_BASE
	if t >= len(m.trie.Base) {
		if inState == ROOT_STATE {
			return ROOT_STATE
		}
		return FAIL_STATE
	}
	if inState == m.trie.Check[t] {
		return m.trie.Base[t]
	}

	if inState == ROOT_STATE {
		return ROOT_STATE
	}

	return FAIL_STATE
}

func (m *Machine) f(index int) (state int) {
	return m.failure[index]
}

func (m *Machine) setF(inState, outState int) {
	m.failure[inState] = outState
}

func (m *Machine) MultiPatternSearch(content []rune, returnImmediately bool) [](*Term) {
	terms := make([](*Term), 0)

	state := ROOT_STATE
	for pos, c := range content {
	start:
		if m.g(state, c) == FAIL_STATE {
			state = m.f(state)
			goto start
		} else {
			state = m.g(state, c)
			if val, ok := m.output[state]; ok != false {
				for _, word := range val {
					term := new(Term)
					term.Pos = pos - len(word) + 1
					term.Word = word
					terms = append(terms, term)
					if returnImmediately {
						return terms
					}
				}
			}
		}
	}

	return terms
}

func (m *Machine) ExactSearch(content []rune) [](*Term) {
	if m.trie.ExactMatchSearch(content, 0) {
		t := new(Term)
		t.Word = content
		t.Pos = 0
		return [](*Term){t}
	}

	return nil
}
<|endoftext|>


<commit_before>package ln

import (
	"math"
	"math/rand"
)

type Sphere struct {
	Center Vector
	Radius float64
	Box    Box
}

func NewSphere(center Vector, radius float64) *Sphere {
	min := Vector{center.X - radius, center.Y - radius, center.Z - radius}
	max := Vector{center.X + radius, center.Y + radius, center.Z + radius}
	box := Box{min, max}
	return &Sphere{center, radius, box}
}

func (s *Sphere) Compile() {
}

func (s *Sphere) BoundingBox() Box {
	return s.Box
}

func (s *Sphere) Contains(v Vector, f float64) bool {
	return v.Sub(s.Center).Length() <= s.Radius+f
}

func (s *Sphere) Intersect(r Ray) Hit {
	radius := s.Radius
	to := r.Origin.Sub(s.Center)
	b := to.Dot(r.Direction)
	c := to.Dot(to) - radius*radius
	d := b*b - c
	if d > 0 {
		d = math.Sqrt(d)
		t1 := -b - d
		if t1 > 1e-3 {
			return Hit{s, t1}
		}
		t2 := -b + d
		if t2 > 1e-3 {
			return Hit{s, t2}
		}
	}
	return NoHit
}

func (s *Sphere) Paths4() Paths {
	var paths Paths
	var seen []Vector
	var radii []float64
	for i := 0; i < 140; i++ {
		var v Vector
		var m float64
		for {
			v = RandomUnitVector()
			m = rand.Float64()*0.25 + 0.05
			ok := true
			for i, other := range seen {
				threshold := m + radii[i] + 0.02
				if other.Sub(v).Length() < threshold {
					ok = false
					break
				}
			}
			if ok {
				seen = append(seen, v)
				radii = append(radii, m)
				break
			}
		}
		p := v.Cross(RandomUnitVector()).Normalize()
		q := p.Cross(v).Normalize()
		n := rand.Intn(4) + 1
		for k := 0; k < n; k++ {
			var path Path
			for j := 0; j <= 360; j += 5 {
				a := Radians(float64(j))
				x := v
				x = x.Add(p.MulScalar(math.Cos(a) * m))
				x = x.Add(q.MulScalar(math.Sin(a) * m))
				x = x.Normalize()
				x = x.MulScalar(s.Radius).Add(s.Center)
				path = append(path, x)
			}
			paths = append(paths, path)
			m *= 0.75
		}
	}
	return paths
}

func (s *Sphere) Paths3() Paths {
	var paths Paths
	for i := 0; i < 20000; i++ {
		v := RandomUnitVector()
		v = v.MulScalar(s.Radius).Add(s.Center)
		paths = append(paths, Path{v, v})
	}
	return paths
}

func (s *Sphere) Paths2() Paths {
	var equator Path
	for lng := 0; lng <= 360; lng++ {
		v := LatLngToXYZ(0, float64(lng), s.Radius)
		equator = append(equator, v)
	}
	var paths Paths
	for i := 0; i < 100; i++ {
		m := Identity()
		for j := 0; j < 3; j++ {
			v := RandomUnitVector()
			m = m.Rotate(v, rand.Float64()*2*math.Pi)
		}
		m = m.Translate(s.Center)
		paths = append(paths, equator.Transform(m))
	}
	return paths
}

func (s *Sphere) Paths() Paths {
	var paths Paths
	n := 10
	o := 10
	for lat := -90 + o; lat <= 90-o; lat += n {
		var path Path
		for lng := 0; lng <= 360; lng++ {
			v := LatLngToXYZ(float64(lat), float64(lng), s.Radius).Add(s.Center)
			path = append(path, v)
		}
		paths = append(paths, path)
	}
	for lng := 0; lng <= 360; lng += n {
		var path Path
		for lat := -90 + o; lat <= 90-o; lat++ {
			v := LatLngToXYZ(float64(lat), float64(lng), s.Radius).Add(s.Center)
			path = append(path, v)
		}
		paths = append(paths, path)
	}
	return paths
}

func LatLngToXYZ(lat, lng, radius float64) Vector {
	lat, lng = Radians(lat), Radians(lng)
	x := radius * math.Cos(lat) * math.Cos(lng)
	y := radius * math.Cos(lat) * math.Sin(lng)
	z := radius * math.Sin(lat)
	return Vector{x, y, z}
}

type OutlineSphere struct {
	Sphere
	Eye Vector
	Up  Vector
}

func NewOutlineSphere(eye, up, center Vector, radius float64) *OutlineSphere {
	sphere := NewSphere(center, radius)
	return &OutlineSphere{*sphere, eye, up}
}

func (s *OutlineSphere) Paths() Paths {
	var path Path
	c := s.Sphere.Center
	r := s.Sphere.Radius + 0.005
	w := c.Sub(s.Eye)
	u := w.Cross(s.Up).Normalize()
	v := w.Cross(u).Normalize()
	for i := 0; i <= 360; i++ {
		a := Radians(float64(i))
		p := c
		p = p.Add(u.MulScalar(math.Cos(a) * r))
		p = p.Add(v.MulScalar(math.Sin(a) * r))
		path = append(path, p)
	}
	return Paths{path}
}
<commit_msg>fix OutlineSphere<commit_after>package ln

import (
	"math"
	"math/rand"
)

type Sphere struct {
	Center Vector
	Radius float64
	Box    Box
}

func NewSphere(center Vector, radius float64) *Sphere {
	min := Vector{center.X - radius, center.Y - radius, center.Z - radius}
	max := Vector{center.X + radius, center.Y + radius, center.Z + radius}
	box := Box{min, max}
	return &Sphere{center, radius, box}
}

func (s *Sphere) Compile() {
}

func (s *Sphere) BoundingBox() Box {
	return s.Box
}

func (s *Sphere) Contains(v Vector, f float64) bool {
	return v.Sub(s.Center).Length() <= s.Radius+f
}

func (s *Sphere) Intersect(r Ray) Hit {
	radius := s.Radius
	to := r.Origin.Sub(s.Center)
	b := to.Dot(r.Direction)
	c := to.Dot(to) - radius*radius
	d := b*b - c
	if d > 0 {
		d = math.Sqrt(d)
		t1 := -b - d
		if t1 > 1e-3 {
			return Hit{s, t1}
		}
		t2 := -b + d
		if t2 > 1e-3 {
			return Hit{s, t2}
		}
	}
	return NoHit
}

func (s *Sphere) Paths4() Paths {
	var paths Paths
	var seen []Vector
	var radii []float64
	for i := 0; i < 140; i++ {
		var v Vector
		var m float64
		for {
			v = RandomUnitVector()
			m = rand.Float64()*0.25 + 0.05
			ok := true
			for i, other := range seen {
				threshold := m + radii[i] + 0.02
				if other.Sub(v).Length() < threshold {
					ok = false
					break
				}
			}
			if ok {
				seen = append(seen, v)
				radii = append(radii, m)
				break
			}
		}
		p := v.Cross(RandomUnitVector()).Normalize()
		q := p.Cross(v).Normalize()
		n := rand.Intn(4) + 1
		for k := 0; k < n; k++ {
			var path Path
			for j := 0; j <= 360; j += 5 {
				a := Radians(float64(j))
				x := v
				x = x.Add(p.MulScalar(math.Cos(a) * m))
				x = x.Add(q.MulScalar(math.Sin(a) * m))
				x = x.Normalize()
				x = x.MulScalar(s.Radius).Add(s.Center)
				path = append(path, x)
			}
			paths = append(paths, path)
			m *= 0.75
		}
	}
	return paths
}

func (s *Sphere) Paths3() Paths {
	var paths Paths
	for i := 0; i < 20000; i++ {
		v := RandomUnitVector()
		v = v.MulScalar(s.Radius).Add(s.Center)
		paths = append(paths, Path{v, v})
	}
	return paths
}

func (s *Sphere) Paths2() Paths {
	var equator Path
	for lng := 0; lng <= 360; lng++ {
		v := LatLngToXYZ(0, float64(lng), s.Radius)
		equator = append(equator, v)
	}
	var paths Paths
	for i := 0; i < 100; i++ {
		m := Identity()
		for j := 0; j < 3; j++ {
			v := RandomUnitVector()
			m = m.Rotate(v, rand.Float64()*2*math.Pi)
		}
		m = m.Translate(s.Center)
		paths = append(paths, equator.Transform(m))
	}
	return paths
}

func (s *Sphere) Paths() Paths {
	var paths Paths
	n := 10
	o := 10
	for lat := -90 + o; lat <= 90-o; lat += n {
		var path Path
		for lng := 0; lng <= 360; lng++ {
			v := LatLngToXYZ(float64(lat), float64(lng), s.Radius).Add(s.Center)
			path = append(path, v)
		}
		paths = append(paths, path)
	}
	for lng := 0; lng <= 360; lng += n {
		var path Path
		for lat := -90 + o; lat <= 90-o; lat++ {
			v := LatLngToXYZ(float64(lat), float64(lng), s.Radius).Add(s.Center)
			path = append(path, v)
		}
		paths = append(paths, path)
	}
	return paths
}

func LatLngToXYZ(lat, lng, radius float64) Vector {
	lat, lng = Radians(lat), Radians(lng)
	x := radius * math.Cos(lat) * math.Cos(lng)
	y := radius * math.Cos(lat) * math.Sin(lng)
	z := radius * math.Sin(lat)
	return Vector{x, y, z}
}

type OutlineSphere struct {
	Sphere
	Eye Vector
	Up  Vector
}

func NewOutlineSphere(eye, up, center Vector, radius float64) *OutlineSphere {
	sphere := NewSphere(center, radius)
	return &OutlineSphere{*sphere, eye, up}
}

func (s *OutlineSphere) Paths() Paths {
	var path Path
	center := s.Sphere.Center
	radius := s.Sphere.Radius

	hyp := center.Sub(s.Eye).Length()
	opp := radius
	theta := math.Asin(opp / hyp)
	adj := opp / math.Tan(theta)
	d := math.Cos(theta) * adj
	r := math.Sin(theta) * adj

	w := center.Sub(s.Eye).Normalize()
	u := w.Cross(s.Up).Normalize()
	v := w.Cross(u).Normalize()
	c := s.Eye.Add(w.MulScalar(d))
	for i := 0; i <= 360; i++ {
		a := Radians(float64(i))
		p := c
		p = p.Add(u.MulScalar(math.Cos(a) * r))
		p = p.Add(v.MulScalar(math.Sin(a) * r))
		path = append(path, p)
	}
	return Paths{path}
}
<|endoftext|>


<commit_before>// Copyright 2019 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 main

import (
	"context"
	"time"

	"golang.org/x/time/rate"

	"google.golang.org/genproto/protobuf/field_mask"
	"google.golang.org/grpc"
	"google.golang.org/grpc/metadata"

	"go.chromium.org/luci/auth"
	"go.chromium.org/luci/buildbucket"
	bbpb "go.chromium.org/luci/buildbucket/proto"
	"go.chromium.org/luci/common/clock"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/lhttp"
	"go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/common/retry"
	"go.chromium.org/luci/common/retry/transient"
	"go.chromium.org/luci/common/sync/dispatcher"
	"go.chromium.org/luci/common/sync/dispatcher/buffer"
	"go.chromium.org/luci/grpc/prpc"
	"go.chromium.org/luci/lucictx"
)

// BuildsClient is a trimmed version of `bbpb.BuildsClient` which only
// contains the required RPC method for bbagent.
type BuildsClient interface {
	// Updates a build. Returns an empty build.
	//
	// RPC metadata must include "X-Build-Token" key with a token
	// generated by the server when scheduling the build.
	UpdateBuild(ctx context.Context, in *bbpb.UpdateBuildRequest, opts ...grpc.CallOption) (*bbpb.Build, error)
}

var _ BuildsClient = dummyBBClient{}

type dummyBBClient struct{}

func (dummyBBClient) UpdateBuild(ctx context.Context, in *bbpb.UpdateBuildRequest, opts ...grpc.CallOption) (*bbpb.Build, error) {
	return nil, nil
}

// retryEnabled allows us to switch retries for this client on and off
func newBuildsClient(ctx context.Context, infraOpts *bbpb.BuildInfra_Buildbucket, retryEnabled *bool) (BuildsClient, *bbpb.BuildSecrets, error) {
	hostname := infraOpts.GetHostname()
	if hostname == "" {
		logging.Infof(ctx, "No buildbucket hostname set; making dummy buildbucket client.")
		return dummyBBClient{}, &bbpb.BuildSecrets{BuildToken: "dummy token"}, nil
	}
	opts := prpc.DefaultOptions()
	opts.Insecure = lhttp.IsLocalHost(hostname)
	originalRetry := opts.Retry
	opts.Retry = func() retry.Iterator {
		if *retryEnabled {
			return originalRetry()
		}
		return nil
	}

	prpcClient := &prpc.Client{
		Host:    hostname,
		Options: opts,
	}
	secrets, err := readBuildSecrets(ctx)
	if err != nil {
		return nil, nil, err
	}

	// Use "system" account to call UpdateBuild RPCs.
	sctx, err := lucictx.SwitchLocalAccount(ctx, "system")
	if err != nil {
		return nil, nil, errors.Annotate(err, "could not switch to 'system' account in LUCI_CONTEXT").Err()
	}
	prpcClient.C, err = auth.NewAuthenticator(sctx, auth.SilentLogin, auth.Options{
		MonitorAs: "bbagent/buildbucket",
	}).Client()
	if err != nil {
		return nil, nil, err
	}
	// TODO(iannucci): Exchange secret build token+nonce for a running build token
	// here to confirm that:
	//   * We're the ONLY ones servicing this build (detect duplicate Swarming
	//     tasks). Failure to exchange the token would let us know that we got
	//     double-booked.
	//   * Auth is properly configured for buildbucket before we start running the
	//     user code.
	return bbpb.NewBuildsPRPCClient(prpcClient), secrets, nil
}

// options for the dispatcher.Channel
func channelOpts(ctx context.Context) (*dispatcher.Options, <-chan error) {
	errorFn, errCh := dispatcher.ErrorFnReport(10, func(failedBatch *buffer.Batch, err error) bool {
		return transient.Tag.In(err)
	})
	opts := &dispatcher.Options{
		QPSLimit: rate.NewLimiter(1, 1),
		Buffer: buffer.Options{
			BatchSize:    1,
			MaxLeases:    1,
			FullBehavior: &buffer.DropOldestBatch{MaxLiveItems: 1},
			Retry: func() retry.Iterator {
				return &retry.ExponentialBackoff{
					Limited: retry.Limited{
						Delay:    200 * time.Millisecond, // initial delay
						Retries:  -1,
						MaxTotal: 5 * time.Minute,
					},
					Multiplier: 1.2,
					MaxDelay:   30 * time.Second,
				}
			},
		},
		DropFn:  dispatcher.DropFnSummarized(ctx, rate.NewLimiter(.1, 1)),
		ErrorFn: errorFn,
	}
	return opts, errCh
}

func mkSendFn(ctx context.Context, secrets *bbpb.BuildSecrets, client BuildsClient) dispatcher.SendFn {
	return func(b *buffer.Batch) error {
		ctx = metadata.NewOutgoingContext(ctx, metadata.Pairs(buildbucket.BuildTokenHeader, secrets.BuildToken))

		var req *bbpb.UpdateBuildRequest

		if b.Meta != nil {
			req = b.Meta.(*bbpb.UpdateBuildRequest)
		} else {
			build := b.Data[0].(*bbpb.Build)
			// We always set status=STARTED; bbagent will do one final send after the
			// dispatcher channel is closed which includes the real final status.
			build.Status = bbpb.Status_STARTED
			req = &bbpb.UpdateBuildRequest{
				Build: build,
				UpdateMask: &field_mask.FieldMask{
					Paths: []string{
						"build.steps",
						"build.status",
						"build.output",
						"build.summary_markdown",
					},
				},
			}
			if len(build.Tags) > 0 {
				req.UpdateMask.Paths = append(req.UpdateMask.Paths, "build.tags")
			}
			b.Meta = req
			b.Data[0] = nil
		}

		// This RPC is currently served by an AppEngine frontend instance which
		// is capped at a 60s request time.
		tctx, cancel := clock.WithTimeout(ctx, 60*time.Second+500*time.Millisecond)
		defer cancel()

		_, err := client.UpdateBuild(tctx, req)
		return err
	}
}
<commit_msg>[bbagent] Copy Build before mutating.<commit_after>// Copyright 2019 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 main

import (
	"context"
	"time"

	"golang.org/x/time/rate"

	"google.golang.org/genproto/protobuf/field_mask"
	"google.golang.org/grpc"
	"google.golang.org/grpc/metadata"
	"google.golang.org/protobuf/proto"

	"go.chromium.org/luci/auth"
	"go.chromium.org/luci/buildbucket"
	bbpb "go.chromium.org/luci/buildbucket/proto"
	"go.chromium.org/luci/common/clock"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/lhttp"
	"go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/common/retry"
	"go.chromium.org/luci/common/retry/transient"
	"go.chromium.org/luci/common/sync/dispatcher"
	"go.chromium.org/luci/common/sync/dispatcher/buffer"
	"go.chromium.org/luci/grpc/prpc"
	"go.chromium.org/luci/lucictx"
)

// BuildsClient is a trimmed version of `bbpb.BuildsClient` which only
// contains the required RPC method for bbagent.
type BuildsClient interface {
	// Updates a build. Returns an empty build.
	//
	// RPC metadata must include "X-Build-Token" key with a token
	// generated by the server when scheduling the build.
	UpdateBuild(ctx context.Context, in *bbpb.UpdateBuildRequest, opts ...grpc.CallOption) (*bbpb.Build, error)
}

var _ BuildsClient = dummyBBClient{}

type dummyBBClient struct{}

func (dummyBBClient) UpdateBuild(ctx context.Context, in *bbpb.UpdateBuildRequest, opts ...grpc.CallOption) (*bbpb.Build, error) {
	return nil, nil
}

// retryEnabled allows us to switch retries for this client on and off
func newBuildsClient(ctx context.Context, infraOpts *bbpb.BuildInfra_Buildbucket, retryEnabled *bool) (BuildsClient, *bbpb.BuildSecrets, error) {
	hostname := infraOpts.GetHostname()
	if hostname == "" {
		logging.Infof(ctx, "No buildbucket hostname set; making dummy buildbucket client.")
		return dummyBBClient{}, &bbpb.BuildSecrets{BuildToken: "dummy token"}, nil
	}
	opts := prpc.DefaultOptions()
	opts.Insecure = lhttp.IsLocalHost(hostname)
	originalRetry := opts.Retry
	opts.Retry = func() retry.Iterator {
		if *retryEnabled {
			return originalRetry()
		}
		return nil
	}

	prpcClient := &prpc.Client{
		Host:    hostname,
		Options: opts,
	}
	secrets, err := readBuildSecrets(ctx)
	if err != nil {
		return nil, nil, err
	}

	// Use "system" account to call UpdateBuild RPCs.
	sctx, err := lucictx.SwitchLocalAccount(ctx, "system")
	if err != nil {
		return nil, nil, errors.Annotate(err, "could not switch to 'system' account in LUCI_CONTEXT").Err()
	}
	prpcClient.C, err = auth.NewAuthenticator(sctx, auth.SilentLogin, auth.Options{
		MonitorAs: "bbagent/buildbucket",
	}).Client()
	if err != nil {
		return nil, nil, err
	}
	// TODO(iannucci): Exchange secret build token+nonce for a running build token
	// here to confirm that:
	//   * We're the ONLY ones servicing this build (detect duplicate Swarming
	//     tasks). Failure to exchange the token would let us know that we got
	//     double-booked.
	//   * Auth is properly configured for buildbucket before we start running the
	//     user code.
	return bbpb.NewBuildsPRPCClient(prpcClient), secrets, nil
}

// options for the dispatcher.Channel
func channelOpts(ctx context.Context) (*dispatcher.Options, <-chan error) {
	errorFn, errCh := dispatcher.ErrorFnReport(10, func(failedBatch *buffer.Batch, err error) bool {
		return transient.Tag.In(err)
	})
	opts := &dispatcher.Options{
		QPSLimit: rate.NewLimiter(1, 1),
		Buffer: buffer.Options{
			BatchSize:    1,
			MaxLeases:    1,
			FullBehavior: &buffer.DropOldestBatch{MaxLiveItems: 1},
			Retry: func() retry.Iterator {
				return &retry.ExponentialBackoff{
					Limited: retry.Limited{
						Delay:    200 * time.Millisecond, // initial delay
						Retries:  -1,
						MaxTotal: 5 * time.Minute,
					},
					Multiplier: 1.2,
					MaxDelay:   30 * time.Second,
				}
			},
		},
		DropFn:  dispatcher.DropFnSummarized(ctx, rate.NewLimiter(.1, 1)),
		ErrorFn: errorFn,
	}
	return opts, errCh
}

func mkSendFn(ctx context.Context, secrets *bbpb.BuildSecrets, client BuildsClient) dispatcher.SendFn {
	return func(b *buffer.Batch) error {
		ctx = metadata.NewOutgoingContext(ctx, metadata.Pairs(buildbucket.BuildTokenHeader, secrets.BuildToken))

		var req *bbpb.UpdateBuildRequest

		if b.Meta != nil {
			req = b.Meta.(*bbpb.UpdateBuildRequest)
		} else {
			build := proto.Clone(b.Data[0].(proto.Message)).(*bbpb.Build)
			// We always set status=STARTED; bbagent will do one final send after the
			// dispatcher channel is closed which includes the real final status.
			build.Status = bbpb.Status_STARTED
			req = &bbpb.UpdateBuildRequest{
				Build: build,
				UpdateMask: &field_mask.FieldMask{
					Paths: []string{
						"build.steps",
						"build.status",
						"build.output",
						"build.summary_markdown",
					},
				},
			}
			if len(build.Tags) > 0 {
				req.UpdateMask.Paths = append(req.UpdateMask.Paths, "build.tags")
			}
			b.Meta = req
			b.Data[0] = nil
		}

		// This RPC is currently served by an AppEngine frontend instance which
		// is capped at a 60s request time.
		tctx, cancel := clock.WithTimeout(ctx, 60*time.Second+500*time.Millisecond)
		defer cancel()

		_, err := client.UpdateBuild(tctx, req)
		return err
	}
}
<|endoftext|>


<commit_before>package common

import (
	"bytes"
	"fmt"
	"log"
	"os/exec"
	"regexp"
	"strings"
	"time"
)

type VBox42Driver struct {
	// This is the path to the "VBoxManage" application.
	VBoxManagePath string
}

func (d *VBox42Driver) CreateSATAController(vmName string, name string) error {
	version, err := d.Version()
	if err != nil {
		return err
	}

	portCountArg := "--sataportcount"
	if strings.HasPrefix(version, "4.3") {
		portCountArg = "--portcount"
	}

	command := []string{
		"storagectl", vmName,
		"--name", name,
		"--add", "sata",
		portCountArg, "1",
	}

	return d.VBoxManage(command...)
}

func (d *VBox42Driver) Delete(name string) error {
	return d.VBoxManage("unregistervm", name, "--delete")
}

func (d *VBox42Driver) Iso() (string, error) {
	var stdout bytes.Buffer

	cmd := exec.Command(d.VBoxManagePath, "list", "systemproperties")
	cmd.Stdout = &stdout
	if err := cmd.Run(); err != nil {
		return "", err
	}

	DefaultGuestAdditionsRe := regexp.MustCompile("Default Guest Additions ISO:(.*)")

	for _, line := range strings.Split(stdout.String(), "\n") {
		// Need to trim off CR character when running in windows
		line = strings.TrimRight(line, "\r")

		matches := DefaultGuestAdditionsRe.FindStringSubmatch(line)
		if matches == nil {
			continue
		}

		isoname := strings.Trim(matches[1], " \r\n")
		log.Printf("Found Default Guest Additions ISO: %s", isoname)

		return isoname, nil
	}

	return "", fmt.Errorf("Cannot find \"Default Guest Additions ISO\" in vboxmanage output")
}

func (d *VBox42Driver) Import(name, path, opts string) error {
	args := []string{
		"import", path,
		"--vsys", "0",
		"--vmname", name,
		"--options", opts,
	}

	return d.VBoxManage(args...)
}

func (d *VBox42Driver) IsRunning(name string) (bool, error) {
	var stdout bytes.Buffer

	cmd := exec.Command(d.VBoxManagePath, "showvminfo", name, "--machinereadable")
	cmd.Stdout = &stdout
	if err := cmd.Run(); err != nil {
		return false, err
	}

	for _, line := range strings.Split(stdout.String(), "\n") {
		// Need to trim off CR character when running in windows
		line = strings.TrimRight(line, "\r")

		if line == `VMState="running"` {
			return true, nil
		}

		// We consider "stopping" to still be running. We wait for it to
		// be completely stopped or some other state.
		if line == `VMState="stopping"` {
			return true, nil
		}

		// We consider "paused" to still be running. We wait for it to
		// be completely stopped or some other state.
		if line == `VMState="paused"` {
			return true, nil
		}
	}

	return false, nil
}

func (d *VBox42Driver) Stop(name string) error {
	if err := d.VBoxManage("controlvm", name, "poweroff"); err != nil {
		return err
	}

	// We sleep here for a little bit to let the session "unlock"
	time.Sleep(2 * time.Second)

	return nil
}

func (d *VBox42Driver) SuppressMessages() error {
	extraData := map[string]string{
		"GUI/RegistrationData": "triesLeft=0",
		"GUI/SuppressMessages": "confirmInputCapture,remindAboutAutoCapture,remindAboutMouseIntegrationOff,remindAboutMouseIntegrationOn,remindAboutWrongColorDepth",
		"GUI/UpdateDate":       fmt.Sprintf("1 d, %d-01-01, stable", time.Now().Year()+1),
		"GUI/UpdateCheckCount": "60",
	}

	for k, v := range extraData {
		if err := d.VBoxManage("setextradata", "global", k, v); err != nil {
			return err
		}
	}

	return nil
}

func (d *VBox42Driver) VBoxManage(args ...string) error {
	var stdout, stderr bytes.Buffer

	log.Printf("Executing VBoxManage: %#v", args)
	cmd := exec.Command(d.VBoxManagePath, args...)
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr
	err := cmd.Run()

	stdoutString := strings.TrimSpace(stdout.String())
	stderrString := strings.TrimSpace(stderr.String())

	if _, ok := err.(*exec.ExitError); ok {
		err = fmt.Errorf("VBoxManage error: %s", stderrString)
	}

	if err == nil {
		m, _ := regexp.MatchString("VBoxManage([.a-z]+?): error:", stderrString)
		if m {
			err = fmt.Errorf("VBoxManage error: %s", stderrString)
		}
	}

	log.Printf("stdout: %s", stdoutString)
	log.Printf("stderr: %s", stderrString)

	return err
}

func (d *VBox42Driver) Verify() error {
	return nil
}

func (d *VBox42Driver) Version() (string, error) {
	var stdout bytes.Buffer

	cmd := exec.Command(d.VBoxManagePath, "--version")
	cmd.Stdout = &stdout
	if err := cmd.Run(); err != nil {
		return "", err
	}

	versionOutput := strings.TrimSpace(stdout.String())
	log.Printf("VBoxManage --version output: %s", versionOutput)

	// If the "--version" output contains vboxdrv, then this is indicative
	// of problems with the VirtualBox setup and we shouldn't really continue,
	// whether or not we can read the version.
	if strings.Contains(versionOutput, "vboxdrv") {
		return "", fmt.Errorf("VirtualBox is not properly setup: %s", versionOutput)
	}

	versionRe := regexp.MustCompile("^[.0-9]+(?:_RC[0-9]+)?")
	matches := versionRe.FindAllString(versionOutput, 1)
	if matches == nil {
		return "", fmt.Errorf("No version found: %s", versionOutput)
	}

	log.Printf("VirtualBox version: %s", matches[0])
	return matches[0], nil
}
<commit_msg>builder/virtualbox: comment<commit_after>package common

import (
	"bytes"
	"fmt"
	"log"
	"os/exec"
	"regexp"
	"strings"
	"time"
)

type VBox42Driver struct {
	// This is the path to the "VBoxManage" application.
	VBoxManagePath string
}

func (d *VBox42Driver) CreateSATAController(vmName string, name string) error {
	version, err := d.Version()
	if err != nil {
		return err
	}

	portCountArg := "--sataportcount"
	if strings.HasPrefix(version, "4.3") {
		portCountArg = "--portcount"
	}

	command := []string{
		"storagectl", vmName,
		"--name", name,
		"--add", "sata",
		portCountArg, "1",
	}

	return d.VBoxManage(command...)
}

func (d *VBox42Driver) Delete(name string) error {
	return d.VBoxManage("unregistervm", name, "--delete")
}

func (d *VBox42Driver) Iso() (string, error) {
	var stdout bytes.Buffer

	cmd := exec.Command(d.VBoxManagePath, "list", "systemproperties")
	cmd.Stdout = &stdout
	if err := cmd.Run(); err != nil {
		return "", err
	}

	DefaultGuestAdditionsRe := regexp.MustCompile("Default Guest Additions ISO:(.*)")

	for _, line := range strings.Split(stdout.String(), "\n") {
		// Need to trim off CR character when running in windows
		line = strings.TrimRight(line, "\r")

		matches := DefaultGuestAdditionsRe.FindStringSubmatch(line)
		if matches == nil {
			continue
		}

		isoname := strings.Trim(matches[1], " \r\n")
		log.Printf("Found Default Guest Additions ISO: %s", isoname)

		return isoname, nil
	}

	return "", fmt.Errorf("Cannot find \"Default Guest Additions ISO\" in vboxmanage output")
}

func (d *VBox42Driver) Import(name, path, opts string) error {
	args := []string{
		"import", path,
		"--vsys", "0",
		"--vmname", name,
		"--options", opts,
	}

	return d.VBoxManage(args...)
}

func (d *VBox42Driver) IsRunning(name string) (bool, error) {
	var stdout bytes.Buffer

	cmd := exec.Command(d.VBoxManagePath, "showvminfo", name, "--machinereadable")
	cmd.Stdout = &stdout
	if err := cmd.Run(); err != nil {
		return false, err
	}

	for _, line := range strings.Split(stdout.String(), "\n") {
		// Need to trim off CR character when running in windows
		line = strings.TrimRight(line, "\r")

		if line == `VMState="running"` {
			return true, nil
		}

		// We consider "stopping" to still be running. We wait for it to
		// be completely stopped or some other state.
		if line == `VMState="stopping"` {
			return true, nil
		}

		// We consider "paused" to still be running. We wait for it to
		// be completely stopped or some other state.
		if line == `VMState="paused"` {
			return true, nil
		}
	}

	return false, nil
}

func (d *VBox42Driver) Stop(name string) error {
	if err := d.VBoxManage("controlvm", name, "poweroff"); err != nil {
		return err
	}

	// We sleep here for a little bit to let the session "unlock"
	time.Sleep(2 * time.Second)

	return nil
}

func (d *VBox42Driver) SuppressMessages() error {
	extraData := map[string]string{
		"GUI/RegistrationData": "triesLeft=0",
		"GUI/SuppressMessages": "confirmInputCapture,remindAboutAutoCapture,remindAboutMouseIntegrationOff,remindAboutMouseIntegrationOn,remindAboutWrongColorDepth",
		"GUI/UpdateDate":       fmt.Sprintf("1 d, %d-01-01, stable", time.Now().Year()+1),
		"GUI/UpdateCheckCount": "60",
	}

	for k, v := range extraData {
		if err := d.VBoxManage("setextradata", "global", k, v); err != nil {
			return err
		}
	}

	return nil
}

func (d *VBox42Driver) VBoxManage(args ...string) error {
	var stdout, stderr bytes.Buffer

	log.Printf("Executing VBoxManage: %#v", args)
	cmd := exec.Command(d.VBoxManagePath, args...)
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr
	err := cmd.Run()

	stdoutString := strings.TrimSpace(stdout.String())
	stderrString := strings.TrimSpace(stderr.String())

	if _, ok := err.(*exec.ExitError); ok {
		err = fmt.Errorf("VBoxManage error: %s", stderrString)
	}

	if err == nil {
		// Sometimes VBoxManage gives us an error with a zero exit code,
		// so we also regexp match an error string.
		m, _ := regexp.MatchString("VBoxManage([.a-z]+?): error:", stderrString)
		if m {
			err = fmt.Errorf("VBoxManage error: %s", stderrString)
		}
	}

	log.Printf("stdout: %s", stdoutString)
	log.Printf("stderr: %s", stderrString)

	return err
}

func (d *VBox42Driver) Verify() error {
	return nil
}

func (d *VBox42Driver) Version() (string, error) {
	var stdout bytes.Buffer

	cmd := exec.Command(d.VBoxManagePath, "--version")
	cmd.Stdout = &stdout
	if err := cmd.Run(); err != nil {
		return "", err
	}

	versionOutput := strings.TrimSpace(stdout.String())
	log.Printf("VBoxManage --version output: %s", versionOutput)

	// If the "--version" output contains vboxdrv, then this is indicative
	// of problems with the VirtualBox setup and we shouldn't really continue,
	// whether or not we can read the version.
	if strings.Contains(versionOutput, "vboxdrv") {
		return "", fmt.Errorf("VirtualBox is not properly setup: %s", versionOutput)
	}

	versionRe := regexp.MustCompile("^[.0-9]+(?:_RC[0-9]+)?")
	matches := versionRe.FindAllString(versionOutput, 1)
	if matches == nil {
		return "", fmt.Errorf("No version found: %s", versionOutput)
	}

	log.Printf("VirtualBox version: %s", matches[0])
	return matches[0], nil
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"log"
	"os"

	"github.com/ViBiOh/funds/mailjet"
	"github.com/ViBiOh/funds/notifier"
	"github.com/ViBiOh/httputils/db"
)

func main() {
	check := flag.Bool(`c`, false, `Healthcheck (check and exit)`)
	recipients := flag.String(`recipients`, ``, `Email of notifications recipients`)
	score := flag.Float64(`score`, 25.0, `Score value to notification when above`)
	hour := flag.Int(`hour`, 8, `Hour of day for sending notifications in Europe/Paris`)
	minute := flag.Int(`minute`, 0, `Minute of hour for sending notifications`)
	dbConfig := db.Flags(`db`)
	flag.Parse()

	fundsDB, err := db.GetDB(dbConfig)
	if err != nil {
		log.Printf(`Error while initializing database: %v`, err)
	} else {
		log.Print(`Database ready`)
	}

	if *check {
		if !(db.Ping(fundsDB) && mailjet.Ping()) {
			os.Exit(1)
		}
		return
	}

	log.Printf(`Notification to %s at %02d:%02d for score above %.2f`, *recipients, *hour, *minute, *score)

	if err := notifier.Init(fundsDB); err != nil {
		log.Printf(`Error while initializing notifier: %v`, err)
	}
	notifier.Start(*recipients, *score, *hour, *minute)
}
<commit_msg>Removing setted flags<commit_after>package main

import (
	"flag"
	"log"
	"os"

	"github.com/ViBiOh/funds/mailjet"
	"github.com/ViBiOh/funds/notifier"
	"github.com/ViBiOh/httputils/db"
)

func main() {
	check := flag.Bool(`c`, false, `Healthcheck (check and exit)`)
	recipients := flag.String(`recipients`, ``, `Email of notifications recipients`)
	score := flag.Float64(`score`, 25.0, `Score value to notification when above`)
	hour := flag.Int(`hour`, 8, `Hour of day for sending notifications in Europe/Paris`)
	minute := flag.Int(`minute`, 0, `Minute of hour for sending notifications`)
	dbConfig := db.Flags(``)
	flag.Parse()

	fundsDB, err := db.GetDB(dbConfig)
	if err != nil {
		log.Printf(`Error while initializing database: %v`, err)
	} else {
		log.Print(`Database ready`)
	}

	if *check {
		if !(db.Ping(fundsDB) && mailjet.Ping()) {
			os.Exit(1)
		}
		return
	}

	log.Printf(`Notification to %s at %02d:%02d for score above %.2f`, *recipients, *hour, *minute, *score)

	if err := notifier.Init(fundsDB); err != nil {
		log.Printf(`Error while initializing notifier: %v`, err)
	}
	notifier.Start(*recipients, *score, *hour, *minute)
}
<|endoftext|>


<commit_before>package plugins

import (
	"fmt"
	"testing"
)

func test(t *testing.T, filename, expectedHash string) {
	hash, _, err := hashFile(filename)
	if err != nil {
		t.Fatalf("Somethings gone fucked! Error:\n%s\n", err)
	}

	if hash != expectedHash {
		t.Fatalf("\nExpected hash:  %s\nGot hash:  %s\n", expectedHash, hash)
	}
}

func TestSimple(t *testing.T) {
	test(t, "breakdance.avi", "8e245d9679d31e12")
}

func TestLarge(t *testing.T) {
	test(t, "dummy.bin", "61f7751fc2a72bfb")
}

func TestRequest(t *testing.T) {
	err := GetSubtitle("C:/Users/Rubys/Videos/Doctor Who/Season 7/doctor.who.2005.s07e01.bdrip.xvid-haggis.avi")
	if err != nil {
		fmt.Printf("ERR:\n%s\n", err)
	}
}
<commit_msg>Irrelevant test<commit_after><|endoftext|>


<commit_before>// Package puppetserver implements a provisioner for Packer that executes
// Puppet on the remote machine connecting to a Puppet master.
package puppetserver

import (
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"github.com/hashicorp/packer/common"
	"github.com/hashicorp/packer/helper/config"
	"github.com/hashicorp/packer/packer"
	"github.com/hashicorp/packer/provisioner"
	"github.com/hashicorp/packer/template/interpolate"
)

type Config struct {
	common.PackerConfig `mapstructure:",squash"`
	ctx                 interpolate.Context

	// If true, staging directory is removed after executing puppet.
	CleanStagingDir bool `mapstructure:"clean_staging_directory"`

	// A path to the client certificate
	ClientCertPath string `mapstructure:"client_cert_path"`

	// A path to a directory containing the client private keys
	ClientPrivateKeyPath string `mapstructure:"client_private_key_path"`

	// The command used to execute Puppet.
	ExecuteCommand string `mapstructure:"execute_command"`

	// Additional argument to pass when executing Puppet.
	ExtraArguments []string `mapstructure:"extra_arguments"`

	// Additional facts to set when executing Puppet
	Facter map[string]string

	// The Guest OS Type (unix or windows)
	GuestOSType string `mapstructure:"guest_os_type"`

	// If true, packer will ignore all exit-codes from a puppet run
	IgnoreExitCodes bool `mapstructure:"ignore_exit_codes"`

	// If true, `sudo` will NOT be used to execute Puppet.
	PreventSudo bool `mapstructure:"prevent_sudo"`

	// The directory that contains the puppet binary.
	// E.g. if it can't be found on the standard path.
	PuppetBinDir string `mapstructure:"puppet_bin_dir"`

	// The hostname of the Puppet node.
	PuppetNode string `mapstructure:"puppet_node"`

	// The hostname of the Puppet server.
	PuppetServer string `mapstructure:"puppet_server"`

	// The directory where files will be uploaded. Packer requires write
	// permissions in this directory.
	StagingDir string `mapstructure:"staging_dir"`

	// The directory from which the command will be executed.
	// Packer requires the directory to exist when running puppet.
	WorkingDir string `mapstructure:"working_directory"`
}

type guestOSTypeConfig struct {
	executeCommand   string
	facterVarsFmt    string
	facterVarsJoiner string
	stagingDir       string
	tempDir          string
}

// FIXME assumes both Packer host and target are same OS
var guestOSTypeConfigs = map[string]guestOSTypeConfig{
	provisioner.UnixOSType: {
		tempDir:    "/tmp",
		stagingDir: "/tmp/packer-puppet-server",
		executeCommand: "cd {{.WorkingDir}} && " +
			`{{if ne .FacterVars ""}}{{.FacterVars}} {{end}}` +
			"{{if .Sudo}}sudo -E {{end}}" +
			`{{if ne .PuppetBinDir ""}}{{.PuppetBinDir}}/{{end}}` +
			"puppet agent --onetime --no-daemonize --detailed-exitcodes " +
			"{{if .Debug}}--debug {{end}}" +
			`{{if ne .PuppetServer ""}}--server='{{.PuppetServer}}' {{end}}` +
			`{{if ne .PuppetNode ""}}--certname={{.PuppetNode}} {{end}}` +
			`{{if ne .ClientCertPath ""}}--certdir='{{.ClientCertPath}}' {{end}}` +
			`{{if ne .ClientPrivateKeyPath ""}}--privatekeydir='{{.ClientPrivateKeyPath}}' {{end}}` +
			`{{if ne .ExtraArguments ""}}{{.ExtraArguments}} {{end}}`,
		facterVarsFmt:    "FACTER_%s='%s'",
		facterVarsJoiner: " ",
	},
	provisioner.WindowsOSType: {
		tempDir:    filepath.ToSlash(os.Getenv("TEMP")),
		stagingDir: filepath.ToSlash(os.Getenv("SYSTEMROOT")) + "/Temp/packer-puppet-server",
		executeCommand: "cd {{.WorkingDir}} && " +
			`{{if ne .FacterVars ""}}{{.FacterVars}} && {{end}}` +
			`{{if ne .PuppetBinDir ""}}{{.PuppetBinDir}}/{{end}}` +
			"puppet agent --onetime --no-daemonize --detailed-exitcodes " +
			"{{if .Debug}}--debug {{end}}" +
			`{{if ne .PuppetServer ""}}--server='{{.PuppetServer}}' {{end}}` +
			`{{if ne .PuppetNode ""}}--certname={{.PuppetNode}} {{end}}` +
			`{{if ne .ClientCertPath ""}}--certdir='{{.ClientCertPath}}' {{end}}` +
			`{{if ne .ClientPrivateKeyPath ""}}--privatekeydir='{{.ClientPrivateKeyPath}}' {{end}}` +
			`{{if ne .ExtraArguments ""}}{{.ExtraArguments}} {{end}}`,
		facterVarsFmt:    `SET "FACTER_%s=%s"`,
		facterVarsJoiner: " & ",
	},
}

type Provisioner struct {
	config            Config
	guestOSTypeConfig guestOSTypeConfig
	guestCommands     *provisioner.GuestCommands
}

type ExecuteTemplate struct {
	ClientCertPath       string
	ClientPrivateKeyPath string
	Debug                bool
	ExtraArguments       string
	FacterVars           string
	PuppetNode           string
	PuppetServer         string
	PuppetBinDir         string
	Sudo                 bool
	WorkingDir           string
}

func (p *Provisioner) Prepare(raws ...interface{}) error {
	err := config.Decode(&p.config, &config.DecodeOpts{
		Interpolate:        true,
		InterpolateContext: &p.config.ctx,
		InterpolateFilter: &interpolate.RenderFilter{
			Exclude: []string{
				"execute_command",
				"extra_arguments",
			},
		},
	}, raws...)
	if err != nil {
		return err
	}

	if p.config.GuestOSType == "" {
		p.config.GuestOSType = provisioner.DefaultOSType
	}
	p.config.GuestOSType = strings.ToLower(p.config.GuestOSType)

	var ok bool
	p.guestOSTypeConfig, ok = guestOSTypeConfigs[p.config.GuestOSType]
	if !ok {
		return fmt.Errorf("Invalid guest_os_type: \"%s\"", p.config.GuestOSType)
	}

	p.guestCommands, err = provisioner.NewGuestCommands(p.config.GuestOSType, !p.config.PreventSudo)
	if err != nil {
		return fmt.Errorf("Invalid guest_os_type: \"%s\"", p.config.GuestOSType)
	}

	if p.config.ExecuteCommand == "" {
		p.config.ExecuteCommand = p.guestOSTypeConfig.executeCommand
	}

	if p.config.StagingDir == "" {
		p.config.StagingDir = p.guestOSTypeConfig.stagingDir
	}

	if p.config.WorkingDir == "" {
		p.config.WorkingDir = p.config.StagingDir
	}

	if p.config.Facter == nil {
		p.config.Facter = make(map[string]string)
	}
	p.config.Facter["packer_build_name"] = p.config.PackerBuildName
	p.config.Facter["packer_builder_type"] = p.config.PackerBuilderType

	var errs *packer.MultiError
	if p.config.ClientCertPath != "" {
		info, err := os.Stat(p.config.ClientCertPath)
		if err != nil {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("client_cert_dir is invalid: %s", err))
		} else if !info.IsDir() {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("client_cert_dir must point to a directory"))
		}
	}

	if p.config.ClientPrivateKeyPath != "" {
		info, err := os.Stat(p.config.ClientPrivateKeyPath)
		if err != nil {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("client_private_key_dir is invalid: %s", err))
		} else if !info.IsDir() {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("client_private_key_dir must point to a directory"))
		}
	}

	if errs != nil && len(errs.Errors) > 0 {
		return errs
	}

	return nil
}

func (p *Provisioner) Provision(ui packer.Ui, comm packer.Communicator) error {
	ui.Say("Provisioning with Puppet...")
	ui.Message("Creating Puppet staging directory...")
	if err := p.createDir(ui, comm, p.config.StagingDir); err != nil {
		return fmt.Errorf("Error creating staging directory: %s", err)
	}

	// Upload client cert dir if set
	remoteClientCertPath := ""
	if p.config.ClientCertPath != "" {
		ui.Message(fmt.Sprintf(
			"Uploading client cert from: %s", p.config.ClientCertPath))
		remoteClientCertPath = fmt.Sprintf("%s/certs", p.config.StagingDir)
		err := p.uploadDirectory(ui, comm, remoteClientCertPath, p.config.ClientCertPath)
		if err != nil {
			return fmt.Errorf("Error uploading client cert: %s", err)
		}
	}

	// Upload client cert dir if set
	remoteClientPrivateKeyPath := ""
	if p.config.ClientPrivateKeyPath != "" {
		ui.Message(fmt.Sprintf(
			"Uploading client private keys from: %s", p.config.ClientPrivateKeyPath))
		remoteClientPrivateKeyPath = fmt.Sprintf("%s/private_keys", p.config.StagingDir)
		err := p.uploadDirectory(ui, comm, remoteClientPrivateKeyPath, p.config.ClientPrivateKeyPath)
		if err != nil {
			return fmt.Errorf("Error uploading client private keys: %s", err)
		}
	}

	// Compile the facter variables
	facterVars := make([]string, 0, len(p.config.Facter))
	for k, v := range p.config.Facter {
		facterVars = append(facterVars, fmt.Sprintf(p.guestOSTypeConfig.facterVarsFmt, k, v))
	}

	data := ExecuteTemplate{
		ClientCertPath:       remoteClientCertPath,
		ClientPrivateKeyPath: remoteClientPrivateKeyPath,
		ExtraArguments:       "",
		FacterVars:           strings.Join(facterVars, p.guestOSTypeConfig.facterVarsJoiner),
		PuppetNode:           p.config.PuppetNode,
		PuppetServer:         p.config.PuppetServer,
		PuppetBinDir:         p.config.PuppetBinDir,
		Sudo:                 !p.config.PreventSudo,
		WorkingDir:           p.config.WorkingDir,
	}

	p.config.ctx.Data = &data
	_ExtraArguments, err := interpolate.Render(strings.Join(p.config.ExtraArguments, " "), &p.config.ctx)
	if err != nil {
		return err
	}
	data.ExtraArguments = _ExtraArguments

	command, err := interpolate.Render(p.config.ExecuteCommand, &p.config.ctx)
	if err != nil {
		return err
	}

	cmd := &packer.RemoteCmd{
		Command: command,
	}

	ui.Message(fmt.Sprintf("Running Puppet: %s", command))
	if err := cmd.StartWithUi(comm, ui); err != nil {
		return err
	}

	if cmd.ExitStatus != 0 && cmd.ExitStatus != 2 && !p.config.IgnoreExitCodes {
		return fmt.Errorf("Puppet exited with a non-zero exit status: %d", cmd.ExitStatus)
	}

	if p.config.CleanStagingDir {
		if err := p.removeDir(ui, comm, p.config.StagingDir); err != nil {
			return fmt.Errorf("Error removing staging directory: %s", err)
		}
	}

	return nil
}

func (p *Provisioner) Cancel() {
	// Just hard quit. It isn't a big deal if what we're doing keeps
	// running on the other side.
	os.Exit(0)
}

func (p *Provisioner) createDir(ui packer.Ui, comm packer.Communicator, dir string) error {
	ui.Message(fmt.Sprintf("Creating directory: %s", dir))

	cmd := &packer.RemoteCmd{Command: p.guestCommands.CreateDir(dir)}
	if err := cmd.StartWithUi(comm, ui); err != nil {
		return err
	}
	if cmd.ExitStatus != 0 {
		return fmt.Errorf("Non-zero exit status. See output above for more info.")
	}

	// Chmod the directory to 0777 just so that we can access it as our user
	cmd = &packer.RemoteCmd{Command: p.guestCommands.Chmod(dir, "0777")}
	if err := cmd.StartWithUi(comm, ui); err != nil {
		return err
	}
	if cmd.ExitStatus != 0 {
		return fmt.Errorf("Non-zero exit status. See output above for more info.")
	}

	return nil
}

func (p *Provisioner) uploadDirectory(ui packer.Ui, comm packer.Communicator, dst string, src string) error {
	if err := p.createDir(ui, comm, dst); err != nil {
		return err
	}

	// Make sure there is a trailing "/" so that the directory isn't
	// created on the other side.
	if src[len(src)-1] != '/' {
		src = src + "/"
	}

	return comm.UploadDir(dst, src, nil)
}
<commit_msg>add missing removeDir() to allow StagingDir clean-up post-run<commit_after>// Package puppetserver implements a provisioner for Packer that executes
// Puppet on the remote machine connecting to a Puppet master.
package puppetserver

import (
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"github.com/hashicorp/packer/common"
	"github.com/hashicorp/packer/helper/config"
	"github.com/hashicorp/packer/packer"
	"github.com/hashicorp/packer/provisioner"
	"github.com/hashicorp/packer/template/interpolate"
)

type Config struct {
	common.PackerConfig `mapstructure:",squash"`
	ctx                 interpolate.Context

	// If true, staging directory is removed after executing puppet.
	CleanStagingDir bool `mapstructure:"clean_staging_directory"`

	// A path to the client certificate
	ClientCertPath string `mapstructure:"client_cert_path"`

	// A path to a directory containing the client private keys
	ClientPrivateKeyPath string `mapstructure:"client_private_key_path"`

	// The command used to execute Puppet.
	ExecuteCommand string `mapstructure:"execute_command"`

	// Additional argument to pass when executing Puppet.
	ExtraArguments []string `mapstructure:"extra_arguments"`

	// Additional facts to set when executing Puppet
	Facter map[string]string

	// The Guest OS Type (unix or windows)
	GuestOSType string `mapstructure:"guest_os_type"`

	// If true, packer will ignore all exit-codes from a puppet run
	IgnoreExitCodes bool `mapstructure:"ignore_exit_codes"`

	// If true, `sudo` will NOT be used to execute Puppet.
	PreventSudo bool `mapstructure:"prevent_sudo"`

	// The directory that contains the puppet binary.
	// E.g. if it can't be found on the standard path.
	PuppetBinDir string `mapstructure:"puppet_bin_dir"`

	// The hostname of the Puppet node.
	PuppetNode string `mapstructure:"puppet_node"`

	// The hostname of the Puppet server.
	PuppetServer string `mapstructure:"puppet_server"`

	// The directory where files will be uploaded. Packer requires write
	// permissions in this directory.
	StagingDir string `mapstructure:"staging_dir"`

	// The directory from which the command will be executed.
	// Packer requires the directory to exist when running puppet.
	WorkingDir string `mapstructure:"working_directory"`
}

type guestOSTypeConfig struct {
	executeCommand   string
	facterVarsFmt    string
	facterVarsJoiner string
	stagingDir       string
	tempDir          string
}

// FIXME assumes both Packer host and target are same OS
var guestOSTypeConfigs = map[string]guestOSTypeConfig{
	provisioner.UnixOSType: {
		tempDir:    "/tmp",
		stagingDir: "/tmp/packer-puppet-server",
		executeCommand: "cd {{.WorkingDir}} && " +
			`{{if ne .FacterVars ""}}{{.FacterVars}} {{end}}` +
			"{{if .Sudo}}sudo -E {{end}}" +
			`{{if ne .PuppetBinDir ""}}{{.PuppetBinDir}}/{{end}}` +
			"puppet agent --onetime --no-daemonize --detailed-exitcodes " +
			"{{if .Debug}}--debug {{end}}" +
			`{{if ne .PuppetServer ""}}--server='{{.PuppetServer}}' {{end}}` +
			`{{if ne .PuppetNode ""}}--certname={{.PuppetNode}} {{end}}` +
			`{{if ne .ClientCertPath ""}}--certdir='{{.ClientCertPath}}' {{end}}` +
			`{{if ne .ClientPrivateKeyPath ""}}--privatekeydir='{{.ClientPrivateKeyPath}}' {{end}}` +
			`{{if ne .ExtraArguments ""}}{{.ExtraArguments}} {{end}}`,
		facterVarsFmt:    "FACTER_%s='%s'",
		facterVarsJoiner: " ",
	},
	provisioner.WindowsOSType: {
		tempDir:    filepath.ToSlash(os.Getenv("TEMP")),
		stagingDir: filepath.ToSlash(os.Getenv("SYSTEMROOT")) + "/Temp/packer-puppet-server",
		executeCommand: "cd {{.WorkingDir}} && " +
			`{{if ne .FacterVars ""}}{{.FacterVars}} && {{end}}` +
			`{{if ne .PuppetBinDir ""}}{{.PuppetBinDir}}/{{end}}` +
			"puppet agent --onetime --no-daemonize --detailed-exitcodes " +
			"{{if .Debug}}--debug {{end}}" +
			`{{if ne .PuppetServer ""}}--server='{{.PuppetServer}}' {{end}}` +
			`{{if ne .PuppetNode ""}}--certname={{.PuppetNode}} {{end}}` +
			`{{if ne .ClientCertPath ""}}--certdir='{{.ClientCertPath}}' {{end}}` +
			`{{if ne .ClientPrivateKeyPath ""}}--privatekeydir='{{.ClientPrivateKeyPath}}' {{end}}` +
			`{{if ne .ExtraArguments ""}}{{.ExtraArguments}} {{end}}`,
		facterVarsFmt:    `SET "FACTER_%s=%s"`,
		facterVarsJoiner: " & ",
	},
}

type Provisioner struct {
	config            Config
	guestOSTypeConfig guestOSTypeConfig
	guestCommands     *provisioner.GuestCommands
}

type ExecuteTemplate struct {
	ClientCertPath       string
	ClientPrivateKeyPath string
	Debug                bool
	ExtraArguments       string
	FacterVars           string
	PuppetNode           string
	PuppetServer         string
	PuppetBinDir         string
	Sudo                 bool
	WorkingDir           string
}

func (p *Provisioner) Prepare(raws ...interface{}) error {
	err := config.Decode(&p.config, &config.DecodeOpts{
		Interpolate:        true,
		InterpolateContext: &p.config.ctx,
		InterpolateFilter: &interpolate.RenderFilter{
			Exclude: []string{
				"execute_command",
				"extra_arguments",
			},
		},
	}, raws...)
	if err != nil {
		return err
	}

	if p.config.GuestOSType == "" {
		p.config.GuestOSType = provisioner.DefaultOSType
	}
	p.config.GuestOSType = strings.ToLower(p.config.GuestOSType)

	var ok bool
	p.guestOSTypeConfig, ok = guestOSTypeConfigs[p.config.GuestOSType]
	if !ok {
		return fmt.Errorf("Invalid guest_os_type: \"%s\"", p.config.GuestOSType)
	}

	p.guestCommands, err = provisioner.NewGuestCommands(p.config.GuestOSType, !p.config.PreventSudo)
	if err != nil {
		return fmt.Errorf("Invalid guest_os_type: \"%s\"", p.config.GuestOSType)
	}

	if p.config.ExecuteCommand == "" {
		p.config.ExecuteCommand = p.guestOSTypeConfig.executeCommand
	}

	if p.config.StagingDir == "" {
		p.config.StagingDir = p.guestOSTypeConfig.stagingDir
	}

	if p.config.WorkingDir == "" {
		p.config.WorkingDir = p.config.StagingDir
	}

	if p.config.Facter == nil {
		p.config.Facter = make(map[string]string)
	}
	p.config.Facter["packer_build_name"] = p.config.PackerBuildName
	p.config.Facter["packer_builder_type"] = p.config.PackerBuilderType

	var errs *packer.MultiError
	if p.config.ClientCertPath != "" {
		info, err := os.Stat(p.config.ClientCertPath)
		if err != nil {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("client_cert_dir is invalid: %s", err))
		} else if !info.IsDir() {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("client_cert_dir must point to a directory"))
		}
	}

	if p.config.ClientPrivateKeyPath != "" {
		info, err := os.Stat(p.config.ClientPrivateKeyPath)
		if err != nil {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("client_private_key_dir is invalid: %s", err))
		} else if !info.IsDir() {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("client_private_key_dir must point to a directory"))
		}
	}

	if errs != nil && len(errs.Errors) > 0 {
		return errs
	}

	return nil
}

func (p *Provisioner) Provision(ui packer.Ui, comm packer.Communicator) error {
	ui.Say("Provisioning with Puppet...")
	ui.Message("Creating Puppet staging directory...")
	if err := p.createDir(ui, comm, p.config.StagingDir); err != nil {
		return fmt.Errorf("Error creating staging directory: %s", err)
	}

	// Upload client cert dir if set
	remoteClientCertPath := ""
	if p.config.ClientCertPath != "" {
		ui.Message(fmt.Sprintf(
			"Uploading client cert from: %s", p.config.ClientCertPath))
		remoteClientCertPath = fmt.Sprintf("%s/certs", p.config.StagingDir)
		err := p.uploadDirectory(ui, comm, remoteClientCertPath, p.config.ClientCertPath)
		if err != nil {
			return fmt.Errorf("Error uploading client cert: %s", err)
		}
	}

	// Upload client cert dir if set
	remoteClientPrivateKeyPath := ""
	if p.config.ClientPrivateKeyPath != "" {
		ui.Message(fmt.Sprintf(
			"Uploading client private keys from: %s", p.config.ClientPrivateKeyPath))
		remoteClientPrivateKeyPath = fmt.Sprintf("%s/private_keys", p.config.StagingDir)
		err := p.uploadDirectory(ui, comm, remoteClientPrivateKeyPath, p.config.ClientPrivateKeyPath)
		if err != nil {
			return fmt.Errorf("Error uploading client private keys: %s", err)
		}
	}

	// Compile the facter variables
	facterVars := make([]string, 0, len(p.config.Facter))
	for k, v := range p.config.Facter {
		facterVars = append(facterVars, fmt.Sprintf(p.guestOSTypeConfig.facterVarsFmt, k, v))
	}

	data := ExecuteTemplate{
		ClientCertPath:       remoteClientCertPath,
		ClientPrivateKeyPath: remoteClientPrivateKeyPath,
		ExtraArguments:       "",
		FacterVars:           strings.Join(facterVars, p.guestOSTypeConfig.facterVarsJoiner),
		PuppetNode:           p.config.PuppetNode,
		PuppetServer:         p.config.PuppetServer,
		PuppetBinDir:         p.config.PuppetBinDir,
		Sudo:                 !p.config.PreventSudo,
		WorkingDir:           p.config.WorkingDir,
	}

	p.config.ctx.Data = &data
	_ExtraArguments, err := interpolate.Render(strings.Join(p.config.ExtraArguments, " "), &p.config.ctx)
	if err != nil {
		return err
	}
	data.ExtraArguments = _ExtraArguments

	command, err := interpolate.Render(p.config.ExecuteCommand, &p.config.ctx)
	if err != nil {
		return err
	}

	cmd := &packer.RemoteCmd{
		Command: command,
	}

	ui.Message(fmt.Sprintf("Running Puppet: %s", command))
	if err := cmd.StartWithUi(comm, ui); err != nil {
		return err
	}

	if cmd.ExitStatus != 0 && cmd.ExitStatus != 2 && !p.config.IgnoreExitCodes {
		return fmt.Errorf("Puppet exited with a non-zero exit status: %d", cmd.ExitStatus)
	}

	if p.config.CleanStagingDir {
		if err := p.removeDir(ui, comm, p.config.StagingDir); err != nil {
			return fmt.Errorf("Error removing staging directory: %s", err)
		}
	}

	return nil
}

func (p *Provisioner) Cancel() {
	// Just hard quit. It isn't a big deal if what we're doing keeps
	// running on the other side.
	os.Exit(0)
}

func (p *Provisioner) createDir(ui packer.Ui, comm packer.Communicator, dir string) error {
	ui.Message(fmt.Sprintf("Creating directory: %s", dir))

	cmd := &packer.RemoteCmd{Command: p.guestCommands.CreateDir(dir)}
	if err := cmd.StartWithUi(comm, ui); err != nil {
		return err
	}
	if cmd.ExitStatus != 0 {
		return fmt.Errorf("Non-zero exit status. See output above for more info.")
	}

	// Chmod the directory to 0777 just so that we can access it as our user
	cmd = &packer.RemoteCmd{Command: p.guestCommands.Chmod(dir, "0777")}
	if err := cmd.StartWithUi(comm, ui); err != nil {
		return err
	}
	if cmd.ExitStatus != 0 {
		return fmt.Errorf("Non-zero exit status. See output above for more info.")
	}

	return nil
}

func (p *Provisioner) removeDir(ui packer.Ui, comm packer.Communicator, dir string) error {
	cmd := &packer.RemoteCmd{
		Command: fmt.Sprintf("rm -fr '%s'", dir),
	}

	if err := cmd.StartWithUi(comm, ui); err != nil {
		return err
	}

	if cmd.ExitStatus != 0 {
		return fmt.Errorf("Non-zero exit status.")
	}

	return nil
}

func (p *Provisioner) uploadDirectory(ui packer.Ui, comm packer.Communicator, dst string, src string) error {
	if err := p.createDir(ui, comm, dst); err != nil {
		return err
	}

	// Make sure there is a trailing "/" so that the directory isn't
	// created on the other side.
	if src[len(src)-1] != '/' {
		src = src + "/"
	}

	return comm.UploadDir(dst, src, nil)
}
<|endoftext|>


<commit_before>/**********************************************************\
|                                                          |
|                          hprose                          |
|                                                          |
| Official WebSite: http://www.hprose.com/                 |
|                   http://www.hprose.org/                 |
|                                                          |
\**********************************************************/
/**********************************************************\
 *                                                        *
 * hprose/raw_reader_test.go                              *
 *                                                        *
 * hprose RawReader Test for Go.                          *
 *                                                        *
 * LastModified: Jan 31, 2014                             *
 * Author: Ma Bingyao <andot@hprose.com>                  *
 *                                                        *
\**********************************************************/

package io

import (
	"bytes"
	"testing"
	"time"
)

func TestRawReaderTag(t *testing.T) {
	b := new(bytes.Buffer)
	writer := NewWriter(b, true)
	writer.Serialize(nil)
	rawReader := NewRawReader(b.Bytes())
	raw, err := rawReader.ReadRaw()
	if err != nil {
		t.Error(err.Error())
	}
	if string(raw) != "n" {
		t.Error("read tag error: ", string(raw))
	}
}

func TestRawReaderNumber(t *testing.T) {
	b := new(bytes.Buffer)
	writer := NewWriter(b, true)
	writer.Serialize(123)
	rawReader := NewRawReader(b.Bytes())
	raw, err := rawReader.ReadRaw()
	if err != nil {
		t.Error(err.Error())
	}
	if string(raw) != "i123;" {
		t.Error("read integer error: ", string(raw))
	}
}

func TestRawReaderChar(t *testing.T) {
	b := new(bytes.Buffer)
	writer := NewWriter(b, true)
	writer.Serialize("我")
	rawReader := NewRawReader(b.Bytes())
	raw, err := rawReader.ReadRaw()
	if err != nil {
		t.Error(err.Error())
	}
	if string(raw) != `u我` {
		t.Error("read char error: ", string(raw))
	}
}

func TestRawReaderBytes(t *testing.T) {
	b := new(bytes.Buffer)
	writer := NewWriter(b, true)
	writer.Serialize([]byte("hello world!"))
	rawReader := NewRawReader(b.Bytes())
	raw, err := rawReader.ReadRaw()
	if err != nil {
		t.Error(err.Error())
	}
	if string(raw) != `b12"hello world!"` {
		t.Error("read bytes error: ", string(raw))
	}
}

func TestRawReaderString(t *testing.T) {
	b := new(bytes.Buffer)
	writer := NewWriter(b, true)
	writer.Serialize("我爱你!")
	rawReader := NewRawReader(b.Bytes())
	raw, err := rawReader.ReadRaw()
	if err != nil {
		t.Error(err.Error())
	}
	if string(raw) != `s4"我爱你!"` {
		t.Error("read string error: ", string(raw))
	}
}

func TestRawReaderGuid(t *testing.T) {
	rawReader := NewRawReader([]byte("g{AFA7F4B1-A64D-46FA-886F-ED7FBCE569B6}"))
	raw, err := rawReader.ReadRaw()
	if err != nil {
		t.Error(err.Error())
	}
	if string(raw) != "g{AFA7F4B1-A64D-46FA-886F-ED7FBCE569B6}" {
		t.Error("read guid error: ", string(raw))
	}
}

func TestRawReaderDateTime(t *testing.T) {
	b := new(bytes.Buffer)
	writer := NewWriter(b, true)
	writer.Serialize(time.Date(2008, 12, 11, 23, 12, 21, 123433453, time.UTC))
	rawReader := NewRawReader(b.Bytes())
	raw, err := rawReader.ReadRaw()
	if err != nil {
		t.Error(err.Error())
	}
	if string(raw) != "D20081211T231221.123433453Z" {
		t.Error("read datetime error: ", string(raw))
	}
}

func TestRawReaderComplex(t *testing.T) {
	b := new(bytes.Buffer)
	writer := NewWriter(b, true)
	writer.Serialize([]int{0, 1, 2, 3, 4, 5, 6, 7, 8, 9})
	rawReader := NewRawReader(b.Bytes())
	raw, err := rawReader.ReadRaw()
	if err != nil {
		t.Error(err.Error())
	}
	if string(raw) != "a10{0123456789}" {
		t.Error("read complex error: ", string(raw))
	}
}

func TestRawReaderStruct(t *testing.T) {
	type Person struct {
		Name string
		Age  int
	}
	b := new(bytes.Buffer)
	writer := NewWriter(b, true)
	writer.Serialize(Person{"Tom", 18})
	rawReader := NewRawReader(b.Bytes())
	raw, err := rawReader.ReadRaw()
	if err != nil {
		t.Error(err.Error())
	}
	if string(raw) != `c6"Person"2{s4"name"s3"age"}o0{s3"Tom"i18;}` {
		t.Error("read object error: ", string(raw))
	}
}

func BenchmarkRawReaderReadUTF8StringEmpty(b *testing.B) {
	buf := new(bytes.Buffer)
	writer := NewWriter(buf, true)
	s := "我爱你🇨🇳"
	for i := 0; i < 20; i++ {
		s += s
	}
	writer.Serialize(s)
	data := buf.Bytes()
	test := func(_ []byte) {}
	for i := 0; i < b.N; i++ {
		test(data)
	}
}

func BenchmarkRawReaderReadUTF8String(b *testing.B) {
	buf := new(bytes.Buffer)
	writer := NewWriter(buf, true)
	s := "我爱你🇨🇳"
	for i := 0; i < 10; i++ {
		s += s
	}
	writer.Serialize(s)
	data := buf.Bytes()
	for i := 0; i < b.N; i++ {
		rawReader := NewRawReader(data)
		rawReader.ReadRaw()
	}
}
<commit_msg>Update comment of raw_reader_test<commit_after>/**********************************************************\
|                                                          |
|                          hprose                          |
|                                                          |
| Official WebSite: http://www.hprose.com/                 |
|                   http://www.hprose.org/                 |
|                                                          |
\**********************************************************/
/**********************************************************\
 *                                                        *
 * io/raw_reader_test.go                                  *
 *                                                        *
 * hprose RawReader Test for Go.                          *
 *                                                        *
 * LastModified: Aug 30, 2016                             *
 * Author: Ma Bingyao <andot@hprose.com>                  *
 *                                                        *
\**********************************************************/

package io

import (
	"bytes"
	"testing"
	"time"
)

func TestRawReaderTag(t *testing.T) {
	b := new(bytes.Buffer)
	writer := NewWriter(b, true)
	writer.Serialize(nil)
	rawReader := NewRawReader(b.Bytes())
	raw, err := rawReader.ReadRaw()
	if err != nil {
		t.Error(err.Error())
	}
	if string(raw) != "n" {
		t.Error("read tag error: ", string(raw))
	}
}

func TestRawReaderNumber(t *testing.T) {
	b := new(bytes.Buffer)
	writer := NewWriter(b, true)
	writer.Serialize(123)
	rawReader := NewRawReader(b.Bytes())
	raw, err := rawReader.ReadRaw()
	if err != nil {
		t.Error(err.Error())
	}
	if string(raw) != "i123;" {
		t.Error("read integer error: ", string(raw))
	}
}

func TestRawReaderChar(t *testing.T) {
	b := new(bytes.Buffer)
	writer := NewWriter(b, true)
	writer.Serialize("我")
	rawReader := NewRawReader(b.Bytes())
	raw, err := rawReader.ReadRaw()
	if err != nil {
		t.Error(err.Error())
	}
	if string(raw) != `u我` {
		t.Error("read char error: ", string(raw))
	}
}

func TestRawReaderBytes(t *testing.T) {
	b := new(bytes.Buffer)
	writer := NewWriter(b, true)
	writer.Serialize([]byte("hello world!"))
	rawReader := NewRawReader(b.Bytes())
	raw, err := rawReader.ReadRaw()
	if err != nil {
		t.Error(err.Error())
	}
	if string(raw) != `b12"hello world!"` {
		t.Error("read bytes error: ", string(raw))
	}
}

func TestRawReaderString(t *testing.T) {
	b := new(bytes.Buffer)
	writer := NewWriter(b, true)
	writer.Serialize("我爱你!")
	rawReader := NewRawReader(b.Bytes())
	raw, err := rawReader.ReadRaw()
	if err != nil {
		t.Error(err.Error())
	}
	if string(raw) != `s4"我爱你!"` {
		t.Error("read string error: ", string(raw))
	}
}

func TestRawReaderGuid(t *testing.T) {
	rawReader := NewRawReader([]byte("g{AFA7F4B1-A64D-46FA-886F-ED7FBCE569B6}"))
	raw, err := rawReader.ReadRaw()
	if err != nil {
		t.Error(err.Error())
	}
	if string(raw) != "g{AFA7F4B1-A64D-46FA-886F-ED7FBCE569B6}" {
		t.Error("read guid error: ", string(raw))
	}
}

func TestRawReaderDateTime(t *testing.T) {
	b := new(bytes.Buffer)
	writer := NewWriter(b, true)
	writer.Serialize(time.Date(2008, 12, 11, 23, 12, 21, 123433453, time.UTC))
	rawReader := NewRawReader(b.Bytes())
	raw, err := rawReader.ReadRaw()
	if err != nil {
		t.Error(err.Error())
	}
	if string(raw) != "D20081211T231221.123433453Z" {
		t.Error("read datetime error: ", string(raw))
	}
}

func TestRawReaderComplex(t *testing.T) {
	b := new(bytes.Buffer)
	writer := NewWriter(b, true)
	writer.Serialize([]int{0, 1, 2, 3, 4, 5, 6, 7, 8, 9})
	rawReader := NewRawReader(b.Bytes())
	raw, err := rawReader.ReadRaw()
	if err != nil {
		t.Error(err.Error())
	}
	if string(raw) != "a10{0123456789}" {
		t.Error("read complex error: ", string(raw))
	}
}

func TestRawReaderStruct(t *testing.T) {
	type Person struct {
		Name string
		Age  int
	}
	b := new(bytes.Buffer)
	writer := NewWriter(b, true)
	writer.Serialize(Person{"Tom", 18})
	rawReader := NewRawReader(b.Bytes())
	raw, err := rawReader.ReadRaw()
	if err != nil {
		t.Error(err.Error())
	}
	if string(raw) != `c6"Person"2{s4"name"s3"age"}o0{s3"Tom"i18;}` {
		t.Error("read object error: ", string(raw))
	}
}

func BenchmarkRawReaderReadUTF8StringEmpty(b *testing.B) {
	buf := new(bytes.Buffer)
	writer := NewWriter(buf, true)
	s := "我爱你🇨🇳"
	for i := 0; i < 20; i++ {
		s += s
	}
	writer.Serialize(s)
	data := buf.Bytes()
	test := func(_ []byte) {}
	for i := 0; i < b.N; i++ {
		test(data)
	}
}

func BenchmarkRawReaderReadUTF8String(b *testing.B) {
	buf := new(bytes.Buffer)
	writer := NewWriter(buf, true)
	s := "我爱你🇨🇳"
	for i := 0; i < 10; i++ {
		s += s
	}
	writer.Serialize(s)
	data := buf.Bytes()
	for i := 0; i < b.N; i++ {
		rawReader := NewRawReader(data)
		rawReader.ReadRaw()
	}
}
<|endoftext|>


<commit_before>package loop // import "fknsrs.biz/p/ottoext/loop"

import (
	"sync"
	"sync/atomic"

	"github.com/robertkrimen/otto"
)

type Task interface {
	SetID(id int64)
	GetID() int64
	Execute(vm *otto.Otto, l *Loop) error
	Cancel()
}

type Loop struct {
	vm     *otto.Otto
	id     int64
	lock   sync.RWMutex
	tasks  map[int64]Task
	ready  chan Task
	closed bool
}

func New(vm *otto.Otto) *Loop {
	return &Loop{
		vm:    vm,
		tasks: make(map[int64]Task),
		ready: make(chan Task),
	}
}

func (l *Loop) Add(t Task) {
	l.lock.Lock()
	t.SetID(atomic.AddInt64(&l.id, 1))
	l.tasks[t.GetID()] = t
	l.lock.Unlock()
}

func (l *Loop) Remove(t Task) {
	l.remove(t)
	go l.Ready(nil)
}

func (l *Loop) remove(t Task) {
	l.removeByID(t.GetID())
}

func (l *Loop) removeByID(id int64) {
	l.lock.Lock()
	delete(l.tasks, id)
	l.lock.Unlock()
}

func (l *Loop) Ready(t Task) {
	if l.closed {
		return
	}

	l.ready <- t
}

func (l *Loop) EvalAndRun(s interface{}) error {
	if err := l.Eval(s); err != nil {
		return err
	}

	return l.Run()
}

func (l *Loop) Eval(s interface{}) error {
	if _, err := l.vm.Run(s); err != nil {
		return err
	}

	return nil
}

func (l *Loop) Run() error {
	i := 0

	for t := range l.ready {
		i++

		if t != nil {
			id := t.GetID()

			if err := t.Execute(l.vm, l); err != nil {
				l.lock.RLock()
				for _, t := range l.tasks {
					t.Cancel()
				}
				l.lock.RUnlock()

				return err
			}

			l.removeByID(id)
		}

		l.lock.RLock()
		if len(l.tasks) == 0 {
			close(l.ready)
			l.closed = true
		}
		l.lock.RUnlock()
	}

	return nil
}

func (l *Loop) Step() (error, bool) {
	var tasks []Task

outer:
	for {
		select {
		case t := <-l.ready:
			tasks = append(tasks, t)

			l.remove(t)
		default:
			break outer
		}
	}

	for _, t := range tasks {
		if err := t.Execute(l.vm, l); err != nil {
			l.lock.RLock()
			for _, t := range l.tasks {
				t.Cancel()
			}
			l.lock.RUnlock()

			return err, false
		}
	}

	return nil, len(l.tasks) == 0
}
<commit_msg>ensure we drain all pending tasks at loop shutdown time<commit_after>package loop // import "fknsrs.biz/p/ottoext/loop"

import (
	"fmt"
	"sync"
	"sync/atomic"

	"github.com/robertkrimen/otto"
)

func formatTask(t Task) string {
	if t == nil {
		return "<nil>"
	}

	return fmt.Sprintf("<%T> %d", t, t.GetID())
}

type Task interface {
	SetID(id int64)
	GetID() int64
	Execute(vm *otto.Otto, l *Loop) error
	Cancel()
}

type Loop struct {
	vm     *otto.Otto
	id     int64
	lock   sync.RWMutex
	tasks  map[int64]Task
	ready  chan Task
	closed bool
}

func New(vm *otto.Otto) *Loop {
	return NewWithBacklog(vm, 0)
}

func NewWithBacklog(vm *otto.Otto, backlog int) *Loop {
	return &Loop{
		vm:    vm,
		tasks: make(map[int64]Task),
		ready: make(chan Task, backlog),
	}
}

func (l *Loop) Add(t Task) {
	l.lock.Lock()
	t.SetID(atomic.AddInt64(&l.id, 1))
	l.tasks[t.GetID()] = t
	l.lock.Unlock()
}

func (l *Loop) Remove(t Task) {
	l.remove(t)
	go l.Ready(nil)
}

func (l *Loop) remove(t Task) {
	l.removeByID(t.GetID())
}

func (l *Loop) removeByID(id int64) {
	l.lock.Lock()
	delete(l.tasks, id)
	l.lock.Unlock()
}

func (l *Loop) Ready(t Task) {
	if l.closed {
		return
	}

	l.ready <- t
}

func (l *Loop) EvalAndRun(s interface{}) error {
	if err := l.Eval(s); err != nil {
		return err
	}

	return l.Run()
}

func (l *Loop) Eval(s interface{}) error {
	if _, err := l.vm.Run(s); err != nil {
		return err
	}

	return nil
}

func (l *Loop) processTask(t Task) error {
	id := t.GetID()

	if err := t.Execute(l.vm, l); err != nil {
		l.lock.RLock()
		for _, t := range l.tasks {
			t.Cancel()
		}
		l.lock.RUnlock()

		return err
	}

	l.removeByID(id)

	return nil
}

func (l *Loop) Run() error {
	for {
		t := <-l.ready

		if t != nil {
			if err := l.processTask(t); err != nil {
				return err
			}
		}

		l.lock.Lock()
		if len(l.tasks) == 0 {
			// prevent any more tasks entering the ready channel
			l.closed = true

			l.lock.Unlock()

			break
		}
		l.lock.Unlock()
	}

	// drain ready channel of any existing tasks
outer:
	for {
		select {
		case t := <-l.ready:
			if err := l.processTask(t); err != nil {
				return err
			}
		default:
			break outer
		}
	}

	close(l.ready)

	return nil
}
<|endoftext|>


<commit_before>package transfer

import (
	"bufio"
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os/exec"
	"regexp"
	"strings"

	"github.com/github/git-lfs/localstorage"
	"github.com/github/git-lfs/tools"

	"github.com/github/git-lfs/api"
	"github.com/github/git-lfs/subprocess"
	"github.com/rubyist/tracerx"

	"github.com/github/git-lfs/config"
)

// Adapter for custom transfer via external process
type customAdapter struct {
	*adapterBase
	path                string
	args                string
	concurrent          bool
	originalConcurrency int
}

// Struct to capture stderr and write to trace
type traceWriter struct {
	buf bytes.Buffer
}

func (t *traceWriter) Write(b []byte) (int, error) {
	n, err := t.buf.Write(b)
	t.Flush()
	return n, err
}
func (t *traceWriter) Flush() {
	var err error
	for err == nil {
		var s string
		s, err = t.buf.ReadString('\n')
		if len(s) > 0 {
			tracerx.Printf("xfer_custom_stderr: %v", strings.TrimSpace(s))
		}
	}
}

type customAdapterWorkerContext struct {
	cmd         *exec.Cmd
	stdout      io.ReadCloser
	bufferedOut *bufio.Reader
	stdin       io.WriteCloser
	errTracer   *traceWriter
}

type customAdapterInitRequest struct {
	Id                  string `json:"id"`
	Operation           string `json:"operation"`
	Concurrent          bool   `json:"concurrent"`
	ConcurrentTransfers int    `json:"concurrenttransfers"`
}

func NewCustomAdapterInitRequest(op string, concurrent bool, concurrentTransfers int) *customAdapterInitRequest {
	return &customAdapterInitRequest{"init", op, concurrent, concurrentTransfers}
}

type customAdapterTransferRequest struct { // common between upload/download
	Id     string            `json:"id"`
	Oid    string            `json:"oid"`
	Size   int64             `json:"size"`
	Path   string            `json:"path,omitempty"`
	Action *api.LinkRelation `json:"action"`
}

func NewCustomAdapterUploadRequest(oid string, size int64, path string, action *api.LinkRelation) *customAdapterTransferRequest {
	return &customAdapterTransferRequest{"upload", oid, size, path, action}
}
func NewCustomAdapterDownloadRequest(oid string, size int64, action *api.LinkRelation) *customAdapterTransferRequest {
	return &customAdapterTransferRequest{"download", oid, size, "", action}
}

type customAdapterTerminateRequest struct {
	MessageType string `json:"type"`
}

func NewCustomAdapterTerminateRequest() *customAdapterTerminateRequest {
	return &customAdapterTerminateRequest{"terminate"}
}

// A common struct that allows all types of response to be identified
type customAdapterResponseMessage struct {
	Id             string           `json:"id"`
	Error          *api.ObjectError `json:"error"`
	Oid            string           `json:"oid"`
	Path           string           `json:"path,omitempty"` // always blank for upload
	BytesSoFar     int64            `json:"bytesSoFar"`
	BytesSinceLast int              `json:"bytesSinceLast"`
}

func (a *customAdapter) Begin(maxConcurrency int, cb TransferProgressCallback, completion chan TransferResult) error {
	// If config says not to launch multiple processes, downgrade incoming value
	useConcurrency := maxConcurrency
	if !a.concurrent {
		useConcurrency = 1
	}
	a.originalConcurrency = maxConcurrency

	tracerx.Printf("xfer: Custom transfer adapter %q using concurrency %d", a.name, useConcurrency)

	// Use common workers impl, but downgrade workers to number of processes
	return a.adapterBase.Begin(useConcurrency, cb, completion)
}

func (a *customAdapter) ClearTempStorage() error {
	// no action requred
	return nil
}

func (a *customAdapter) WorkerStarting(workerNum int) (interface{}, error) {

	// Start a process per worker
	// If concurrent = false we have already dialled back workers to 1
	tracerx.Printf("xfer: starting up custom transfer process %q for worker %d", a.name, workerNum)
	cmd := subprocess.ExecCommand(a.path, a.args)
	outp, err := cmd.StdoutPipe()
	if err != nil {
		return nil, fmt.Errorf("Failed to get stdout for custom transfer command %q remote: %v", a.path, err)
	}
	inp, err := cmd.StdinPipe()
	if err != nil {
		return nil, fmt.Errorf("Failed to get stdin for custom transfer command %q remote: %v", a.path, err)
	}
	// Capture stderr to trace
	tracer := &traceWriter{}
	cmd.Stderr = tracer
	err = cmd.Start()
	if err != nil {
		return nil, fmt.Errorf("Failed to start custom transfer command %q remote: %v", a.path, err)
	}
	// Set up buffered reader/writer since we operate on lines
	ctx := &customAdapterWorkerContext{cmd, outp, bufio.NewReader(outp), inp, tracer}

	// send initiate message
	initReq := NewCustomAdapterInitRequest(a.getOperationName(), a.concurrent, a.originalConcurrency)
	resp, err := a.exchangeMessage(ctx, initReq)
	if err != nil {
		a.abortWorkerProcess(ctx)
		return nil, err
	}
	if resp.Error != nil {
		a.abortWorkerProcess(ctx)
		return nil, fmt.Errorf("Error initializing custom adapter %q: %v", a.name, resp.Error.Error())
	}

	tracerx.Printf("xfer: %q for worker %d started OK", a.name, workerNum)

	// Save this process context and use in future callbacks
	return ctx, nil
}

func (a *customAdapter) getOperationName() string {
	if a.direction == Download {
		return "download"
	}
	return "upload"
}

// sendMessage sends a JSON message to the custom adapter process
func (a *customAdapter) sendMessage(ctx *customAdapterWorkerContext, req interface{}) error {
	b, err := json.Marshal(req)
	if err != nil {
		return err
	}
	// Line oriented JSON
	b = append(b, '\n')
	_, err = ctx.stdin.Write(b)
	if err != nil {
		return err
	}
	return nil
}

func (a *customAdapter) readResponse(ctx *customAdapterWorkerContext) (*customAdapterResponseMessage, error) {
	line, err := ctx.bufferedOut.ReadString('\n')
	if err != nil {
		return nil, err
	}
	resp := &customAdapterResponseMessage{}
	err = json.Unmarshal([]byte(line), resp)
	return resp, err
}

// exchangeMessage sends a message to a process and reads a response if resp != nil
// Only fatal errors to communicate return an error, errors may be embedded in reply
func (a *customAdapter) exchangeMessage(ctx *customAdapterWorkerContext, req interface{}) (*customAdapterResponseMessage, error) {

	err := a.sendMessage(ctx, req)
	if err != nil {
		return nil, err
	}
	return a.readResponse(ctx)
}

// shutdownWorkerProcess terminates gracefully a custom adapter process
// returns an error if it couldn't shut down gracefully (caller may abortWorkerProcess)
func (a *customAdapter) shutdownWorkerProcess(ctx *customAdapterWorkerContext) error {
	defer ctx.errTracer.Flush()

	termReq := NewCustomAdapterTerminateRequest()
	err := a.sendMessage(ctx, termReq)
	if err != nil {
		return err
	}
	ctx.stdin.Close()
	ctx.stdout.Close()
	return ctx.cmd.Wait()
}

// abortWorkerProcess terminates & aborts untidily, most probably breakdown of comms or internal error
func (a *customAdapter) abortWorkerProcess(ctx *customAdapterWorkerContext) {
	ctx.stdin.Close()
	ctx.stdout.Close()
	ctx.cmd.Process.Kill()
}
func (a *customAdapter) WorkerEnding(workerNum int, ctx interface{}) {
	customCtx, ok := ctx.(*customAdapterWorkerContext)
	if !ok {
		tracerx.Printf("Context object for custom transfer %q was of the wrong type", a.name)
		return
	}

	err := a.shutdownWorkerProcess(customCtx)
	if err != nil {
		tracerx.Printf("xfer: error finishing up custom transfer process %q, aborting: %v", a.name, err)
		a.abortWorkerProcess(customCtx)
	}
}

func (a *customAdapter) DoTransfer(ctx interface{}, t *Transfer, cb TransferProgressCallback, authOkFunc func()) error {
	if ctx == nil {
		return fmt.Errorf("Custom transfer %q was not properly initialized, see previous errors", a.name)
	}

	customCtx, ok := ctx.(*customAdapterWorkerContext)
	if !ok {
		return fmt.Errorf("Context object for custom transfer %q was of the wrong type", a.name)
	}
	var authCalled bool

	rel, ok := t.Object.Rel(a.getOperationName())
	if !ok {
		return errors.New("Object not found on the server.")
	}
	var req *customAdapterTransferRequest
	if a.direction == Upload {
		req = NewCustomAdapterUploadRequest(t.Object.Oid, t.Object.Size, localstorage.Objects().ObjectPath(t.Object.Oid), rel)
	} else {
		req = NewCustomAdapterDownloadRequest(t.Object.Oid, t.Object.Size, rel)
	}
	err := a.sendMessage(customCtx, req)
	if err != nil {
		return err
	}

	// 1..N replies (including progress & one of download / upload)
	var complete bool
	for !complete {
		resp, err := a.readResponse(customCtx)
		if err != nil {
			return err
		}
		var wasAuthOk bool
		switch resp.Id {
		case "progress":
			// Progress
			if resp.Oid != t.Object.Oid {
				return fmt.Errorf("Unexpected oid %q in response, expecting %q", resp.Oid, t.Object.Oid)
			}
			if cb != nil {
				cb(t.Name, t.Object.Size, resp.BytesSoFar, resp.BytesSinceLast)
			}
			wasAuthOk = resp.BytesSoFar > 0
		case "complete":
			// Download/Upload complete
			if resp.Oid != t.Object.Oid {
				return fmt.Errorf("Unexpected oid %q in response, expecting %q", resp.Oid, t.Object.Oid)
			}
			if resp.Error != nil {
				return fmt.Errorf("Error transferring %q: %v", t.Object.Oid, resp.Error.Error())
			}
			if a.direction == Download {
				// So we don't have to blindly trust external providers, check SHA
				if err = tools.VerifyFileHash(t.Object.Oid, resp.Path); err != nil {
					return fmt.Errorf("Downloaded file failed checks: %v", err)
				}
				// Move file to final location
				if err = tools.RenameFileCopyPermissions(resp.Path, t.Path); err != nil {
					return fmt.Errorf("Failed to copy downloaded file: %v", err)
				}
			} else if a.direction == Upload {
				if err = api.VerifyUpload(t.Object); err != nil {
					return err
				}
			}
			wasAuthOk = true
			complete = true
		}
		// Fall through from both progress and completion messages
		// Call auth on first progress or success to free up other workers
		if wasAuthOk && authOkFunc != nil && !authCalled {
			authOkFunc()
			authCalled = true
		}
	}

	return nil
}

func newCustomAdapter(name string, dir Direction, path, args string, concurrent bool) *customAdapter {
	c := &customAdapter{newAdapterBase(name, dir, nil), path, args, concurrent, 3}
	// self implements impl
	c.transferImpl = c
	return c
}

// Initialise custom adapters based on current config
func ConfigureCustomAdapters() {
	pathRegex := regexp.MustCompile(`lfs.customtransfer.([^.]+).path`)
	for k, v := range config.Config.AllGitConfig() {
		if match := pathRegex.FindStringSubmatch(k); match != nil {
			name := match[1]
			path := v
			var args string
			var concurrent bool
			var direction string
			// retrieve other values
			args, _ = config.Config.GitConfig(fmt.Sprintf("lfs.customtransfer.%s.args", name))
			concurrent = config.Config.GitConfigBool(fmt.Sprintf("lfs.customtransfer.%s.concurrent", name), true)
			direction, _ = config.Config.GitConfig(fmt.Sprintf("lfs.customtransfer.%s.direction", name))
			if len(direction) == 0 {
				direction = "both"
			} else {
				direction = strings.ToLower(direction)
			}

			// Separate closure for each since we need to capture vars above
			newfunc := func(name string, dir Direction) TransferAdapter {
				return newCustomAdapter(name, dir, path, args, concurrent)
			}

			if direction == "download" || direction == "both" {
				RegisterNewTransferAdapterFunc(name, Download, newfunc)
			}
			if direction == "upload" || direction == "both" {
				RegisterNewTransferAdapterFunc(name, Upload, newfunc)
			}

		}
	}

}

func init() {
	ConfigureCustomAdapters()
}
<commit_msg>Improve tracing<commit_after>package transfer

import (
	"bufio"
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"

	"github.com/github/git-lfs/localstorage"
	"github.com/github/git-lfs/tools"

	"github.com/github/git-lfs/api"
	"github.com/github/git-lfs/subprocess"
	"github.com/rubyist/tracerx"

	"github.com/github/git-lfs/config"
)

// Adapter for custom transfer via external process
type customAdapter struct {
	*adapterBase
	path                string
	args                string
	concurrent          bool
	originalConcurrency int
}

// Struct to capture stderr and write to trace
type traceWriter struct {
	buf         bytes.Buffer
	processName string
}

func (t *traceWriter) Write(b []byte) (int, error) {
	n, err := t.buf.Write(b)
	t.Flush()
	return n, err
}
func (t *traceWriter) Flush() {
	var err error
	for err == nil {
		var s string
		s, err = t.buf.ReadString('\n')
		if len(s) > 0 {
			tracerx.Printf("xfer[%v]: %v", t.processName, strings.TrimSpace(s))
		}
	}
}

type customAdapterWorkerContext struct {
	workerNum   int
	cmd         *exec.Cmd
	stdout      io.ReadCloser
	bufferedOut *bufio.Reader
	stdin       io.WriteCloser
	errTracer   *traceWriter
}

type customAdapterInitRequest struct {
	Id                  string `json:"id"`
	Operation           string `json:"operation"`
	Concurrent          bool   `json:"concurrent"`
	ConcurrentTransfers int    `json:"concurrenttransfers"`
}

func NewCustomAdapterInitRequest(op string, concurrent bool, concurrentTransfers int) *customAdapterInitRequest {
	return &customAdapterInitRequest{"init", op, concurrent, concurrentTransfers}
}

type customAdapterTransferRequest struct { // common between upload/download
	Id     string            `json:"id"`
	Oid    string            `json:"oid"`
	Size   int64             `json:"size"`
	Path   string            `json:"path,omitempty"`
	Action *api.LinkRelation `json:"action"`
}

func NewCustomAdapterUploadRequest(oid string, size int64, path string, action *api.LinkRelation) *customAdapterTransferRequest {
	return &customAdapterTransferRequest{"upload", oid, size, path, action}
}
func NewCustomAdapterDownloadRequest(oid string, size int64, action *api.LinkRelation) *customAdapterTransferRequest {
	return &customAdapterTransferRequest{"download", oid, size, "", action}
}

type customAdapterTerminateRequest struct {
	MessageType string `json:"type"`
}

func NewCustomAdapterTerminateRequest() *customAdapterTerminateRequest {
	return &customAdapterTerminateRequest{"terminate"}
}

// A common struct that allows all types of response to be identified
type customAdapterResponseMessage struct {
	Id             string           `json:"id"`
	Error          *api.ObjectError `json:"error"`
	Oid            string           `json:"oid"`
	Path           string           `json:"path,omitempty"` // always blank for upload
	BytesSoFar     int64            `json:"bytesSoFar"`
	BytesSinceLast int              `json:"bytesSinceLast"`
}

func (a *customAdapter) Begin(maxConcurrency int, cb TransferProgressCallback, completion chan TransferResult) error {
	// If config says not to launch multiple processes, downgrade incoming value
	useConcurrency := maxConcurrency
	if !a.concurrent {
		useConcurrency = 1
	}
	a.originalConcurrency = maxConcurrency

	tracerx.Printf("xfer: Custom transfer adapter %q using concurrency %d", a.name, useConcurrency)

	// Use common workers impl, but downgrade workers to number of processes
	return a.adapterBase.Begin(useConcurrency, cb, completion)
}

func (a *customAdapter) ClearTempStorage() error {
	// no action requred
	return nil
}

func (a *customAdapter) WorkerStarting(workerNum int) (interface{}, error) {

	// Start a process per worker
	// If concurrent = false we have already dialled back workers to 1
	tracerx.Printf("xfer: starting up custom transfer process %q for worker %d", a.name, workerNum)
	cmd := subprocess.ExecCommand(a.path, a.args)
	outp, err := cmd.StdoutPipe()
	if err != nil {
		return nil, fmt.Errorf("Failed to get stdout for custom transfer command %q remote: %v", a.path, err)
	}
	inp, err := cmd.StdinPipe()
	if err != nil {
		return nil, fmt.Errorf("Failed to get stdin for custom transfer command %q remote: %v", a.path, err)
	}
	// Capture stderr to trace
	tracer := &traceWriter{}
	tracer.processName = filepath.Base(a.path)
	cmd.Stderr = tracer
	err = cmd.Start()
	if err != nil {
		return nil, fmt.Errorf("Failed to start custom transfer command %q remote: %v", a.path, err)
	}
	// Set up buffered reader/writer since we operate on lines
	ctx := &customAdapterWorkerContext{workerNum, cmd, outp, bufio.NewReader(outp), inp, tracer}

	// send initiate message
	initReq := NewCustomAdapterInitRequest(a.getOperationName(), a.concurrent, a.originalConcurrency)
	resp, err := a.exchangeMessage(ctx, initReq)
	if err != nil {
		a.abortWorkerProcess(ctx)
		return nil, err
	}
	if resp.Error != nil {
		a.abortWorkerProcess(ctx)
		return nil, fmt.Errorf("Error initializing custom adapter %q worker %d: %v", a.name, workerNum, resp.Error.Error())
	}

	tracerx.Printf("xfer: started custom adapter process %q for worker %d OK", a.path, workerNum)

	// Save this process context and use in future callbacks
	return ctx, nil
}

func (a *customAdapter) getOperationName() string {
	if a.direction == Download {
		return "download"
	}
	return "upload"
}

// sendMessage sends a JSON message to the custom adapter process
func (a *customAdapter) sendMessage(ctx *customAdapterWorkerContext, req interface{}) error {
	b, err := json.Marshal(req)
	if err != nil {
		return err
	}
	tracerx.Printf("xfer: Custom adapter worker %d sending message: %v", ctx.workerNum, string(b))
	// Line oriented JSON
	b = append(b, '\n')
	_, err = ctx.stdin.Write(b)
	if err != nil {
		return err
	}

	return nil
}

func (a *customAdapter) readResponse(ctx *customAdapterWorkerContext) (*customAdapterResponseMessage, error) {
	line, err := ctx.bufferedOut.ReadString('\n')
	if err != nil {
		return nil, err
	}
	tracerx.Printf("xfer: Custom adapter worker %d received response: %v", ctx.workerNum, line)
	resp := &customAdapterResponseMessage{}
	err = json.Unmarshal([]byte(line), resp)
	return resp, err
}

// exchangeMessage sends a message to a process and reads a response if resp != nil
// Only fatal errors to communicate return an error, errors may be embedded in reply
func (a *customAdapter) exchangeMessage(ctx *customAdapterWorkerContext, req interface{}) (*customAdapterResponseMessage, error) {

	err := a.sendMessage(ctx, req)
	if err != nil {
		return nil, err
	}
	return a.readResponse(ctx)
}

// shutdownWorkerProcess terminates gracefully a custom adapter process
// returns an error if it couldn't shut down gracefully (caller may abortWorkerProcess)
func (a *customAdapter) shutdownWorkerProcess(ctx *customAdapterWorkerContext) error {
	defer ctx.errTracer.Flush()

	tracerx.Printf("xfer: Shutting down adapter worker %d", ctx.workerNum)
	termReq := NewCustomAdapterTerminateRequest()
	err := a.sendMessage(ctx, termReq)
	if err != nil {
		return err
	}
	ctx.stdin.Close()
	ctx.stdout.Close()
	return ctx.cmd.Wait()
}

// abortWorkerProcess terminates & aborts untidily, most probably breakdown of comms or internal error
func (a *customAdapter) abortWorkerProcess(ctx *customAdapterWorkerContext) {
	tracerx.Printf("xfer: Aborting worker process: %d", ctx.workerNum)
	ctx.stdin.Close()
	ctx.stdout.Close()
	ctx.cmd.Process.Kill()
}
func (a *customAdapter) WorkerEnding(workerNum int, ctx interface{}) {
	customCtx, ok := ctx.(*customAdapterWorkerContext)
	if !ok {
		tracerx.Printf("Context object for custom transfer %q was of the wrong type", a.name)
		return
	}

	err := a.shutdownWorkerProcess(customCtx)
	if err != nil {
		tracerx.Printf("xfer: error finishing up custom transfer process %q worker %d, aborting: %v", a.path, customCtx.workerNum, err)
		a.abortWorkerProcess(customCtx)
	}
}

func (a *customAdapter) DoTransfer(ctx interface{}, t *Transfer, cb TransferProgressCallback, authOkFunc func()) error {
	if ctx == nil {
		return fmt.Errorf("Custom transfer %q was not properly initialized, see previous errors", a.name)
	}

	customCtx, ok := ctx.(*customAdapterWorkerContext)
	if !ok {
		return fmt.Errorf("Context object for custom transfer %q was of the wrong type", a.name)
	}
	var authCalled bool

	rel, ok := t.Object.Rel(a.getOperationName())
	if !ok {
		return errors.New("Object not found on the server.")
	}
	var req *customAdapterTransferRequest
	if a.direction == Upload {
		req = NewCustomAdapterUploadRequest(t.Object.Oid, t.Object.Size, localstorage.Objects().ObjectPath(t.Object.Oid), rel)
	} else {
		req = NewCustomAdapterDownloadRequest(t.Object.Oid, t.Object.Size, rel)
	}
	err := a.sendMessage(customCtx, req)
	if err != nil {
		return err
	}

	// 1..N replies (including progress & one of download / upload)
	var complete bool
	for !complete {
		resp, err := a.readResponse(customCtx)
		if err != nil {
			return err
		}
		var wasAuthOk bool
		switch resp.Id {
		case "progress":
			// Progress
			if resp.Oid != t.Object.Oid {
				return fmt.Errorf("Unexpected oid %q in response, expecting %q", resp.Oid, t.Object.Oid)
			}
			if cb != nil {
				cb(t.Name, t.Object.Size, resp.BytesSoFar, resp.BytesSinceLast)
			}
			wasAuthOk = resp.BytesSoFar > 0
		case "complete":
			// Download/Upload complete
			if resp.Oid != t.Object.Oid {
				return fmt.Errorf("Unexpected oid %q in response, expecting %q", resp.Oid, t.Object.Oid)
			}
			if resp.Error != nil {
				return fmt.Errorf("Error transferring %q: %v", t.Object.Oid, resp.Error.Error())
			}
			if a.direction == Download {
				// So we don't have to blindly trust external providers, check SHA
				if err = tools.VerifyFileHash(t.Object.Oid, resp.Path); err != nil {
					return fmt.Errorf("Downloaded file failed checks: %v", err)
				}
				// Move file to final location
				if err = tools.RenameFileCopyPermissions(resp.Path, t.Path); err != nil {
					return fmt.Errorf("Failed to copy downloaded file: %v", err)
				}
			} else if a.direction == Upload {
				if err = api.VerifyUpload(t.Object); err != nil {
					return err
				}
			}
			wasAuthOk = true
			complete = true
		}
		// Fall through from both progress and completion messages
		// Call auth on first progress or success to free up other workers
		if wasAuthOk && authOkFunc != nil && !authCalled {
			authOkFunc()
			authCalled = true
		}
	}

	return nil
}

func newCustomAdapter(name string, dir Direction, path, args string, concurrent bool) *customAdapter {
	c := &customAdapter{newAdapterBase(name, dir, nil), path, args, concurrent, 3}
	// self implements impl
	c.transferImpl = c
	return c
}

// Initialise custom adapters based on current config
func ConfigureCustomAdapters() {
	pathRegex := regexp.MustCompile(`lfs.customtransfer.([^.]+).path`)
	for k, v := range config.Config.AllGitConfig() {
		if match := pathRegex.FindStringSubmatch(k); match != nil {
			name := match[1]
			path := v
			var args string
			var concurrent bool
			var direction string
			// retrieve other values
			args, _ = config.Config.GitConfig(fmt.Sprintf("lfs.customtransfer.%s.args", name))
			concurrent = config.Config.GitConfigBool(fmt.Sprintf("lfs.customtransfer.%s.concurrent", name), true)
			direction, _ = config.Config.GitConfig(fmt.Sprintf("lfs.customtransfer.%s.direction", name))
			if len(direction) == 0 {
				direction = "both"
			} else {
				direction = strings.ToLower(direction)
			}

			// Separate closure for each since we need to capture vars above
			newfunc := func(name string, dir Direction) TransferAdapter {
				return newCustomAdapter(name, dir, path, args, concurrent)
			}

			if direction == "download" || direction == "both" {
				RegisterNewTransferAdapterFunc(name, Download, newfunc)
			}
			if direction == "upload" || direction == "both" {
				RegisterNewTransferAdapterFunc(name, Upload, newfunc)
			}

		}
	}

}

func init() {
	ConfigureCustomAdapters()
}
<|endoftext|>