<commit_before>package kdtree

import (
	"alg"
	"ellipsoid"
	"errors"
	"fmt"
	"geo"
	"graph"
	"log"
	"math"
	"mm"
	"path"
)

var (
	e             ellipsoid.Ellipsoid
	clusterKdTree ClusterKdTree
)

func init() {
	e = ellipsoid.Init("WGS84", ellipsoid.Degrees, ellipsoid.Meter, ellipsoid.Longitude_is_symmetric, ellipsoid.Bearing_is_symmetric)
}

func LoadKdTree(clusterGraph *graph.ClusterGraph, base string) error {
	// Load the k-d tree of all clusters
	clusterKdTrees := make([]*KdTree, len(clusterGraph.Cluster))
	for i, g := range clusterGraph.Cluster {
		clusterDir := path.Join(base, fmt.Sprintf("/cluster%d", i+1))
		var encodedSteps []uint64
		var encodedCoordinates []int32
		err := mm.Open(path.Join(clusterDir, "kdtree.ftf"), &encodedSteps)
		if err != nil {
			return err
		}
		err = mm.Open(path.Join(clusterDir, "coordinates.ftf"), &encodedCoordinates)
		if err != nil {
			return err
		}
		clusterKdTrees[i] = &KdTree{Graph: g, EncodedSteps: encodedSteps, EncodedCoordinates: encodedCoordinates}
	}

	// Load the k-d tree of the overlay graph
	var encodedSteps []uint64
	var encodedCoordinates []int32
	err := mm.Open(path.Join(base, "/overlay/kdtree.ftf"), &encodedSteps)
	if err != nil {
		return err
	}
	err = mm.Open(path.Join(base, "/overlay/coordinates.ftf"), &encodedCoordinates)
	if err != nil {
		return err
	}
	overlayKdTree := &KdTree{Graph: clusterGraph.Overlay, EncodedSteps: encodedSteps, EncodedCoordinates: encodedCoordinates}

	// Load the bounding boxes of the clusters
	var bboxesFile []int32
	err = mm.Open(path.Join(base, "bboxes.ftf"), &bboxesFile)
	if err != nil {
		return err
	}
	if len(bboxesFile)/4 != clusterGraph.Overlay.ClusterCount() {
		return errors.New("size of bboxes file does not match cluster count")
	}
	bboxes := make([]geo.BBox, len(bboxesFile)/4)
	for i, _ := range bboxes {
		bboxes[i] = geo.DecodeBBox(bboxesFile[4*i : 4*i+4])
	}

	clusterKdTree = ClusterKdTree{Overlay: overlayKdTree, Cluster: clusterKdTrees, BBoxes: bboxes}
	return nil
}

// NearestNeighbor returns -1 if the way is on the overlay graph
// No fail strategy: a nearest point on the overlay graph is always returned if no point
// is found in the clusters.
func NearestNeighbor(x geo.Coordinate, forward bool, trans graph.Transport) (int, []graph.Way) {
	edges := []graph.Edge(nil)
	steps := []geo.Coordinate(nil)

	// first search on the overlay graph
	overlay := clusterKdTree.Overlay
	bestStepIndex, coordOverlay, foundPoint := binarySearch(overlay, x, 0, overlay.EncodedStepLen()-1,
		true /* compareLat */, trans, &edges)
	minDistance, _ := e.To(x.Lat, x.Lng, coordOverlay.Lat, coordOverlay.Lng)

	// for debugging
	bestCoord := coordOverlay

	// then search on all clusters where the point is inside the bounding box of the cluster
	clusterIndex := -1
	for i, b := range clusterKdTree.BBoxes {
		if b.Contains(x) {
			kdTree := clusterKdTree.Cluster[i]
			stepIndex, coord, ok := binarySearch(kdTree, x, 0, kdTree.EncodedStepLen()-1, true /* compareLat */, trans, &edges)
			dist, _ := e.To(x.Lat, x.Lng, coord.Lat, coord.Lng)

			if ok && (!foundPoint || dist < minDistance) {
				foundPoint = true
				minDistance = dist
				bestStepIndex = stepIndex
				clusterIndex = i

				// for debugging
				bestCoord = coord
			}
		}
	}

	// for debugging: linear search
	minD := math.Inf(1)
	minI := -1
	minCluster := -1
	var minC geo.Coordinate
	for i, b := range clusterKdTree.BBoxes {
		if b.Contains(x) {
			kdTree := clusterKdTree.Cluster[i]
			for j := 0; j < kdTree.EncodedStepLen(); j++ {
				c, oki := decodeCoordinate(kdTree, j, trans, &edges)
				dd, _ := e.To(x.Lat, x.Lng, c.Lat, c.Lng)
				if oki && dd < minD {
					minD = dd
					minI = j
					minC = c
					minCluster = i
				}
			}
		}
	}
	fmt.Printf("nearest neighbor for %v\n", x)
	fmt.Printf("binary search: min distance: %v at %v in cluster %v  %v\n", minDistance, bestStepIndex, clusterIndex, bestCoord)
	fmt.Printf("linear search: min distance: %v at %v in cluster %v  %v\n", minD, minI, minCluster, minC)

	if clusterIndex >= 0 {
		kdTree := clusterKdTree.Cluster[clusterIndex]
		return clusterIndex, decodeWays(kdTree.Graph, kdTree.EncodedStep(bestStepIndex), forward, trans, &edges, &steps)
	}
	log.Printf("no matching bounding box found for (%v, %v)", x.Lat, x.Lng)
	return clusterIndex, decodeWays(overlay.Graph, overlay.EncodedStep(bestStepIndex), forward, trans, &edges, &steps)
}

// binarySearch in one k-d tree. The index, the coordinate, and if the returned step/vertex
// is accessible are returned.
func binarySearch(kdTree *KdTree, x geo.Coordinate, start, end int, compareLat bool,
	trans graph.Transport, edges *[]graph.Edge) (int, geo.Coordinate, bool) {

	if end-start <= 0 {
		startCoord, startAccessible := decodeCoordinate(kdTree, start, trans, edges)
		return start, startCoord, startAccessible
	}

	middle := (end-start)/2 + start

	// exact hit
	middleCoord, middleAccessible := decodeCoordinate(kdTree, middle, trans, edges)
	if middleAccessible && x.Lat == middleCoord.Lat && x.Lng == middleCoord.Lng {
		return middle, middleCoord, middleAccessible
	}

	// corner case where the nearest point can be on both sides of the middle
	if !middleAccessible || (compareLat && x.Lat == middleCoord.Lat) || (!compareLat && x.Lng == middleCoord.Lng) {
		// recursion on both halfs
		leftRecIndex, leftCoord, leftAccessible := binarySearch(kdTree, x, start, middle-1, !compareLat, trans, edges)
		rightRecIndex, rightCoord, rightAccessible := binarySearch(kdTree, x, middle+1, end, !compareLat, trans, edges)

		if !middleAccessible && !leftAccessible && !rightAccessible {
			return middle, middleCoord, middleAccessible
		}

		// Infinity is used if a vertex/step it is not accessible as we know that at least one is accessible.
		distMiddle := math.Inf(1)
		distRecursionLeft := math.Inf(1)
		distRecursionRight := math.Inf(1)
		if middleAccessible {
			distMiddle, _ = e.To(x.Lat, x.Lng, middleCoord.Lat, middleCoord.Lng)
		}
		if leftAccessible {
			distRecursionLeft, _ = e.To(x.Lat, x.Lng, leftCoord.Lat, leftCoord.Lng)
		}
		if rightAccessible {
			distRecursionRight, _ = e.To(x.Lat, x.Lng, rightCoord.Lat, rightCoord.Lng)
		}

		if distRecursionLeft < distRecursionRight {
			if distRecursionLeft < distMiddle {
				return leftRecIndex, leftCoord, leftAccessible
			}
			return middle, middleCoord, middleAccessible
		}
		if distRecursionRight < distMiddle {
			return rightRecIndex, rightCoord, rightAccessible
		}
		return middle, middleCoord, middleAccessible
	}

	var left bool
	if compareLat {
		left = x.Lat < middleCoord.Lat
	} else {
		left = x.Lng < middleCoord.Lng
	}
	if left {
		// stop if there is nothing left of the middle
		if middle == start {
			return middle, middleCoord, middleAccessible
		}

		// recursion on the left half
		recIndex, recCoord, recAccessible := binarySearch(kdTree, x, start, middle-1, !compareLat, trans, edges)

		// compare middle and result from the left
		distMiddle, _ := e.To(x.Lat, x.Lng, middleCoord.Lat, middleCoord.Lng)
		distRecursion, _ := e.To(x.Lat, x.Lng, recCoord.Lat, recCoord.Lng)
		if !recAccessible || distMiddle < distRecursion {
			return middle, middleCoord, middleAccessible
		}
		return recIndex, recCoord, recAccessible
	}
	// stop if there is nothing right of the middle
	if middle == end {
		return middle, middleCoord, middleAccessible
	}

	// recursion on the right half
	recIndex, recCoord, recAccessible := binarySearch(kdTree, x, middle+1, end, !compareLat, trans, edges)

	// compare middle and result from the right
	distMiddle, _ := e.To(x.Lat, x.Lng, middleCoord.Lat, middleCoord.Lng)
	distRecursion, _ := e.To(x.Lat, x.Lng, recCoord.Lat, recCoord.Lng)
	if !recAccessible || distMiddle < distRecursion {
		return middle, middleCoord, middleAccessible
	}
	return recIndex, recCoord, recAccessible
}

// decodeCoordinate returns the coordinate of the encoded vertex/step and if it is accessible by the
// given transport mode
func decodeCoordinate(t *KdTree, i int, trans graph.Transport, edges *[]graph.Edge) (geo.Coordinate, bool) {
	g := t.Graph
	ec := t.EncodedStep(i)
	coord := geo.DecodeCoordinate(t.EncodedCoordinates[2*i], t.EncodedCoordinates[2*i+1])

	vertexIndex := ec >> (EdgeOffsetBits + StepOffsetBits)
	edgeOffset := (ec >> StepOffsetBits) & MaxEdgeOffset
	stepOffset := ec & MaxStepOffset
	vertex := graph.Vertex(vertexIndex)

	if edgeOffset == MaxEdgeOffset && stepOffset == MaxStepOffset {
		// it is a vertex and not a step
		return coord, g.VertexAccessible(vertex, trans)
	}

	var edge graph.Edge
	var edgeAccessible bool
	switch t := g.(type) {
	case *graph.GraphFile:
		(*edges) = t.VertexRawEdges(vertex, *edges)
		edge = (*edges)[edgeOffset]
		edgeAccessible = t.EdgeAccessible(edge, trans)
	case *graph.OverlayGraphFile:
		(*edges) = t.VertexRawEdges(vertex, *edges)
		edge = (*edges)[edgeOffset]
		edgeAccessible = t.EdgeAccessible(edge, trans)
	default:
		panic("unexpected graph implementation")
	}

	return coord, edgeAccessible
}

// decodeWays returns one or two partital edges that are therefore called ways. Even for a vertex a
// way is returned so that later code only has to consider ways. The other cases are explained below.
func decodeWays(g graph.Graph, ec uint64, forward bool, trans graph.Transport, edges *[]graph.Edge, steps *[]geo.Coordinate) []graph.Way {
	vertexIndex := ec >> (EdgeOffsetBits + StepOffsetBits)
	edgeOffset := (ec >> StepOffsetBits) & MaxEdgeOffset
	offset := ec & MaxStepOffset
	vertex := graph.Vertex(vertexIndex)

	if edgeOffset == MaxEdgeOffset && offset == MaxStepOffset {
		// The easy case, where we hit some vertex exactly.
		w := make([]graph.Way, 1)
		target := g.VertexCoordinate(vertex)
		w[0] = graph.Way{Length: 0, Vertex: vertex, Steps: nil, Target: target}
		return w
	}

	var edge graph.Edge
	oneway := false
	switch t := g.(type) {
	case *graph.GraphFile:
		(*edges) = t.VertexRawEdges(vertex, *edges)
		edge = (*edges)[edgeOffset]
		oneway = alg.GetBit(t.Oneway, uint(edge))
	case *graph.OverlayGraphFile:
		(*edges) = t.VertexRawEdges(vertex, *edges)
		edge = (*edges)[edgeOffset]
		oneway = alg.GetBit(t.Oneway, uint(edge))
	default:
		panic("unexpected graph implementation")
	}
	if trans == graph.Foot {
		oneway = false
	}

	t1 := vertex                       // start vertex
	t2 := g.EdgeOpposite(edge, vertex) // end vertex

	// now we can allocate the way corresponding to (edge,offset),
	// but there are three cases to consider:
	// - if the way is bidirectional we have to compute both directions,
	//   if forward == true the from the offset two both endpoints,
	//   and the reverse otherwise
	// - if the way is unidirectional then we have to compute the way
	//   from the StartPoint to offset if forward == false
	// - otherwise we have to compute the way from offset to the EndPoint
	// Strictly speaking only the second case needs an additional binary
	// search in the form of edge.StartPoint, but let's keep this simple
	// for now.
	(*steps) = g.EdgeSteps(edge, vertex, *steps)
	s := *steps

	b1 := make([]geo.Coordinate, len(s[:offset]))
	b2 := make([]geo.Coordinate, len(s[offset+1:]))
	copy(b1, s[:offset])
	copy(b2, s[offset+1:])
	l1 := geo.StepLength(s[:offset+1])
	l2 := geo.StepLength(s[offset:])
	t1Coord := g.VertexCoordinate(t1)
	t2Coord := g.VertexCoordinate(t2)
	d1, _ := e.To(t1Coord.Lat, t1Coord.Lng, s[0].Lat, s[0].Lng)
	d2, _ := e.To(t2Coord.Lat, t2Coord.Lng, s[len(s)-1].Lat, s[len(s)-1].Lng)
	l1 += d1
	l2 += d2
	target := s[offset]

	if !forward {
		reverse(b2)
	} else {
		reverse(b1)
	}

	var w []graph.Way
	if !oneway {
		w = make([]graph.Way, 2) // bidirectional
		w[0] = graph.Way{Length: l1, Vertex: t1, Steps: b1, Forward: forward, Target: target}
		w[1] = graph.Way{Length: l2, Vertex: t2, Steps: b2, Forward: forward, Target: target}
	} else {
		w = make([]graph.Way, 1) // one way
		if forward {
			w[0] = graph.Way{Length: l2, Vertex: t2, Steps: b2, Forward: forward, Target: target}
		} else {
			w[0] = graph.Way{Length: l1, Vertex: t1, Steps: b1, Forward: forward, Target: target}
		}
	}
	return w
}

func reverse(steps []geo.Coordinate) {
	for i, j := 0, len(steps)-1; i < j; i, j = i+1, j-1 {
		steps[i], steps[j] = steps[j], steps[i]
	}
}
<commit_msg>better nearest neighbor search<commit_after>package kdtree

import (
	"alg"
	"ellipsoid"
	"errors"
	"fmt"
	"geo"
	"graph"
	"log"
	"math"
	"mm"
	"path"
)

var (
	e             ellipsoid.Ellipsoid
	clusterKdTree ClusterKdTree
)

func init() {
	e = ellipsoid.Init("WGS84", ellipsoid.Degrees, ellipsoid.Meter, ellipsoid.Longitude_is_symmetric, ellipsoid.Bearing_is_symmetric)
}

func LoadKdTree(clusterGraph *graph.ClusterGraph, base string) error {
	// Load the k-d tree of all clusters
	clusterKdTrees := make([]*KdTree, len(clusterGraph.Cluster))
	for i, g := range clusterGraph.Cluster {
		clusterDir := path.Join(base, fmt.Sprintf("/cluster%d", i+1))
		var encodedSteps []uint64
		var encodedCoordinates []int32
		err := mm.Open(path.Join(clusterDir, "kdtree.ftf"), &encodedSteps)
		if err != nil {
			return err
		}
		err = mm.Open(path.Join(clusterDir, "coordinates.ftf"), &encodedCoordinates)
		if err != nil {
			return err
		}
		clusterKdTrees[i] = &KdTree{Graph: g, EncodedSteps: encodedSteps, EncodedCoordinates: encodedCoordinates}
	}

	// Load the k-d tree of the overlay graph
	var encodedSteps []uint64
	var encodedCoordinates []int32
	err := mm.Open(path.Join(base, "/overlay/kdtree.ftf"), &encodedSteps)
	if err != nil {
		return err
	}
	err = mm.Open(path.Join(base, "/overlay/coordinates.ftf"), &encodedCoordinates)
	if err != nil {
		return err
	}
	overlayKdTree := &KdTree{Graph: clusterGraph.Overlay, EncodedSteps: encodedSteps, EncodedCoordinates: encodedCoordinates}

	// Load the bounding boxes of the clusters
	var bboxesFile []int32
	err = mm.Open(path.Join(base, "bboxes.ftf"), &bboxesFile)
	if err != nil {
		return err
	}
	if len(bboxesFile)/4 != clusterGraph.Overlay.ClusterCount() {
		return errors.New("size of bboxes file does not match cluster count")
	}
	bboxes := make([]geo.BBox, len(bboxesFile)/4)
	for i, _ := range bboxes {
		bboxes[i] = geo.DecodeBBox(bboxesFile[4*i : 4*i+4])
	}

	clusterKdTree = ClusterKdTree{Overlay: overlayKdTree, Cluster: clusterKdTrees, BBoxes: bboxes}
	return nil
}

// NearestNeighbor returns -1 if the way is on the overlay graph
// No fail strategy: a nearest point on the overlay graph is always returned if no point
// is found in the clusters.
func NearestNeighbor(x geo.Coordinate, forward bool, trans graph.Transport) (int, []graph.Way) {
	edges := []graph.Edge(nil)
	steps := []geo.Coordinate(nil)

	// first search on the overlay graph
	overlay := clusterKdTree.Overlay
	bestStepIndex, coordOverlay, foundPoint := binarySearch(overlay, x, 0, overlay.EncodedStepLen()-1,
		true /* compareLat */, trans, &edges)
	minDistance, _ := e.To(x.Lat, x.Lng, coordOverlay.Lat, coordOverlay.Lng)

	// for debugging
	//bestCoord := coordOverlay

	// then search on all clusters where the point is inside the bounding box of the cluster
	clusterIndex := -1
	for i, b := range clusterKdTree.BBoxes {
		if b.Contains(x) {
			kdTree := clusterKdTree.Cluster[i]
			stepIndex, coord, ok := binarySearch(kdTree, x, 0, kdTree.EncodedStepLen()-1, true /* compareLat */, trans, &edges)
			dist, _ := e.To(x.Lat, x.Lng, coord.Lat, coord.Lng)

			if ok && (!foundPoint || dist < minDistance) {
				foundPoint = true
				minDistance = dist
				bestStepIndex = stepIndex
				clusterIndex = i

				// for debugging
				//bestCoord = coord
			}
		}
	}

	// for debugging: linear search
	/*minD := math.Inf(1)
	minI := -1
	minCluster := -1
	var minC geo.Coordinate
	for i, b := range clusterKdTree.BBoxes {
		if b.Contains(x) {
			kdTree := clusterKdTree.Cluster[i]
			for j := 0; j < kdTree.EncodedStepLen(); j++ {
				c, oki := decodeCoordinate(kdTree, j, trans, &edges)
				dd, _ := e.To(x.Lat, x.Lng, c.Lat, c.Lng)
				if oki && dd < minD {
					minD = dd
					minI = j
					minC = c
					minCluster = i
				}
			}
		}
	}
	fmt.Printf("nearest neighbor for %v\n", x)
	fmt.Printf("binary search: min distance: %v at %v in cluster %v  %v\n", minDistance, bestStepIndex, clusterIndex, bestCoord)
	fmt.Printf("linear search: min distance: %v at %v in cluster %v  %v\n", minD, minI, minCluster, minC)*/

	if clusterIndex >= 0 {
		kdTree := clusterKdTree.Cluster[clusterIndex]
		return clusterIndex, decodeWays(kdTree.Graph, kdTree.EncodedStep(bestStepIndex), forward, trans, &edges, &steps)
	}
	log.Printf("no matching bounding box found for (%v, %v)", x.Lat, x.Lng)
	return clusterIndex, decodeWays(overlay.Graph, overlay.EncodedStep(bestStepIndex), forward, trans, &edges, &steps)
}

// binarySearch in one k-d tree. The index, the coordinate, and if the returned step/vertex
// is accessible are returned.
func binarySearch(kdTree *KdTree, x geo.Coordinate, start, end int, compareLat bool,
	trans graph.Transport, edges *[]graph.Edge) (int, geo.Coordinate, bool) {

	if end <= start {
		// end of the recursion
		startCoord, startAccessible := decodeCoordinate(kdTree, start, trans, edges)
		return start, startCoord, startAccessible
	}

	middle := (end-start)/2 + start
	middleCoord, middleAccessible := decodeCoordinate(kdTree, middle, trans, edges)

	// exact hit
	if middleAccessible && x.Lat == middleCoord.Lat && x.Lng == middleCoord.Lng {
		return middle, middleCoord, middleAccessible
	}

	middleDist := math.Inf(1)
	if middleAccessible {
		middleDist, _ = e.To(x.Lat, x.Lng, middleCoord.Lat, middleCoord.Lng)
	}

	// recursion one half and if no accessible point is returned also on the other half
	var left bool
	if compareLat {
		left = x.Lat < middleCoord.Lat
	} else {
		left = x.Lng < middleCoord.Lng
	}
	var recIndex int
	var recCoord geo.Coordinate
	var recAccessible bool
	bothHalfs := false
	if left {
		// left
		recIndex, recCoord, recAccessible = binarySearch(kdTree, x, start, middle-1, !compareLat, trans, edges)
		if !recAccessible {
			// other half -> right
			recIndex, recCoord, recAccessible = binarySearch(kdTree, x, middle+1, end, !compareLat, trans, edges)
			bothHalfs = true
		}
	} else {
		// right
		recIndex, recCoord, recAccessible = binarySearch(kdTree, x, middle+1, end, !compareLat, trans, edges)
		if !recAccessible {
			// other half -> left
			recIndex, recCoord, recAccessible = binarySearch(kdTree, x, start, middle-1, !compareLat, trans, edges)
			bothHalfs = true
		}
	}
	bestDistance, _ := e.To(x.Lat, x.Lng, recCoord.Lat, recCoord.Lng)

	// we are finished if both have already been searched
	if bothHalfs {
		if middleDist < bestDistance {
			return middle, middleCoord, middleAccessible
		} else {
			return recIndex, recCoord, recAccessible
		}
	}

	distToPlane := 0.0
	if compareLat {
		distToPlane, _ = e.To(middleCoord.Lat, x.Lng, x.Lat, x.Lng)
	} else {
		distToPlane, _ = e.To(x.Lat, middleCoord.Lng, x.Lat, x.Lng)
	}

	var recIndex2 int
	var recCoord2 geo.Coordinate
	recAccessible2 := false
	// test whether the current best distance circle crosses the plane
	//distToPlane -= 0.002
	if bestDistance >= distToPlane {
		// search on the other half
		if !left {
			// left
			recIndex2, recCoord2, recAccessible2 = binarySearch(kdTree, x, start, middle-1, !compareLat, trans, edges)
		} else {
			// right
			recIndex2, recCoord2, recAccessible2 = binarySearch(kdTree, x, middle+1, end, !compareLat, trans, edges)
		}
	}

	bestDistance2 := math.Inf(1)
	if recAccessible2 {
		bestDistance2, _ = e.To(x.Lat, x.Lng, recCoord2.Lat, recCoord2.Lng)
	}

	if bestDistance < bestDistance2 {
		if middleDist < bestDistance {
			return middle, middleCoord, middleAccessible
		} else {
			return recIndex, recCoord, recAccessible
		}
	}
	if middleDist < bestDistance2 {
		return middle, middleCoord, middleAccessible
	}
	return recIndex2, recCoord2, recAccessible2
}

// decodeCoordinate returns the coordinate of the encoded vertex/step and if it is accessible by the
// given transport mode
func decodeCoordinate(t *KdTree, i int, trans graph.Transport, edges *[]graph.Edge) (geo.Coordinate, bool) {
	g := t.Graph
	ec := t.EncodedStep(i)
	coord := geo.DecodeCoordinate(t.EncodedCoordinates[2*i], t.EncodedCoordinates[2*i+1])

	vertexIndex := ec >> (EdgeOffsetBits + StepOffsetBits)
	edgeOffset := (ec >> StepOffsetBits) & MaxEdgeOffset
	stepOffset := ec & MaxStepOffset
	vertex := graph.Vertex(vertexIndex)

	if edgeOffset == MaxEdgeOffset && stepOffset == MaxStepOffset {
		// it is a vertex and not a step
		return coord, g.VertexAccessible(vertex, trans)
	}

	var edge graph.Edge
	var edgeAccessible bool
	switch t := g.(type) {
	case *graph.GraphFile:
		(*edges) = t.VertexRawEdges(vertex, *edges)
		edge = (*edges)[edgeOffset]
		edgeAccessible = t.EdgeAccessible(edge, trans)
	case *graph.OverlayGraphFile:
		(*edges) = t.VertexRawEdges(vertex, *edges)
		edge = (*edges)[edgeOffset]
		edgeAccessible = t.EdgeAccessible(edge, trans)
	default:
		panic("unexpected graph implementation")
	}

	return coord, edgeAccessible
}

// decodeWays returns one or two partital edges that are therefore called ways. Even for a vertex a
// way is returned so that later code only has to consider ways. The other cases are explained below.
func decodeWays(g graph.Graph, ec uint64, forward bool, trans graph.Transport, edges *[]graph.Edge, steps *[]geo.Coordinate) []graph.Way {
	vertexIndex := ec >> (EdgeOffsetBits + StepOffsetBits)
	edgeOffset := (ec >> StepOffsetBits) & MaxEdgeOffset
	offset := ec & MaxStepOffset
	vertex := graph.Vertex(vertexIndex)

	if edgeOffset == MaxEdgeOffset && offset == MaxStepOffset {
		// The easy case, where we hit some vertex exactly.
		w := make([]graph.Way, 1)
		target := g.VertexCoordinate(vertex)
		w[0] = graph.Way{Length: 0, Vertex: vertex, Steps: nil, Target: target}
		return w
	}

	var edge graph.Edge
	oneway := false
	switch t := g.(type) {
	case *graph.GraphFile:
		(*edges) = t.VertexRawEdges(vertex, *edges)
		edge = (*edges)[edgeOffset]
		oneway = alg.GetBit(t.Oneway, uint(edge))
	case *graph.OverlayGraphFile:
		(*edges) = t.VertexRawEdges(vertex, *edges)
		edge = (*edges)[edgeOffset]
		oneway = alg.GetBit(t.Oneway, uint(edge))
	default:
		panic("unexpected graph implementation")
	}
	if trans == graph.Foot {
		oneway = false
	}

	t1 := vertex                       // start vertex
	t2 := g.EdgeOpposite(edge, vertex) // end vertex

	// now we can allocate the way corresponding to (edge,offset),
	// but there are three cases to consider:
	// - if the way is bidirectional we have to compute both directions,
	//   if forward == true the from the offset two both endpoints,
	//   and the reverse otherwise
	// - if the way is unidirectional then we have to compute the way
	//   from the StartPoint to offset if forward == false
	// - otherwise we have to compute the way from offset to the EndPoint
	// Strictly speaking only the second case needs an additional binary
	// search in the form of edge.StartPoint, but let's keep this simple
	// for now.
	(*steps) = g.EdgeSteps(edge, vertex, *steps)
	s := *steps

	b1 := make([]geo.Coordinate, len(s[:offset]))
	b2 := make([]geo.Coordinate, len(s[offset+1:]))
	copy(b1, s[:offset])
	copy(b2, s[offset+1:])
	l1 := geo.StepLength(s[:offset+1])
	l2 := geo.StepLength(s[offset:])
	t1Coord := g.VertexCoordinate(t1)
	t2Coord := g.VertexCoordinate(t2)
	d1, _ := e.To(t1Coord.Lat, t1Coord.Lng, s[0].Lat, s[0].Lng)
	d2, _ := e.To(t2Coord.Lat, t2Coord.Lng, s[len(s)-1].Lat, s[len(s)-1].Lng)
	l1 += d1
	l2 += d2
	target := s[offset]

	if !forward {
		reverse(b2)
	} else {
		reverse(b1)
	}

	var w []graph.Way
	if !oneway {
		w = make([]graph.Way, 2) // bidirectional
		w[0] = graph.Way{Length: l1, Vertex: t1, Steps: b1, Forward: forward, Target: target}
		w[1] = graph.Way{Length: l2, Vertex: t2, Steps: b2, Forward: forward, Target: target}
	} else {
		w = make([]graph.Way, 1) // one way
		if forward {
			w[0] = graph.Way{Length: l2, Vertex: t2, Steps: b2, Forward: forward, Target: target}
		} else {
			w[0] = graph.Way{Length: l1, Vertex: t1, Steps: b1, Forward: forward, Target: target}
		}
	}
	return w
}

func reverse(steps []geo.Coordinate) {
	for i, j := 0, len(steps)-1; i < j; i, j = i+1, j-1 {
		steps[i], steps[j] = steps[j], steps[i]
	}
}
<|endoftext|>


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

package net

import (
	"net";
	"flag";
	"io";
	"os";
	"testing";
)

// If an IPv6 tunnel is running (see go/stubl), we can try dialing a real IPv6 address.
var ipv6 = flag.Bool("ipv6", false, "assume ipv6 tunnel is present")

// fd is already connected to www.google.com port 80.
// Run an HTTP request to fetch the main page.
func fetchGoogle(t *testing.T, fd net.Conn, network, addr string) {
	req := io.StringBytes("GET / HTTP/1.0\r\nHost: www.google.com\r\n\r\n");
	n, errno := fd.Write(req);

	buf := make([]byte, 1000);
	n, errno = io.Readn(fd, buf);

	if n < 1000 {
		t.Errorf("fetchGoogle: short HTTP read from %s %s", network, addr);
		return
	}
}

func doDial(t *testing.T, network, addr string) {
	fd, err := net.Dial(network, "", addr);
	if err != nil {
		t.Errorf("net.Dial(%q, %q, %q) = _, %v", network, "", addr, err);
		return
	}
	fetchGoogle(t, fd, network, addr);
	fd.Close()
}

func doDialTCP(t *testing.T, network, addr string) {
	fd, err := net.DialTCP(network, "", addr);
	if err != nil {
		t.Errorf("net.DialTCP(%q, %q, %q) = _, %v", network, "", addr, err);
	} else {
		fetchGoogle(t, fd, network, addr);
	}
	fd.Close()
}

var googleaddrs = []string (
	"74.125.19.99:80",
	"www.google.com:80",
	"74.125.19.99:http",
	"www.google.com:http",
	"074.125.019.099:0080",
	"[::ffff:74.125.19.99]:80",
	"[::ffff:4a7d:1363]:80",
	"[0:0:0:0:0000:ffff:74.125.19.99]:80",
	"[0:0:0:0:000000:ffff:74.125.19.99]:80",
	"[0:0:0:0:0:ffff::74.125.19.99]:80",
	"[2001:4860:0:2001::68]:80"	// ipv6.google.com; removed if ipv6 flag not set
)

func TestDialGoogle(t *testing.T) {
	// If no ipv6 tunnel, don't try the last address.
	if !*ipv6 {
		googleaddrs[len(googleaddrs)-1] = ""
	}

	for i := 0; i < len(googleaddrs); i++ {
		addr := googleaddrs[i];
		if addr == "" {
			continue
		}
		t.Logf("-- %s --", addr);
		doDial(t, "tcp", addr);
		doDialTCP(t, "tcp", addr);
		if addr[0] != '[' {
			doDial(t, "tcp4", addr);
			doDialTCP(t, "tcp4", addr)
		}
		doDial(t, "tcp6", addr);
		doDialTCP(t, "tcp6", addr)
	}
}
<commit_msg>change the URL in the test to avoid a redirection that breaks it in sydney.<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package net

import (
	"net";
	"flag";
	"io";
	"os";
	"testing";
)

// If an IPv6 tunnel is running (see go/stubl), we can try dialing a real IPv6 address.
var ipv6 = flag.Bool("ipv6", false, "assume ipv6 tunnel is present")

// fd is already connected to the destination, port 80.
// Run an HTTP request to fetch the appropriate page.
func fetchGoogle(t *testing.T, fd net.Conn, network, addr string) {
	req := io.StringBytes("GET /intl/en/privacy.html HTTP/1.0\r\nHost: www.google.com\r\n\r\n");
	n, errno := fd.Write(req);

	buf := make([]byte, 1000);
	n, errno = io.Readn(fd, buf);

	if n < 1000 {
		t.Errorf("fetchGoogle: short HTTP read from %s %s", network, addr);
		return
	}
}

func doDial(t *testing.T, network, addr string) {
	fd, err := net.Dial(network, "", addr);
	if err != nil {
		t.Errorf("net.Dial(%q, %q, %q) = _, %v", network, "", addr, err);
		return
	}
	fetchGoogle(t, fd, network, addr);
	fd.Close()
}

func doDialTCP(t *testing.T, network, addr string) {
	fd, err := net.DialTCP(network, "", addr);
	if err != nil {
		t.Errorf("net.DialTCP(%q, %q, %q) = _, %v", network, "", addr, err);
	} else {
		fetchGoogle(t, fd, network, addr);
	}
	fd.Close()
}

var googleaddrs = []string (
	"74.125.19.99:80",
	"www.google.com:80",
	"74.125.19.99:http",
	"www.google.com:http",
	"074.125.019.099:0080",
	"[::ffff:74.125.19.99]:80",
	"[::ffff:4a7d:1363]:80",
	"[0:0:0:0:0000:ffff:74.125.19.99]:80",
	"[0:0:0:0:000000:ffff:74.125.19.99]:80",
	"[0:0:0:0:0:ffff::74.125.19.99]:80",
	"[2001:4860:0:2001::68]:80"	// ipv6.google.com; removed if ipv6 flag not set
)

func TestDialGoogle(t *testing.T) {
	// If no ipv6 tunnel, don't try the last address.
	if !*ipv6 {
		googleaddrs[len(googleaddrs)-1] = ""
	}

	for i := 0; i < len(googleaddrs); i++ {
		addr := googleaddrs[i];
		if addr == "" {
			continue
		}
		t.Logf("-- %s --", addr);
		doDial(t, "tcp", addr);
		doDialTCP(t, "tcp", addr);
		if addr[0] != '[' {
			doDial(t, "tcp4", addr);
			doDialTCP(t, "tcp4", addr)
		}
		doDial(t, "tcp6", addr);
		doDialTCP(t, "tcp6", addr)
	}
}
<|endoftext|>


<commit_before>package liblumberjack

import (
  "bytes"
  "encoding/json"
  zmq "github.com/alecthomas/gozmq"
  "log"
  "math/big"
  "syscall"
  "time"
  //"syscall"
  "compress/zlib"
  //"compress/flate"
  "crypto/rand"
)

var context *zmq.Context

func init() {
  context, _ = zmq.NewContext()
}

// Forever Faithful Socket
type FFS struct {
  Endpoints []string // set of endpoints available to ship to

  // Socket type; zmq.REQ, etc
  SocketType zmq.SocketType

  // Various timeout values
  SendTimeout time.Duration
  RecvTimeout time.Duration

  endpoint  string      // the current endpoint in use
  socket    *zmq.Socket // the current zmq socket
  connected bool        // are we connected?
}

func (s *FFS) Send(data []byte, flags zmq.SendRecvOption) (err error) {
  for {
    s.ensure_connect()

    pi := zmq.PollItems{zmq.PollItem{Socket: s.socket, Events: zmq.POLLOUT}}
    count, err := zmq.Poll(pi, s.SendTimeout)
    if count == 0 {
      // not ready in time, fail the socket and try again.
      log.Printf("%s: timed out waiting to Send(): %s\n",
        s.endpoint, err)
      s.fail_socket()
    } else {
      //log.Printf("%s: sending %d payload\n", s.endpoint, len(data))
      err = s.socket.Send(data, flags)
      if err != nil {
        log.Printf("%s: Failed to Send() %d byte message: %s\n",
          s.endpoint, len(data), err)
        s.fail_socket()
      } else {
        // Success!
        break
      }
    }
  }
  return
}

func (s *FFS) Recv(flags zmq.SendRecvOption) (data []byte, err error) {
  s.ensure_connect()

  pi := zmq.PollItems{zmq.PollItem{Socket: s.socket, Events: zmq.POLLIN}}
  count, err := zmq.Poll(pi, s.RecvTimeout)
  if count == 0 {
    // not ready in time, fail the socket and try again.
    s.fail_socket()

    err = syscall.ETIMEDOUT
    log.Printf("%s: timed out waiting to Recv(): %s\n",
      s.endpoint, err)
    return nil, err
  } else {
    data, err = s.socket.Recv(flags)
    if err != nil {
      log.Printf("%s: Failed to Recv() %d byte message: %s\n",
        s.endpoint, len(data), err)
      s.fail_socket()
      return nil, err
    } else {
      // Success!
    }
  }
  return
}

func (s *FFS) Close() (err error) {
  err = s.socket.Close()
  if err != nil {
    return
  }

  s.socket = nil
  s.connected = false
  return nil
}

func (s *FFS) ensure_connect() {
  if s.connected {
    return
  }

  if s.SendTimeout == 0 {
    s.SendTimeout = 1 * time.Second
  }
  if s.RecvTimeout == 0 {
    s.RecvTimeout = 1 * time.Second
  }

  if s.SocketType == 0 {
    log.Panicf("No socket type set on zmq socket")
  }
  if s.socket != nil {
    s.socket.Close()
    s.socket = nil
  }

  var err error
  s.socket, err = context.NewSocket(s.SocketType)
  if err != nil {
    log.Panicf("zmq.NewSocket(%d) failed: %s\n", s.SocketType, err)
  }

  //s.socket.SetSockOptUInt64(zmq.HWM, 1)
  //s.socket.SetSockOptInt(zmq.RCVTIMEO, int(s.RecvTimeout.Nanoseconds() / 1000000))
  //s.socket.SetSockOptInt(zmq.SNDTIMEO, int(s.SendTimeout.Nanoseconds() / 1000000))

  // Abort anything in-flight on a socket that's closed.
  s.socket.SetSockOptInt(zmq.LINGER, 0)

  for !s.connected {
    var max *big.Int = big.NewInt(int64(len(s.Endpoints)))
    i, _ := rand.Int(rand.Reader, max)
    s.endpoint = s.Endpoints[i.Int64()]
    err := s.socket.Connect(s.endpoint)
    if err != nil {
      log.Printf("%s: Error connecting: %s\n", s.endpoint, err)
      time.Sleep(500 * time.Millisecond)
      continue
    }

    // No error, we're connected.
    s.connected = true
  }
}

func (s *FFS) fail_socket() {
  if !s.connected {
    return
  }
  s.Close()
}

func Publish(input chan []*FileEvent, server_list []string,
  server_timeout time.Duration) {
  var buffer bytes.Buffer
  //key := "abcdefghijklmnop"
  //cipher, err := aes.NewCipher([]byte(key))

  socket := FFS{
    Endpoints:   server_list,
    SocketType:  zmq.REQ,
    RecvTimeout: server_timeout,
    SendTimeout: server_timeout,
  }
  //defer socket.Close()

  for events := range input {
    // got a bunch of events, ship them out.
    log.Printf("Spooler gave me %d events\n", len(events))

    // Serialize with msgpack
    data, err := json.Marshal(events)
    // TODO(sissel): chefk error
    _ = err
    //log.Printf("json serialized %d bytes\n", len(data))

    // Compress it
    // A new compressor is used for every payload of events so
    // that any individual payload can be decompressed alone.
    // TODO(sissel): Make compression level tunable
    compressor, _ := zlib.NewWriterLevel(&buffer, 3)
    buffer.Truncate(0)
    _, err = compressor.Write(data)
    err = compressor.Flush()
    compressor.Close()
    //log.Printf("compressed %d bytes\n", buffer.Len())
    // TODO(sissel): check err

    // TODO(sissel): implement security/encryption/etc

    // Send full payload over zeromq REQ/REP
    // TODO(sissel): check error
    //buffer.Write(data)

    // Loop forever trying to send.
    // This will cause reconnects/etc on failures automatically
    for {
      err = socket.Send(buffer.Bytes(), 0)
      data, err = socket.Recv(0)
      if err == nil {
        // success!
        break
      }
    }
    // TODO(sissel): Check data value of reply?

    // TODO(sissel): retry on failure or timeout
    // TODO(sissel): notify registrar of success
  } /* for each event payload */
} // Publish
<commit_msg>- get the publisher using sodium for crypto<commit_after>package liblumberjack

import (
  "bytes"
  "encoding/json"
  zmq "github.com/alecthomas/gozmq"
  "log"
  "math/big"
  "syscall"
  "time"
  "compress/zlib"
  "crypto/rand"
  "sodium"
)

var context *zmq.Context

func init() {
  context, _ = zmq.NewContext()
}

// Forever Faithful Socket
type FFS struct {
  Endpoints []string // set of endpoints available to ship to

  // Socket type; zmq.REQ, etc
  SocketType zmq.SocketType

  // Various timeout values
  SendTimeout time.Duration
  RecvTimeout time.Duration

  endpoint  string      // the current endpoint in use
  socket    *zmq.Socket // the current zmq socket
  connected bool        // are we connected?
}

func (s *FFS) Send(data []byte, flags zmq.SendRecvOption) (err error) {
  for {
    s.ensure_connect()

    pi := zmq.PollItems{zmq.PollItem{Socket: s.socket, Events: zmq.POLLOUT}}
    count, err := zmq.Poll(pi, s.SendTimeout)
    if count == 0 {
      // not ready in time, fail the socket and try again.
      log.Printf("%s: timed out waiting to Send(): %s\n",
        s.endpoint, err)
      s.fail_socket()
    } else {
      //log.Printf("%s: sending %d payload\n", s.endpoint, len(data))
      err = s.socket.Send(data, flags)
      if err != nil {
        log.Printf("%s: Failed to Send() %d byte message: %s\n",
          s.endpoint, len(data), err)
        s.fail_socket()
      } else {
        // Success!
        break
      }
    }
  }
  return
}

func (s *FFS) Recv(flags zmq.SendRecvOption) (data []byte, err error) {
  s.ensure_connect()

  pi := zmq.PollItems{zmq.PollItem{Socket: s.socket, Events: zmq.POLLIN}}
  count, err := zmq.Poll(pi, s.RecvTimeout)
  if count == 0 {
    // not ready in time, fail the socket and try again.
    s.fail_socket()

    err = syscall.ETIMEDOUT
    log.Printf("%s: timed out waiting to Recv(): %s\n",
      s.endpoint, err)
    return nil, err
  } else {
    data, err = s.socket.Recv(flags)
    if err != nil {
      log.Printf("%s: Failed to Recv() %d byte message: %s\n",
        s.endpoint, len(data), err)
      s.fail_socket()
      return nil, err
    } else {
      // Success!
    }
  }
  return
}

func (s *FFS) Close() (err error) {
  err = s.socket.Close()
  if err != nil {
    return
  }

  s.socket = nil
  s.connected = false
  return nil
}

func (s *FFS) ensure_connect() {
  if s.connected {
    return
  }

  if s.SendTimeout == 0 {
    s.SendTimeout = 1 * time.Second
  }
  if s.RecvTimeout == 0 {
    s.RecvTimeout = 1 * time.Second
  }

  if s.SocketType == 0 {
    log.Panicf("No socket type set on zmq socket")
  }
  if s.socket != nil {
    s.socket.Close()
    s.socket = nil
  }

  var err error
  s.socket, err = context.NewSocket(s.SocketType)
  if err != nil {
    log.Panicf("zmq.NewSocket(%d) failed: %s\n", s.SocketType, err)
  }

  //s.socket.SetSockOptUInt64(zmq.HWM, 1)
  //s.socket.SetSockOptInt(zmq.RCVTIMEO, int(s.RecvTimeout.Nanoseconds() / 1000000))
  //s.socket.SetSockOptInt(zmq.SNDTIMEO, int(s.SendTimeout.Nanoseconds() / 1000000))

  // Abort anything in-flight on a socket that's closed.
  s.socket.SetSockOptInt(zmq.LINGER, 0)

  for !s.connected {
    var max *big.Int = big.NewInt(int64(len(s.Endpoints)))
    i, _ := rand.Int(rand.Reader, max)
    s.endpoint = s.Endpoints[i.Int64()]
    err := s.socket.Connect(s.endpoint)
    if err != nil {
      log.Printf("%s: Error connecting: %s\n", s.endpoint, err)
      time.Sleep(500 * time.Millisecond)
      continue
    }

    // No error, we're connected.
    s.connected = true
  }
}

func (s *FFS) fail_socket() {
  if !s.connected {
    return
  }
  s.Close()
}

func Publish(input chan []*FileEvent,
             server_list []string,
             public_key [sodium.PUBLICKEYBYTES]byte,
             secret_key [sodium.SECRETKEYBYTES]byte,
             server_timeout time.Duration) {
  var buffer bytes.Buffer
  session := sodium.NewSession(public_key, secret_key)

  socket := FFS{
    Endpoints:   server_list,
    SocketType:  zmq.REQ,
    RecvTimeout: server_timeout,
    SendTimeout: server_timeout,
  }
  //defer socket.Close()

  for events := range input {
    // got a bunch of events, ship them out.
    log.Printf("Publisher received %d events\n", len(events))

    data, err := json.Marshal(events)
    // TODO(sissel): check error
    _ = err

    // Compress it
    // A new compressor is used for every payload of events so
    // that any individual payload can be decompressed alone.
    // TODO(sissel): Make compression level tunable
    compressor, _ := zlib.NewWriterLevel(&buffer, 3)
    buffer.Truncate(0)
    _, err = compressor.Write(data)
    err = compressor.Flush()
    compressor.Close()

    //log.Printf("compressed %d bytes\n", buffer.Len())
    // TODO(sissel): check err
    // TODO(sissel): implement security/encryption/etc

    // Send full payload over zeromq REQ/REP
    // TODO(sissel): check error
    //buffer.Write(data)
    ciphertext, nonce := session.Box(buffer.Bytes())

    log.Printf("plaintext: %d\n", len(data))
    log.Printf("compressed: %d\n", buffer.Len())
    log.Printf("ciphertext: %d\n", len(ciphertext))
    log.Printf("nonce: %d\n", len(nonce))

    // TODO(sissel): figure out encoding for ciphertext + nonce
    // TODO(sissel): Figure out the protocol for envelopes
    //   - compression envelope:
    //     - what compression system was used
    //     - any parameters?
    //     - payload
    //   - encryption envelope:
    //     - what encryption system was used
    //     - any public parameters necessary to decrypt (public key, nonce)
    //     - payload

    // Loop forever trying to send.
    // This will cause reconnects/etc on failures automatically
    for {
      err = socket.Send(buffer.Bytes(), 0)
      data, err = socket.Recv(0)
      if err == nil {
        // success!
        break
      }
    }
    // TODO(sissel): Check data value of reply?

    // TODO(sissel): retry on failure or timeout
    // TODO(sissel): notify registrar of success
  } /* for each event payload */
} // Publish
<|endoftext|>


<commit_before>package mcstore

import (
	"github.com/emicklei/go-restful"
	"github.com/materials-commons/mcstore/pkg/app"
	"github.com/materials-commons/mcstore/pkg/db/schema"
	"github.com/materials-commons/mcstore/pkg/ws/rest"
	"github.com/materials-commons/mcstore/server/mcstore/pkg/filters"
	"gopkg.in/olivere/elastic.v2"
)

type searchResource struct {
	log *app.Logger
}

type Query struct {
	QString string `json:"query_string"`
}

func newSearchResource() rest.Service {
	return &searchResource{
		log: app.NewLog("resource", "search"),
	}
}

func (r *searchResource) WebService() *restful.WebService {
	ws := new(restful.WebService)

	ws.Path("/search").Produces(restful.MIME_JSON).Consumes(restful.MIME_JSON)

	ws.Route(ws.POST("/project/{project}").
		Filter(filters.ProjectAccess).
		Filter(filters.SearchClient).
		To(rest.RouteHandler(r.searchProject)).
		Param(ws.PathParameter("project", "project id").DataType("string")).
		Doc("Searches all project items - files, processes, samples, notes, users").
		Reads(Query{}).
		Writes(elastic.SearchHits{}))

	ws.Route(ws.POST("/project/{project}/files").
		Filter(filters.ProjectAccess).
		Filter(filters.SearchClient).
		To(rest.RouteHandler(r.searchProjectFiles)).
		Param(ws.PathParameter("project", "project id").DataType("string")).
		Doc("Searches files in project").
		Reads(Query{}).
		Writes(elastic.SearchHits{}))

	return ws
}

func (r *searchResource) searchProject(request *restful.Request, response *restful.Response, user schema.User) (interface{}, error) {
	var query Query
	if err := request.ReadEntity(&query); err != nil {
		r.log.Debugf("Failed reading query: %s", err)
		return nil, err
	}

	project := request.Attribute("project").(schema.Project)
	client := request.Attribute("searchclient").(*elastic.Client)
	q := createQuery(query, project.ID)
	results, err := client.Search().Index("mc").Query(q).Size(20).Do()
	if err != nil {
		r.log.Infof("Query failed: %s", err)
		return nil, err
	}
	return results.Hits, nil
}

//func (r *searchResource) searchUserProjects(request *restful.Request, response *restful.Response, user schema.User) (interface{}, error) {
//	return nil, nil
//}

func (r *searchResource) searchProjectFiles(request *restful.Request, response *restful.Response, user schema.User) (interface{}, error) {
	var query Query

	if err := request.ReadEntity(&query); err != nil {
		r.log.Debugf("Failed reading query: %s", err)
		return nil, err
	}

	project := request.Attribute("project").(schema.Project)
	client := request.Attribute("searchclient").(*elastic.Client)
	q := createQuery(query, project.ID)
	results, err := client.Search().Index("mc").Type("files").Query(q).Size(20).Do()
	if err != nil {
		r.log.Infof("Query failed: %s", err)
		return nil, err
	}

	return results.Hits, nil
}

func createQuery(query Query, projectID string) elastic.Query {
	termQueryProj := elastic.NewTermQuery("project_id", projectID)
	userQuery := elastic.NewQueryStringQuery(query.QString)
	boolQuery := elastic.NewBoolQuery()
	boolQuery = boolQuery.Must(termQueryProj, userQuery)
	return &boolQuery
}
<commit_msg>Send back top 100.<commit_after>package mcstore

import (
	"github.com/emicklei/go-restful"
	"github.com/materials-commons/mcstore/pkg/app"
	"github.com/materials-commons/mcstore/pkg/db/schema"
	"github.com/materials-commons/mcstore/pkg/ws/rest"
	"github.com/materials-commons/mcstore/server/mcstore/pkg/filters"
	"gopkg.in/olivere/elastic.v2"
)

type searchResource struct {
	log *app.Logger
}

type Query struct {
	QString string `json:"query_string"`
}

func newSearchResource() rest.Service {
	return &searchResource{
		log: app.NewLog("resource", "search"),
	}
}

func (r *searchResource) WebService() *restful.WebService {
	ws := new(restful.WebService)

	ws.Path("/search").Produces(restful.MIME_JSON).Consumes(restful.MIME_JSON)

	ws.Route(ws.POST("/project/{project}").
		Filter(filters.ProjectAccess).
		Filter(filters.SearchClient).
		To(rest.RouteHandler(r.searchProject)).
		Param(ws.PathParameter("project", "project id").DataType("string")).
		Doc("Searches all project items - files, processes, samples, notes, users").
		Reads(Query{}).
		Writes(elastic.SearchHits{}))

	ws.Route(ws.POST("/project/{project}/files").
		Filter(filters.ProjectAccess).
		Filter(filters.SearchClient).
		To(rest.RouteHandler(r.searchProjectFiles)).
		Param(ws.PathParameter("project", "project id").DataType("string")).
		Doc("Searches files in project").
		Reads(Query{}).
		Writes(elastic.SearchHits{}))

	return ws
}

func (r *searchResource) searchProject(request *restful.Request, response *restful.Response, user schema.User) (interface{}, error) {
	var query Query
	if err := request.ReadEntity(&query); err != nil {
		r.log.Debugf("Failed reading query: %s", err)
		return nil, err
	}

	project := request.Attribute("project").(schema.Project)
	client := request.Attribute("searchclient").(*elastic.Client)
	q := createQuery(query, project.ID)
	results, err := client.Search().Index("mc").Query(q).Size(100).Do()
	if err != nil {
		r.log.Infof("Query failed: %s", err)
		return nil, err
	}
	return results.Hits, nil
}

//func (r *searchResource) searchUserProjects(request *restful.Request, response *restful.Response, user schema.User) (interface{}, error) {
//	return nil, nil
//}

func (r *searchResource) searchProjectFiles(request *restful.Request, response *restful.Response, user schema.User) (interface{}, error) {
	var query Query

	if err := request.ReadEntity(&query); err != nil {
		r.log.Debugf("Failed reading query: %s", err)
		return nil, err
	}

	project := request.Attribute("project").(schema.Project)
	client := request.Attribute("searchclient").(*elastic.Client)
	q := createQuery(query, project.ID)
	results, err := client.Search().Index("mc").Type("files").Query(q).Size(100).Do()
	if err != nil {
		r.log.Infof("Query failed: %s", err)
		return nil, err
	}

	return results.Hits, nil
}

func createQuery(query Query, projectID string) elastic.Query {
	termQueryProj := elastic.NewTermQuery("project_id", projectID)
	userQuery := elastic.NewQueryStringQuery(query.QString)
	boolQuery := elastic.NewBoolQuery()
	boolQuery = boolQuery.Must(termQueryProj, userQuery)
	return &boolQuery
}
<|endoftext|>


<commit_before>/*
Copyright 2016 - 2017 Huawei Technologies Co., Ltd. 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 template

import (
	"regexp"
)

//getTemplateContent get the template content by given version, if no content
//for the version, try its parenet version(like, 1.7.8 -> 1.7)
func getTemplateContent(templates map[string]string, version string) string {
	tplContent := templates[version]
	if tplContent == "" {
		parentVersion := parentVersionPattern.FindString(version)
		// TODO Log the Warnning
		// objects.WriteLog(fmt.Sprintf("No ca template for version %s, trying to use parent version %s", version, parentVersion), stdout, timestamp, d)

		tplContent = templates[parentVersion]
		if tplContent == "" {
			// return files, fmt.Errorf("No template for version %s or %s", version, parentVersion)
			// TODO Log the error
		}
	}
	return tplContent
}

var parentVersionPattern *regexp.Regexp = regexp.MustCompile(`[a-z]+-\d{1}\.\d{1}`)

func getParentVersion(version string) string {
	parentVersion := parentVersionPattern.FindString(version)
	return parentVersion
}
<commit_msg>Update the regex to check k8s version. Fix the bug that number greather than 10 cannot be parsed correctly.<commit_after>/*
Copyright 2016 - 2017 Huawei Technologies Co., Ltd. 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 template

import (
	"regexp"
)

//getTemplateContent get the template content by given version, if no content
//for the version, try its parenet version(like, 1.7.8 -> 1.7)
func getTemplateContent(templates map[string]string, version string) string {
	tplContent := templates[version]
	if tplContent == "" {
		parentVersion := parentVersionPattern.FindString(version)
		// TODO Log the Warnning
		// objects.WriteLog(fmt.Sprintf("No ca template for version %s, trying to use parent version %s", version, parentVersion), stdout, timestamp, d)

		tplContent = templates[parentVersion]
		if tplContent == "" {
			// return files, fmt.Errorf("No template for version %s or %s", version, parentVersion)
			// TODO Log the error
		}
	}
	return tplContent
}

var parentVersionPattern *regexp.Regexp = regexp.MustCompile(`[a-z]+-\d+\.\d+`)

func getParentVersion(version string) string {
	parentVersion := parentVersionPattern.FindString(version)
	return parentVersion
}
<|endoftext|>


<commit_before>package main

import (
        "./comms"
        "./EncryptionEngine"

        "flag"

        "net"
        "fmt"
        "io"

        "net/http"
        "net/url"
        "html"
        "log"
        "encoding/binary"
       )

type State struct {
    Kc [16]byte
    clk [4]byte
    BD_ADDR [6]byte
}

func server_main(s *State) net.Conn {
    ln, err := net.Listen("tcp", ":8080")

    if err != nil {
        fmt.Println("net.Listen failed:", err)
    }

    conn, err := ln.Accept()

    fmt.Printf("Accepted Connection\n")

    if err != nil {
        fmt.Println("ln.Accept failed:", err)
    }

    s.BD_ADDR = [6]byte{}
    s.Kc = [16]byte{0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0}
    s.clk = [4]byte{}

    comms.Send_neg(conn, [6]byte{})
    comms.Send_init(conn, 0, [16]byte{}, [16]byte{0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0})

    return conn
}

func client_main() net.Conn {
    conn, err := net.Dial("tcp", "127.0.0.1:8080")

    if err != nil {
        fmt.Println("net.Dial failed:", err)
    }

    return conn
}

func receiver(conn io.Reader, s *State) {
  LOOP:
  for {
        packet_type := comms.Recv_packet(conn)

        switch packet_type {
            case 0:
                s.BD_ADDR = comms.Recv_neg(conn)
            case 1:
                var clock uint32
                clock, _, s.Kc = comms.Recv_init(conn)
                binary.BigEndian.PutUint32(s.clk[:], clock)
            case 2:
                var clock uint32
                clock, msg := comms.Recv_data(conn)
                binary.BigEndian.PutUint32(s.clk[:], clock)

                keyStream := EncryptionEngine.GetKeyStream(s.Kc, s.BD_ADDR, s.clk, len(msg))
                msg = EncryptionEngine.Encrypt(msg, keyStream)

                _, err := http.PostForm("http://127.0.0.1:9999/bar",
                    url.Values{"msg": {string(msg)}})

                if err != nil {
                    fmt.Println("There was an http error: ", err)
                }

            case 99: break LOOP
        }
    }
}

func main() {
    isServerPtr := flag.Bool("server", false, "Run in server mode?")
    flag.Parse()
    fmt.Println("Is Server: ", *isServerPtr)

    var conn net.Conn
    var p string

    var state State

    if *isServerPtr {
        conn = server_main(&state)
        p = ":8888"
    } else {
        conn = client_main()
        p = ":6666"
    }

    go receiver(conn, &state)

    http.HandleFunc("/bar", func(w http.ResponseWriter, r *http.Request) {

        if r.Method == "POST" {
            r.ParseForm()

            fmt.Println("Requested to send: ", r.PostForm["msg"][0])

            fmt.Fprintf(w, html.EscapeString("Sending packet"))

            pt := []byte(r.PostForm["msg"][0])

            keyStream := EncryptionEngine.GetKeyStream(
                state.Kc, state.BD_ADDR, state.clk, len(pt))
            msg := EncryptionEngine.Encrypt(pt, keyStream)

            clock := binary.BigEndian.Uint32(state.clk[:])

            comms.Send_data(conn, clock, msg)
        }
    })

    log.Fatal(http.ListenAndServe(p, nil))

    conn.Close()
}
<commit_msg>Added some prints to show what's going on<commit_after>package main

import (
        "./comms"
        "./EncryptionEngine"

        "flag"

        "net"
        "fmt"
        "io"

        "net/http"
        "net/url"
        "html"
        "log"
        "encoding/binary"
       )

type State struct {
    Kc [16]byte
    clk [4]byte
    BD_ADDR [6]byte
}

func server_main(s *State) net.Conn {
    ln, err := net.Listen("tcp", ":8080")

    if err != nil {
        fmt.Println("net.Listen failed:", err)
    }

    conn, err := ln.Accept()

    fmt.Printf("Accepted Connection\n")

    if err != nil {
        fmt.Println("ln.Accept failed:", err)
    }

    s.BD_ADDR = [6]byte{}
    s.Kc = [16]byte{0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0}
    s.clk = [4]byte{}

    comms.Send_neg(conn, [6]byte{})
    comms.Send_init(conn, 0, [16]byte{}, [16]byte{0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0})

    return conn
}

func client_main() net.Conn {
    conn, err := net.Dial("tcp", "127.0.0.1:8080")

    if err != nil {
        fmt.Println("net.Dial failed:", err)
    }

    return conn
}

func receiver(conn io.Reader, s *State) {
  LOOP:
  for {
        packet_type := comms.Recv_packet(conn)

        switch packet_type {
            case 0:
                s.BD_ADDR = comms.Recv_neg(conn)
            case 1:
                var clock uint32
                clock, _, s.Kc = comms.Recv_init(conn)
                binary.BigEndian.PutUint32(s.clk[:], clock)
            case 2:
                var clock uint32
                clock, msg := comms.Recv_data(conn)
                binary.BigEndian.PutUint32(s.clk[:], clock)

                fmt.Println("Recieved: ", string(msg))

                keyStream := EncryptionEngine.GetKeyStream(s.Kc, s.BD_ADDR, s.clk, len(msg))
                msg = EncryptionEngine.Encrypt(msg, keyStream)

                fmt.Println("Decypted as: ", string(msg))

                _, err := http.PostForm("http://127.0.0.1:9999/bar",
                    url.Values{"msg": {string(msg)}})

                if err != nil {
                    fmt.Println("There was an http error: ", err)
                }

            case 99: break LOOP
        }
    }
}

func main() {
    isServerPtr := flag.Bool("server", false, "Run in server mode?")
    flag.Parse()
    fmt.Println("Is Server: ", *isServerPtr)

    var conn net.Conn
    var p string

    var state State

    if *isServerPtr {
        conn = server_main(&state)
        p = ":8888"
    } else {
        conn = client_main()
        p = ":6666"
    }

    go receiver(conn, &state)

    http.HandleFunc("/bar", func(w http.ResponseWriter, r *http.Request) {

        if r.Method == "POST" {
            r.ParseForm()

            fmt.Println("Sending: ", r.PostForm["msg"][0])

            fmt.Fprintf(w, html.EscapeString("Sending packet"))

            pt := []byte(r.PostForm["msg"][0])

            keyStream := EncryptionEngine.GetKeyStream(
                state.Kc, state.BD_ADDR, state.clk, len(pt))
            msg := EncryptionEngine.Encrypt(pt, keyStream)

            clock := binary.BigEndian.Uint32(state.clk[:])

            comms.Send_data(conn, clock, msg)

            fmt.Println("Encrypted as: ", string(msg))
        }
    })

    log.Fatal(http.ListenAndServe(p, nil))

    conn.Close()
}
<|endoftext|>


<commit_before>// Copyright 2015 Luke Shumaker

package store

import (
	"encoding/json"
	"github.com/jinzhu/gorm"
	"io/ioutil"
	"log"
	"net/http"
	"os"
)

type TwilioNumber struct {
	Id     int64
	Number string
	// TODO
}

func (o TwilioNumber) dbSchema(db *gorm.DB) error {
	return db.CreateTable(&o).Error
}

type TwilioPool struct {
	UserId   string
	GroupId  string
	NumberId string
}

type Incoming_numbers struct {
	Phone_numbers []Incoming_number `json:"incoming_phone_numbers"`
}

type Incoming_number struct {
	Number string `json:"phone_number"`
}

func (o TwilioPool) dbSchema(db *gorm.DB) error {
	return db.CreateTable(&o).
		AddForeignKey("user_id", "users(id)", "CASCADE", "RESTRICT").
		AddForeignKey("group_id", "groups(id)", "CASCADE", "RESTRICT").
		AddForeignKey("number_id", "twilio_numbers(id)", "RESTRICT", "RESTRICT").
		Error
}

func GetAllTwilioNumbers(db *gorm.DB) (ret []TwilioNumber) {
	panic("TODO")
}

func GetAllExistingTwilioNumbers() []string {

	// account SID for Twilio account
	account_sid := os.Getenv("TWILIO_ACCOUNTID")

	// Authorization token for Twilio account
	auth_token := os.Getenv("TWILIO_TOKEN")

	// gets url for the numbers we own in the Twilio Account
	incoming_num_url := "https://api.twilio.com/2010-04-01/Accounts/" + account_sid + "/IncomingPhoneNumbers.json"

	client := &http.Client{}

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

	req.SetBasicAuth(account_sid, auth_token)

	resp, err := client.Do(req)

	if err != nil {
		log.Println(err)
		return nil
	}

	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Println(err)
		return nil
	}

	if resp.StatusCode != 200 {
		log.Println(resp.Status)
		return nil
	}

	numbers := Incoming_numbers{}
	if err := json.Unmarshal(body, &numbers); err != nil {
		log.Println(err)
		return nil
	}

	if len(numbers.Phone_numbers) > 0 {

		existing_numbers := make([]string, len(numbers.Phone_numbers))

		for i, num := range numbers.Phone_numbers {
			existing_numbers[i] = num.Number
		}

		return existing_numbers

	} else {
		log.Println("You do not have a number in your Twilio account")
		return nil
	}

}
<commit_msg>added functions for twilopool<commit_after>// Copyright 2015 Luke Shumaker
// Copyright 2015 Zhandos Suleimenov
package store

import (
	"encoding/json"
	"github.com/jinzhu/gorm"
	"io/ioutil"
	"log"
	"net/http"
	"os"
)

type TwilioNumber struct {
	Id     int64  `json:"number_id"`
	Number string `json:"-"`
	// TODO
}

func (o TwilioNumber) dbSchema(db *gorm.DB) error {
	return db.CreateTable(&o).Error
}

type TwilioPool struct {
	UserId   string `json:"user_id"`
	GroupId  string `json:"group_id"`
	NumberId int64  `json:"number_id"`
}

type Incoming_numbers struct {
	Phone_numbers []Incoming_number `json:"incoming_phone_numbers"`
}

type Incoming_number struct {
	Number string `json:"phone_number"`
}

func (o TwilioPool) dbSchema(db *gorm.DB) error {
	return db.CreateTable(&o).
		AddForeignKey("user_id", "users(id)", "CASCADE", "RESTRICT").
		AddForeignKey("group_id", "groups(id)", "CASCADE", "RESTRICT").
		AddForeignKey("number_id", "twilio_numbers(id)", "RESTRICT", "RESTRICT").
		Error
}

func GetAllTwilioNumbers(db *gorm.DB) (ret []TwilioNumber) {
	var twilio_num []TwilioNumber
	if result := db.Find(&twilio_num); result.Error != nil {
		if result.RecordNotFound() {
			return nil
		}
		panic(result.Error)
	}
	return twilio_num

}

func GetTwilioNumberByUserAndGroup(db *gorm.DB, username string, groupname string) string {

	var o TwilioPool
	if result := db.Where(&TwilioPool{UserId: username, GroupId: groupname}).First(&o); result.Error != nil {
		if result.RecordNotFound() {
			log.Println("RecordNotFound")
			return ""
		}
		panic(result.Error)
	}

	var twilio_num TwilioNumber
	if result := db.Where("number_id = ?", o.NumberId).First(&twilio_num); result.Error != nil {
		if result.RecordNotFound() {
			log.Println("RecordNotFound")
			return ""
		}
		panic(result.Error)
	}

	return twilio_num.Number
}

func GetAllExistingTwilioNumbers() []string {

	// account SID for Twilio account
	account_sid := os.Getenv("TWILIO_ACCOUNTID")

	// Authorization token for Twilio account
	auth_token := os.Getenv("TWILIO_TOKEN")

	// gets url for the numbers we own in the Twilio Account
	incoming_num_url := "https://api.twilio.com/2010-04-01/Accounts/" + account_sid + "/IncomingPhoneNumbers.json"

	client := &http.Client{}

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

	req.SetBasicAuth(account_sid, auth_token)

	resp, err := client.Do(req)

	if err != nil {
		log.Println(err)
		return nil
	}

	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Println(err)
		return nil
	}

	if resp.StatusCode != 200 {
		log.Println(resp.Status)
		return nil
	}

	numbers := Incoming_numbers{}
	if err := json.Unmarshal(body, &numbers); err != nil {
		log.Println(err)
		return nil
	}

	if len(numbers.Phone_numbers) > 0 {

		existing_numbers := make([]string, len(numbers.Phone_numbers))

		for i, num := range numbers.Phone_numbers {
			existing_numbers[i] = num.Number
		}

		return existing_numbers

	} else {
		log.Println("You do not have a number in your Twilio account")
		return nil
	}

}
<|endoftext|>


<commit_before>package discovery

import (
	"errors"
)

var (
	ErrNetwork       = errors.New("pachyderm: network failure")
	ErrNotFound      = errors.New("pachyderm: Not found")
	ErrIsDir         = errors.New("pachyderm: Found dir")
	ErrWrongOldValue = errors.New("pachyderm: Wrong old value")
	ErrExists        = errors.New("pachyderm: Key already exists")
)

type Client interface {
	// Close closes the underlying connection
	// Errors: ErrNetwork
	Close() error
	// Get gets the value of a key
	// Errors: ErrNetwork, ErrNotFound, ErrIsDir
	Get(key string) (string, error)
	// GetAll returns all of the keys in a directory and its subdirectories as
	// a map from absolute keys to values.
	// Errors: ErrNetwork, ErrNotFound
	GetAll(key string) (map[string]string, error)
	// Set sets the value for a key
	// Errors: ErrNetwork, ErrIsDir
	Set(key string, value string) error
	// Delete deletes a key.
	// Errors: ErrNetwork, ErrNotFound, ErrIsDir
	Delete(key string) error
	// Create is like Set but only succeeds if the key doesn't already exist
	// Errors: ErrNetwork, ErrExists
	Create(key string, value string) error
	// CheckAndSet is like Set but only succeeds if the key is already set to oldValue
	// Errors: ErrNetwork, ErrNotFound, ErrIsDir, ErrWrongOldValue
	CheckAndSet(key string, value string, oldValue string) error
}

func NewEtcdClient(addresses ...string) Client {
	return newEtcdClient(addresses...)
}

func NewMockClient() *mockClient {
	return newMockClient()
}
<commit_msg>Adds a mock implementation of discovery.mock_client<commit_after>package discovery

import (
	"errors"
)

var (
	ErrNetwork      = errors.New("pachyderm: network failure")
	ErrNotFound     = errors.New("pachyderm: Not found")
	ErrDirectory    = errors.New("pachyderm: Found directory instead of value")
	ErrExists       = errors.New("pachyderm: Key already exists")
	ErrPrecondition = errors.New("pachyderm: Precondition not met")
)

type Client interface {
	// Close closes the underlying connection
	// Errors: ErrNetwork
	Close() error
	// Get gets the value of a key
	// Errors: ErrNetwork, ErrNotFound, ErrDirectory
	Get(key string) (string, error)
	// GetAll returns all of the keys in a directory and its subdirectories as
	// a map from absolute keys to values.
	// Errors: ErrNetwork
	GetAll(key string) (map[string]string, error)
	// Set sets the value for a key
	// Errors: ErrNetwork, ErrDirectory
	Set(key string, value string) error
	// Delete deletes a key.
	// Errors: ErrNetwork, ErrNotFound, ErrDirectory
	Delete(key string) error
	// Create is like Set but only succeeds if the key doesn't already exist
	// Errors: ErrNetwork, ErrExists
	Create(key string, value string) error
	// CheckAndSet is like Set but only succeeds if the key is already set to oldValue
	// Errors: ErrNetwork, ErrNotFound, ErrDirectory, ErrPrecondition
	CheckAndSet(key string, value string, oldValue string) error
}

func NewEtcdClient(addresses ...string) Client {
	return newEtcdClient(addresses...)
}

func NewMockClient() *mockClient {
	return newMockClient()
}
<|endoftext|>


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

package build

import (
	"exec"
	"path/filepath"
	"reflect"
	"runtime"
	"sort"
	"testing"
)

func sortstr(x []string) []string {
	sort.Strings(x)
	return x
}

var buildPkgs = []struct {
	dir  string
	info *DirInfo
}{
	{
		"go/build/pkgtest",
		&DirInfo{
			GoFiles:      []string{"pkgtest.go"},
			SFiles:       []string{"sqrt_" + runtime.GOARCH + ".s"},
			PkgName:      "pkgtest",
			Imports:      []string{"fmt", "os", "pkgtest"},
			TestGoFiles:  sortstr([]string{"sqrt_test.go", "sqrt_" + runtime.GOARCH + "_test.go"}),
			XTestGoFiles: []string{"xsqrt_test.go"},
		},
	},
	{
		"go/build/cmdtest",
		&DirInfo{
			GoFiles: []string{"main.go"},
			PkgName: "main",
			Imports: []string{"go/build/pkgtest"},
		},
	},
	{
		"go/build/cgotest",
		&DirInfo{
			CgoFiles: []string{"cgotest.go"},
			CFiles:   []string{"cgotest.c"},
			Imports:  []string{"C", "unsafe"},
			PkgName:  "cgotest",
		},
	},
}

const cmdtestOutput = "3"

func TestBuild(t *testing.T) {
	for _, tt := range buildPkgs {
		tree := Path[0] // Goroot
		dir := filepath.Join(tree.SrcDir(), tt.dir)

		info, err := ScanDir(dir, true)
		if err != nil {
			t.Errorf("ScanDir(%#q): %v", tt.dir, err)
			continue
		}
		if !reflect.DeepEqual(info, tt.info) {
			t.Errorf("ScanDir(%#q) = %#v, want %#v\n", tt.dir, info, tt.info)
			continue
		}

		s, err := Build(tree, tt.dir, info)
		if err != nil {
			t.Errorf("Build(%#q): %v", tt.dir, err)
			continue
		}

		if err := s.Run(); err != nil {
			t.Errorf("Run(%#q): %v", tt.dir, err)
			continue
		}

		if tt.dir == "go/build/cmdtest" {
			bin := s.Output[0]
			b, err := exec.Command(bin).CombinedOutput()
			if err != nil {
				t.Errorf("exec %s: %v", bin, err)
				continue
			}
			if string(b) != cmdtestOutput {
				t.Errorf("cmdtest output: %s want: %s", b, cmdtestOutput)
			}
		}

		// Deferred because cmdtest depends on pkgtest.
		defer func(s *Script) {
			if err := s.Nuke(); err != nil {
				t.Errorf("nuking: %v", err)
			}
		}(s)
	}
}
<commit_msg>Changed test to the behavior we want.<commit_after>// Copyright 2011 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package build

import (
	"exec"
	"path/filepath"
	"reflect"
	"runtime"
	"sort"
	"testing"
)

func sortstr(x []string) []string {
	sort.Strings(x)
	return x
}

var buildPkgs = []struct {
	dir  string
	info *DirInfo
}{
	{
		"go/build/pkgtest",
		&DirInfo{
			GoFiles:      []string{"pkgtest.go"},
			SFiles:       []string{"sqrt_" + runtime.GOARCH + ".s"},
			PkgName:      "pkgtest",
			Imports:      []string{"pkgtest"},
			TestImports:  []string{"fmt", "os"},
			TestGoFiles:  sortstr([]string{"sqrt_test.go", "sqrt_" + runtime.GOARCH + "_test.go"}),
			XTestGoFiles: []string{"xsqrt_test.go"},
		},
	},
	{
		"go/build/cmdtest",
		&DirInfo{
			GoFiles: []string{"main.go"},
			PkgName: "main",
			Imports: []string{"go/build/pkgtest"},
		},
	},
	{
		"go/build/cgotest",
		&DirInfo{
			CgoFiles: []string{"cgotest.go"},
			CFiles:   []string{"cgotest.c"},
			Imports:  []string{"C", "unsafe"},
			PkgName:  "cgotest",
		},
	},
}

const cmdtestOutput = "3"

func TestBuild(t *testing.T) {
	for _, tt := range buildPkgs {
		tree := Path[0] // Goroot
		dir := filepath.Join(tree.SrcDir(), tt.dir)

		info, err := ScanDir(dir, true)
		if err != nil {
			t.Errorf("ScanDir(%#q): %v", tt.dir, err)
			continue
		}
		if !reflect.DeepEqual(info, tt.info) {
			t.Errorf("ScanDir(%#q) = %#v, want %#v\n", tt.dir, info, tt.info)
			continue
		}

		s, err := Build(tree, tt.dir, info)
		if err != nil {
			t.Errorf("Build(%#q): %v", tt.dir, err)
			continue
		}

		if err := s.Run(); err != nil {
			t.Errorf("Run(%#q): %v", tt.dir, err)
			continue
		}

		if tt.dir == "go/build/cmdtest" {
			bin := s.Output[0]
			b, err := exec.Command(bin).CombinedOutput()
			if err != nil {
				t.Errorf("exec %s: %v", bin, err)
				continue
			}
			if string(b) != cmdtestOutput {
				t.Errorf("cmdtest output: %s want: %s", b, cmdtestOutput)
			}
		}

		// Deferred because cmdtest depends on pkgtest.
		defer func(s *Script) {
			if err := s.Nuke(); err != nil {
				t.Errorf("nuking: %v", err)
			}
		}(s)
	}
}
<|endoftext|>


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

// Package debug contains facilities for programs to debug themselves while
// they are running.
package debug

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
	"runtime"
)

var (
	dunno     = []byte("???")
	centerDot = []byte("·")
	dot       = []byte(".")
)

// PrintStack prints to standard error the stack trace returned by Stack.
func PrintStack() {
	os.Stderr.Write(stack())
}

// Stack returns a formatted stack trace of the goroutine that calls it.
// For each routine, it includes the source line information and PC value,
// then attempts to discover, for Go functions, the calling function or
// method and the text of the line containing the invocation.
func Stack() []byte {
	return stack()
}

// stack implements Stack, skipping 2 frames
func stack() []byte {
	buf := new(bytes.Buffer) // the returned data
	// As we loop, we open files and read them. These variables record the currently
	// loaded file.
	var lines [][]byte
	var lastFile string
	for i := 2; ; i++ { // Caller we care about is the user, 2 frames up
		pc, file, line, ok := runtime.Caller(i)
		if !ok {
			break
		}
		// Print this much at least.  If we can't find the source, it won't show.
		fmt.Fprintf(buf, "%s:%d (0x%x)\n", file, line, pc)
		if file != lastFile {
			data, err := ioutil.ReadFile(file)
			if err != nil {
				continue
			}
			lines = bytes.Split(data, []byte{'\n'})
			lastFile = file
		}
		line-- // in stack trace, lines are 1-indexed but our array is 0-indexed
		fmt.Fprintf(buf, "\t%s: %s\n", function(pc), source(lines, line))
	}
	return buf.Bytes()
}

// source returns a space-trimmed slice of the n'th line.
func source(lines [][]byte, n int) []byte {
	if n < 0 || n >= len(lines) {
		return dunno
	}
	return bytes.Trim(lines[n], " \t")
}

// function returns, if possible, the name of the function containing the PC.
func function(pc uintptr) []byte {
	fn := runtime.FuncForPC(pc)
	if fn == nil {
		return dunno
	}
	name := []byte(fn.Name())
	// The name includes the path name to the package, which is unnecessary
	// since the file name is already included.  Plus, it has center dots.
	// That is, we see
	//	runtime/debug.*T·ptrmethod
	// and want
	//	*T.ptrmethod
	if period := bytes.Index(name, dot); period >= 0 {
		name = name[period+1:]
	}
	name = bytes.Replace(name, centerDot, dot, -1)
	return name
}
<commit_msg>runtime/debug: document that Stack is deprecated<commit_after>// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package debug contains facilities for programs to debug themselves while
// they are running.
package debug

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
	"runtime"
)

var (
	dunno     = []byte("???")
	centerDot = []byte("·")
	dot       = []byte(".")
)

// PrintStack prints to standard error the stack trace returned by Stack.
func PrintStack() {
	os.Stderr.Write(stack())
}

// Stack returns a formatted stack trace of the goroutine that calls it.
// For each routine, it includes the source line information and PC value,
// then attempts to discover, for Go functions, the calling function or
// method and the text of the line containing the invocation.
//
// This function is deprecated. Use package runtime's Stack instead.
func Stack() []byte {
	return stack()
}

// stack implements Stack, skipping 2 frames
func stack() []byte {
	buf := new(bytes.Buffer) // the returned data
	// As we loop, we open files and read them. These variables record the currently
	// loaded file.
	var lines [][]byte
	var lastFile string
	for i := 2; ; i++ { // Caller we care about is the user, 2 frames up
		pc, file, line, ok := runtime.Caller(i)
		if !ok {
			break
		}
		// Print this much at least.  If we can't find the source, it won't show.
		fmt.Fprintf(buf, "%s:%d (0x%x)\n", file, line, pc)
		if file != lastFile {
			data, err := ioutil.ReadFile(file)
			if err != nil {
				continue
			}
			lines = bytes.Split(data, []byte{'\n'})
			lastFile = file
		}
		line-- // in stack trace, lines are 1-indexed but our array is 0-indexed
		fmt.Fprintf(buf, "\t%s: %s\n", function(pc), source(lines, line))
	}
	return buf.Bytes()
}

// source returns a space-trimmed slice of the n'th line.
func source(lines [][]byte, n int) []byte {
	if n < 0 || n >= len(lines) {
		return dunno
	}
	return bytes.Trim(lines[n], " \t")
}

// function returns, if possible, the name of the function containing the PC.
func function(pc uintptr) []byte {
	fn := runtime.FuncForPC(pc)
	if fn == nil {
		return dunno
	}
	name := []byte(fn.Name())
	// The name includes the path name to the package, which is unnecessary
	// since the file name is already included.  Plus, it has center dots.
	// That is, we see
	//	runtime/debug.*T·ptrmethod
	// and want
	//	*T.ptrmethod
	if period := bytes.Index(name, dot); period >= 0 {
		name = name[period+1:]
	}
	name = bytes.Replace(name, centerDot, dot, -1)
	return name
}
<|endoftext|>


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

// Package websocket implements a client and server for the WebSocket protocol.
// The protocol is defined at http://tools.ietf.org/html/draft-ietf-hybi-thewebsocketprotocol
package websocket

import (
	"bufio"
	"crypto/tls"
	"http"
	"io"
	"io/ioutil"
	"json"
	"net"
	"os"
	"url"
)

const (
	ProtocolVersionHixie75 = -75
	ProtocolVersionHixie76 = -76
	ProtocolVersionHybi00  = 0
	ProtocolVersionHybi    = 8

	ContinuationFrame = 0
	TextFrame         = 1
	BinaryFrame       = 2
	CloseFrame        = 8
	PingFrame         = 9
	PongFrame         = 10
	UnknownFrame      = 255
)

// WebSocket protocol errors.
type ProtocolError struct {
	ErrorString string
}

func (err ProtocolError) String() string { return err.ErrorString }

var (
	ErrBadProtocolVersion   = ProtocolError{"bad protocol version"}
	ErrBadScheme            = ProtocolError{"bad scheme"}
	ErrBadStatus            = ProtocolError{"bad status"}
	ErrBadUpgrade           = ProtocolError{"missing or bad upgrade"}
	ErrBadWebSocketOrigin   = ProtocolError{"missing or bad WebSocket-Origin"}
	ErrBadWebSocketLocation = ProtocolError{"missing or bad WebSocket-Location"}
	ErrBadWebSocketProtocol = ProtocolError{"missing or bad WebSocket-Protocol"}
	ErrBadWebSocketVersion  = ProtocolError{"missing or bad WebSocket Version"}
	ErrChallengeResponse    = ProtocolError{"mismatch challenge/response"}
	ErrBadFrame             = ProtocolError{"bad frame"}
	ErrBadFrameBoundary     = ProtocolError{"not on frame boundary"}
	ErrNotWebSocket         = ProtocolError{"not websocket protocol"}
	ErrBadRequestMethod     = ProtocolError{"bad method"}
	ErrNotSupported         = ProtocolError{"not supported"}
)

// WebSocketAddr is an implementation of net.Addr for WebSocket.
type WebSocketAddr struct {
	*url.URL
}

// Network returns the network type for a WebSocket, "websocket".
func (addr WebSocketAddr) Network() string { return "websocket" }

// Config is a WebSocket configuration
type Config struct {
	// A WebSocket server address.
	Location *url.URL

	// A Websocket client origin.
	Origin *url.URL

	// WebSocket subprotocols.
	Protocol []string

	// WebSocket protocol version.
	Version int

	// TLS config for secure WebSocket (wss).
	TlsConfig *tls.Config

	handshakeData map[string]string
}

// serverHandshaker is an interface to handle WebSocket server side handshake.
type serverHandshaker interface {
	// ReadHandshake reads handshake request message from client.
	// Returns http response code and error if any.
	ReadHandshake(buf *bufio.Reader, req *http.Request) (code int, err os.Error)

	// AcceptHandshake accepts the client handshake request and sends
	// handshake response back to client.
	AcceptHandshake(buf *bufio.Writer) (err os.Error)

	// NewServerConn creates a new WebSocket connection.
	NewServerConn(buf *bufio.ReadWriter, rwc io.ReadWriteCloser, request *http.Request) (conn *Conn)
}

// frameReader is an interface to read a WebSocket frame.
type frameReader interface {
	// Reader is to read payload of the frame.
	io.Reader

	// PayloadType returns payload type.
	PayloadType() byte

	// HeaderReader returns a reader to read header of the frame.
	HeaderReader() io.Reader

	// TrailerReader returns a reader to read trailer of the frame.
	// If it returns nil, there is no trailer in the frame.
	TrailerReader() io.Reader

	// Len returns total length of the frame, including header and trailer.
	Len() int
}

// frameReaderFactory is an interface to creates new frame reader.
type frameReaderFactory interface {
	NewFrameReader() (r frameReader, err os.Error)
}

// frameWriter is an interface to write a WebSocket frame.
type frameWriter interface {
	// Writer is to write playload of the frame.
	io.WriteCloser
}

// frameWriterFactory is an interface to create new frame writer.
type frameWriterFactory interface {
	NewFrameWriter(payloadType byte) (w frameWriter, err os.Error)
}

type frameHandler interface {
	HandleFrame(frame frameReader) (r frameReader, err os.Error)
	WriteClose(status int) (err os.Error)
}

// Conn represents a WebSocket connection.
type Conn struct {
	config  *Config
	request *http.Request

	buf *bufio.ReadWriter
	rwc io.ReadWriteCloser

	frameReaderFactory
	frameReader

	frameWriterFactory

	frameHandler
	PayloadType        byte
	defaultCloseStatus int
}

// Read implements the io.Reader interface:
// it reads data of a frame from the WebSocket connection.
// if msg is not large enough for the frame data, it fills the msg and next Read
// will read the rest of the frame data.
// it reads Text frame or Binary frame.
func (ws *Conn) Read(msg []byte) (n int, err os.Error) {
again:
	if ws.frameReader == nil {
		frame, err := ws.frameReaderFactory.NewFrameReader()
		if err != nil {
			return 0, err
		}
		ws.frameReader, err = ws.frameHandler.HandleFrame(frame)
		if err != nil {
			return 0, err
		}
		if ws.frameReader == nil {
			goto again
		}
	}
	n, err = ws.frameReader.Read(msg)
	if err == os.EOF {
		if trailer := ws.frameReader.TrailerReader(); trailer != nil {
			io.Copy(ioutil.Discard, trailer)
		}
		ws.frameReader = nil
		goto again
	}
	return n, err
}

// Write implements the io.Writer interface:
// it writes data as a frame to the WebSocket connection.
func (ws *Conn) Write(msg []byte) (n int, err os.Error) {
	w, err := ws.frameWriterFactory.NewFrameWriter(ws.PayloadType)
	if err != nil {
		return 0, err
	}
	n, err = w.Write(msg)
	w.Close()
	if err != nil {
		return n, err
	}
	return n, err
}

// Close implements the io.Closer interface.
func (ws *Conn) Close() os.Error {
	err := ws.frameHandler.WriteClose(ws.defaultCloseStatus)
	if err != nil {
		return err
	}
	return ws.rwc.Close()
}

func (ws *Conn) IsClientConn() bool { return ws.request == nil }
func (ws *Conn) IsServerConn() bool { return ws.request != nil }

// LocalAddr returns the WebSocket Origin for the connection for client, or
// the WebSocket location for server.
func (ws *Conn) LocalAddr() net.Addr {
	if ws.IsClientConn() {
		return WebSocketAddr{ws.config.Origin}
	}
	return WebSocketAddr{ws.config.Location}
}

// RemoteAddr returns the WebSocket location for the connection for client, or
// the Websocket Origin for server.
func (ws *Conn) RemoteAddr() net.Addr {
	if ws.IsClientConn() {
		return WebSocketAddr{ws.config.Location}
	}
	return WebSocketAddr{ws.config.Origin}
}

// SetTimeout sets the connection's network timeout in nanoseconds.
func (ws *Conn) SetTimeout(nsec int64) os.Error {
	if conn, ok := ws.rwc.(net.Conn); ok {
		return conn.SetTimeout(nsec)
	}
	return os.EINVAL
}

// SetReadTimeout sets the connection's network read timeout in nanoseconds.
func (ws *Conn) SetReadTimeout(nsec int64) os.Error {
	if conn, ok := ws.rwc.(net.Conn); ok {
		return conn.SetReadTimeout(nsec)
	}
	return os.EINVAL
}

// SetWriteTimeout sets the connection's network write timeout in nanoseconds.
func (ws *Conn) SetWriteTimeout(nsec int64) os.Error {
	if conn, ok := ws.rwc.(net.Conn); ok {
		return conn.SetWriteTimeout(nsec)
	}
	return os.EINVAL
}

// Config returns the WebSocket config.
func (ws *Conn) Config() *Config { return ws.config }

// Request returns the http request upgraded to the WebSocket.
// It is nil for client side.
func (ws *Conn) Request() *http.Request { return ws.request }

// Codec represents a symmetric pair of functions that implement a codec.
type Codec struct {
	Marshal   func(v interface{}) (data []byte, payloadType byte, err os.Error)
	Unmarshal func(data []byte, payloadType byte, v interface{}) (err os.Error)
}

// Send sends v marshaled by cd.Marshal as single frame to ws.
func (cd Codec) Send(ws *Conn, v interface{}) (err os.Error) {
	if err != nil {
		return err
	}
	data, payloadType, err := cd.Marshal(v)
	if err != nil {
		return err
	}
	w, err := ws.frameWriterFactory.NewFrameWriter(payloadType)
	_, err = w.Write(data)
	w.Close()
	return err
}

// Receive receives single frame from ws, unmarshaled by cd.Unmarshal and stores in v.
func (cd Codec) Receive(ws *Conn, v interface{}) (err os.Error) {
	if ws.frameReader != nil {
		_, err = io.Copy(ioutil.Discard, ws.frameReader)
		if err != nil {
			return err
		}
		ws.frameReader = nil
	}
again:
	frame, err := ws.frameReaderFactory.NewFrameReader()
	if err != nil {
		return err
	}
	frame, err = ws.frameHandler.HandleFrame(frame)
	if err != nil {
		return err
	}
	if frame == nil {
		goto again
	}
	payloadType := frame.PayloadType()
	data, err := ioutil.ReadAll(frame)
	if err != nil {
		return err
	}
	return cd.Unmarshal(data, payloadType, v)
}

func marshal(v interface{}) (msg []byte, payloadType byte, err os.Error) {
	switch data := v.(type) {
	case string:
		return []byte(data), TextFrame, nil
	case []byte:
		return data, BinaryFrame, nil
	}
	return nil, UnknownFrame, ErrNotSupported
}

func unmarshal(msg []byte, payloadType byte, v interface{}) (err os.Error) {
	switch data := v.(type) {
	case *string:
		*data = string(msg)
		return nil
	case *[]byte:
		*data = msg
		return nil
	}
	return ErrNotSupported
}

/*
Message is a codec to send/receive text/binary data in a frame on WebSocket connection.
To send/receive text frame, use string type.
To send/receive binary frame, use []byte type.

Trivial usage:

	import "websocket"

	// receive text frame
	var message string
	websocket.Message.Receive(ws, &message)

	// send text frame
	message = "hello"
	websocket.Message.Send(ws, message)

	// receive binary frame
	var data []byte
	websocket.Message.Receive(ws, &data)

	// send binary frame
	data = []byte{0, 1, 2}
	websocket.Message.Send(ws, data)

*/
var Message = Codec{marshal, unmarshal}

func jsonMarshal(v interface{}) (msg []byte, payloadType byte, err os.Error) {
	msg, err = json.Marshal(v)
	return msg, TextFrame, err
}

func jsonUnmarshal(msg []byte, payloadType byte, v interface{}) (err os.Error) {
	return json.Unmarshal(msg, v)
}

/*
JSON is a codec to send/receive JSON data in a frame from a WebSocket connection.

Trival usage:

	import "websocket"

	type T struct {
		Msg string
		Count int
	}

	// receive JSON type T
	var data T
	websocket.JSON.Receive(ws, &data)

	// send JSON type T
	websocket.JSON.Send(ws, data)
*/
var JSON = Codec{jsonMarshal, jsonUnmarshal}
<commit_msg>websocket: rename websocket.WebSocketAddr to *websocket.Addr.<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package websocket implements a client and server for the WebSocket protocol.
// The protocol is defined at http://tools.ietf.org/html/draft-ietf-hybi-thewebsocketprotocol
package websocket

import (
	"bufio"
	"crypto/tls"
	"http"
	"io"
	"io/ioutil"
	"json"
	"net"
	"os"
	"url"
)

const (
	ProtocolVersionHixie75 = -75
	ProtocolVersionHixie76 = -76
	ProtocolVersionHybi00  = 0
	ProtocolVersionHybi    = 8

	ContinuationFrame = 0
	TextFrame         = 1
	BinaryFrame       = 2
	CloseFrame        = 8
	PingFrame         = 9
	PongFrame         = 10
	UnknownFrame      = 255
)

// WebSocket protocol errors.
type ProtocolError struct {
	ErrorString string
}

func (err ProtocolError) String() string { return err.ErrorString }

var (
	ErrBadProtocolVersion   = ProtocolError{"bad protocol version"}
	ErrBadScheme            = ProtocolError{"bad scheme"}
	ErrBadStatus            = ProtocolError{"bad status"}
	ErrBadUpgrade           = ProtocolError{"missing or bad upgrade"}
	ErrBadWebSocketOrigin   = ProtocolError{"missing or bad WebSocket-Origin"}
	ErrBadWebSocketLocation = ProtocolError{"missing or bad WebSocket-Location"}
	ErrBadWebSocketProtocol = ProtocolError{"missing or bad WebSocket-Protocol"}
	ErrBadWebSocketVersion  = ProtocolError{"missing or bad WebSocket Version"}
	ErrChallengeResponse    = ProtocolError{"mismatch challenge/response"}
	ErrBadFrame             = ProtocolError{"bad frame"}
	ErrBadFrameBoundary     = ProtocolError{"not on frame boundary"}
	ErrNotWebSocket         = ProtocolError{"not websocket protocol"}
	ErrBadRequestMethod     = ProtocolError{"bad method"}
	ErrNotSupported         = ProtocolError{"not supported"}
)

// Addr is an implementation of net.Addr for WebSocket.
type Addr struct {
	*url.URL
}

// Network returns the network type for a WebSocket, "websocket".
func (addr *Addr) Network() string { return "websocket" }

// Config is a WebSocket configuration
type Config struct {
	// A WebSocket server address.
	Location *url.URL

	// A Websocket client origin.
	Origin *url.URL

	// WebSocket subprotocols.
	Protocol []string

	// WebSocket protocol version.
	Version int

	// TLS config for secure WebSocket (wss).
	TlsConfig *tls.Config

	handshakeData map[string]string
}

// serverHandshaker is an interface to handle WebSocket server side handshake.
type serverHandshaker interface {
	// ReadHandshake reads handshake request message from client.
	// Returns http response code and error if any.
	ReadHandshake(buf *bufio.Reader, req *http.Request) (code int, err os.Error)

	// AcceptHandshake accepts the client handshake request and sends
	// handshake response back to client.
	AcceptHandshake(buf *bufio.Writer) (err os.Error)

	// NewServerConn creates a new WebSocket connection.
	NewServerConn(buf *bufio.ReadWriter, rwc io.ReadWriteCloser, request *http.Request) (conn *Conn)
}

// frameReader is an interface to read a WebSocket frame.
type frameReader interface {
	// Reader is to read payload of the frame.
	io.Reader

	// PayloadType returns payload type.
	PayloadType() byte

	// HeaderReader returns a reader to read header of the frame.
	HeaderReader() io.Reader

	// TrailerReader returns a reader to read trailer of the frame.
	// If it returns nil, there is no trailer in the frame.
	TrailerReader() io.Reader

	// Len returns total length of the frame, including header and trailer.
	Len() int
}

// frameReaderFactory is an interface to creates new frame reader.
type frameReaderFactory interface {
	NewFrameReader() (r frameReader, err os.Error)
}

// frameWriter is an interface to write a WebSocket frame.
type frameWriter interface {
	// Writer is to write playload of the frame.
	io.WriteCloser
}

// frameWriterFactory is an interface to create new frame writer.
type frameWriterFactory interface {
	NewFrameWriter(payloadType byte) (w frameWriter, err os.Error)
}

type frameHandler interface {
	HandleFrame(frame frameReader) (r frameReader, err os.Error)
	WriteClose(status int) (err os.Error)
}

// Conn represents a WebSocket connection.
type Conn struct {
	config  *Config
	request *http.Request

	buf *bufio.ReadWriter
	rwc io.ReadWriteCloser

	frameReaderFactory
	frameReader

	frameWriterFactory

	frameHandler
	PayloadType        byte
	defaultCloseStatus int
}

// Read implements the io.Reader interface:
// it reads data of a frame from the WebSocket connection.
// if msg is not large enough for the frame data, it fills the msg and next Read
// will read the rest of the frame data.
// it reads Text frame or Binary frame.
func (ws *Conn) Read(msg []byte) (n int, err os.Error) {
again:
	if ws.frameReader == nil {
		frame, err := ws.frameReaderFactory.NewFrameReader()
		if err != nil {
			return 0, err
		}
		ws.frameReader, err = ws.frameHandler.HandleFrame(frame)
		if err != nil {
			return 0, err
		}
		if ws.frameReader == nil {
			goto again
		}
	}
	n, err = ws.frameReader.Read(msg)
	if err == os.EOF {
		if trailer := ws.frameReader.TrailerReader(); trailer != nil {
			io.Copy(ioutil.Discard, trailer)
		}
		ws.frameReader = nil
		goto again
	}
	return n, err
}

// Write implements the io.Writer interface:
// it writes data as a frame to the WebSocket connection.
func (ws *Conn) Write(msg []byte) (n int, err os.Error) {
	w, err := ws.frameWriterFactory.NewFrameWriter(ws.PayloadType)
	if err != nil {
		return 0, err
	}
	n, err = w.Write(msg)
	w.Close()
	if err != nil {
		return n, err
	}
	return n, err
}

// Close implements the io.Closer interface.
func (ws *Conn) Close() os.Error {
	err := ws.frameHandler.WriteClose(ws.defaultCloseStatus)
	if err != nil {
		return err
	}
	return ws.rwc.Close()
}

func (ws *Conn) IsClientConn() bool { return ws.request == nil }
func (ws *Conn) IsServerConn() bool { return ws.request != nil }

// LocalAddr returns the WebSocket Origin for the connection for client, or
// the WebSocket location for server.
func (ws *Conn) LocalAddr() net.Addr {
	if ws.IsClientConn() {
		return &Addr{ws.config.Origin}
	}
	return &Addr{ws.config.Location}
}

// RemoteAddr returns the WebSocket location for the connection for client, or
// the Websocket Origin for server.
func (ws *Conn) RemoteAddr() net.Addr {
	if ws.IsClientConn() {
		return &Addr{ws.config.Location}
	}
	return &Addr{ws.config.Origin}
}

// SetTimeout sets the connection's network timeout in nanoseconds.
func (ws *Conn) SetTimeout(nsec int64) os.Error {
	if conn, ok := ws.rwc.(net.Conn); ok {
		return conn.SetTimeout(nsec)
	}
	return os.EINVAL
}

// SetReadTimeout sets the connection's network read timeout in nanoseconds.
func (ws *Conn) SetReadTimeout(nsec int64) os.Error {
	if conn, ok := ws.rwc.(net.Conn); ok {
		return conn.SetReadTimeout(nsec)
	}
	return os.EINVAL
}

// SetWriteTimeout sets the connection's network write timeout in nanoseconds.
func (ws *Conn) SetWriteTimeout(nsec int64) os.Error {
	if conn, ok := ws.rwc.(net.Conn); ok {
		return conn.SetWriteTimeout(nsec)
	}
	return os.EINVAL
}

// Config returns the WebSocket config.
func (ws *Conn) Config() *Config { return ws.config }

// Request returns the http request upgraded to the WebSocket.
// It is nil for client side.
func (ws *Conn) Request() *http.Request { return ws.request }

// Codec represents a symmetric pair of functions that implement a codec.
type Codec struct {
	Marshal   func(v interface{}) (data []byte, payloadType byte, err os.Error)
	Unmarshal func(data []byte, payloadType byte, v interface{}) (err os.Error)
}

// Send sends v marshaled by cd.Marshal as single frame to ws.
func (cd Codec) Send(ws *Conn, v interface{}) (err os.Error) {
	if err != nil {
		return err
	}
	data, payloadType, err := cd.Marshal(v)
	if err != nil {
		return err
	}
	w, err := ws.frameWriterFactory.NewFrameWriter(payloadType)
	_, err = w.Write(data)
	w.Close()
	return err
}

// Receive receives single frame from ws, unmarshaled by cd.Unmarshal and stores in v.
func (cd Codec) Receive(ws *Conn, v interface{}) (err os.Error) {
	if ws.frameReader != nil {
		_, err = io.Copy(ioutil.Discard, ws.frameReader)
		if err != nil {
			return err
		}
		ws.frameReader = nil
	}
again:
	frame, err := ws.frameReaderFactory.NewFrameReader()
	if err != nil {
		return err
	}
	frame, err = ws.frameHandler.HandleFrame(frame)
	if err != nil {
		return err
	}
	if frame == nil {
		goto again
	}
	payloadType := frame.PayloadType()
	data, err := ioutil.ReadAll(frame)
	if err != nil {
		return err
	}
	return cd.Unmarshal(data, payloadType, v)
}

func marshal(v interface{}) (msg []byte, payloadType byte, err os.Error) {
	switch data := v.(type) {
	case string:
		return []byte(data), TextFrame, nil
	case []byte:
		return data, BinaryFrame, nil
	}
	return nil, UnknownFrame, ErrNotSupported
}

func unmarshal(msg []byte, payloadType byte, v interface{}) (err os.Error) {
	switch data := v.(type) {
	case *string:
		*data = string(msg)
		return nil
	case *[]byte:
		*data = msg
		return nil
	}
	return ErrNotSupported
}

/*
Message is a codec to send/receive text/binary data in a frame on WebSocket connection.
To send/receive text frame, use string type.
To send/receive binary frame, use []byte type.

Trivial usage:

	import "websocket"

	// receive text frame
	var message string
	websocket.Message.Receive(ws, &message)

	// send text frame
	message = "hello"
	websocket.Message.Send(ws, message)

	// receive binary frame
	var data []byte
	websocket.Message.Receive(ws, &data)

	// send binary frame
	data = []byte{0, 1, 2}
	websocket.Message.Send(ws, data)

*/
var Message = Codec{marshal, unmarshal}

func jsonMarshal(v interface{}) (msg []byte, payloadType byte, err os.Error) {
	msg, err = json.Marshal(v)
	return msg, TextFrame, err
}

func jsonUnmarshal(msg []byte, payloadType byte, v interface{}) (err os.Error) {
	return json.Unmarshal(msg, v)
}

/*
JSON is a codec to send/receive JSON data in a frame from a WebSocket connection.

Trival usage:

	import "websocket"

	type T struct {
		Msg string
		Count int
	}

	// receive JSON type T
	var data T
	websocket.JSON.Receive(ws, &data)

	// send JSON type T
	websocket.JSON.Send(ws, data)
*/
var JSON = Codec{jsonMarshal, jsonUnmarshal}
<|endoftext|>