<commit_before>package quad_test

import (
	"github.com/ghthor/engine/rpg2d/coord"
	"github.com/ghthor/engine/rpg2d/entity"
	"github.com/ghthor/engine/rpg2d/quad"
	"github.com/ghthor/engine/sim/stime"

	"github.com/ghthor/gospec"
	. "github.com/ghthor/gospec"
)

var quadBounds = coord.Bounds{
	coord.Cell{-8, 8},
	coord.Cell{7, -7},
}

// Creates a set of entities in collision groups
// used for testing the broad phase.
func cgEntitiesDataSet() ([]MockEntityWithBounds, []quad.Collision, []quad.CollisionGroup) {
	entities := func() []MockEntityWithBounds {
		c := func(x, y int) coord.Cell { return coord.Cell{x, y} }
		b := func(tl, br coord.Cell) coord.Bounds { return coord.Bounds{tl, br} }

		return []MockEntityWithBounds{
			{ // CollisionGroup 0
				0, c(0, 0),
				b(c(0, 0), c(1, 0)),
			}, {
				1, c(1, 0),
				b(c(1, 0), c(2, 0)),
			},

			{ // CollisionGroup 1
				2, c(1, 1),
				b(c(1, 2), c(1, 1)),
			}, {
				3, c(1, 3),
				b(c(1, 3), c(1, 2)),
			}, {
				4, c(2, 2),
				b(c(1, 2), c(2, 2)),
			},

			{ // CollisionGroup 2
				5, c(-1, 0),
				b(c(-2, 0), c(-2, 0)),
			}, {
				6, c(-2, 0),
				b(c(-2, 0), c(-2, -1)),
			}, {
				7, c(-2, -1),
				b(c(-2, -1), c(-1, -1)),
			}, {
				8, c(-1, -1),
				b(c(-1, -1), c(0, -1)),
			}, {
				9, c(0, -1),
				b(c(0, -1), c(1, -1)),
			}, {
				10, c(1, -1),
				b(c(1, -1), c(1, -2)),
			}, {
				11, c(1, -2),
				b(c(-2, -2), c(1, -2)),
			},

			{ // CollisionGroup 3
				12, c(0, 5),
				b(c(0, 5), c(1, 5)),
			}, {
				13, c(1, 5),
				b(c(1, 5), c(2, 5)),
			}, {
				14, c(2, 5),
				b(c(2, 5), c(3, 5)),
			}, {
				15, c(0, 6),
				b(c(0, 6), c(1, 6)),
			}, {
				16, c(1, 6),
				b(c(1, 6), c(2, 6)),
			}, {
				17, c(2, 6),
				b(c(2, 6), c(3, 6)),
			}, {
				18, c(3, 6),
				b(c(3, 6), c(3, 5)),
			},
		}
	}()

	collisions := func(e []MockEntityWithBounds) []quad.Collision {
		c := func(a, b entity.Entity) quad.Collision { return quad.Collision{a, b} }

		return []quad.Collision{
			// Group 0
			c(e[0], e[1]),

			// Group 1
			c(e[2], e[3]),
			c(e[2], e[4]),
			c(e[3], e[4]),

			// Group 2
			c(e[5], e[6]),
			c(e[6], e[7]),
			c(e[7], e[8]),
			c(e[8], e[9]),
			c(e[9], e[10]),
			c(e[10], e[11]),

			// Group 3,
			c(e[12], e[13]),
			c(e[13], e[14]),
			c(e[14], e[18]),
			c(e[15], e[16]),
			c(e[16], e[17]),
			c(e[17], e[18]),
		}
	}(entities)

	cgroups := func(c []quad.Collision) []quad.CollisionGroup {
		cg := func(collisions ...quad.Collision) (cg quad.CollisionGroup) {
			for _, c := range collisions {
				cg = cg.AddCollision(c)
			}
			return
		}

		return []quad.CollisionGroup{
			cg(c[0]),
			cg(c[1:4]...),
			cg(c[4:10]...),
			cg(c[10:16]...),
		}
	}(collisions)

	return entities, collisions, cgroups
}

func DescribePhase(c gospec.Context) {
	c.Specify("the input phase", func() {
		c.Specify("will remove any entites that move out of bounds", func() {
			q, err := quad.New(coord.Bounds{
				TopL: coord.Cell{-16, 16},
				BotR: coord.Cell{15, -15},
			}, 3, nil)
			c.Assume(err, IsNil)

			// A Single Entity
			q = q.Insert(MockEntity{0, coord.Cell{-16, 16}})
			c.Assume(len(q.QueryBounds(q.Bounds())), Equals, 1)

			q, outOfBounds := quad.RunInputPhaseOn(q, quad.InputPhaseHandlerFn(func(chunk quad.Chunk, now stime.Time) quad.Chunk {
				c.Assume(len(chunk.Entities), Equals, 1)

				// Move the entity out of bounds
				chunk.Entities[0] = MockEntity{0, coord.Cell{-17, 16}}

				return chunk
			}), stime.Time(0))

			c.Expect(len(outOfBounds), Equals, 1)
			c.Expect(len(q.QueryBounds(q.Bounds())), Equals, 0)

			// Multiple entities
			q = q.Insert(MockEntity{0, coord.Cell{-16, 16}})
			q = q.Insert(MockEntity{1, coord.Cell{15, -15}})
			q = q.Insert(MockEntity{2, coord.Cell{-1, 1}})
			q = q.Insert(MockEntity{3, coord.Cell{0, 0}})
			q = q.Insert(MockEntity{4, coord.Cell{5, -2}})

			q, outOfBounds = quad.RunInputPhaseOn(q, quad.InputPhaseHandlerFn(func(chunk quad.Chunk, now stime.Time) quad.Chunk {
				// Move the entity out of bounds
				for i, e := range chunk.Entities {
					switch e.Id() {
					case 1:
						// Move out of quadtree's bounds
						chunk.Entities[i] = MockEntity{1, coord.Cell{16, -15}}
					case 4:
						// Move from SE to NE quadrant
						chunk.Entities[i] = MockEntity{4, coord.Cell{5, 5}}
					}
				}

				return chunk
			}), stime.Time(0))

			c.Expect(len(q.QueryBounds(q.Bounds())), Equals, 4)
			c.Expect(len(outOfBounds), Equals, 1)

			c.Expect(q.QueryCell(coord.Cell{-16, 16})[0].Id(), Equals, int64(0))
			c.Expect(outOfBounds[0].Id(), Equals, int64(1))
			c.Expect(q.QueryCell(coord.Cell{-1, 1})[0].Id(), Equals, int64(2))
			c.Expect(q.QueryCell(coord.Cell{0, 0})[0].Id(), Equals, int64(3))
			c.Expect(q.QueryCell(coord.Cell{5, 5})[0].Id(), Equals, int64(4))

		})
	})

	c.Specify("the broad phase", func() {
		cgEntities, _, cgroups := cgEntitiesDataSet()
		// Sanity check the data set is what I expect it
		// to be. Have these check because this is the first
		// time I've used slice operations extensively and I
		// want to make sure I'm using the right indices in
		// range expressions.
		c.Assume(len(cgroups[0].Entities), Equals, 2)
		c.Assume(len(cgroups[0].Collisions), Equals, 1)
		c.Assume(cgroups[0].Entities, ContainsAll, cgEntities[0:2])
		c.Assume(cgroups[0].Entities, Not(ContainsAny), cgEntities[2:])

		c.Assume(len(cgroups[1].Entities), Equals, 3)
		c.Assume(len(cgroups[1].Collisions), Equals, 3)
		c.Assume(cgroups[1].Entities, Not(ContainsAny), cgEntities[0:2])
		c.Assume(cgroups[1].Entities, ContainsAll, cgEntities[2:5])
		c.Assume(cgroups[1].Entities, Not(ContainsAny), cgEntities[5:])

		c.Assume(len(cgroups[2].Entities), Equals, 7)
		c.Assume(len(cgroups[2].Collisions), Equals, 6)
		c.Assume(cgroups[2].Entities, Not(ContainsAny), cgEntities[0:5])
		c.Assume(cgroups[2].Entities, ContainsAll, cgEntities[5:12])
		c.Assume(cgroups[2].Entities, Not(ContainsAny), cgEntities[12:])

		c.Assume(len(cgroups[3].Entities), Equals, 7)
		c.Assume(len(cgroups[3].Collisions), Equals, 6)
		c.Assume(cgroups[3].Entities, Not(ContainsAny), cgEntities[0:12])
		c.Assume(cgroups[3].Entities, ContainsAll, cgEntities[12:19])
		c.Assume(cgroups[3].Entities, Not(ContainsAny), cgEntities[19:])

		makeQuad := func(entities []entity.Entity, quadMaxSize int) quad.Quad {
			q, err := quad.New(quadBounds, quadMaxSize, nil)
			c.Assume(err, IsNil)

			for _, e := range entities {
				q = q.Insert(e)
			}

			return q
		}

		c.Specify("will create collision groups", func() {
			type testCase struct {
				entities []entity.Entity
				cgroups  []quad.CollisionGroup
			}

			testCases := func(cg []quad.CollisionGroup) []testCase {
				tc := func(cgroups ...quad.CollisionGroup) testCase {
					var entities []entity.Entity
					for _, cg := range cgroups {
						entities = append(entities, cg.Entities...)
					}
					return testCase{entities, cgroups}
				}

				return []testCase{
					// Make test cases that only contain collision groups
					tc(cg[0]),
					tc(cg[0:2]...),
					tc(cg[0:3]...),
					tc(cg[0:4]...),
					tc(cg[1]),
					tc(cg[1:3]...),
					tc(cg[1:4]...),
					tc(cg[2]),
					tc(cg[2:4]...),
					tc(cg[3]),
					tc(cg...),
				}
			}(cgroups)

			for _, testCase := range testCases {
				for i := 4; i < len(testCase.entities)+1; i++ {
					q := makeQuad(testCase.entities, i)

					var cgroups []*quad.CollisionGroup
					q, cgroups, _, _ = quad.RunBroadPhaseOn(q, stime.Time(0))

					c.Expect(len(cgroups), Equals, len(testCase.cgroups))
					c.Expect(cgroups, ContainsAll, testCase.cgroups)
				}
			}
		})
	})

	c.Specify("the narrow phase", func() {
		c.Specify("will realize all future potentials", func() {
		})
	})
}
<commit_msg>Sort the entities by id before inserting them into the tree<commit_after>package quad_test

import (
	"sort"

	"github.com/ghthor/engine/rpg2d/coord"
	"github.com/ghthor/engine/rpg2d/entity"
	"github.com/ghthor/engine/rpg2d/quad"
	"github.com/ghthor/engine/sim/stime"

	"github.com/ghthor/gospec"
	. "github.com/ghthor/gospec"
)

var quadBounds = coord.Bounds{
	coord.Cell{-8, 8},
	coord.Cell{7, -7},
}

// Creates a set of entities in collision groups
// used for testing the broad phase.
func cgEntitiesDataSet() ([]MockEntityWithBounds, []quad.Collision, []quad.CollisionGroup) {
	entities := func() []MockEntityWithBounds {
		c := func(x, y int) coord.Cell { return coord.Cell{x, y} }
		b := func(tl, br coord.Cell) coord.Bounds { return coord.Bounds{tl, br} }

		return []MockEntityWithBounds{
			{ // CollisionGroup 0
				0, c(0, 0),
				b(c(0, 0), c(1, 0)),
			}, {
				1, c(1, 0),
				b(c(1, 0), c(2, 0)),
			},

			{ // CollisionGroup 1
				2, c(1, 1),
				b(c(1, 2), c(1, 1)),
			}, {
				3, c(1, 3),
				b(c(1, 3), c(1, 2)),
			}, {
				4, c(2, 2),
				b(c(1, 2), c(2, 2)),
			},

			{ // CollisionGroup 2
				5, c(-1, 0),
				b(c(-2, 0), c(-2, 0)),
			}, {
				6, c(-2, 0),
				b(c(-2, 0), c(-2, -1)),
			}, {
				7, c(-2, -1),
				b(c(-2, -1), c(-1, -1)),
			}, {
				8, c(-1, -1),
				b(c(-1, -1), c(0, -1)),
			}, {
				9, c(0, -1),
				b(c(0, -1), c(1, -1)),
			}, {
				10, c(1, -1),
				b(c(1, -1), c(1, -2)),
			}, {
				11, c(1, -2),
				b(c(-2, -2), c(1, -2)),
			},

			{ // CollisionGroup 3
				12, c(0, 5),
				b(c(0, 5), c(1, 5)),
			}, {
				13, c(1, 5),
				b(c(1, 5), c(2, 5)),
			}, {
				14, c(2, 5),
				b(c(2, 5), c(3, 5)),
			}, {
				15, c(0, 6),
				b(c(0, 6), c(1, 6)),
			}, {
				16, c(1, 6),
				b(c(1, 6), c(2, 6)),
			}, {
				17, c(2, 6),
				b(c(2, 6), c(3, 6)),
			}, {
				18, c(3, 6),
				b(c(3, 6), c(3, 5)),
			},
		}
	}()

	collisions := func(e []MockEntityWithBounds) []quad.Collision {
		c := func(a, b entity.Entity) quad.Collision { return quad.Collision{a, b} }

		return []quad.Collision{
			// Group 0
			c(e[0], e[1]),

			// Group 1
			c(e[2], e[3]),
			c(e[2], e[4]),
			c(e[3], e[4]),

			// Group 2
			c(e[5], e[6]),
			c(e[6], e[7]),
			c(e[7], e[8]),
			c(e[8], e[9]),
			c(e[9], e[10]),
			c(e[10], e[11]),

			// Group 3,
			c(e[12], e[13]),
			c(e[13], e[14]),
			c(e[14], e[18]),
			c(e[15], e[16]),
			c(e[16], e[17]),
			c(e[17], e[18]),
		}
	}(entities)

	cgroups := func(c []quad.Collision) []quad.CollisionGroup {
		cg := func(collisions ...quad.Collision) (cg quad.CollisionGroup) {
			for _, c := range collisions {
				cg = cg.AddCollision(c)
			}
			return
		}

		return []quad.CollisionGroup{
			cg(c[0]),
			cg(c[1:4]...),
			cg(c[4:10]...),
			cg(c[10:16]...),
			cg(c[16:18]...),
		}
	}(collisions)

	return entities, collisions, cgroups
}

type byId []entity.Entity

func (e byId) Len() int      { return len(e) }
func (e byId) Swap(i, j int) { e[i], e[j] = e[j], e[i] }
func (e byId) Less(i, j int) bool {
	return e[i].Id() < e[j].Id()
}

func DescribePhase(c gospec.Context) {
	c.Specify("the input phase", func() {
		c.Specify("will remove any entites that move out of bounds", func() {
			q, err := quad.New(coord.Bounds{
				TopL: coord.Cell{-16, 16},
				BotR: coord.Cell{15, -15},
			}, 3, nil)
			c.Assume(err, IsNil)

			// A Single Entity
			q = q.Insert(MockEntity{0, coord.Cell{-16, 16}})
			c.Assume(len(q.QueryBounds(q.Bounds())), Equals, 1)

			q, outOfBounds := quad.RunInputPhaseOn(q, quad.InputPhaseHandlerFn(func(chunk quad.Chunk, now stime.Time) quad.Chunk {
				c.Assume(len(chunk.Entities), Equals, 1)

				// Move the entity out of bounds
				chunk.Entities[0] = MockEntity{0, coord.Cell{-17, 16}}

				return chunk
			}), stime.Time(0))

			c.Expect(len(outOfBounds), Equals, 1)
			c.Expect(len(q.QueryBounds(q.Bounds())), Equals, 0)

			// Multiple entities
			q = q.Insert(MockEntity{0, coord.Cell{-16, 16}})
			q = q.Insert(MockEntity{1, coord.Cell{15, -15}})
			q = q.Insert(MockEntity{2, coord.Cell{-1, 1}})
			q = q.Insert(MockEntity{3, coord.Cell{0, 0}})
			q = q.Insert(MockEntity{4, coord.Cell{5, -2}})

			q, outOfBounds = quad.RunInputPhaseOn(q, quad.InputPhaseHandlerFn(func(chunk quad.Chunk, now stime.Time) quad.Chunk {
				// Move the entity out of bounds
				for i, e := range chunk.Entities {
					switch e.Id() {
					case 1:
						// Move out of quadtree's bounds
						chunk.Entities[i] = MockEntity{1, coord.Cell{16, -15}}
					case 4:
						// Move from SE to NE quadrant
						chunk.Entities[i] = MockEntity{4, coord.Cell{5, 5}}
					}
				}

				return chunk
			}), stime.Time(0))

			c.Expect(len(q.QueryBounds(q.Bounds())), Equals, 4)
			c.Expect(len(outOfBounds), Equals, 1)

			c.Expect(q.QueryCell(coord.Cell{-16, 16})[0].Id(), Equals, int64(0))
			c.Expect(outOfBounds[0].Id(), Equals, int64(1))
			c.Expect(q.QueryCell(coord.Cell{-1, 1})[0].Id(), Equals, int64(2))
			c.Expect(q.QueryCell(coord.Cell{0, 0})[0].Id(), Equals, int64(3))
			c.Expect(q.QueryCell(coord.Cell{5, 5})[0].Id(), Equals, int64(4))

		})
	})

	c.Specify("the broad phase", func() {
		cgEntities, _, cgroups := cgEntitiesDataSet()
		// Sanity check the data set is what I expect it
		// to be. Have these check because this is the first
		// time I've used slice operations extensively and I
		// want to make sure I'm using the right indices in
		// range expressions.
		c.Assume(len(cgroups[0].Entities), Equals, 2)
		c.Assume(len(cgroups[0].Collisions), Equals, 1)
		c.Assume(cgroups[0].Entities, ContainsAll, cgEntities[0:2])
		c.Assume(cgroups[0].Entities, Not(ContainsAny), cgEntities[2:])

		c.Assume(len(cgroups[1].Entities), Equals, 3)
		c.Assume(len(cgroups[1].Collisions), Equals, 3)
		c.Assume(cgroups[1].Entities, Not(ContainsAny), cgEntities[0:2])
		c.Assume(cgroups[1].Entities, ContainsAll, cgEntities[2:5])
		c.Assume(cgroups[1].Entities, Not(ContainsAny), cgEntities[5:])

		c.Assume(len(cgroups[2].Entities), Equals, 7)
		c.Assume(len(cgroups[2].Collisions), Equals, 6)
		c.Assume(cgroups[2].Entities, Not(ContainsAny), cgEntities[0:5])
		c.Assume(cgroups[2].Entities, ContainsAll, cgEntities[5:12])
		c.Assume(cgroups[2].Entities, Not(ContainsAny), cgEntities[12:])

		c.Assume(len(cgroups[3].Entities), Equals, 7)
		c.Assume(len(cgroups[3].Collisions), Equals, 6)
		c.Assume(cgroups[3].Entities, Not(ContainsAny), cgEntities[0:12])
		c.Assume(cgroups[3].Entities, ContainsAll, cgEntities[12:19])
		c.Assume(cgroups[3].Entities, Not(ContainsAny), cgEntities[19:])

		makeQuad := func(entities []entity.Entity, quadMaxSize int) quad.Quad {
			q, err := quad.New(quadBounds, quadMaxSize, nil)
			c.Assume(err, IsNil)

			for _, e := range entities {
				q = q.Insert(e)
			}

			return q
		}

		c.Specify("will create collision groups", func() {
			type testCase struct {
				entities []entity.Entity
				cgroups  []quad.CollisionGroup
			}

			testCases := func(cg []quad.CollisionGroup) []testCase {
				tc := func(cgroups ...quad.CollisionGroup) testCase {
					var entities []entity.Entity
					for _, cg := range cgroups {
						entities = append(entities, cg.Entities...)
					}
					sort.Sort(byId(entities))
					return testCase{entities, cgroups}
				}

				return []testCase{
					// Make test cases that only contain collision groups
					tc(cg[0]),
					tc(cg[0:2]...),
					tc(cg[0:3]...),
					tc(cg[0:4]...),
					tc(cg[1]),
					tc(cg[1:3]...),
					tc(cg[1:4]...),
					tc(cg[2]),
					tc(cg[2:4]...),
					tc(cg[3]),
					tc(cg...),
				}
			}(cgroups)

			for _, testCase := range testCases {
				for i := 4; i < len(testCase.entities)+1; i++ {
					q := makeQuad(testCase.entities, i)

					var cgroups []*quad.CollisionGroup
					q, cgroups, _, _ = quad.RunBroadPhaseOn(q, stime.Time(0))

					c.Expect(len(cgroups), Equals, len(testCase.cgroups))
					c.Expect(cgroups, ContainsAll, testCase.cgroups)
				}
			}
		})
	})

	c.Specify("the narrow phase", func() {
		c.Specify("will realize all future potentials", func() {
		})
	})
}
<|endoftext|>


<commit_before>package rt

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

	"veyron.io/veyron/veyron2"
	"veyron.io/veyron/veyron2/config"
	"veyron.io/veyron/veyron2/i18n"
	"veyron.io/veyron/veyron2/ipc"
	"veyron.io/veyron/veyron2/ipc/stream"
	"veyron.io/veyron/veyron2/naming"
	"veyron.io/veyron/veyron2/rt"

	"veyron.io/veyron/veyron/lib/exec"
	"veyron.io/veyron/veyron/lib/flags"
	_ "veyron.io/veyron/veyron/lib/stats/sysstats"
	"veyron.io/veyron/veyron/runtimes/google/naming/namespace"
	"veyron.io/veyron/veyron2/options"
	"veyron.io/veyron/veyron2/security"
	"veyron.io/veyron/veyron2/vlog"
)

// TODO(caprita): Verrorize this, and maybe move it in the API.
var errCleaningUp = fmt.Errorf("operation rejected: runtime is being cleaned up")

type vrt struct {
	mu sync.Mutex

	profile        veyron2.Profile
	publisher      *config.Publisher
	sm             []stream.Manager // GUARDED_BY(mu)
	ns             naming.Namespace
	signals        chan os.Signal
	principal      security.Principal
	client         ipc.Client
	mgmt           *mgmtImpl
	flags          flags.RuntimeFlags
	reservedDisp   ipc.Dispatcher
	reservedPrefix string
	reservedOpts   []ipc.ServerOpt
	nServers       int  // GUARDED_BY(mu)
	cleaningUp     bool // GUARDED_BY(mu)

	lang    i18n.LangID // Language, from environment variables.
	program string      // Program name, from os.Args[0].
}

var _ veyron2.Runtime = (*vrt)(nil)

var flagsOnce sync.Once
var runtimeFlags *flags.Flags

func init() {
	rt.RegisterRuntime(veyron2.GoogleRuntimeName, New)
	runtimeFlags = flags.CreateAndRegister(flag.CommandLine, flags.Runtime)
}

// Implements veyron2/rt.New
func New(opts ...veyron2.ROpt) (veyron2.Runtime, error) {
	rt := &vrt{mgmt: new(mgmtImpl), lang: i18n.LangIDFromEnv(), program: filepath.Base(os.Args[0])}
	flagsOnce.Do(func() {
		runtimeFlags.Parse(os.Args[1:])
	})
	rt.flags = runtimeFlags.RuntimeFlags()
	nsRoots := []string{}
	for _, o := range opts {
		switch v := o.(type) {
		case options.RuntimePrincipal:
			rt.principal = v.Principal
		case options.Profile:
			rt.profile = v.Profile
		default:
			return nil, fmt.Errorf("option has wrong type %T", o)
		}
	}

	rt.initLogging()
	rt.initSignalHandling()

	if rt.profile == nil {
		return nil, fmt.Errorf("No profile configured!")
	} else {
		vlog.VI(1).Infof("Using profile %q", rt.profile.Name())
	}

	if ns, err := namespace.New(rt, rt.flags.NamespaceRoots...); err != nil {
		return nil, fmt.Errorf("Couldn't create mount table: %v", err)
	} else {
		rt.ns = ns
	}
	vlog.VI(2).Infof("Namespace Roots: %s", nsRoots)

	// Create the default stream manager.
	if _, err := rt.NewStreamManager(); err != nil {
		return nil, fmt.Errorf("failed to create stream manager: %s", err)
	}

	if err := rt.initSecurity(rt.flags.Credentials); err != nil {
		return nil, fmt.Errorf("failed to init sercurity: %s", err)
	}

	var err error
	if rt.client, err = rt.NewClient(); err != nil {
		return nil, fmt.Errorf("failed to create new client: %s", err)
	}

	handle, err := exec.GetChildHandle()
	switch err {
	case exec.ErrNoVersion:
		// The process has not been started through the veyron exec
		// library. No further action is needed.
	case nil:
		// The process has been started through the veyron exec
		// library. Signal the parent the the child is ready.
		handle.SetReady()
	default:
		return nil, fmt.Errorf("failed to get child handle: %s", err)
	}

	rt.publisher = config.NewPublisher()
	if err := rt.profile.Init(rt, rt.publisher); err != nil {
		return nil, err
	}

	// TODO(caprita, cnicolaou): how is this to be configured?
	// Can it ever be anything other than a localhost/loopback address?
	listenSpec := ipc.ListenSpec{Protocol: "tcp", Address: "127.0.0.1:0"}
	if err := rt.mgmt.initMgmt(rt, listenSpec); err != nil {
		return nil, err
	}

	vlog.VI(2).Infof("rt.Init done")
	return rt, nil
}

func (rt *vrt) Publisher() *config.Publisher {
	return rt.publisher
}

func (rt *vrt) Profile() veyron2.Profile {
	return rt.profile
}

func (rt *vrt) ConfigureReservedName(prefix string, server ipc.Dispatcher, opts ...ipc.ServerOpt) {
	rt.mu.Lock()
	defer rt.mu.Unlock()
	rt.reservedDisp = server
	prefix = strings.TrimLeft(prefix, "_")
	rt.reservedPrefix = naming.ReservedNamePrefix + prefix
	rt.reservedOpts = make([]ipc.ServerOpt, 0, len(opts))
	copy(rt.reservedOpts, opts)
}

func (rt *vrt) Cleanup() {
	rt.mu.Lock()
	if rt.cleaningUp {
		rt.mu.Unlock()
		// TODO(caprita): Should we actually treat extra Cleanups as
		// silent no-ops?  For now, it's better not to, in order to
		// expose programming errors.
		vlog.Errorf("rt.Cleanup done")
		return
	}
	rt.cleaningUp = true
	rt.mu.Unlock()

	// TODO(caprita): Consider shutting down mgmt later in the runtime's
	// shutdown sequence, to capture some of the runtime internal shutdown
	// tasks in the task tracker.
	rt.mgmt.shutdown()
	// It's ok to access rt.sm out of lock below, since a Mutex acts as a
	// barrier in Go and hence we're guaranteed that cleaningUp is true at
	// this point.  The only code that mutates rt.sm is NewStreamManager in
	// ipc.go, which respects the value of cleaningUp.
	//
	// TODO(caprita): It would be cleaner if we do something like this
	// inside the lock (and then iterate over smgrs below):
	//   smgrs := rt.sm
	//   rt.sm = nil
	// However, to make that work we need to prevent NewClient and NewServer
	// from using rt.sm after cleaningUp has been set.  Hence the TODO.
	for _, sm := range rt.sm {
		sm.Shutdown()
	}
	rt.shutdownSignalHandling()
	rt.shutdownLogging()
}
<commit_msg>TBR Fix nsRoots bug in rt.go<commit_after>package rt

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

	"veyron.io/veyron/veyron2"
	"veyron.io/veyron/veyron2/config"
	"veyron.io/veyron/veyron2/i18n"
	"veyron.io/veyron/veyron2/ipc"
	"veyron.io/veyron/veyron2/ipc/stream"
	"veyron.io/veyron/veyron2/naming"
	"veyron.io/veyron/veyron2/rt"

	"veyron.io/veyron/veyron/lib/exec"
	"veyron.io/veyron/veyron/lib/flags"
	_ "veyron.io/veyron/veyron/lib/stats/sysstats"
	"veyron.io/veyron/veyron/runtimes/google/naming/namespace"
	"veyron.io/veyron/veyron2/options"
	"veyron.io/veyron/veyron2/security"
	"veyron.io/veyron/veyron2/vlog"
)

// TODO(caprita): Verrorize this, and maybe move it in the API.
var errCleaningUp = fmt.Errorf("operation rejected: runtime is being cleaned up")

type vrt struct {
	mu sync.Mutex

	profile        veyron2.Profile
	publisher      *config.Publisher
	sm             []stream.Manager // GUARDED_BY(mu)
	ns             naming.Namespace
	signals        chan os.Signal
	principal      security.Principal
	client         ipc.Client
	mgmt           *mgmtImpl
	flags          flags.RuntimeFlags
	reservedDisp   ipc.Dispatcher
	reservedPrefix string
	reservedOpts   []ipc.ServerOpt
	nServers       int  // GUARDED_BY(mu)
	cleaningUp     bool // GUARDED_BY(mu)

	lang    i18n.LangID // Language, from environment variables.
	program string      // Program name, from os.Args[0].
}

var _ veyron2.Runtime = (*vrt)(nil)

var flagsOnce sync.Once
var runtimeFlags *flags.Flags

func init() {
	rt.RegisterRuntime(veyron2.GoogleRuntimeName, New)
	runtimeFlags = flags.CreateAndRegister(flag.CommandLine, flags.Runtime)
}

// Implements veyron2/rt.New
func New(opts ...veyron2.ROpt) (veyron2.Runtime, error) {
	rt := &vrt{mgmt: new(mgmtImpl), lang: i18n.LangIDFromEnv(), program: filepath.Base(os.Args[0])}
	flagsOnce.Do(func() {
		runtimeFlags.Parse(os.Args[1:])
	})
	rt.flags = runtimeFlags.RuntimeFlags()
	for _, o := range opts {
		switch v := o.(type) {
		case options.RuntimePrincipal:
			rt.principal = v.Principal
		case options.Profile:
			rt.profile = v.Profile
		default:
			return nil, fmt.Errorf("option has wrong type %T", o)
		}
	}

	rt.initLogging()
	rt.initSignalHandling()

	if rt.profile == nil {
		return nil, fmt.Errorf("No profile configured!")
	} else {
		vlog.VI(1).Infof("Using profile %q", rt.profile.Name())
	}

	if ns, err := namespace.New(rt, rt.flags.NamespaceRoots...); err != nil {
		return nil, fmt.Errorf("Couldn't create mount table: %v", err)
	} else {
		rt.ns = ns
	}
	vlog.VI(2).Infof("Namespace Roots: %s", rt.ns.Roots())

	// Create the default stream manager.
	if _, err := rt.NewStreamManager(); err != nil {
		return nil, fmt.Errorf("failed to create stream manager: %s", err)
	}

	if err := rt.initSecurity(rt.flags.Credentials); err != nil {
		return nil, fmt.Errorf("failed to init sercurity: %s", err)
	}

	var err error
	if rt.client, err = rt.NewClient(); err != nil {
		return nil, fmt.Errorf("failed to create new client: %s", err)
	}

	handle, err := exec.GetChildHandle()
	switch err {
	case exec.ErrNoVersion:
		// The process has not been started through the veyron exec
		// library. No further action is needed.
	case nil:
		// The process has been started through the veyron exec
		// library. Signal the parent the the child is ready.
		handle.SetReady()
	default:
		return nil, fmt.Errorf("failed to get child handle: %s", err)
	}

	rt.publisher = config.NewPublisher()
	if err := rt.profile.Init(rt, rt.publisher); err != nil {
		return nil, err
	}

	// TODO(caprita, cnicolaou): how is this to be configured?
	// Can it ever be anything other than a localhost/loopback address?
	listenSpec := ipc.ListenSpec{Protocol: "tcp", Address: "127.0.0.1:0"}
	if err := rt.mgmt.initMgmt(rt, listenSpec); err != nil {
		return nil, err
	}

	vlog.VI(2).Infof("rt.Init done")
	return rt, nil
}

func (rt *vrt) Publisher() *config.Publisher {
	return rt.publisher
}

func (rt *vrt) Profile() veyron2.Profile {
	return rt.profile
}

func (rt *vrt) ConfigureReservedName(prefix string, server ipc.Dispatcher, opts ...ipc.ServerOpt) {
	rt.mu.Lock()
	defer rt.mu.Unlock()
	rt.reservedDisp = server
	prefix = strings.TrimLeft(prefix, "_")
	rt.reservedPrefix = naming.ReservedNamePrefix + prefix
	rt.reservedOpts = make([]ipc.ServerOpt, 0, len(opts))
	copy(rt.reservedOpts, opts)
}

func (rt *vrt) Cleanup() {
	rt.mu.Lock()
	if rt.cleaningUp {
		rt.mu.Unlock()
		// TODO(caprita): Should we actually treat extra Cleanups as
		// silent no-ops?  For now, it's better not to, in order to
		// expose programming errors.
		vlog.Errorf("rt.Cleanup done")
		return
	}
	rt.cleaningUp = true
	rt.mu.Unlock()

	// TODO(caprita): Consider shutting down mgmt later in the runtime's
	// shutdown sequence, to capture some of the runtime internal shutdown
	// tasks in the task tracker.
	rt.mgmt.shutdown()
	// It's ok to access rt.sm out of lock below, since a Mutex acts as a
	// barrier in Go and hence we're guaranteed that cleaningUp is true at
	// this point.  The only code that mutates rt.sm is NewStreamManager in
	// ipc.go, which respects the value of cleaningUp.
	//
	// TODO(caprita): It would be cleaner if we do something like this
	// inside the lock (and then iterate over smgrs below):
	//   smgrs := rt.sm
	//   rt.sm = nil
	// However, to make that work we need to prevent NewClient and NewServer
	// from using rt.sm after cleaningUp has been set.  Hence the TODO.
	for _, sm := range rt.sm {
		sm.Shutdown()
	}
	rt.shutdownSignalHandling()
	rt.shutdownLogging()
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"github.com/bitly/go-nsq"
	"io"
	"log"
	"net"
	"strings"
)

// udp_server
func run_udp_server(port string, w *nsq.Writer, exitchan chan int) {
	udp_addr, err := net.ResolveUDPAddr("udp", port)
	if err != nil {
		log.Fatal("udp:", err)
	}
	server, err := net.ListenUDP("udp", udp_addr)
	if err != nil {
		log.Fatal("server bind failed:", err)
	}
	defer server.Close()
	rbuf := bufio.NewReader(server)
	for {
		select {
		case <-exitchan:
			return
		default:
			msg, err := rbuf.ReadString('\n')
			if err != nil && strings.Contains(err.Error(), "use of closed network connection") {
				return
			}
			if err == io.EOF {
				return
			}
			if err != nil {
				log.Fatal("read log failed", err)
				continue
			}
			w.Publish(topic, []byte(msg))
		}
	}
}
<commit_msg>fix udp<commit_after>package main

import (
	"github.com/bitly/go-nsq"
	"io"
	"log"
	"net"
	"strings"
)

// udp_server
func run_udp_server(port string, w *nsq.Writer, exitchan chan int) {
	udp_addr, err := net.ResolveUDPAddr("udp", port)
	if err != nil {
		log.Fatal("udp:", err)
	}
	server, err := net.ListenUDP("udp", udp_addr)
	if err != nil {
		log.Fatal("server bind failed:", err)
	}
	defer server.Close()
	buf := make([]byte, 8192)
	for {
		select {
		case <-exitchan:
			return
		default:
			size, _, err := server.ReadFromUDP(buf)
			if err != nil && strings.Contains(err.Error(), "use of closed network connection") {
				return
			}
			if err == io.EOF {
				return
			}
			if err != nil {
				log.Println("read log failed", err)
				continue
			}
			w.Publish(topic, []byte(buf[:size]))
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2018 The Ebiten 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.

// +build !js

package devicescale

import (
	"fmt"
	"syscall"
	"unsafe"
)

const (
	logPixelSx              = 88
	monitorDefaultToNearest = 2
	mdtEffectiveDpi         = 0
)

var (
	user32 = syscall.NewLazyDLL("user32")
	gdi32  = syscall.NewLazyDLL("gdi32")
	shcore = syscall.NewLazyDLL("shcore")
)

var (
	procSetProcessDPIAware = user32.NewProc("SetProcessDPIAware")
	procGetActiveWindow    = user32.NewProc("GetActiveWindow")
	procGetWindowDC        = user32.NewProc("GetWindowDC")
	procReleaseDC          = user32.NewProc("ReleaseDC")
	procMonitorFromWindow  = user32.NewProc("MonitorFromWindow")
	procGetMonitorInfo     = user32.NewProc("GetMonitorInfoW")

	procGetDeviceCaps = gdi32.NewProc("GetDeviceCaps")

	// GetScaleFactorForMonitor function can return unrelaiavle value (e.g. returning 180
	// for 200% scale). Use GetDpiForMonitor instead.
	procGetDpiForMonitor = shcore.NewProc("GetDpiForMonitor")
)

var shcoreAvailable = false

func init() {
	if shcore.Load() == nil {
		shcoreAvailable = true
	}
}

func setProcessDPIAware() error {
	r, _, e := syscall.Syscall(procSetProcessDPIAware.Addr(), 0, 0, 0, 0)
	if e != 0 {
		return fmt.Errorf("devicescale: SetProcessDPIAware failed: error code: %d", e)
	}
	if r == 0 {
		return fmt.Errorf("devicescale: SetProcessDPIAware failed: returned value: %d", r)
	}
	return nil
}

func getActiveWindow() (uintptr, error) {
	r, _, e := syscall.Syscall(procGetActiveWindow.Addr(), 0, 0, 0, 0)
	if e != 0 {
		return 0, fmt.Errorf("devicescale: GetActiveWindow failed: error code: %d", e)
	}
	return r, nil
}

func getWindowDC(hwnd uintptr) (uintptr, error) {
	r, _, e := syscall.Syscall(procGetWindowDC.Addr(), 1, hwnd, 0, 0)
	if e != 0 {
		return 0, fmt.Errorf("devicescale: GetWindowDC failed: error code: %d", e)
	}
	if r == 0 {
		return 0, fmt.Errorf("devicescale: GetWindowDC failed: returned value: %d", r)
	}
	return r, nil
}

func releaseDC(hwnd, hdc uintptr) error {
	r, _, e := syscall.Syscall(procReleaseDC.Addr(), 2, hwnd, hdc, 0)
	if e != 0 {
		return fmt.Errorf("devicescale: ReleaseDC failed: error code: %d", e)
	}
	if r == 0 {
		return fmt.Errorf("devicescale: ReleaseDC failed: returned value: %d", r)
	}
	return nil
}

func getDeviceCaps(hdc uintptr, nindex int) (int, error) {
	r, _, e := syscall.Syscall(procGetDeviceCaps.Addr(), 2, hdc, uintptr(nindex), 0)
	if e != 0 {
		return 0, fmt.Errorf("devicescale: GetDeviceCaps failed: error code: %d", e)
	}
	return int(r), nil
}

func monitorFromWindow(hwnd uintptr, dwFlags int) (uintptr, error) {
	r, _, e := syscall.Syscall(procMonitorFromWindow.Addr(), 2, hwnd, uintptr(dwFlags), 0)
	if e != 0 {
		return 0, fmt.Errorf("devicescale: MonitorFromWindow failed: error code: %d", e)
	}
	if r == 0 {
		return 0, fmt.Errorf("devicescale: MonitorFromWindow failed: returned value: %d", r)
	}
	return r, nil
}

func getMonitorInfo(hMonitor uintptr, lpMonitorInfo uintptr) error {
	r, _, e := syscall.Syscall(procGetMonitorInfo.Addr(), 2, hMonitor, lpMonitorInfo, 0)
	if e != 0 {
		return fmt.Errorf("devicescale: GetMonitorInfo failed: error code: %d", e)
	}
	if r == 0 {
		return fmt.Errorf("devicescale: GetMonitorInfo failed: returned value: %d", r)
	}
	return nil
}

func getDpiForMonitor(hMonitor uintptr, dpiType uintptr, dpiX, dpiY uintptr) error {
	r, _, e := syscall.Syscall6(procGetDpiForMonitor.Addr(), 4,
		hMonitor, dpiType, dpiX, dpiY, 0, 0)
	if e != 0 {
		return fmt.Errorf("devicescale: GetDpiForMonitor failed: error code: %d", e)
	}
	if r != 0 {
		return fmt.Errorf("devicescale: GetDpiForMonitor failed: returned value: %d", r)
	}
	return nil
}

func getFromLogPixelSx() float64 {
	dc, err := getWindowDC(0)
	if err != nil {
		panic(err)
	}
	defer releaseDC(0, dc)

	// Note that GetDeviceCaps with LOGPIXELSX always returns a same value for any monitors
	// even if multiple monitors are used.
	dpi, err := getDeviceCaps(dc, logPixelSx)
	if err != nil {
		panic(err)
	}

	return float64(dpi) / 96
}

func impl() float64 {
	if err := setProcessDPIAware(); err != nil {
		panic(err)
	}

	// On Windows 7 or older, shcore.dll is not available.
	if !shcoreAvailable {
		return getFromLogPixelSx()
	}

	w, err := getActiveWindow()
	if err != nil {
		panic(err)
	}
	// The window is not initialized yet when w == 0.
	if w == 0 {
		return getFromLogPixelSx()
	}
	defer releaseDC(0, w)

	m, err := monitorFromWindow(w, monitorDefaultToNearest)
	if err != nil {
		panic(err)
	}

	dpiX := uint32(0)
	dpiY := uint32(0) // Passing dpiY is needed even though this is not used.
	if err := getDpiForMonitor(m, mdtEffectiveDpi, uintptr(unsafe.Pointer(&dpiX)), uintptr(unsafe.Pointer(&dpiY))); err != nil {
		panic(err)
	}
	return float64(dpiX) / 96
}
<commit_msg>devicescale: Rename logPixelSx -> logPixelsX<commit_after>// Copyright 2018 The Ebiten 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.

// +build !js

package devicescale

import (
	"fmt"
	"syscall"
	"unsafe"
)

const (
	logPixelsX              = 88
	monitorDefaultToNearest = 2
	mdtEffectiveDpi         = 0
)

var (
	user32 = syscall.NewLazyDLL("user32")
	gdi32  = syscall.NewLazyDLL("gdi32")
	shcore = syscall.NewLazyDLL("shcore")
)

var (
	procSetProcessDPIAware = user32.NewProc("SetProcessDPIAware")
	procGetActiveWindow    = user32.NewProc("GetActiveWindow")
	procGetWindowDC        = user32.NewProc("GetWindowDC")
	procReleaseDC          = user32.NewProc("ReleaseDC")
	procMonitorFromWindow  = user32.NewProc("MonitorFromWindow")
	procGetMonitorInfo     = user32.NewProc("GetMonitorInfoW")

	procGetDeviceCaps = gdi32.NewProc("GetDeviceCaps")

	// GetScaleFactorForMonitor function can return unrelaiavle value (e.g. returning 180
	// for 200% scale). Use GetDpiForMonitor instead.
	procGetDpiForMonitor = shcore.NewProc("GetDpiForMonitor")
)

var shcoreAvailable = false

func init() {
	if shcore.Load() == nil {
		shcoreAvailable = true
	}
}

func setProcessDPIAware() error {
	r, _, e := syscall.Syscall(procSetProcessDPIAware.Addr(), 0, 0, 0, 0)
	if e != 0 {
		return fmt.Errorf("devicescale: SetProcessDPIAware failed: error code: %d", e)
	}
	if r == 0 {
		return fmt.Errorf("devicescale: SetProcessDPIAware failed: returned value: %d", r)
	}
	return nil
}

func getActiveWindow() (uintptr, error) {
	r, _, e := syscall.Syscall(procGetActiveWindow.Addr(), 0, 0, 0, 0)
	if e != 0 {
		return 0, fmt.Errorf("devicescale: GetActiveWindow failed: error code: %d", e)
	}
	return r, nil
}

func getWindowDC(hwnd uintptr) (uintptr, error) {
	r, _, e := syscall.Syscall(procGetWindowDC.Addr(), 1, hwnd, 0, 0)
	if e != 0 {
		return 0, fmt.Errorf("devicescale: GetWindowDC failed: error code: %d", e)
	}
	if r == 0 {
		return 0, fmt.Errorf("devicescale: GetWindowDC failed: returned value: %d", r)
	}
	return r, nil
}

func releaseDC(hwnd, hdc uintptr) error {
	r, _, e := syscall.Syscall(procReleaseDC.Addr(), 2, hwnd, hdc, 0)
	if e != 0 {
		return fmt.Errorf("devicescale: ReleaseDC failed: error code: %d", e)
	}
	if r == 0 {
		return fmt.Errorf("devicescale: ReleaseDC failed: returned value: %d", r)
	}
	return nil
}

func getDeviceCaps(hdc uintptr, nindex int) (int, error) {
	r, _, e := syscall.Syscall(procGetDeviceCaps.Addr(), 2, hdc, uintptr(nindex), 0)
	if e != 0 {
		return 0, fmt.Errorf("devicescale: GetDeviceCaps failed: error code: %d", e)
	}
	return int(r), nil
}

func monitorFromWindow(hwnd uintptr, dwFlags int) (uintptr, error) {
	r, _, e := syscall.Syscall(procMonitorFromWindow.Addr(), 2, hwnd, uintptr(dwFlags), 0)
	if e != 0 {
		return 0, fmt.Errorf("devicescale: MonitorFromWindow failed: error code: %d", e)
	}
	if r == 0 {
		return 0, fmt.Errorf("devicescale: MonitorFromWindow failed: returned value: %d", r)
	}
	return r, nil
}

func getMonitorInfo(hMonitor uintptr, lpMonitorInfo uintptr) error {
	r, _, e := syscall.Syscall(procGetMonitorInfo.Addr(), 2, hMonitor, lpMonitorInfo, 0)
	if e != 0 {
		return fmt.Errorf("devicescale: GetMonitorInfo failed: error code: %d", e)
	}
	if r == 0 {
		return fmt.Errorf("devicescale: GetMonitorInfo failed: returned value: %d", r)
	}
	return nil
}

func getDpiForMonitor(hMonitor uintptr, dpiType uintptr, dpiX, dpiY uintptr) error {
	r, _, e := syscall.Syscall6(procGetDpiForMonitor.Addr(), 4,
		hMonitor, dpiType, dpiX, dpiY, 0, 0)
	if e != 0 {
		return fmt.Errorf("devicescale: GetDpiForMonitor failed: error code: %d", e)
	}
	if r != 0 {
		return fmt.Errorf("devicescale: GetDpiForMonitor failed: returned value: %d", r)
	}
	return nil
}

func getFromLogPixelSx() float64 {
	dc, err := getWindowDC(0)
	if err != nil {
		panic(err)
	}
	defer releaseDC(0, dc)

	// Note that GetDeviceCaps with LOGPIXELSX always returns a same value for any monitors
	// even if multiple monitors are used.
	dpi, err := getDeviceCaps(dc, logPixelsX)
	if err != nil {
		panic(err)
	}

	return float64(dpi) / 96
}

func impl() float64 {
	if err := setProcessDPIAware(); err != nil {
		panic(err)
	}

	// On Windows 7 or older, shcore.dll is not available.
	if !shcoreAvailable {
		return getFromLogPixelSx()
	}

	w, err := getActiveWindow()
	if err != nil {
		panic(err)
	}
	// The window is not initialized yet when w == 0.
	if w == 0 {
		return getFromLogPixelSx()
	}
	defer releaseDC(0, w)

	m, err := monitorFromWindow(w, monitorDefaultToNearest)
	if err != nil {
		panic(err)
	}

	dpiX := uint32(0)
	dpiY := uint32(0) // Passing dpiY is needed even though this is not used.
	if err := getDpiForMonitor(m, mdtEffectiveDpi, uintptr(unsafe.Pointer(&dpiX)), uintptr(unsafe.Pointer(&dpiY))); err != nil {
		panic(err)
	}
	return float64(dpiX) / 96
}
<|endoftext|>


<commit_before>package jsonrpc

import (
	"encoding/json"
	"io/ioutil"

	"github.com/awslabs/aws-sdk-go/aws"
)

var emptyJSON = []byte("{}")

func Build(req *aws.Request) {
	var buf []byte
	var err error
	if req.ParamsFilled() {
		buf, err = json.Marshal(req.Params)
		if err != nil {
			req.Error = err
			return
		}
	} else {
		buf = emptyJSON
	}

	if req.Service.TargetPrefix != "" || string(buf) != "{}" {
		req.SetBufferBody(buf)
	}

	if req.Service.TargetPrefix != "" {
		target := req.Service.TargetPrefix + "." + req.Operation.Name
		req.HTTPRequest.Header.Add("X-Amz-Target", target)
	}
	if req.Service.JSONVersion != "" {
		jsonVersion := req.Service.JSONVersion
		req.HTTPRequest.Header.Add("Content-Type", "application/x-amz-json-"+jsonVersion)
	}
}

func Unmarshal(req *aws.Request) {
	defer req.HTTPResponse.Body.Close()
	if req.Data != nil {
		json.NewDecoder(req.HTTPResponse.Body).Decode(req.Data)
	}
	return
}

func UnmarshalMeta(req *aws.Request) {
	req.RequestID = req.HTTPResponse.Header.Get("x-amzn-requestid")
}

func UnmarshalError(req *aws.Request) {
	bodyBytes, err := ioutil.ReadAll(req.HTTPResponse.Body)
	if err != nil {
		req.Error = err
		return
	}
	if len(bodyBytes) == 0 {
		req.Error = aws.APIError{
			StatusCode: req.HTTPResponse.StatusCode,
			Message:    req.HTTPResponse.Status,
		}
		return
	}
	var jsonErr jsonErrorResponse
	if err := json.Unmarshal(bodyBytes, &jsonErr); err != nil {
		req.Error = err
		return
	}
	req.Error = jsonErr.Err(req.HTTPResponse.StatusCode)
}

type jsonErrorResponse struct {
	Code    string `json:"__type"`
	Message string `json:"message"`
}

func (e jsonErrorResponse) Err(StatusCode int) error {
	return aws.APIError{
		StatusCode: StatusCode,
		Code:       e.Code,
		Message:    e.Message,
	}
}
<commit_msg>Strip prefix off of error messages in JSON protocol<commit_after>package jsonrpc

import (
	"encoding/json"
	"io/ioutil"
	"strings"

	"github.com/awslabs/aws-sdk-go/aws"
)

var emptyJSON = []byte("{}")

func Build(req *aws.Request) {
	var buf []byte
	var err error
	if req.ParamsFilled() {
		buf, err = json.Marshal(req.Params)
		if err != nil {
			req.Error = err
			return
		}
	} else {
		buf = emptyJSON
	}

	if req.Service.TargetPrefix != "" || string(buf) != "{}" {
		req.SetBufferBody(buf)
	}

	if req.Service.TargetPrefix != "" {
		target := req.Service.TargetPrefix + "." + req.Operation.Name
		req.HTTPRequest.Header.Add("X-Amz-Target", target)
	}
	if req.Service.JSONVersion != "" {
		jsonVersion := req.Service.JSONVersion
		req.HTTPRequest.Header.Add("Content-Type", "application/x-amz-json-"+jsonVersion)
	}
}

func Unmarshal(req *aws.Request) {
	defer req.HTTPResponse.Body.Close()
	if req.Data != nil {
		json.NewDecoder(req.HTTPResponse.Body).Decode(req.Data)
	}
	return
}

func UnmarshalMeta(req *aws.Request) {
	req.RequestID = req.HTTPResponse.Header.Get("x-amzn-requestid")
}

func UnmarshalError(req *aws.Request) {
	bodyBytes, err := ioutil.ReadAll(req.HTTPResponse.Body)
	if err != nil {
		req.Error = err
		return
	}
	if len(bodyBytes) == 0 {
		req.Error = aws.APIError{
			StatusCode: req.HTTPResponse.StatusCode,
			Message:    req.HTTPResponse.Status,
		}
		return
	}
	var jsonErr jsonErrorResponse
	if err := json.Unmarshal(bodyBytes, &jsonErr); err != nil {
		req.Error = err
		return
	}

	codes := strings.SplitN(jsonErr.Code, "#", 2)
	req.Error = aws.APIError{
		StatusCode: req.HTTPResponse.StatusCode,
		Code:       codes[len(codes)-1],
		Message:    jsonErr.Message,
	}
}

type jsonErrorResponse struct {
	Code    string `json:"__type"`
	Message string `json:"message"`
}
<|endoftext|>


<commit_before>// Copyright 2015 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 internal contains code used internally by common/auth.
package internal

import (
	"bytes"
	"reflect"
	"strings"
	"sync"
	"time"

	"golang.org/x/net/context"
	"golang.org/x/oauth2"

	"go.chromium.org/luci/common/clock"
	"go.chromium.org/luci/common/data/rand/mathrand"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/retry/transient"
)

// expiryRandInterval is used by TokenExpiresInRnd.
const expiryRandInterval = 30 * time.Second

const (
	// NoEmail indicates an OAuth2 token is not associated with an email.
	//
	// See Token below. We need this special value to distinguish "an email can not
	// possibly be fetched ever" from "the cached token doesn't have an email yet"
	// cases.
	NoEmail = "-"

	// UnknownEmail indicates an OAuth2 token may potentially be associated with
	// an email, but we haven't tried to fetch the email yet.
	UnknownEmail = ""
)

var (
	// ErrInsufficientAccess is returned by MintToken() if token can't be minted
	// for given OAuth scopes. For example, if GCE instance wasn't granted access
	// to requested scopes when it was created.
	ErrInsufficientAccess = errors.New("can't get access token for the given account and scopes")

	// ErrBadRefreshToken is returned by RefreshToken if refresh token was revoked
	// or otherwise invalid. It means MintToken must be used to get a new refresh
	// token.
	ErrBadRefreshToken = errors.New("refresh_token is not valid")

	// ErrBadCredentials is returned by MintToken or RefreshToken if provided
	// offline credentials (like service account key) are invalid.
	ErrBadCredentials = errors.New("invalid service account credentials")
)

// Token is an oauth2.Token with an email that corresponds to it.
//
// Email may be an empty string, in which case we assume the email hasn't been
// fetched yet. It can also be a special NoEmail string, which means the token
// is not associated with an email (happens for tokens without 'userinfo.email'
// scope).
type Token struct {
	oauth2.Token

	Email string // an email or NoEmail or empty string (aka UnknownEmail)
}

// TokenProvider knows how to mint new tokens or refresh existing ones.
type TokenProvider interface {
	// RequiresInteraction is true if provider may start user interaction
	// in MintToken.
	RequiresInteraction() bool

	// Lightweight is true if MintToken is very cheap to call.
	//
	// In this case the token is not being cached on disk (only in memory), since
	// it's easy to get a new one each time the process starts.
	//
	// By avoiding the disk cache, we reduce the chance of a leak.
	Lightweight() bool

	// Email is email associated with tokens produced by the provider, if known.
	//
	// May return UnknownEmail, which means the provider doesn't know the email
	// in advance and RefreshToken must be used to get the token and the email.
	// This happens, for example, for interactive providers before user has
	// logged in.
	//
	// It can also be NoEmail which means the email is not available, even if
	// caller is using RefreshToken.
	Email() string

	// CacheKey identifies a slot in the token cache to store the token in.
	//
	// Note: CacheKey MAY change during lifetime of a TokenProvider. It happens,
	// for example, for ServiceAccount token provider if the underlying service
	// account key is replaced while the process is still running.
	CacheKey(ctx context.Context) (*CacheKey, error)

	// MintToken launches authentication flow (possibly interactive) and returns
	// a new refreshable token (or error). It must never return (nil, nil).
	//
	// In actor mode 'base' is an IAM-scoped sufficiently fresh oauth token. It's
	// nil otherwise. Used by IAM-based token provider.
	MintToken(ctx context.Context, base *Token) (*Token, error)

	// RefreshToken takes existing token (probably expired, but not necessarily)
	// and returns a new refreshed token. It should never do any user interaction.
	// If a user interaction is required, a error should be returned instead.
	//
	// In actor mode 'base' is an IAM-scoped sufficiently fresh oauth token. It's
	// nil otherwise. Used by IAM-based token provider.
	RefreshToken(ctx context.Context, prev, base *Token) (*Token, error)
}

// TokenCache stores access and refresh tokens to avoid requesting them all
// the time.
type TokenCache interface {
	// GetToken reads the token from cache.
	//
	// Returns (nil, nil) if requested token is not in the cache.
	GetToken(key *CacheKey) (*Token, error)

	// PutToken writes the token to cache.
	PutToken(key *CacheKey, tok *Token) error

	// DeleteToken removes the token from cache.
	DeleteToken(key *CacheKey) error
}

// CacheKey identifies a slot in the token cache to store the token in.
type CacheKey struct {
	// Key identifies an auth method being used to get the token and its
	// parameters.
	//
	// Its exact form is not important, since it is used only for string matching
	// when searching for a token inside the cache.
	//
	// The following forms are being used currently:
	//  * user/<client_id> when using UserCredentialsMethod with some ClientID.
	//  * service_account/<email>/<key_id> when using ServiceAccountMethod.
	//  * gce/<account> when using GCEMetadataMethod.
	//  * iam/<account> when using actor mode with ActAsServiceAccount != "".
	//  * luci_ctx/<digest> when using LUCIContextMethod.
	Key string `json:"key"`

	// Scopes is the list of requested OAuth scopes.
	Scopes []string `json:"scopes"`
}

var bufPool = sync.Pool{}

// ToMapKey returns a string that can be used as map[string] key.
//
// This string IS NOT PRINTABLE. It's a merely a string-looking []byte.
func (k *CacheKey) ToMapKey() string {
	b, _ := bufPool.Get().(*bytes.Buffer)
	if b == nil {
		b = &bytes.Buffer{}
	} else {
		b.Reset()
	}
	defer bufPool.Put(b)
	b.WriteString(k.Key)
	b.WriteByte(0)
	for _, s := range k.Scopes {
		b.WriteString(s)
		b.WriteByte(0)
	}
	return b.String()
}

// EqualCacheKeys returns true if keys are equal.
func EqualCacheKeys(a, b *CacheKey) bool {
	return reflect.DeepEqual(a, b)
}

// TokenExpiresIn returns True if the token is not valid or expires within given
// duration.
//
// The function returns True in any of the following conditions:
//  * The token is not valid.
//  * The token expires before now+lifetime.
//
// In all other cases it returns False.
func TokenExpiresIn(ctx context.Context, t *Token, lifetime time.Duration) bool {
	if t == nil || t.AccessToken == "" {
		return true
	}
	if t.Expiry.IsZero() {
		return false
	}
	return t.Expiry.Before(clock.Now(ctx).Add(lifetime))
}

// TokenExpiresInRnd is like TokenExpiresIn, except it slightly randomizes the
// token expiration time.
//
// If the function returns False, the token expires past now+lifetime. In other
// words, it is totally safe to use the token until now+lifetime. The inverse of
// this statement is not correct though: if the function returns True, it
// doesn't necessarily imply the token will expire before now+lifetime.
//
// The function returns True in any of the following conditions:
//  * The token is not valid.
//  * The token expires before now+lifetime.
//  * The token expiration time is between (now+lifetime, now+lifetime+rnd),
//    where rnd is a uniformly distributed random number between 0 and
//    expiryRandInterval sec (which is set to 30 sec).
//
// This is useful for processes that use multiple service account keys at
// around the same time. Without randomization, access tokens for such keys
// expire at the same time (strictly 1h after process startup, where 1h is
// the default token lifetime). This causes unnecessary contention on the token
// cache file.
func TokenExpiresInRnd(ctx context.Context, t *Token, lifetime time.Duration) bool {
	if t == nil || t.AccessToken == "" {
		return true
	}
	if t.Expiry.IsZero() {
		return false
	}
	deadline := clock.Now(ctx).Add(lifetime)
	if t.Expiry.Before(deadline) {
		// Definitely expires within 'lifetime'.
		return true
	}
	if t.Expiry.After(deadline.Add(expiryRandInterval)) {
		// Definitely expires much later than 'lifetime', no need to involve RNG.
		return false
	}
	// Semi-randomly declare it as expired.
	rnd := time.Duration(mathrand.Int63n(ctx, int64(expiryRandInterval)))
	return t.Expiry.Before(deadline.Add(rnd))
}

// EqualTokens returns true if tokens are equal.
//
// 'nil' token corresponds to an empty access token.
func EqualTokens(a, b *Token) bool {
	if a == b {
		return true
	}
	if a == nil {
		a = &Token{}
	}
	if b == nil {
		b = &Token{}
	}
	return a.AccessToken == b.AccessToken && a.Expiry.Equal(b.Expiry) && a.Email == b.Email
}

// isBadTokenError sniffs out HTTP 400/401 from token source errors.
func isBadTokenError(err error) bool {
	// See https://github.com/golang/oauth2/blob/master/internal/token.go.
	// Unfortunately, fmt.Errorf is used there, so there's no other way to
	// differentiate between bad tokens/creds and transient errors.
	// Note that mis-categorization is not a big deal: we'll unnecessarily retry
	// fatal error a bunch of times, thinking it is transient.
	if err == nil {
		return false
	}
	s := strings.ToLower(err.Error())
	return strings.Contains(s, "400 bad request") || strings.Contains(s, "401 unauthorized")
}

// isBadKeyError sniffs out errors related to malformed private keys.
func isBadKeyError(err error) bool {
	// See https://github.com/golang/oauth2/blob/master/internal/oauth2.go.
	// Unfortunately, if also uses fmt.Errorf.
	// Note that mis-categorization is not a big deal: we'll unnecessarily retry
	// fatal error a bunch of times, thinking it is transient.
	if err == nil {
		return false
	}
	s := err.Error()
	return strings.Contains(s, "private key should be a PEM or plain PKSC1 or PKCS8") || s == "private key is invalid"
}

// grabToken uses token source to create a new *oauth2.Token.
//
// It recognizes transient errors.
func grabToken(src oauth2.TokenSource) (*oauth2.Token, error) {
	switch tok, err := src.Token(); {
	case isBadTokenError(err):
		return nil, err
	case isBadKeyError(err):
		return nil, err
	case err != nil:
		// More often than not errors here are transient (network connectivity
		// errors, HTTP 500 responses, etc). It is difficult to categorize them,
		// since oauth2 library uses fmt.Errorf(...) for errors. Retrying a fatal
		// error a bunch of times is not very bad, so pick safer approach and assume
		// any error is transient. Revoked refresh token or bad credentials (most
		// common source of fatal errors) is already handled above.
		return nil, transient.Tag.Apply(err)
	default:
		return tok, nil
	}
}
<commit_msg>[auth] Use oauth2.RetrieveError instead of string parsing, where possible.<commit_after>// Copyright 2015 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 internal contains code used internally by common/auth.
package internal

import (
	"bytes"
	"reflect"
	"strings"
	"sync"
	"time"

	"golang.org/x/net/context"
	"golang.org/x/oauth2"

	"go.chromium.org/luci/common/clock"
	"go.chromium.org/luci/common/data/rand/mathrand"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/retry/transient"
)

// expiryRandInterval is used by TokenExpiresInRnd.
const expiryRandInterval = 30 * time.Second

const (
	// NoEmail indicates an OAuth2 token is not associated with an email.
	//
	// See Token below. We need this special value to distinguish "an email can not
	// possibly be fetched ever" from "the cached token doesn't have an email yet"
	// cases.
	NoEmail = "-"

	// UnknownEmail indicates an OAuth2 token may potentially be associated with
	// an email, but we haven't tried to fetch the email yet.
	UnknownEmail = ""
)

var (
	// ErrInsufficientAccess is returned by MintToken() if token can't be minted
	// for given OAuth scopes. For example, if GCE instance wasn't granted access
	// to requested scopes when it was created.
	ErrInsufficientAccess = errors.New("can't get access token for the given account and scopes")

	// ErrBadRefreshToken is returned by RefreshToken if refresh token was revoked
	// or otherwise invalid. It means MintToken must be used to get a new refresh
	// token.
	ErrBadRefreshToken = errors.New("refresh_token is not valid")

	// ErrBadCredentials is returned by MintToken or RefreshToken if provided
	// offline credentials (like service account key) are invalid.
	ErrBadCredentials = errors.New("invalid service account credentials")
)

// Token is an oauth2.Token with an email that corresponds to it.
//
// Email may be an empty string, in which case we assume the email hasn't been
// fetched yet. It can also be a special NoEmail string, which means the token
// is not associated with an email (happens for tokens without 'userinfo.email'
// scope).
type Token struct {
	oauth2.Token

	Email string // an email or NoEmail or empty string (aka UnknownEmail)
}

// TokenProvider knows how to mint new tokens or refresh existing ones.
type TokenProvider interface {
	// RequiresInteraction is true if provider may start user interaction
	// in MintToken.
	RequiresInteraction() bool

	// Lightweight is true if MintToken is very cheap to call.
	//
	// In this case the token is not being cached on disk (only in memory), since
	// it's easy to get a new one each time the process starts.
	//
	// By avoiding the disk cache, we reduce the chance of a leak.
	Lightweight() bool

	// Email is email associated with tokens produced by the provider, if known.
	//
	// May return UnknownEmail, which means the provider doesn't know the email
	// in advance and RefreshToken must be used to get the token and the email.
	// This happens, for example, for interactive providers before user has
	// logged in.
	//
	// It can also be NoEmail which means the email is not available, even if
	// caller is using RefreshToken.
	Email() string

	// CacheKey identifies a slot in the token cache to store the token in.
	//
	// Note: CacheKey MAY change during lifetime of a TokenProvider. It happens,
	// for example, for ServiceAccount token provider if the underlying service
	// account key is replaced while the process is still running.
	CacheKey(ctx context.Context) (*CacheKey, error)

	// MintToken launches authentication flow (possibly interactive) and returns
	// a new refreshable token (or error). It must never return (nil, nil).
	//
	// In actor mode 'base' is an IAM-scoped sufficiently fresh oauth token. It's
	// nil otherwise. Used by IAM-based token provider.
	MintToken(ctx context.Context, base *Token) (*Token, error)

	// RefreshToken takes existing token (probably expired, but not necessarily)
	// and returns a new refreshed token. It should never do any user interaction.
	// If a user interaction is required, a error should be returned instead.
	//
	// In actor mode 'base' is an IAM-scoped sufficiently fresh oauth token. It's
	// nil otherwise. Used by IAM-based token provider.
	RefreshToken(ctx context.Context, prev, base *Token) (*Token, error)
}

// TokenCache stores access and refresh tokens to avoid requesting them all
// the time.
type TokenCache interface {
	// GetToken reads the token from cache.
	//
	// Returns (nil, nil) if requested token is not in the cache.
	GetToken(key *CacheKey) (*Token, error)

	// PutToken writes the token to cache.
	PutToken(key *CacheKey, tok *Token) error

	// DeleteToken removes the token from cache.
	DeleteToken(key *CacheKey) error
}

// CacheKey identifies a slot in the token cache to store the token in.
type CacheKey struct {
	// Key identifies an auth method being used to get the token and its
	// parameters.
	//
	// Its exact form is not important, since it is used only for string matching
	// when searching for a token inside the cache.
	//
	// The following forms are being used currently:
	//  * user/<client_id> when using UserCredentialsMethod with some ClientID.
	//  * service_account/<email>/<key_id> when using ServiceAccountMethod.
	//  * gce/<account> when using GCEMetadataMethod.
	//  * iam/<account> when using actor mode with ActAsServiceAccount != "".
	//  * luci_ctx/<digest> when using LUCIContextMethod.
	Key string `json:"key"`

	// Scopes is the list of requested OAuth scopes.
	Scopes []string `json:"scopes"`
}

var bufPool = sync.Pool{}

// ToMapKey returns a string that can be used as map[string] key.
//
// This string IS NOT PRINTABLE. It's a merely a string-looking []byte.
func (k *CacheKey) ToMapKey() string {
	b, _ := bufPool.Get().(*bytes.Buffer)
	if b == nil {
		b = &bytes.Buffer{}
	} else {
		b.Reset()
	}
	defer bufPool.Put(b)
	b.WriteString(k.Key)
	b.WriteByte(0)
	for _, s := range k.Scopes {
		b.WriteString(s)
		b.WriteByte(0)
	}
	return b.String()
}

// EqualCacheKeys returns true if keys are equal.
func EqualCacheKeys(a, b *CacheKey) bool {
	return reflect.DeepEqual(a, b)
}

// TokenExpiresIn returns True if the token is not valid or expires within given
// duration.
//
// The function returns True in any of the following conditions:
//  * The token is not valid.
//  * The token expires before now+lifetime.
//
// In all other cases it returns False.
func TokenExpiresIn(ctx context.Context, t *Token, lifetime time.Duration) bool {
	if t == nil || t.AccessToken == "" {
		return true
	}
	if t.Expiry.IsZero() {
		return false
	}
	return t.Expiry.Before(clock.Now(ctx).Add(lifetime))
}

// TokenExpiresInRnd is like TokenExpiresIn, except it slightly randomizes the
// token expiration time.
//
// If the function returns False, the token expires past now+lifetime. In other
// words, it is totally safe to use the token until now+lifetime. The inverse of
// this statement is not correct though: if the function returns True, it
// doesn't necessarily imply the token will expire before now+lifetime.
//
// The function returns True in any of the following conditions:
//  * The token is not valid.
//  * The token expires before now+lifetime.
//  * The token expiration time is between (now+lifetime, now+lifetime+rnd),
//    where rnd is a uniformly distributed random number between 0 and
//    expiryRandInterval sec (which is set to 30 sec).
//
// This is useful for processes that use multiple service account keys at
// around the same time. Without randomization, access tokens for such keys
// expire at the same time (strictly 1h after process startup, where 1h is
// the default token lifetime). This causes unnecessary contention on the token
// cache file.
func TokenExpiresInRnd(ctx context.Context, t *Token, lifetime time.Duration) bool {
	if t == nil || t.AccessToken == "" {
		return true
	}
	if t.Expiry.IsZero() {
		return false
	}
	deadline := clock.Now(ctx).Add(lifetime)
	if t.Expiry.Before(deadline) {
		// Definitely expires within 'lifetime'.
		return true
	}
	if t.Expiry.After(deadline.Add(expiryRandInterval)) {
		// Definitely expires much later than 'lifetime', no need to involve RNG.
		return false
	}
	// Semi-randomly declare it as expired.
	rnd := time.Duration(mathrand.Int63n(ctx, int64(expiryRandInterval)))
	return t.Expiry.Before(deadline.Add(rnd))
}

// EqualTokens returns true if tokens are equal.
//
// 'nil' token corresponds to an empty access token.
func EqualTokens(a, b *Token) bool {
	if a == b {
		return true
	}
	if a == nil {
		a = &Token{}
	}
	if b == nil {
		b = &Token{}
	}
	return a.AccessToken == b.AccessToken && a.Expiry.Equal(b.Expiry) && a.Email == b.Email
}

// isBadTokenError sniffs out HTTP 400/401 from token source errors.
func isBadTokenError(err error) bool {
	if rerr, _ := err.(*oauth2.RetrieveError); rerr != nil {
		return rerr.Response.StatusCode == 400 || rerr.Response.StatusCode == 401
	}
	return false
}

// isBadKeyError sniffs out errors related to malformed private keys.
func isBadKeyError(err error) bool {
	if err == nil {
		return false
	}
	// See https://go.googlesource.com/oauth2.git/+/197281d4/internal/oauth2.go#32
	// Unfortunately, if uses fmt.Errorf.
	s := err.Error()
	return strings.Contains(s, "private key should be a PEM or plain PKSC1 or PKCS8") ||
		s == "private key is invalid"
}

// grabToken uses token source to create a new *oauth2.Token.
//
// It recognizes transient errors.
func grabToken(src oauth2.TokenSource) (*oauth2.Token, error) {
	switch tok, err := src.Token(); {
	case isBadTokenError(err):
		return nil, err
	case isBadKeyError(err):
		return nil, err
	case err != nil:
		// More often than not errors here are transient (network connectivity
		// errors, HTTP 500 responses, etc). Retrying a fatal error a bunch of times
		// is not very bad, so pick safer approach and assume any error is
		// transient. Revoked refresh token or bad credentials (most common source
		// of fatal errors) is already handled above.
		return nil, transient.Tag.Apply(err)
	default:
		return tok, nil
	}
}
<|endoftext|>


<commit_before>/*
Copyright IBM Corp. 2017 All Rights Reserved.

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 (
	"flag"
	"fmt"
	"io"
	"os"
	"path/filepath"

	"github.com/hyperledger/fabric/common/tools/cryptogen/ca"
	"github.com/hyperledger/fabric/common/tools/cryptogen/msp"
)

const (
	peerOrgBaseName  = "peerOrg"
	peerBaseName     = "Peer"
	userBaseName     = "User"
	adminBaseName    = "Admin"
	orderOrgBaseName = "ordererOrg"
	ordererBaseName  = "Orderer"
)

//command line flags
var (
	numPeerOrgs = flag.Int("peerOrgs", 2,
		"number of unique organizations with peers")
	numPeers = flag.Int("peersPerOrg", 1,
		"number of peers per organization")
	numPeerOrgUsers = flag.Int("peerOrgUsers", 1,
		"number of users per peer organization")
	numOrderers = flag.Int("ordererNodes", 1,
		"number of ordering service nodes")
	baseDir = flag.String("baseDir", ".",
		"directory in which to place artifacts")
)

var numOrdererOrgs = 1

func main() {
	flag.Parse()

	if flag.NFlag() == 0 {
		fmt.Println("\nYou must specify at least one parameter")
		flag.Usage()
		os.Exit(1)
	}

	genDir := filepath.Join(*baseDir, "crypto-config")
	if *numPeerOrgs > 0 {
		fmt.Printf("Generating %d peer organization(s) each with %d peer(s) ...\n",
			*numPeerOrgs, *numPeers)

		// TODO: add ability to specify peer org names
		// for name just use default base name
		peerOrgNames := []string{}
		for i := 1; i <= *numPeerOrgs; i++ {
			peerOrgNames = append(peerOrgNames, fmt.Sprintf("%s%d", peerOrgBaseName, i))
		}
		generatePeerOrgs(genDir, peerOrgNames)

	}

	if *numOrderers > 0 {
		fmt.Printf("Generating %d orderer organization(s) and %d ordering node(s) ...\n",
			numOrdererOrgs, *numOrderers)
		generateOrdererOrg(genDir, fmt.Sprintf("%s1", orderOrgBaseName))
	}

}

func generatePeerOrgs(baseDir string, orgNames []string) {

	for _, orgName := range orgNames {
		fmt.Println(orgName)
		// generate CA
		orgDir := filepath.Join(baseDir, "peerOrganizations", orgName)
		caDir := filepath.Join(orgDir, "ca")
		mspDir := filepath.Join(orgDir, "msp")
		peersDir := filepath.Join(orgDir, "peers")
		usersDir := filepath.Join(orgDir, "users")
		adminCertsDir := filepath.Join(mspDir, "admincerts")
		rootCA, err := ca.NewCA(caDir, orgName)
		if err != nil {
			fmt.Printf("Error generating CA for org %s:\n%v\n", orgName, err)
			os.Exit(1)
		}
		err = msp.GenerateVerifyingMSP(mspDir, rootCA)
		if err != nil {
			fmt.Printf("Error generating MSP for org %s:\n%v\n", orgName, err)
			os.Exit(1)
		}

		// TODO: add ability to specify peer names
		// for name just use default base name
		peerNames := []string{}
		for i := 1; i <= *numPeers; i++ {
			peerNames = append(peerNames, fmt.Sprintf("%s%s%d",
				orgName, peerBaseName, i))
		}
		generateNodes(peersDir, peerNames, rootCA)

		// TODO: add ability to specify usernames
		usernames := []string{}
		for j := 1; j <= *numPeerOrgUsers; j++ {
			usernames = append(usernames, fmt.Sprintf("%s%s%d",
				orgName, userBaseName, j))
		}
		// add an admin user
		usernames = append(usernames, fmt.Sprintf("%s%s",
			orgName, adminBaseName))
		generateNodes(usersDir, usernames, rootCA)

		// copy the admin cert to the org's MSP admincerts
		adminUserName := fmt.Sprintf("%s%s",
			orgName, adminBaseName)
		err = copyAdminCert(usersDir, adminCertsDir, adminUserName)
		if err != nil {
			fmt.Printf("Error copying admin cert for org %s:\n%v\n",
				orgName, err)
			os.Exit(1)
		}

	}
}

func copyAdminCert(usersDir, adminCertsDir, adminUserName string) error {
	// delete the contents of admincerts
	err := os.RemoveAll(adminCertsDir)
	if err != nil {
		return err
	}
	// recreate the admincerts directory
	err = os.MkdirAll(adminCertsDir, 0755)
	if err != nil {
		return err
	}
	err = copyFile(filepath.Join(usersDir, adminUserName, "signcerts",
		adminUserName+"-cert.pem"), filepath.Join(adminCertsDir,
		adminUserName+"-cert.pem"))
	if err != nil {
		return err
	}
	return nil

}

func generateNodes(baseDir string, nodeNames []string, rootCA *ca.CA) {

	for _, nodeName := range nodeNames {
		nodeDir := filepath.Join(baseDir, nodeName)
		err := msp.GenerateLocalMSP(nodeDir, nodeName, rootCA)
		if err != nil {
			fmt.Printf("Error generating local MSP for %s:\n%v\n", nodeName, err)
			os.Exit(1)
		}
	}

}

func generateOrdererOrg(baseDir, orgName string) {

	// generate CA
	orgDir := filepath.Join(baseDir, "ordererOrganizations", orgName)
	caDir := filepath.Join(orgDir, "ca")
	mspDir := filepath.Join(orgDir, "msp")
	orderersDir := filepath.Join(orgDir, "orderers")
	usersDir := filepath.Join(orgDir, "users")
	adminCertsDir := filepath.Join(mspDir, "admincerts")
	rootCA, err := ca.NewCA(caDir, orgName)
	if err != nil {
		fmt.Printf("Error generating CA for org %s:\n%v\n", orgName, err)
		os.Exit(1)
	}
	err = msp.GenerateVerifyingMSP(mspDir, rootCA)
	if err != nil {
		fmt.Printf("Error generating MSP for org %s:\n%v\n", orgName, err)
		os.Exit(1)
	}

	// TODO: add ability to specify orderer names
	// for name just use default base name
	ordererNames := []string{}
	for i := 1; i <= *numOrderers; i++ {
		ordererNames = append(ordererNames, fmt.Sprintf("%s%s%d",
			orgName, ordererBaseName, i))
	}
	generateNodes(orderersDir, ordererNames, rootCA)

	// generate an admin for the orderer org
	usernames := []string{}
	// add an admin user
	usernames = append(usernames, fmt.Sprintf("%s%s",
		orgName, adminBaseName))
	generateNodes(usersDir, usernames, rootCA)

	// copy the admin cert to the org's MSP admincerts
	adminUserName := fmt.Sprintf("%s%s",
		orgName, adminBaseName)
	err = copyAdminCert(usersDir, adminCertsDir, adminUserName)
	if err != nil {
		fmt.Printf("Error copying admin cert for org %s:\n%v\n",
			orgName, err)
		os.Exit(1)
	}

}

func copyFile(src, dst string) error {
	in, err := os.Open(src)
	if err != nil {
		return err
	}
	defer in.Close()
	out, err := os.Create(dst)
	if err != nil {
		return err
	}
	defer out.Close()
	_, err = io.Copy(out, in)
	cerr := out.Close()
	if err != nil {
		return err
	}
	return cerr
}
<commit_msg>[FAB-3098] Add org name options to cryptogen<commit_after>/*
Copyright IBM Corp. 2017 All Rights Reserved.

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 (
	"flag"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"

	"github.com/hyperledger/fabric/common/tools/cryptogen/ca"
	"github.com/hyperledger/fabric/common/tools/cryptogen/msp"
)

const (
	peerOrgBaseName    = "peerOrg"
	peerBaseName       = "Peer"
	userBaseName       = "User"
	adminBaseName      = "Admin"
	ordererOrgBaseName = "ordererOrg"
	ordererBaseName    = "Orderer"
	numOrdererOrgs     = 1
)

//command line flags
var (
	numPeerOrgs = flag.Int("peerOrgs", 2,
		"number of unique organizations with peers, used if peerOrgNames is not specified")
	peerOrgNames = flag.String("peerOrgNames", "",
		"comma-separated list of peer organization names")
	numPeers = flag.Int("peersPerOrg", 1,
		"number of peers per organization")
	numPeerOrgUsers = flag.Int("peerOrgUsers", 1,
		"number of users per peer organization")

	numOrderers = flag.Int("ordererNodes", 1,
		"number of ordering service nodes per organization")
	ordererOrgName = flag.String("ordererOrgName", "",
		"orderer organization name")

	baseDir = flag.String("baseDir", ".",
		"directory in which to place artifacts")
)

func main() {

	flag.Parse()

	if flag.NFlag() == 0 {
		fmt.Println("\nYou must specify at least one parameter")
		flag.Usage()
		os.Exit(1)
	}

	genDir := filepath.Join(*baseDir, "crypto-config")
	if *peerOrgNames != "" {
		generatePeerOrgs(genDir, strings.Split(*peerOrgNames, ","))
	} else if *numPeerOrgs > 0 {
		fmt.Printf("Generating %d peer organization(s) each with %d peer(s) ...\n",
			*numPeerOrgs, *numPeers)
		generatePeerOrgs(genDir, getOrgNames(peerOrgBaseName, *numPeerOrgs))
	}

	if *ordererOrgName != "" {
		fmt.Printf("Generating orderer organization %s with %d ordering node(s) ...\n",
			*ordererOrgName, *numOrderers)
		generateOrdererOrg(genDir, *ordererOrgName)
	} else if numOrdererOrgs > 0 {
		fmt.Printf("Generating %d orderer organization(s) and %d ordering node(s) ...\n",
			numOrdererOrgs, *numOrderers)
		generateOrdererOrg(genDir, fmt.Sprintf("%s1", ordererOrgBaseName))
	}

}

func generatePeerOrgs(baseDir string, orgNames []string) {

	for _, orgName := range orgNames {
		fmt.Println(orgName)
		// generate CA
		orgDir := filepath.Join(baseDir, "peerOrganizations", orgName)
		caDir := filepath.Join(orgDir, "ca")
		mspDir := filepath.Join(orgDir, "msp")
		peersDir := filepath.Join(orgDir, "peers")
		usersDir := filepath.Join(orgDir, "users")
		adminCertsDir := filepath.Join(mspDir, "admincerts")
		rootCA, err := ca.NewCA(caDir, orgName)
		if err != nil {
			fmt.Printf("Error generating CA for org %s:\n%v\n", orgName, err)
			os.Exit(1)
		}
		err = msp.GenerateVerifyingMSP(mspDir, rootCA)
		if err != nil {
			fmt.Printf("Error generating MSP for org %s:\n%v\n", orgName, err)
			os.Exit(1)
		}

		// TODO: add ability to specify peer names
		// for name just use default base name
		peerNames := []string{}
		for i := 1; i <= *numPeers; i++ {
			peerNames = append(peerNames, fmt.Sprintf("%s%s%d",
				orgName, peerBaseName, i))
		}
		generateNodes(peersDir, peerNames, rootCA)

		// TODO: add ability to specify usernames
		usernames := []string{}
		for j := 1; j <= *numPeerOrgUsers; j++ {
			usernames = append(usernames, fmt.Sprintf("%s%s%d",
				orgName, userBaseName, j))
		}
		// add an admin user
		usernames = append(usernames, fmt.Sprintf("%s%s",
			orgName, adminBaseName))
		generateNodes(usersDir, usernames, rootCA)

		// copy the admin cert to the org's MSP admincerts
		adminUserName := fmt.Sprintf("%s%s",
			orgName, adminBaseName)
		err = copyAdminCert(usersDir, adminCertsDir, adminUserName)
		if err != nil {
			fmt.Printf("Error copying admin cert for org %s:\n%v\n",
				orgName, err)
			os.Exit(1)
		}

	}
}

func getOrgNames(baseName string, count int) []string {
	orgNames := []string{}
	for i := 1; i <= count; i++ {
		orgNames = append(orgNames, fmt.Sprintf("%s%d", baseName, i))
	}
	return orgNames
}

func copyAdminCert(usersDir, adminCertsDir, adminUserName string) error {
	// delete the contents of admincerts
	err := os.RemoveAll(adminCertsDir)
	if err != nil {
		return err
	}
	// recreate the admincerts directory
	err = os.MkdirAll(adminCertsDir, 0755)
	if err != nil {
		return err
	}
	err = copyFile(filepath.Join(usersDir, adminUserName, "signcerts",
		adminUserName+"-cert.pem"), filepath.Join(adminCertsDir,
		adminUserName+"-cert.pem"))
	if err != nil {
		return err
	}
	return nil

}

func generateNodes(baseDir string, nodeNames []string, rootCA *ca.CA) {

	for _, nodeName := range nodeNames {
		nodeDir := filepath.Join(baseDir, nodeName)
		err := msp.GenerateLocalMSP(nodeDir, nodeName, rootCA)
		if err != nil {
			fmt.Printf("Error generating local MSP for %s:\n%v\n", nodeName, err)
			os.Exit(1)
		}
	}

}

func generateOrdererOrg(baseDir, orgName string) {

	// generate CA
	orgDir := filepath.Join(baseDir, "ordererOrganizations", orgName)
	caDir := filepath.Join(orgDir, "ca")
	mspDir := filepath.Join(orgDir, "msp")
	orderersDir := filepath.Join(orgDir, "orderers")
	usersDir := filepath.Join(orgDir, "users")
	adminCertsDir := filepath.Join(mspDir, "admincerts")
	rootCA, err := ca.NewCA(caDir, orgName)
	if err != nil {
		fmt.Printf("Error generating CA for org %s:\n%v\n", orgName, err)
		os.Exit(1)
	}
	err = msp.GenerateVerifyingMSP(mspDir, rootCA)
	if err != nil {
		fmt.Printf("Error generating MSP for org %s:\n%v\n", orgName, err)
		os.Exit(1)
	}

	// TODO: add ability to specify orderer names
	// for name just use default base name
	ordererNames := []string{}
	for i := 1; i <= *numOrderers; i++ {
		ordererNames = append(ordererNames, fmt.Sprintf("%s%s%d",
			orgName, ordererBaseName, i))
	}
	generateNodes(orderersDir, ordererNames, rootCA)

	// generate an admin for the orderer org
	usernames := []string{}
	// add an admin user
	usernames = append(usernames, fmt.Sprintf("%s%s",
		orgName, adminBaseName))
	generateNodes(usersDir, usernames, rootCA)

	// copy the admin cert to the org's MSP admincerts
	adminUserName := fmt.Sprintf("%s%s",
		orgName, adminBaseName)
	err = copyAdminCert(usersDir, adminCertsDir, adminUserName)
	if err != nil {
		fmt.Printf("Error copying admin cert for org %s:\n%v\n",
			orgName, err)
		os.Exit(1)
	}

}

func copyFile(src, dst string) error {
	in, err := os.Open(src)
	if err != nil {
		return err
	}
	defer in.Close()
	out, err := os.Create(dst)
	if err != nil {
		return err
	}
	defer out.Close()
	_, err = io.Copy(out, in)
	cerr := out.Close()
	if err != nil {
		return err
	}
	return cerr
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"net"
	"os"
	"strings"

	"github.com/godbus/dbus"
	"github.com/oleksandr/bonjour"
)

const avahiDbusName string = "org.freedesktop.Avahi"
const avahiInterfaceServer string = avahiDbusName + ".Server"
const avahiInterfaceEntrygroup string = avahiDbusName + ".EntryGroup"
const avahiEntrygroupAddService string = avahiInterfaceEntrygroup + ".AddService"
const avahiEntrygroupCommit = avahiInterfaceEntrygroup + ".Commit"

const avahiIfUnspec = int32(-1)
const protoUnspec = int32(-1)
const protoInet = int32(0)
const protoInet6 = int32(1)

func registerAvahiService(name, serviceType string, serverPort uint16, txt []string) bool {
	conn, err := dbus.SystemBus()

	if err != nil {
		fmt.Fprintln(os.Stderr, "Failed to connect to D-BUS session bus:", err)
		return false
	}

	obj := conn.Object(avahiDbusName, "/")

	var groupPath dbus.ObjectPath
	call := obj.Call("org.freedesktop.Avahi.Server.EntryGroupNew", 0)
	if call.Err != nil {
		fmt.Fprintln(os.Stderr, "Avahi not found on D-BUS")
		return false
	}
	call.Store(&groupPath)

	group := conn.Object(avahiDbusName, groupPath)

	var txtBytes [][]byte
	for _, s := range txt {
		txtBytes = append(txtBytes, []byte(s))
	}

	var flags = uint32(0)
	var domain = "local"
	var host = string("")
	call = group.Call(avahiEntrygroupAddService, 0, avahiIfUnspec, protoUnspec, flags, name, serviceType, domain, host, serverPort, txtBytes)
	if call.Err != nil {
		fmt.Println(os.Stderr, "Failed to call", avahiEntrygroupAddService)
		return false
	}
	call = group.Call(avahiEntrygroupCommit, 0)
	if call.Err != nil {
		fmt.Println(os.Stderr, "Failed to call", avahiEntrygroupCommit)
		return false
	}

	return true
}

func registerViaBuiltin(hostname, name, serviceType string, serverPort uint16, txt []string) {
	addrs, err := net.LookupIP(hostname + ".")
	if err != nil {
		log.Fatalln("Could not determine host IP addresses for %s", hostname)
	}

	var ip4 string
	var ip6 string
	for i := 0; i < len(addrs); i++ {
		if ipv4 := addrs[i].To4(); ipv4 != nil {
			if !addrs[i].IsLoopback() {
				ip4 = addrs[i].String()
			}
		} else if ipv6 := addrs[i].To16(); ipv6 != nil {
			if !addrs[i].IsLoopback() {
				ip6 = addrs[i].String()
			}
		}
	}

	var ip string
	if len(ip4) > 0 {
		ip = ip4
	} else if len(ip6) > 0 {
		ip = ip6
	} else {
		log.Fatalln("Neither ipv4 or ipv6 address resolved to hostname")
	}

	_, err = bonjour.RegisterProxy(name, serviceType, "local", int(serverPort), hostname, ip, txt, nil)
	if err != nil {
		log.Fatalln(err.Error())
	}
}

func installZeroconfListener(name, serviceType string, serverPort uint16) {
	hostname, err := os.Hostname()
	if err != nil {
		log.Fatalln("Failed to fetch hostname")
	}

	nameHostname := strings.ToLower(name + "@" + hostname)

	txt := []string{"host=" + hostname}
	registered := registerAvahiService(nameHostname, serviceType, serverPort, txt)
	if !registered {
		registerViaBuiltin(hostname, nameHostname, serviceType, serverPort, txt)
		log.Println("Registered via internal bonjour handler")
	} else {
		log.Println("Registered via AVAHI")
	}
}
<commit_msg>don't append final .<commit_after>package main

import (
	"fmt"
	"log"
	"net"
	"os"
	"strings"

	"github.com/godbus/dbus"
	"github.com/oleksandr/bonjour"
)

const avahiDbusName string = "org.freedesktop.Avahi"
const avahiInterfaceServer string = avahiDbusName + ".Server"
const avahiInterfaceEntrygroup string = avahiDbusName + ".EntryGroup"
const avahiEntrygroupAddService string = avahiInterfaceEntrygroup + ".AddService"
const avahiEntrygroupCommit = avahiInterfaceEntrygroup + ".Commit"

const avahiIfUnspec = int32(-1)
const protoUnspec = int32(-1)
const protoInet = int32(0)
const protoInet6 = int32(1)

func registerAvahiService(name, serviceType string, serverPort uint16, txt []string) bool {
	conn, err := dbus.SystemBus()

	if err != nil {
		fmt.Fprintln(os.Stderr, "Failed to connect to D-BUS session bus:", err)
		return false
	}

	obj := conn.Object(avahiDbusName, "/")

	var groupPath dbus.ObjectPath
	call := obj.Call("org.freedesktop.Avahi.Server.EntryGroupNew", 0)
	if call.Err != nil {
		fmt.Fprintln(os.Stderr, "Avahi not found on D-BUS")
		return false
	}
	call.Store(&groupPath)

	group := conn.Object(avahiDbusName, groupPath)

	var txtBytes [][]byte
	for _, s := range txt {
		txtBytes = append(txtBytes, []byte(s))
	}

	var flags = uint32(0)
	var domain = "local"
	var host = string("")
	call = group.Call(avahiEntrygroupAddService, 0, avahiIfUnspec, protoUnspec, flags, name, serviceType, domain, host, serverPort, txtBytes)
	if call.Err != nil {
		fmt.Println(os.Stderr, "Failed to call", avahiEntrygroupAddService)
		return false
	}
	call = group.Call(avahiEntrygroupCommit, 0)
	if call.Err != nil {
		fmt.Println(os.Stderr, "Failed to call", avahiEntrygroupCommit)
		return false
	}

	return true
}

func registerViaBuiltin(hostname, name, serviceType string, serverPort uint16, txt []string) {
	addrs, err := net.LookupIP(hostname)
	if err != nil {
		log.Fatalln("Could not determine host IP addresses for %s", hostname)
	}

	var ip4 string
	var ip6 string
	for i := 0; i < len(addrs); i++ {
		if ipv4 := addrs[i].To4(); ipv4 != nil {
			if !addrs[i].IsLoopback() {
				ip4 = addrs[i].String()
			}
		} else if ipv6 := addrs[i].To16(); ipv6 != nil {
			if !addrs[i].IsLoopback() {
				ip6 = addrs[i].String()
			}
		}
	}

	var ip string
	if len(ip4) > 0 {
		ip = ip4
	} else if len(ip6) > 0 {
		ip = ip6
	} else {
		log.Fatalln("Neither ipv4 or ipv6 address resolved to hostname")
	}

	_, err = bonjour.RegisterProxy(name, serviceType, "local", int(serverPort), hostname, ip, txt, nil)
	if err != nil {
		log.Fatalln(err.Error())
	}
}

func installZeroconfListener(name, serviceType string, serverPort uint16) {
	hostname, err := os.Hostname()
	if err != nil {
		log.Fatalln("Failed to fetch hostname")
	}

	nameHostname := strings.ToLower(name + "@" + hostname)

	txt := []string{"host=" + hostname}
	registered := registerAvahiService(nameHostname, serviceType, serverPort, txt)
	if !registered {
		registerViaBuiltin(hostname, nameHostname, serviceType, serverPort, txt)
		log.Println("Registered via internal bonjour handler")
	} else {
		log.Println("Registered via AVAHI")
	}
}
<|endoftext|>


<commit_before>package main

import (
	"os"
	"testing"

	. "github.com/smartystreets/goconvey/convey"
)

func TestMutilatePluginLaunch(t *testing.T) {
	Convey("Ensure mutilate plugin can be launched", t, func() {
		os.Args = []string{"", "{\"NoDaemon\": true}"}
		So(func() { main() }, ShouldNotPanic)
	})
}

<commit_msg>Minor style fixes<commit_after>package main

import (
	"os"
	"testing"

	. "github.com/smartystreets/goconvey/convey"
)

func TestMutilatePluginLaunch(t *testing.T) {
	Convey("Ensure mutilate plugin can be launched", t, func() {
		os.Args = []string{"", "{\"NoDaemon\": true}"}
		So(func() { main() }, ShouldNotPanic)
	})
}
<|endoftext|>


<commit_before>// Copyright 2016 Albert Nigmatzianov. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package track

import (
	"net/url"
	"runtime"
	"strings"

	"github.com/bogem/nehm/util"
)

const clientID = "11a37feb6ccc034d5975f3f803928a32"

type Track struct {
	artist string
	title  string

	// Fields needed for JSON unmarshalling.
	JArtworkURL string `json:"artwork_url"`
	JCreatedAt  string `json:"created_at"`
	JDuration   int    `json:"duration"`
	JID         int    `json:"id"`
	JTitle      string `json:"title"`
	JURL        string `json:"stream_url"`
	JAuthor     struct {
		AvatarURL string `json:"avatar_url"`
		Username  string `json:"username"`
	} `json:"user"`
}

func (t *Track) Artist() string {
	t.setArtistAndTitle()
	return t.artist
}

func (t Track) ArtworkURL() string {
	artworkURL := t.JArtworkURL
	if artworkURL == "" {
		artworkURL = t.JAuthor.AvatarURL
	}
	return strings.Replace(artworkURL, "large", "t500x500", 1)
}

func (t Track) Duration() string {
	return util.DurationString(util.ParseDuration(t.JDuration))
}

func (t Track) Filename() string {
	// Replace all filesystem non-friendly runes with the underscore.
	var toReplace string
	if runtime.GOOS == "windows" {
		toReplace = "<>:\"\\/|?*" // https://msdn.microsoft.com/en-us/library/windows/desktop/aa365247(v=vs.85).aspx
	} else {
		toReplace = ":/\\"
	}

	replaceRunes := func(r rune) rune {
		if strings.ContainsRune(toReplace, r) {
			return '_'
		}
		return r
	}

	return strings.Map(replaceRunes, t.Fullname()) + ".mp3"
}

func (t Track) Fullname() string {
	return t.Artist() + " — " + t.Title()
}

func (t Track) ID() int {
	return t.JID
}

// name splits track's title to artist and title if there is one of separators
// and sets to t.artist and to t.title respectively.
// If t.artist or t.title are not blank, it will do nothing.
// If there is no separator in title, it willt use t.JAuthor.Username and
// t.JTitle.
//
// E.g. if track has title "Michael Jackson - Thriller", then it will use
// "Michael Jackson" as artist and "Thriller" as title.
func (t *Track) setArtistAndTitle() {
	if t.artist != "" || t.title != "" {
		return
	}

	separators := [...]string{" - ", " ~ ", " – "}
	for _, sep := range separators {
		if strings.Contains(t.JTitle, sep) {
			splitted := strings.SplitN(t.JTitle, sep, 2)
			t.artist = strings.TrimSpace(splitted[0])
			t.title = strings.TrimSpace(splitted[1])
			return
		}
	}

	t.artist = strings.TrimSpace(t.JAuthor.Username)
	t.title = strings.TrimSpace(t.JTitle)
}

func (t *Track) Title() string {
	t.setArtistAndTitle()
	return t.title
}

func (t Track) URL() string {
	u, err := url.Parse(t.JURL)
	if err != nil {
		return ""
	}

	q := u.Query()
	q.Add("client_id", clientID)
	u.RawQuery = q.Encode()

	return u.String()
}

func (t Track) Year() string {
	return t.JCreatedAt[0:4]
}
<commit_msg>Change clientID<commit_after>// Copyright 2016 Albert Nigmatzianov. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package track

import (
	"net/url"
	"runtime"
	"strings"

	"github.com/bogem/nehm/util"
)

const clientID = "a3e059563d7fd3372b49b37f00a00bcf"

type Track struct {
	artist string
	title  string

	// Fields needed for JSON unmarshalling.
	JArtworkURL string `json:"artwork_url"`
	JCreatedAt  string `json:"created_at"`
	JDuration   int    `json:"duration"`
	JID         int    `json:"id"`
	JTitle      string `json:"title"`
	JURL        string `json:"stream_url"`
	JAuthor     struct {
		AvatarURL string `json:"avatar_url"`
		Username  string `json:"username"`
	} `json:"user"`
}

func (t *Track) Artist() string {
	t.setArtistAndTitle()
	return t.artist
}

func (t Track) ArtworkURL() string {
	artworkURL := t.JArtworkURL
	if artworkURL == "" {
		artworkURL = t.JAuthor.AvatarURL
	}
	return strings.Replace(artworkURL, "large", "t500x500", 1)
}

func (t Track) Duration() string {
	return util.DurationString(util.ParseDuration(t.JDuration))
}

func (t Track) Filename() string {
	// Replace all filesystem non-friendly runes with the underscore.
	var toReplace string
	if runtime.GOOS == "windows" {
		toReplace = "<>:\"\\/|?*" // https://msdn.microsoft.com/en-us/library/windows/desktop/aa365247(v=vs.85).aspx
	} else {
		toReplace = ":/\\"
	}

	replaceRunes := func(r rune) rune {
		if strings.ContainsRune(toReplace, r) {
			return '_'
		}
		return r
	}

	return strings.Map(replaceRunes, t.Fullname()) + ".mp3"
}

func (t Track) Fullname() string {
	return t.Artist() + " — " + t.Title()
}

func (t Track) ID() int {
	return t.JID
}

// name splits track's title to artist and title if there is one of separators
// and sets to t.artist and to t.title respectively.
// If t.artist or t.title are not blank, it will do nothing.
// If there is no separator in title, it willt use t.JAuthor.Username and
// t.JTitle.
//
// E.g. if track has title "Michael Jackson - Thriller", then it will use
// "Michael Jackson" as artist and "Thriller" as title.
func (t *Track) setArtistAndTitle() {
	if t.artist != "" || t.title != "" {
		return
	}

	separators := [...]string{" - ", " ~ ", " – "}
	for _, sep := range separators {
		if strings.Contains(t.JTitle, sep) {
			splitted := strings.SplitN(t.JTitle, sep, 2)
			t.artist = strings.TrimSpace(splitted[0])
			t.title = strings.TrimSpace(splitted[1])
			return
		}
	}

	t.artist = strings.TrimSpace(t.JAuthor.Username)
	t.title = strings.TrimSpace(t.JTitle)
}

func (t *Track) Title() string {
	t.setArtistAndTitle()
	return t.title
}

func (t Track) URL() string {
	u, err := url.Parse(t.JURL)
	if err != nil {
		return ""
	}

	q := u.Query()
	q.Add("client_id", clientID)
	u.RawQuery = q.Encode()

	return u.String()
}

func (t Track) Year() string {
	return t.JCreatedAt[0:4]
}
<|endoftext|>


<commit_before>package cas

import (
	"bytes"
	"compress/bzip2"
	"compress/gzip"
	"crypto/sha256"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"

	"github.com/coreos-inc/rkt/app-container/aci"
)

func NewRemote(name string, mirrors []string) *Remote {
	// TODO: don't assume the name is a mirror?
	r := &Remote{}
	r.Name = name
	r.Mirrors = []string{name}
	return r
}

type Remote struct {
	Name    string
	Mirrors []string
	ETag    string
	File    string
}

func (r Remote) Marshal() []byte {
	m, _ := json.Marshal(r)
	return m
}

func (r *Remote) Unmarshal(data []byte) {
	err := json.Unmarshal(data, r)
	if err != nil {
		panic(err)
	}
}

func (r Remote) Hash() string {
	return sha256sum(r.Name)
}

func (r Remote) Type() int64 {
	return remoteType
}

func decompress(rs io.Reader, typ aci.FileType) (io.Reader, error) {
	var (
		dr  io.Reader
		err error
	)
	switch typ {
	case aci.TypeGzip:
		dr, err = gzip.NewReader(rs)
		if err != nil {
			return nil, err
		}
	case aci.TypeBzip2:
		dr = bzip2.NewReader(rs)
	case aci.TypeXz:
		dr = aci.XzReader(rs)
	case aci.TypeUnknown:
		fmt.Fprintf(os.Stderr, "error: unknown image filetype\n")
	default:
		// should never happen
		panic("no type returned from DetectFileType?")
	}
	return dr, nil
}

// TODO: add locking
func (r Remote) Download(ds Store) (*Remote, error) {
	var b bytes.Buffer

	res, err := http.Get(r.Name)
	if err != nil {
		return nil, err
	}
	defer res.Body.Close()

	// TODO(jonboulle): handle http more robustly (redirects?)
	if res.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("bad HTTP status code: %d", res.StatusCode)
	}

	// TODO(philips): use go routines to parallelize this pipeline and make
	// the file type detection happen without a second stream
	_, err = io.Copy(&b, res.Body)
	if err != nil {
		return nil, err
	}
	err = ds.stores[downloadType].WriteStream(r.Hash(), &b, true)
	if err != nil {
		return nil, err
	}

	// Detect the filetype
	rs, err := ds.stores[downloadType].ReadStream(r.Hash(), false)
	if err != nil {
		return nil, err
	}
	defer rs.Close()
	typ, err := aci.DetectFileType(rs)
	if err != nil {
		return nil, err
	}
	rs, err = ds.stores[downloadType].ReadStream(r.Hash(), false)
	if err != nil {
		return nil, err
	}
	defer rs.Close()

	// Generate the hash of the decompressed tar
	dr, err := decompress(rs, typ)
	if err != nil {
		return nil, err
	}
	hash := sha256.New()
	_, err = io.Copy(hash, dr)
	if err != nil {
		return nil, err
	}

	// Store the decompressed tar
	rs, err = ds.stores[downloadType].ReadStream(r.Hash(), false)
	if err != nil {
		return nil, err
	}
	defer rs.Close()
	dr, err = decompress(rs, typ)
	if err != nil {
		return nil, err
	}

	key := fmt.Sprintf("sha256-%x", hash.Sum(nil))
	err = ds.stores[objectType].WriteStream(key, dr, true)
	if err != nil {
		return nil, err
	}

	ds.stores[downloadType].Erase(r.Hash())
	r.File = key
	ds.stores[remoteType].Write(r.Hash(), r.Marshal())

	return &r, nil
}
<commit_msg>cas: minor refactor, split download into download & import<commit_after>package cas

import (
	"bytes"
	"compress/bzip2"
	"compress/gzip"
	"crypto/sha256"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"

	"github.com/coreos-inc/rkt/app-container/aci"
)

func NewRemote(name string, mirrors []string) *Remote {
	// TODO: don't assume the name is a mirror?
	r := &Remote{}
	r.Name = name
	r.Mirrors = []string{name}
	return r
}

type Remote struct {
	Name    string
	Mirrors []string
	ETag    string
	File    string
}

func (r Remote) Marshal() []byte {
	m, _ := json.Marshal(r)
	return m
}

func (r *Remote) Unmarshal(data []byte) {
	err := json.Unmarshal(data, r)
	if err != nil {
		panic(err)
	}
}

func (r Remote) Hash() string {
	return sha256sum(r.Name)
}

func (r Remote) Type() int64 {
	return remoteType
}

func decompress(rs io.Reader, typ aci.FileType) (io.Reader, error) {
	var (
		dr  io.Reader
		err error
	)
	switch typ {
	case aci.TypeGzip:
		dr, err = gzip.NewReader(rs)
		if err != nil {
			return nil, err
		}
	case aci.TypeBzip2:
		dr = bzip2.NewReader(rs)
	case aci.TypeXz:
		dr = aci.XzReader(rs)
	case aci.TypeUnknown:
		fmt.Fprintf(os.Stderr, "error: unknown image filetype\n")
	default:
		// should never happen
		panic("no type returned from DetectFileType?")
	}
	return dr, nil
}

func (r Remote) Import(ds Store, orig io.Reader) (*Remote, error) {
	var b bytes.Buffer

	// TODO(philips): use go routines to parallelize this pipeline and make
	// the file type detection happen without a second stream
	_, err := io.Copy(&b, orig)
	if err != nil {
		return nil, err
	}
	err = ds.stores[downloadType].WriteStream(r.Hash(), &b, true)
	if err != nil {
		return nil, err
	}

	// Detect the filetype
	rs, err := ds.stores[downloadType].ReadStream(r.Hash(), false)
	if err != nil {
		return nil, err
	}
	defer rs.Close()
	typ, err := aci.DetectFileType(rs)
	if err != nil {
		return nil, err
	}
	rs, err = ds.stores[downloadType].ReadStream(r.Hash(), false)
	if err != nil {
		return nil, err
	}
	defer rs.Close()

	// Generate the hash of the decompressed tar
	dr, err := decompress(rs, typ)
	if err != nil {
		return nil, err
	}
	hash := sha256.New()
	_, err = io.Copy(hash, dr)
	if err != nil {
		return nil, err
	}

	// Store the decompressed tar
	rs, err = ds.stores[downloadType].ReadStream(r.Hash(), false)
	if err != nil {
		return nil, err
	}
	defer rs.Close()
	dr, err = decompress(rs, typ)
	if err != nil {
		return nil, err
	}

	key := fmt.Sprintf("sha256-%x", hash.Sum(nil))
	err = ds.stores[objectType].WriteStream(key, dr, true)
	if err != nil {
		return nil, err
	}

	ds.stores[downloadType].Erase(r.Hash())
	r.File = key
	ds.stores[remoteType].Write(r.Hash(), r.Marshal())

	return &r, nil
}

// TODO: add locking
func (r Remote) Download(ds Store) (*Remote, error) {
	res, err := http.Get(r.Name)
	if err != nil {
		return nil, err
	}
	defer res.Body.Close()

	// TODO(jonboulle): handle http more robustly (redirects?)
	if res.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("bad HTTP status code: %d", res.StatusCode)
	}

	return r.Import(ds, res.Body)
}
<|endoftext|>