<commit_before>package aero

import (
	"fmt"
	"io"
	"strings"

	"github.com/akyoto/color"
)

// node types
const (
	separator = '/'
	parameter = ':'
	wildcard  = '*'
)

// treeNode represents a radix tree node.
type treeNode struct {
	prefix     string
	startIndex uint8
	endIndex   uint8
	indices    []uint8
	children   []*treeNode
	data       dataType
	parameter  *treeNode
	wildcard   *treeNode
	kind       byte
}

// split splits the node at the given index and inserts
// a new child node with the given path and data.
// If path is empty, it will not create another child node
// and instead assign the data directly to the node.
func (node *treeNode) split(index int, path string, data dataType) {
	// Create split node with the remaining string
	splitNode := node.clone(node.prefix[index:])

	/// The existing data must be removed
	node.reset(node.prefix[:index])

	// If the path is empty, it means we don't create a 2nd child node.
	// Just assign the data for the existing node and store a single child node.
	if path == "" {
		node.data = data
		node.addChild(splitNode)
		return
	}

	node.addChild(splitNode)

	// Create new nodes with the remaining path
	node.append(path, data)
}

// clone clones the node with a new prefix.
func (node *treeNode) clone(prefix string) *treeNode {
	return &treeNode{
		prefix:     prefix,
		data:       node.data,
		indices:    node.indices,
		startIndex: node.startIndex,
		endIndex:   node.endIndex,
		children:   node.children,
		parameter:  node.parameter,
		wildcard:   node.wildcard,
		kind:       node.kind,
	}
}

// reset resets the existing node data.
func (node *treeNode) reset(prefix string) {
	node.prefix = prefix
	node.data = nil
	node.parameter = nil
	node.wildcard = nil
	node.kind = 0
	node.startIndex = 0
	node.endIndex = 0
	node.indices = nil
	node.children = nil
}

// addChild adds a child tree.
func (node *treeNode) addChild(child *treeNode) {
	if len(node.children) == 0 {
		node.children = append(node.children, nil)
	}

	firstChar := child.prefix[0]

	switch {
	case node.startIndex == 0:
		node.startIndex = firstChar
		node.indices = []uint8{0}
		node.endIndex = node.startIndex + uint8(len(node.indices))

	case firstChar < node.startIndex:
		diff := node.startIndex - firstChar
		newIndices := make([]uint8, diff+uint8(len(node.indices)))
		copy(newIndices[diff:], node.indices)
		node.startIndex = firstChar
		node.indices = newIndices
		node.endIndex = node.startIndex + uint8(len(node.indices))

	case firstChar >= node.endIndex:
		diff := firstChar - node.endIndex + 1
		newIndices := make([]uint8, diff+uint8(len(node.indices)))
		copy(newIndices, node.indices)
		node.indices = newIndices
		node.endIndex = node.startIndex + uint8(len(node.indices))
	}

	index := node.indices[firstChar-node.startIndex]

	if index == 0 {
		node.indices[firstChar-node.startIndex] = uint8(len(node.children))
		node.children = append(node.children, child)
		return
	}

	node.children[index] = child
}

// addTrailingSlash adds a trailing slash with the same data.
func (node *treeNode) addTrailingSlash(data dataType) {
	if strings.HasSuffix(node.prefix, "/") || node.kind == wildcard || (separator >= node.startIndex && separator < node.endIndex && node.indices[separator-node.startIndex] != 0) {
		return
	}

	node.addChild(&treeNode{
		prefix: "/",
		data:   data,
	})
}

// append appends the given path to the tree.
func (node *treeNode) append(path string, data dataType) {
	// At this point, all we know is that somewhere
	// in the remaining string we have parameters.
	// node: /user|
	// path: /user|/:userid
	for {
		if path == "" {
			node.data = data
			return
		}

		paramStart := strings.IndexByte(path, parameter)

		if paramStart == -1 {
			paramStart = strings.IndexByte(path, wildcard)
		}

		// If it's a static route we are adding,
		// just add the remainder as a normal node.
		if paramStart == -1 {
			// If the node itself doesn't have a prefix (root node),
			// don't add a child and use the node itself.
			if node.prefix == "" {
				node.prefix = path
				node.data = data
				return
			}

			child := &treeNode{
				prefix: path,
				data:   data,
			}

			node.addChild(child)
			child.addTrailingSlash(data)
			return
		}

		// If we're directly in front of a parameter,
		// add a parameter node.
		if paramStart == 0 {
			paramEnd := strings.IndexByte(path, separator)

			if paramEnd == -1 {
				paramEnd = len(path)
			}

			child := &treeNode{
				prefix: path[1:paramEnd],
				kind:   path[paramStart],
			}

			switch child.kind {
			case parameter:
				child.addTrailingSlash(data)
				node.parameter = child
				node = child
				path = path[paramEnd:]
				continue

			case wildcard:
				child.data = data
				node.wildcard = child
				return
			}
		}

		// We know there's a parameter, but not directly at the start.

		// If the node itself doesn't have a prefix (root node),
		// don't add a child and use the node itself.
		if node.prefix == "" {
			node.prefix = path[:paramStart]
			path = path[paramStart:]
			continue
		}

		// Add a normal node with the path before the parameter start.
		child := &treeNode{
			prefix: path[:paramStart],
		}

		// Allow trailing slashes to return
		// the same content as their parent node.
		if child.prefix == "/" {
			child.data = node.data
		}

		node.addChild(child)
		node = child
		path = path[paramStart:]
	}
}

// end is called when the node was fully parsed
// and needs to decide the next control flow.
func (node *treeNode) end(path string, data dataType, i int, offset int) (*treeNode, int, controlFlow) {
	char := path[i]

	if char >= node.startIndex && char < node.endIndex {
		index := node.indices[char-node.startIndex]

		if index != 0 {
			node = node.children[index]
			offset = i
			return node, offset, controlNext
		}
	}

	// No fitting children found, does this node even contain a prefix yet?
	// If no prefix is set, this is the starting node.
	if node.prefix == "" {
		node.append(path[i:], data)
		return node, offset, controlStop
	}

	// node: /user/|:id
	// path: /user/|:id/profile
	if node.parameter != nil {
		node = node.parameter
		offset = i
		return node, offset, controlBegin
	}

	node.append(path[i:], data)
	return node, offset, controlStop
}

// each traverses the tree and calls the given function on every node.
func (node *treeNode) each(callback func(*treeNode)) {
	callback(node)

	for _, child := range node.children {
		if child == nil {
			continue
		}

		child.each(callback)
	}

	if node.parameter != nil {
		node.parameter.each(callback)
	}

	if node.wildcard != nil {
		node.wildcard.each(callback)
	}
}

// PrettyPrint prints a human-readable form of the tree to the given writer.
func (node *treeNode) PrettyPrint(writer io.Writer) {
	node.prettyPrint(writer, -1)
}

// prettyPrint
func (node *treeNode) prettyPrint(writer io.Writer, level int) {
	prefix := ""

	if level >= 0 {
		prefix = strings.Repeat("  ", level) + "|_ "
	}

	colorFunc := color.CyanString

	switch node.kind {
	case ':':
		prefix += ":"
		colorFunc = color.YellowString
	case '*':
		prefix += "*"
		colorFunc = color.GreenString
	}

	fmt.Fprintf(writer, "%s%s [%t]\n", prefix, colorFunc(node.prefix), node.data != nil)

	for _, child := range node.children {
		if child == nil {
			continue
		}

		child.prettyPrint(writer, level+1)
	}

	if node.parameter != nil {
		node.parameter.prettyPrint(writer, level+1)
	}

	if node.wildcard != nil {
		node.wildcard.prettyPrint(writer, level+1)
	}
}
<commit_msg>Minor typo<commit_after>package aero

import (
	"fmt"
	"io"
	"strings"

	"github.com/akyoto/color"
)

// node types
const (
	separator = '/'
	parameter = ':'
	wildcard  = '*'
)

// treeNode represents a radix tree node.
type treeNode struct {
	prefix     string
	startIndex uint8
	endIndex   uint8
	indices    []uint8
	children   []*treeNode
	data       dataType
	parameter  *treeNode
	wildcard   *treeNode
	kind       byte
}

// split splits the node at the given index and inserts
// a new child node with the given path and data.
// If path is empty, it will not create another child node
// and instead assign the data directly to the node.
func (node *treeNode) split(index int, path string, data dataType) {
	// Create split node with the remaining string
	splitNode := node.clone(node.prefix[index:])

	// The existing data must be removed
	node.reset(node.prefix[:index])

	// If the path is empty, it means we don't create a 2nd child node.
	// Just assign the data for the existing node and store a single child node.
	if path == "" {
		node.data = data
		node.addChild(splitNode)
		return
	}

	node.addChild(splitNode)

	// Create new nodes with the remaining path
	node.append(path, data)
}

// clone clones the node with a new prefix.
func (node *treeNode) clone(prefix string) *treeNode {
	return &treeNode{
		prefix:     prefix,
		data:       node.data,
		indices:    node.indices,
		startIndex: node.startIndex,
		endIndex:   node.endIndex,
		children:   node.children,
		parameter:  node.parameter,
		wildcard:   node.wildcard,
		kind:       node.kind,
	}
}

// reset resets the existing node data.
func (node *treeNode) reset(prefix string) {
	node.prefix = prefix
	node.data = nil
	node.parameter = nil
	node.wildcard = nil
	node.kind = 0
	node.startIndex = 0
	node.endIndex = 0
	node.indices = nil
	node.children = nil
}

// addChild adds a child tree.
func (node *treeNode) addChild(child *treeNode) {
	if len(node.children) == 0 {
		node.children = append(node.children, nil)
	}

	firstChar := child.prefix[0]

	switch {
	case node.startIndex == 0:
		node.startIndex = firstChar
		node.indices = []uint8{0}
		node.endIndex = node.startIndex + uint8(len(node.indices))

	case firstChar < node.startIndex:
		diff := node.startIndex - firstChar
		newIndices := make([]uint8, diff+uint8(len(node.indices)))
		copy(newIndices[diff:], node.indices)
		node.startIndex = firstChar
		node.indices = newIndices
		node.endIndex = node.startIndex + uint8(len(node.indices))

	case firstChar >= node.endIndex:
		diff := firstChar - node.endIndex + 1
		newIndices := make([]uint8, diff+uint8(len(node.indices)))
		copy(newIndices, node.indices)
		node.indices = newIndices
		node.endIndex = node.startIndex + uint8(len(node.indices))
	}

	index := node.indices[firstChar-node.startIndex]

	if index == 0 {
		node.indices[firstChar-node.startIndex] = uint8(len(node.children))
		node.children = append(node.children, child)
		return
	}

	node.children[index] = child
}

// addTrailingSlash adds a trailing slash with the same data.
func (node *treeNode) addTrailingSlash(data dataType) {
	if strings.HasSuffix(node.prefix, "/") || node.kind == wildcard || (separator >= node.startIndex && separator < node.endIndex && node.indices[separator-node.startIndex] != 0) {
		return
	}

	node.addChild(&treeNode{
		prefix: "/",
		data:   data,
	})
}

// append appends the given path to the tree.
func (node *treeNode) append(path string, data dataType) {
	// At this point, all we know is that somewhere
	// in the remaining string we have parameters.
	// node: /user|
	// path: /user|/:userid
	for {
		if path == "" {
			node.data = data
			return
		}

		paramStart := strings.IndexByte(path, parameter)

		if paramStart == -1 {
			paramStart = strings.IndexByte(path, wildcard)
		}

		// If it's a static route we are adding,
		// just add the remainder as a normal node.
		if paramStart == -1 {
			// If the node itself doesn't have a prefix (root node),
			// don't add a child and use the node itself.
			if node.prefix == "" {
				node.prefix = path
				node.data = data
				return
			}

			child := &treeNode{
				prefix: path,
				data:   data,
			}

			node.addChild(child)
			child.addTrailingSlash(data)
			return
		}

		// If we're directly in front of a parameter,
		// add a parameter node.
		if paramStart == 0 {
			paramEnd := strings.IndexByte(path, separator)

			if paramEnd == -1 {
				paramEnd = len(path)
			}

			child := &treeNode{
				prefix: path[1:paramEnd],
				kind:   path[paramStart],
			}

			switch child.kind {
			case parameter:
				child.addTrailingSlash(data)
				node.parameter = child
				node = child
				path = path[paramEnd:]
				continue

			case wildcard:
				child.data = data
				node.wildcard = child
				return
			}
		}

		// We know there's a parameter, but not directly at the start.

		// If the node itself doesn't have a prefix (root node),
		// don't add a child and use the node itself.
		if node.prefix == "" {
			node.prefix = path[:paramStart]
			path = path[paramStart:]
			continue
		}

		// Add a normal node with the path before the parameter start.
		child := &treeNode{
			prefix: path[:paramStart],
		}

		// Allow trailing slashes to return
		// the same content as their parent node.
		if child.prefix == "/" {
			child.data = node.data
		}

		node.addChild(child)
		node = child
		path = path[paramStart:]
	}
}

// end is called when the node was fully parsed
// and needs to decide the next control flow.
func (node *treeNode) end(path string, data dataType, i int, offset int) (*treeNode, int, controlFlow) {
	char := path[i]

	if char >= node.startIndex && char < node.endIndex {
		index := node.indices[char-node.startIndex]

		if index != 0 {
			node = node.children[index]
			offset = i
			return node, offset, controlNext
		}
	}

	// No fitting children found, does this node even contain a prefix yet?
	// If no prefix is set, this is the starting node.
	if node.prefix == "" {
		node.append(path[i:], data)
		return node, offset, controlStop
	}

	// node: /user/|:id
	// path: /user/|:id/profile
	if node.parameter != nil {
		node = node.parameter
		offset = i
		return node, offset, controlBegin
	}

	node.append(path[i:], data)
	return node, offset, controlStop
}

// each traverses the tree and calls the given function on every node.
func (node *treeNode) each(callback func(*treeNode)) {
	callback(node)

	for _, child := range node.children {
		if child == nil {
			continue
		}

		child.each(callback)
	}

	if node.parameter != nil {
		node.parameter.each(callback)
	}

	if node.wildcard != nil {
		node.wildcard.each(callback)
	}
}

// PrettyPrint prints a human-readable form of the tree to the given writer.
func (node *treeNode) PrettyPrint(writer io.Writer) {
	node.prettyPrint(writer, -1)
}

// prettyPrint
func (node *treeNode) prettyPrint(writer io.Writer, level int) {
	prefix := ""

	if level >= 0 {
		prefix = strings.Repeat("  ", level) + "|_ "
	}

	colorFunc := color.CyanString

	switch node.kind {
	case ':':
		prefix += ":"
		colorFunc = color.YellowString
	case '*':
		prefix += "*"
		colorFunc = color.GreenString
	}

	fmt.Fprintf(writer, "%s%s [%t]\n", prefix, colorFunc(node.prefix), node.data != nil)

	for _, child := range node.children {
		if child == nil {
			continue
		}

		child.prettyPrint(writer, level+1)
	}

	if node.parameter != nil {
		node.parameter.prettyPrint(writer, level+1)
	}

	if node.wildcard != nil {
		node.wildcard.prettyPrint(writer, level+1)
	}
}
<|endoftext|>


<commit_before>package store

import (
	"encoding/json"
	"errors"
	"net/http"
	"sort"

	"github.com/manishrjain/gocrud/x"
)

var (
	ErrNoParent = errors.New("No parent found")
)

// Query stores the read instrutions, storing the instruction set
// for the entities Query relates to.
type Query struct {
	kind      string
	id        string
	filterOut map[string]bool
	maxDepth  int
	children  []*Query
	parent    *Query
}

type Object struct {
	Value  interface{}
	Source string
	NanoTs int64
}

type Versions struct {
	versions []Object
}

// Result stores the final entity state retrieved from Store upon running
// the instructions provided by Query.
type Result struct {
	Id       string
	Kind     string
	Columns  map[string]*Versions
	Children []*Result
}

type runResult struct {
	Result *Result
	Err    error
}

func (v *Versions) add(o Object) {
	if len(v.versions) > 0 {
		i := len(v.versions) - 1
		if v.versions[i].NanoTs > o.NanoTs {
			// unsorted list. Gocrud code is doing something wrong.
			log.Fatal("Appending an object with lower ts to a sorted list")
		}
	}
	v.versions = append(v.versions, o)
}

func (v Versions) Latest() Object {
	if len(v.versions) == 0 {
		return Object{}
	}
	i := len(v.versions) - 1
	return v.versions[i]
}

func (v Versions) Oldest() Object {
	if len(v.versions) == 0 {
		return Object{}
	}
	return v.versions[0]
}

func (v Versions) Count() int {
	return len(v.versions)
}

// Retrieve the parent id for given entity id. Return ErrNoParent if parent is
// not present. Otherwise, if an error occurs during retrieval, returns that.
func Parent(id string) (parentid string, rerr error) {
	its, err := Get().GetEntity(id)
	if err != nil {
		x.LogErr(log, err).WithField("id", id).Error("While retrieving entity")
		return "", err
	}
	for _, it := range its {
		if it.Predicate == "_parent_" {
			return it.ObjectId, nil
		}
	}
	return "", ErrNoParent
}

// NewQuery is the main entrypoint to data store queries. Returns back a Query
// object pointer, to run read instructions on.
func NewQuery(id string) *Query {
	q := new(Query)
	q.id = id
	return q
}

// UptoDepth specifies the number of levels of descendants that would be
// retrieved for the entity Query points to.
//
// You can think of this in terms of a tree structure, where the Entity pointed
// to by Query points to other Entities, which in turn point to more Entities,
// and so on. For e.g. Post -> Comments -> Likes.
func (q *Query) UptoDepth(level int) *Query {
	q.maxDepth = level
	return q
}

// Collect specifies the kind of child entities to retrieve. Returns back
// a new Query pointer pointing to those children entities as a collective.
//
// Any further operations on this returned pointer would attribute to those
// children entities, and not the caller query entity.
func (q *Query) Collect(kind string) *Query {
	for _, child := range q.children {
		if child.kind == kind {
			return child
		}
	}
	child := new(Query)
	child.parent = q
	child.kind = kind
	q.children = append(q.children, child)
	return child
}

// FilterOut provides a way to well, filter out, any entities which have
// the given property.
func (q *Query) FilterOut(property string) *Query {
	if len(q.filterOut) == 0 {
		q.filterOut = make(map[string]bool)
	}
	q.filterOut[property] = true
	return q
}

func (q *Query) root() *Query {
	for q.parent != nil {
		q = q.parent
	}
	return q
}

func (q *Query) doRun(level, max int, ch chan runResult) {
	log.Debugf("Query: %+v", q)
	its, err := Get().GetEntity(q.id)
	if err != nil {
		x.LogErr(log, err).Error("While retrieving: ", q.id)
		ch <- runResult{Result: nil, Err: err}
		return
	}
	if len(its) == 0 {
		ch <- runResult{Result: new(Result), Err: nil}
		return
	}
	sort.Sort(x.Its(its))

	follow := make(map[string]*Query)
	for _, child := range q.children {
		follow[child.kind] = child
		log.WithField("kind", child.kind).WithField("child", child).Debug("Following")
	}

	result := new(Result)
	result.Columns = make(map[string]*Versions)
	it := its[0]
	result.Id = it.SubjectId
	result.Kind = it.SubjectType

	waitTimes := 0
	childChan := make(chan runResult)
	for _, it := range its {
		if it.Predicate == "_delete_" {
			// If marked as deleted, don't return this node.
			log.WithField("id", result.Id).
				WithField("kind", result.Kind).
				WithField("_delete_", true).
				Debug("Discarding due to delete bit")
			ch <- runResult{Result: new(Result), Err: nil}
			return
		}

		if it.Predicate == "_parent_" {
			log.WithField("id", result.Id).
				WithField("kind", result.Kind).
				WithField("parent", it.ObjectId).
				Debug("Not following edge back to parent")
			continue
		}

		if _, fout := q.filterOut[it.Predicate]; fout {
			log.WithField("id", result.Id).
				WithField("kind", result.Kind).
				WithField("predicate", it.Predicate).
				Debug("Discarding due to predicate filter")
			ch <- runResult{Result: new(Result), Err: nil}
			return
		}

		if len(it.ObjectId) == 0 {
			o := Object{NanoTs: it.NanoTs, Source: it.Source}
			if err := json.Unmarshal(it.Object, &o.Value); err != nil {
				x.LogErr(log, err).Error("While unmarshal")
				ch <- runResult{Result: nil, Err: err}
				return
			}

			if _, vok := result.Columns[it.Predicate]; !vok {
				result.Columns[it.Predicate] = new(Versions)
			}
			result.Columns[it.Predicate].add(o)
			continue
		}

		if childq, fw := follow[it.Predicate]; fw {
			nchildq := new(Query)
			*nchildq = *childq // This is important, otherwise id gets overwritten
			nchildq.id = it.ObjectId

			// Use child's maxDepth here, instead of parent's.
			waitTimes += 1
			log.WithField("child_id", nchildq.id).
				WithField("child_kind", nchildq.kind).Debug("Go routine for child")
			go nchildq.doRun(0, nchildq.maxDepth, childChan)
			continue
		}

		if len(it.ObjectId) > 0 && level < max {
			child := new(Query)
			child.id = it.ObjectId

			waitTimes += 1
			log.WithField("child_id", child.id).WithField("level", level+1).
				Debug("Go routine for child one level deeper")
			go child.doRun(level+1, max, childChan)
		}
	}

	// Wait for all those subroutines
	for i := 0; i < waitTimes; i++ {
		log.Debugf("Waiting for children subroutines: %v/%v", i, waitTimes-1)
		rr := <-childChan
		log.Debugf("Waiting done")
		if rr.Err != nil {
			x.LogErr(log, err).Error("While child doRun")
		} else {
			if len(rr.Result.Id) > 0 && len(rr.Result.Kind) > 0 {
				log.WithField("result", *rr.Result).Debug("Appending child")
				result.Children = append(result.Children, rr.Result)
			}
		}
	}

	ch <- runResult{Result: result, Err: nil}
	return
}

// Run finds the root from the given Query pointer, recursively executes
// the read operations, and returns back pointer to Result object.
// Any errors encountered during these stpeps is returned as well.
func (q *Query) Run() (result *Result, rerr error) {
	q = q.root()

	ch := make(chan runResult)
	go q.doRun(0, q.maxDepth, ch)
	rr := <-ch // Blocking wait
	return rr.Result, rr.Err
}

func (r *Result) Drop(pred string) {
	delete(r.Columns, pred)
}

func (r *Result) Debug(level int) {
	log.Debugf("Result level %v: %+v", level, r)
	for _, child := range r.Children {
		child.Debug(level + 1)
	}
}

func (r *Result) ToMap() (data map[string]interface{}) {
	data = make(map[string]interface{})
	data["id"] = r.Id
	data["kind"] = r.Kind
	var ts int64
	for pred, versions := range r.Columns {
		// During conversion to JSON, to keep things simple,
		// we're dropping older versions of predicates, and
		// source and ts information across all the predicates,
		// keeping only the latest one.
		data[pred] = versions.Latest().Value
		data["source"] = versions.Latest().Source // Loss of information.
		if versions.Latest().NanoTs > ts {
			ts = versions.Latest().NanoTs
		}
	}
	data["ts_millis"] = int(ts / 1000000) // Loss of information. Picking up latest mod time.
	kinds := make(map[string]bool)
	for _, child := range r.Children {
		kinds[child.Kind] = true
	}

	for kind := range kinds {
		var l []map[string]interface{}
		for _, child := range r.Children {
			if kind != child.Kind {
				continue
			}
			l = append(l, child.ToMap())
		}
		data[kind] = l
	}

	return
}

// ToJson creates the JSON for the data pointed by the Result pointer.
//
// Note that this doesn't find the "root" from the Result pointer, instead
// doing the processing only from the current Result pointer.
func (r *Result) ToJson() ([]byte, error) {
	data := r.ToMap()
	return json.Marshal(data)
}

// WriteJsonResponse does the same as ToJson. But also writes the JSON
// generated to http.ResponseWriter. In case of error, writes that error
// instead, in this format:
//
//  {
//   "code":    "E_ERROR",
//   "message": "error message"
//  }
func (r *Result) WriteJsonResponse(w http.ResponseWriter) {
	data, err := r.ToJson()
	if err != nil {
		x.SetStatus(w, x.E_ERROR, err.Error())
		return
	}
	_, err = w.Write(data)
	if err != nil {
		x.SetStatus(w, x.E_ERROR, err.Error())
		return
	}
}
<commit_msg>Added a TODO about unbuffered channel<commit_after>package store

import (
	"encoding/json"
	"errors"
	"net/http"
	"sort"

	"github.com/manishrjain/gocrud/x"
)

var (
	ErrNoParent = errors.New("No parent found")
)

// Query stores the read instrutions, storing the instruction set
// for the entities Query relates to.
type Query struct {
	kind      string
	id        string
	filterOut map[string]bool
	maxDepth  int
	children  []*Query
	parent    *Query
}

type Object struct {
	Value  interface{}
	Source string
	NanoTs int64
}

type Versions struct {
	versions []Object
}

// Result stores the final entity state retrieved from Store upon running
// the instructions provided by Query.
type Result struct {
	Id       string
	Kind     string
	Columns  map[string]*Versions
	Children []*Result
}

type runResult struct {
	Result *Result
	Err    error
}

func (v *Versions) add(o Object) {
	if len(v.versions) > 0 {
		i := len(v.versions) - 1
		if v.versions[i].NanoTs > o.NanoTs {
			// unsorted list. Gocrud code is doing something wrong.
			log.Fatal("Appending an object with lower ts to a sorted list")
		}
	}
	v.versions = append(v.versions, o)
}

func (v Versions) Latest() Object {
	if len(v.versions) == 0 {
		return Object{}
	}
	i := len(v.versions) - 1
	return v.versions[i]
}

func (v Versions) Oldest() Object {
	if len(v.versions) == 0 {
		return Object{}
	}
	return v.versions[0]
}

func (v Versions) Count() int {
	return len(v.versions)
}

// Retrieve the parent id for given entity id. Return ErrNoParent if parent is
// not present. Otherwise, if an error occurs during retrieval, returns that.
func Parent(id string) (parentid string, rerr error) {
	its, err := Get().GetEntity(id)
	if err != nil {
		x.LogErr(log, err).WithField("id", id).Error("While retrieving entity")
		return "", err
	}
	for _, it := range its {
		if it.Predicate == "_parent_" {
			return it.ObjectId, nil
		}
	}
	return "", ErrNoParent
}

// NewQuery is the main entrypoint to data store queries. Returns back a Query
// object pointer, to run read instructions on.
func NewQuery(id string) *Query {
	q := new(Query)
	q.id = id
	return q
}

// UptoDepth specifies the number of levels of descendants that would be
// retrieved for the entity Query points to.
//
// You can think of this in terms of a tree structure, where the Entity pointed
// to by Query points to other Entities, which in turn point to more Entities,
// and so on. For e.g. Post -> Comments -> Likes.
func (q *Query) UptoDepth(level int) *Query {
	q.maxDepth = level
	return q
}

// Collect specifies the kind of child entities to retrieve. Returns back
// a new Query pointer pointing to those children entities as a collective.
//
// Any further operations on this returned pointer would attribute to those
// children entities, and not the caller query entity.
func (q *Query) Collect(kind string) *Query {
	for _, child := range q.children {
		if child.kind == kind {
			return child
		}
	}
	child := new(Query)
	child.parent = q
	child.kind = kind
	q.children = append(q.children, child)
	return child
}

// FilterOut provides a way to well, filter out, any entities which have
// the given property.
func (q *Query) FilterOut(property string) *Query {
	if len(q.filterOut) == 0 {
		q.filterOut = make(map[string]bool)
	}
	q.filterOut[property] = true
	return q
}

func (q *Query) root() *Query {
	for q.parent != nil {
		q = q.parent
	}
	return q
}

func (q *Query) doRun(level, max int, ch chan runResult) {
	log.Debugf("Query: %+v", q)
	its, err := Get().GetEntity(q.id)
	if err != nil {
		x.LogErr(log, err).Error("While retrieving: ", q.id)
		ch <- runResult{Result: nil, Err: err}
		return
	}
	if len(its) == 0 {
		ch <- runResult{Result: new(Result), Err: nil}
		return
	}
	sort.Sort(x.Its(its))

	follow := make(map[string]*Query)
	for _, child := range q.children {
		follow[child.kind] = child
		log.WithField("kind", child.kind).WithField("child", child).Debug("Following")
	}

	result := new(Result)
	result.Columns = make(map[string]*Versions)
	it := its[0]
	result.Id = it.SubjectId
	result.Kind = it.SubjectType

	waitTimes := 0
	// TODO: This is an unbuffered channel. Why am I doing this? Look into it!
	childChan := make(chan runResult)
	for _, it := range its {
		if it.Predicate == "_delete_" {
			// If marked as deleted, don't return this node.
			log.WithField("id", result.Id).
				WithField("kind", result.Kind).
				WithField("_delete_", true).
				Debug("Discarding due to delete bit")
			ch <- runResult{Result: new(Result), Err: nil}
			return
		}

		if it.Predicate == "_parent_" {
			log.WithField("id", result.Id).
				WithField("kind", result.Kind).
				WithField("parent", it.ObjectId).
				Debug("Not following edge back to parent")
			continue
		}

		if _, fout := q.filterOut[it.Predicate]; fout {
			log.WithField("id", result.Id).
				WithField("kind", result.Kind).
				WithField("predicate", it.Predicate).
				Debug("Discarding due to predicate filter")
			ch <- runResult{Result: new(Result), Err: nil}
			return
		}

		if len(it.ObjectId) == 0 {
			o := Object{NanoTs: it.NanoTs, Source: it.Source}
			if err := json.Unmarshal(it.Object, &o.Value); err != nil {
				x.LogErr(log, err).Error("While unmarshal")
				ch <- runResult{Result: nil, Err: err}
				return
			}

			if _, vok := result.Columns[it.Predicate]; !vok {
				result.Columns[it.Predicate] = new(Versions)
			}
			result.Columns[it.Predicate].add(o)
			continue
		}

		if childq, fw := follow[it.Predicate]; fw {
			nchildq := new(Query)
			*nchildq = *childq // This is important, otherwise id gets overwritten
			nchildq.id = it.ObjectId

			// Use child's maxDepth here, instead of parent's.
			waitTimes += 1
			log.WithField("child_id", nchildq.id).
				WithField("child_kind", nchildq.kind).Debug("Go routine for child")
			go nchildq.doRun(0, nchildq.maxDepth, childChan)
			continue
		}

		if len(it.ObjectId) > 0 && level < max {
			child := new(Query)
			child.id = it.ObjectId

			waitTimes += 1
			log.WithField("child_id", child.id).WithField("level", level+1).
				Debug("Go routine for child one level deeper")
			go child.doRun(level+1, max, childChan)
		}
	}

	// Wait for all those subroutines
	for i := 0; i < waitTimes; i++ {
		log.Debugf("Waiting for children subroutines: %v/%v", i, waitTimes-1)
		rr := <-childChan
		log.Debugf("Waiting done")
		if rr.Err != nil {
			x.LogErr(log, err).Error("While child doRun")
		} else {
			if len(rr.Result.Id) > 0 && len(rr.Result.Kind) > 0 {
				log.WithField("result", *rr.Result).Debug("Appending child")
				result.Children = append(result.Children, rr.Result)
			}
		}
	}

	ch <- runResult{Result: result, Err: nil}
	return
}

// Run finds the root from the given Query pointer, recursively executes
// the read operations, and returns back pointer to Result object.
// Any errors encountered during these stpeps is returned as well.
func (q *Query) Run() (result *Result, rerr error) {
	q = q.root()

	ch := make(chan runResult)
	go q.doRun(0, q.maxDepth, ch)
	rr := <-ch // Blocking wait
	return rr.Result, rr.Err
}

func (r *Result) Drop(pred string) {
	delete(r.Columns, pred)
}

func (r *Result) Debug(level int) {
	log.Debugf("Result level %v: %+v", level, r)
	for _, child := range r.Children {
		child.Debug(level + 1)
	}
}

func (r *Result) ToMap() (data map[string]interface{}) {
	data = make(map[string]interface{})
	data["id"] = r.Id
	data["kind"] = r.Kind
	var ts int64
	for pred, versions := range r.Columns {
		// During conversion to JSON, to keep things simple,
		// we're dropping older versions of predicates, and
		// source and ts information across all the predicates,
		// keeping only the latest one.
		data[pred] = versions.Latest().Value
		data["source"] = versions.Latest().Source // Loss of information.
		if versions.Latest().NanoTs > ts {
			ts = versions.Latest().NanoTs
		}
	}
	data["ts_millis"] = int(ts / 1000000) // Loss of information. Picking up latest mod time.
	kinds := make(map[string]bool)
	for _, child := range r.Children {
		kinds[child.Kind] = true
	}

	for kind := range kinds {
		var l []map[string]interface{}
		for _, child := range r.Children {
			if kind != child.Kind {
				continue
			}
			l = append(l, child.ToMap())
		}
		data[kind] = l
	}

	return
}

// ToJson creates the JSON for the data pointed by the Result pointer.
//
// Note that this doesn't find the "root" from the Result pointer, instead
// doing the processing only from the current Result pointer.
func (r *Result) ToJson() ([]byte, error) {
	data := r.ToMap()
	return json.Marshal(data)
}

// WriteJsonResponse does the same as ToJson. But also writes the JSON
// generated to http.ResponseWriter. In case of error, writes that error
// instead, in this format:
//
//  {
//   "code":    "E_ERROR",
//   "message": "error message"
//  }
func (r *Result) WriteJsonResponse(w http.ResponseWriter) {
	data, err := r.ToJson()
	if err != nil {
		x.SetStatus(w, x.E_ERROR, err.Error())
		return
	}
	_, err = w.Write(data)
	if err != nil {
		x.SetStatus(w, x.E_ERROR, err.Error())
		return
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2015 The Kubernetes Authors.

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

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

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

package e2e

import (
	"fmt"
	"math"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"k8s.io/kubernetes/pkg/api/v1"
	metav1 "k8s.io/kubernetes/pkg/apis/meta/v1"
	clientset "k8s.io/kubernetes/pkg/client/clientset_generated/clientset"
	"k8s.io/kubernetes/pkg/labels"
	"k8s.io/kubernetes/pkg/util/intstr"
	"k8s.io/kubernetes/pkg/util/sets"
	"k8s.io/kubernetes/pkg/util/uuid"
	"k8s.io/kubernetes/test/e2e/framework"
	testutils "k8s.io/kubernetes/test/utils"
)

var _ = framework.KubeDescribe("Multi-AZ Clusters", func() {
	f := framework.NewDefaultFramework("multi-az")
	var zoneCount int
	var err error
	image := "gcr.io/google_containers/serve_hostname:v1.4"
	BeforeEach(func() {
		framework.SkipUnlessProviderIs("gce", "gke", "aws")
		if zoneCount <= 0 {
			zoneCount, err = getZoneCount(f.ClientSet)
			Expect(err).NotTo(HaveOccurred())
		}
		By(fmt.Sprintf("Checking for multi-zone cluster.  Zone count = %d", zoneCount))
		framework.SkipUnlessAtLeast(zoneCount, 2, "Zone count is %d, only run for multi-zone clusters, skipping test")
		// TODO: SkipUnlessDefaultScheduler() // Non-default schedulers might not spread
	})
	It("should spread the pods of a service across zones", func() {
		SpreadServiceOrFail(f, (2*zoneCount)+1, image)
	})

	It("should spread the pods of a replication controller across zones", func() {
		SpreadRCOrFail(f, int32((2*zoneCount)+1), image)
	})
})

// Check that the pods comprising a service get spread evenly across available zones
func SpreadServiceOrFail(f *framework.Framework, replicaCount int, image string) {
	// First create the service
	serviceName := "test-service"
	serviceSpec := &v1.Service{
		ObjectMeta: v1.ObjectMeta{
			Name:      serviceName,
			Namespace: f.Namespace.Name,
		},
		Spec: v1.ServiceSpec{
			Selector: map[string]string{
				"service": serviceName,
			},
			Ports: []v1.ServicePort{{
				Port:       80,
				TargetPort: intstr.FromInt(80),
			}},
		},
	}
	_, err := f.ClientSet.Core().Services(f.Namespace.Name).Create(serviceSpec)
	Expect(err).NotTo(HaveOccurred())

	// Now create some pods behind the service
	podSpec := &v1.Pod{
		ObjectMeta: v1.ObjectMeta{
			Name:   serviceName,
			Labels: map[string]string{"service": serviceName},
		},
		Spec: v1.PodSpec{
			Containers: []v1.Container{
				{
					Name:  "test",
					Image: framework.GetPauseImageName(f.ClientSet),
				},
			},
		},
	}

	// Caution: StartPods requires at least one pod to replicate.
	// Based on the callers, replicas is always positive number: zoneCount >= 0 implies (2*zoneCount)+1 > 0.
	// Thus, no need to test for it. Once the precondition changes to zero number of replicas,
	// test for replicaCount > 0. Otherwise, StartPods panics.
	framework.ExpectNoError(testutils.StartPods(f.ClientSet, replicaCount, f.Namespace.Name, serviceName, *podSpec, false, framework.Logf))

	// Wait for all of them to be scheduled
	selector := labels.SelectorFromSet(labels.Set(map[string]string{"service": serviceName}))
	pods, err := framework.WaitForPodsWithLabelScheduled(f.ClientSet, f.Namespace.Name, selector)
	Expect(err).NotTo(HaveOccurred())

	// Now make sure they're spread across zones
	zoneNames, err := getZoneNames(f.ClientSet)
	Expect(err).NotTo(HaveOccurred())
	Expect(checkZoneSpreading(f.ClientSet, pods, zoneNames)).To(Equal(true))
}

// Find the name of the zone in which a Node is running
func getZoneNameForNode(node v1.Node) (string, error) {
	for key, value := range node.Labels {
		if key == metav1.LabelZoneFailureDomain {
			return value, nil
		}
	}
	return "", fmt.Errorf("Zone name for node %s not found. No label with key %s",
		node.Name, metav1.LabelZoneFailureDomain)
}

// Find the names of all zones in which we have nodes in this cluster.
func getZoneNames(c clientset.Interface) ([]string, error) {
	zoneNames := sets.NewString()
	nodes, err := c.Core().Nodes().List(v1.ListOptions{})
	if err != nil {
		return nil, err
	}
	for _, node := range nodes.Items {
		zoneName, err := getZoneNameForNode(node)
		Expect(err).NotTo(HaveOccurred())
		zoneNames.Insert(zoneName)
	}
	return zoneNames.List(), nil
}

// Return the number of zones in which we have nodes in this cluster.
func getZoneCount(c clientset.Interface) (int, error) {
	zoneNames, err := getZoneNames(c)
	if err != nil {
		return -1, err
	}
	return len(zoneNames), nil
}

// Find the name of the zone in which the pod is scheduled
func getZoneNameForPod(c clientset.Interface, pod v1.Pod) (string, error) {
	By(fmt.Sprintf("Getting zone name for pod %s, on node %s", pod.Name, pod.Spec.NodeName))
	node, err := c.Core().Nodes().Get(pod.Spec.NodeName, metav1.GetOptions{})
	Expect(err).NotTo(HaveOccurred())
	return getZoneNameForNode(*node)
}

// Determine whether a set of pods are approximately evenly spread
// across a given set of zones
func checkZoneSpreading(c clientset.Interface, pods *v1.PodList, zoneNames []string) (bool, error) {
	podsPerZone := make(map[string]int)
	for _, zoneName := range zoneNames {
		podsPerZone[zoneName] = 0
	}
	for _, pod := range pods.Items {
		if pod.DeletionTimestamp != nil {
			continue
		}
		zoneName, err := getZoneNameForPod(c, pod)
		Expect(err).NotTo(HaveOccurred())
		podsPerZone[zoneName] = podsPerZone[zoneName] + 1
	}
	minPodsPerZone := math.MaxInt32
	maxPodsPerZone := 0
	for _, podCount := range podsPerZone {
		if podCount < minPodsPerZone {
			minPodsPerZone = podCount
		}
		if podCount > maxPodsPerZone {
			maxPodsPerZone = podCount
		}
	}
	Expect(minPodsPerZone).To(BeNumerically("~", maxPodsPerZone, 1),
		"Pods were not evenly spread across zones.  %d in one zone and %d in another zone",
		minPodsPerZone, maxPodsPerZone)
	return true, nil
}

// Check that the pods comprising a replication controller get spread evenly across available zones
func SpreadRCOrFail(f *framework.Framework, replicaCount int32, image string) {
	name := "ubelite-spread-rc-" + string(uuid.NewUUID())
	By(fmt.Sprintf("Creating replication controller %s", name))
	controller, err := f.ClientSet.Core().ReplicationControllers(f.Namespace.Name).Create(&v1.ReplicationController{
		ObjectMeta: v1.ObjectMeta{
			Namespace: f.Namespace.Name,
			Name:      name,
		},
		Spec: v1.ReplicationControllerSpec{
			Replicas: &replicaCount,
			Selector: map[string]string{
				"name": name,
			},
			Template: &v1.PodTemplateSpec{
				ObjectMeta: v1.ObjectMeta{
					Labels: map[string]string{"name": name},
				},
				Spec: v1.PodSpec{
					Containers: []v1.Container{
						{
							Name:  name,
							Image: image,
							Ports: []v1.ContainerPort{{ContainerPort: 9376}},
						},
					},
				},
			},
		},
	})
	Expect(err).NotTo(HaveOccurred())
	// Cleanup the replication controller when we are done.
	defer func() {
		// Resize the replication controller to zero to get rid of pods.
		if err := framework.DeleteRCAndPods(f.ClientSet, f.InternalClientset, f.Namespace.Name, controller.Name); err != nil {
			framework.Logf("Failed to cleanup replication controller %v: %v.", controller.Name, err)
		}
	}()
	// List the pods, making sure we observe all the replicas.
	selector := labels.SelectorFromSet(labels.Set(map[string]string{"name": name}))
	pods, err := framework.PodsCreated(f.ClientSet, f.Namespace.Name, name, replicaCount)
	Expect(err).NotTo(HaveOccurred())

	// Wait for all of them to be scheduled
	By(fmt.Sprintf("Waiting for %d replicas of %s to be scheduled.  Selector: %v", replicaCount, name, selector))
	pods, err = framework.WaitForPodsWithLabelScheduled(f.ClientSet, f.Namespace.Name, selector)
	Expect(err).NotTo(HaveOccurred())

	// Now make sure they're spread across zones
	zoneNames, err := getZoneNames(f.ClientSet)
	Expect(err).NotTo(HaveOccurred())
	Expect(checkZoneSpreading(f.ClientSet, pods, zoneNames)).To(Equal(true))
}
<commit_msg>Fix wrong skipf parameter<commit_after>/*
Copyright 2015 The Kubernetes Authors.

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

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

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

package e2e

import (
	"fmt"
	"math"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"k8s.io/kubernetes/pkg/api/v1"
	metav1 "k8s.io/kubernetes/pkg/apis/meta/v1"
	clientset "k8s.io/kubernetes/pkg/client/clientset_generated/clientset"
	"k8s.io/kubernetes/pkg/labels"
	"k8s.io/kubernetes/pkg/util/intstr"
	"k8s.io/kubernetes/pkg/util/sets"
	"k8s.io/kubernetes/pkg/util/uuid"
	"k8s.io/kubernetes/test/e2e/framework"
	testutils "k8s.io/kubernetes/test/utils"
)

var _ = framework.KubeDescribe("Multi-AZ Clusters", func() {
	f := framework.NewDefaultFramework("multi-az")
	var zoneCount int
	var err error
	image := "gcr.io/google_containers/serve_hostname:v1.4"
	BeforeEach(func() {
		framework.SkipUnlessProviderIs("gce", "gke", "aws")
		if zoneCount <= 0 {
			zoneCount, err = getZoneCount(f.ClientSet)
			Expect(err).NotTo(HaveOccurred())
		}
		By(fmt.Sprintf("Checking for multi-zone cluster.  Zone count = %d", zoneCount))
		msg := fmt.Sprintf("Zone count is %d, only run for multi-zone clusters, skipping test", zoneCount)
		framework.SkipUnlessAtLeast(zoneCount, 2, msg)
		// TODO: SkipUnlessDefaultScheduler() // Non-default schedulers might not spread
	})
	It("should spread the pods of a service across zones", func() {
		SpreadServiceOrFail(f, (2*zoneCount)+1, image)
	})

	It("should spread the pods of a replication controller across zones", func() {
		SpreadRCOrFail(f, int32((2*zoneCount)+1), image)
	})
})

// Check that the pods comprising a service get spread evenly across available zones
func SpreadServiceOrFail(f *framework.Framework, replicaCount int, image string) {
	// First create the service
	serviceName := "test-service"
	serviceSpec := &v1.Service{
		ObjectMeta: v1.ObjectMeta{
			Name:      serviceName,
			Namespace: f.Namespace.Name,
		},
		Spec: v1.ServiceSpec{
			Selector: map[string]string{
				"service": serviceName,
			},
			Ports: []v1.ServicePort{{
				Port:       80,
				TargetPort: intstr.FromInt(80),
			}},
		},
	}
	_, err := f.ClientSet.Core().Services(f.Namespace.Name).Create(serviceSpec)
	Expect(err).NotTo(HaveOccurred())

	// Now create some pods behind the service
	podSpec := &v1.Pod{
		ObjectMeta: v1.ObjectMeta{
			Name:   serviceName,
			Labels: map[string]string{"service": serviceName},
		},
		Spec: v1.PodSpec{
			Containers: []v1.Container{
				{
					Name:  "test",
					Image: framework.GetPauseImageName(f.ClientSet),
				},
			},
		},
	}

	// Caution: StartPods requires at least one pod to replicate.
	// Based on the callers, replicas is always positive number: zoneCount >= 0 implies (2*zoneCount)+1 > 0.
	// Thus, no need to test for it. Once the precondition changes to zero number of replicas,
	// test for replicaCount > 0. Otherwise, StartPods panics.
	framework.ExpectNoError(testutils.StartPods(f.ClientSet, replicaCount, f.Namespace.Name, serviceName, *podSpec, false, framework.Logf))

	// Wait for all of them to be scheduled
	selector := labels.SelectorFromSet(labels.Set(map[string]string{"service": serviceName}))
	pods, err := framework.WaitForPodsWithLabelScheduled(f.ClientSet, f.Namespace.Name, selector)
	Expect(err).NotTo(HaveOccurred())

	// Now make sure they're spread across zones
	zoneNames, err := getZoneNames(f.ClientSet)
	Expect(err).NotTo(HaveOccurred())
	Expect(checkZoneSpreading(f.ClientSet, pods, zoneNames)).To(Equal(true))
}

// Find the name of the zone in which a Node is running
func getZoneNameForNode(node v1.Node) (string, error) {
	for key, value := range node.Labels {
		if key == metav1.LabelZoneFailureDomain {
			return value, nil
		}
	}
	return "", fmt.Errorf("Zone name for node %s not found. No label with key %s",
		node.Name, metav1.LabelZoneFailureDomain)
}

// Find the names of all zones in which we have nodes in this cluster.
func getZoneNames(c clientset.Interface) ([]string, error) {
	zoneNames := sets.NewString()
	nodes, err := c.Core().Nodes().List(v1.ListOptions{})
	if err != nil {
		return nil, err
	}
	for _, node := range nodes.Items {
		zoneName, err := getZoneNameForNode(node)
		Expect(err).NotTo(HaveOccurred())
		zoneNames.Insert(zoneName)
	}
	return zoneNames.List(), nil
}

// Return the number of zones in which we have nodes in this cluster.
func getZoneCount(c clientset.Interface) (int, error) {
	zoneNames, err := getZoneNames(c)
	if err != nil {
		return -1, err
	}
	return len(zoneNames), nil
}

// Find the name of the zone in which the pod is scheduled
func getZoneNameForPod(c clientset.Interface, pod v1.Pod) (string, error) {
	By(fmt.Sprintf("Getting zone name for pod %s, on node %s", pod.Name, pod.Spec.NodeName))
	node, err := c.Core().Nodes().Get(pod.Spec.NodeName, metav1.GetOptions{})
	Expect(err).NotTo(HaveOccurred())
	return getZoneNameForNode(*node)
}

// Determine whether a set of pods are approximately evenly spread
// across a given set of zones
func checkZoneSpreading(c clientset.Interface, pods *v1.PodList, zoneNames []string) (bool, error) {
	podsPerZone := make(map[string]int)
	for _, zoneName := range zoneNames {
		podsPerZone[zoneName] = 0
	}
	for _, pod := range pods.Items {
		if pod.DeletionTimestamp != nil {
			continue
		}
		zoneName, err := getZoneNameForPod(c, pod)
		Expect(err).NotTo(HaveOccurred())
		podsPerZone[zoneName] = podsPerZone[zoneName] + 1
	}
	minPodsPerZone := math.MaxInt32
	maxPodsPerZone := 0
	for _, podCount := range podsPerZone {
		if podCount < minPodsPerZone {
			minPodsPerZone = podCount
		}
		if podCount > maxPodsPerZone {
			maxPodsPerZone = podCount
		}
	}
	Expect(minPodsPerZone).To(BeNumerically("~", maxPodsPerZone, 1),
		"Pods were not evenly spread across zones.  %d in one zone and %d in another zone",
		minPodsPerZone, maxPodsPerZone)
	return true, nil
}

// Check that the pods comprising a replication controller get spread evenly across available zones
func SpreadRCOrFail(f *framework.Framework, replicaCount int32, image string) {
	name := "ubelite-spread-rc-" + string(uuid.NewUUID())
	By(fmt.Sprintf("Creating replication controller %s", name))
	controller, err := f.ClientSet.Core().ReplicationControllers(f.Namespace.Name).Create(&v1.ReplicationController{
		ObjectMeta: v1.ObjectMeta{
			Namespace: f.Namespace.Name,
			Name:      name,
		},
		Spec: v1.ReplicationControllerSpec{
			Replicas: &replicaCount,
			Selector: map[string]string{
				"name": name,
			},
			Template: &v1.PodTemplateSpec{
				ObjectMeta: v1.ObjectMeta{
					Labels: map[string]string{"name": name},
				},
				Spec: v1.PodSpec{
					Containers: []v1.Container{
						{
							Name:  name,
							Image: image,
							Ports: []v1.ContainerPort{{ContainerPort: 9376}},
						},
					},
				},
			},
		},
	})
	Expect(err).NotTo(HaveOccurred())
	// Cleanup the replication controller when we are done.
	defer func() {
		// Resize the replication controller to zero to get rid of pods.
		if err := framework.DeleteRCAndPods(f.ClientSet, f.InternalClientset, f.Namespace.Name, controller.Name); err != nil {
			framework.Logf("Failed to cleanup replication controller %v: %v.", controller.Name, err)
		}
	}()
	// List the pods, making sure we observe all the replicas.
	selector := labels.SelectorFromSet(labels.Set(map[string]string{"name": name}))
	pods, err := framework.PodsCreated(f.ClientSet, f.Namespace.Name, name, replicaCount)
	Expect(err).NotTo(HaveOccurred())

	// Wait for all of them to be scheduled
	By(fmt.Sprintf("Waiting for %d replicas of %s to be scheduled.  Selector: %v", replicaCount, name, selector))
	pods, err = framework.WaitForPodsWithLabelScheduled(f.ClientSet, f.Namespace.Name, selector)
	Expect(err).NotTo(HaveOccurred())

	// Now make sure they're spread across zones
	zoneNames, err := getZoneNames(f.ClientSet)
	Expect(err).NotTo(HaveOccurred())
	Expect(checkZoneSpreading(f.ClientSet, pods, zoneNames)).To(Equal(true))
}
<|endoftext|>


<commit_before>package store

import (
	"encoding/json"
	"errors"
	"net/http"
	"sort"
	"time"

	"github.com/manishrjain/gocrud/x"
)

var (
	ErrNoParent = errors.New("No parent found")
)

// Query stores the read instrutions, storing the instruction set
// for the entities Query relates to.
type Query struct {
	kind      string
	id        string
	filterOut map[string]bool
	maxDepth  int
	children  []*Query
	parent    *Query
}

type Object struct {
	Value  interface{}
	Source string
	NanoTs int64
}

type Versions struct {
	versions []Object
}

// Result stores the final entity state retrieved from Store upon running
// the instructions provided by Query.
type Result struct {
	Id       string
	Kind     string
	Columns  map[string]*Versions
	Children []*Result
}

type runResult struct {
	Result *Result
	Err    error
}

func (v *Versions) add(o Object) {
	if len(v.versions) > 0 {
		i := len(v.versions) - 1
		if v.versions[i].NanoTs > o.NanoTs {
			// unsorted list. Gocrud code is doing something wrong.
			log.Fatal("Appending an object with lower ts to a sorted list")
		}
	}
	v.versions = append(v.versions, o)
}

func (v Versions) Latest() Object {
	if len(v.versions) == 0 {
		return Object{}
	}
	i := len(v.versions) - 1
	return v.versions[i]
}

func (v Versions) Oldest() Object {
	if len(v.versions) == 0 {
		return Object{}
	}
	return v.versions[0]
}

func (v Versions) Count() int {
	return len(v.versions)
}

// Retrieve the parent id for given entity id. Return ErrNoParent if parent is
// not present. Otherwise, if an error occurs during retrieval, returns that.
func Parent(id string) (parentid string, rerr error) {
	its, err := Get().GetEntity(id)
	if err != nil {
		x.LogErr(log, err).WithField("id", id).Error("While retrieving entity")
		return "", err
	}
	for _, it := range its {
		if it.Predicate == "_parent_" {
			return it.ObjectId, nil
		}
	}
	return "", ErrNoParent
}

// NewQuery is the main entrypoint to data store queries. Returns back a Query
// object pointer, to run read instructions on.
func NewQuery(id string) *Query {
	q := new(Query)
	q.id = id
	return q
}

// UptoDepth specifies the number of levels of descendants that would be
// retrieved for the entity Query points to.
//
// You can think of this in terms of a tree structure, where the Entity pointed
// to by Query points to other Entities, which in turn point to more Entities,
// and so on. For e.g. Post -> Comments -> Likes.
func (q *Query) UptoDepth(level int) *Query {
	q.maxDepth = level
	return q
}

// Collect specifies the kind of child entities to retrieve. Returns back
// a new Query pointer pointing to those children entities as a collective.
//
// Any further operations on this returned pointer would attribute to those
// children entities, and not the caller query entity.
func (q *Query) Collect(kind string) *Query {
	for _, child := range q.children {
		if child.kind == kind {
			return child
		}
	}
	child := new(Query)
	child.parent = q
	child.kind = kind
	q.children = append(q.children, child)
	return child
}

// FilterOut provides a way to well, filter out, any entities which have
// the given property.
func (q *Query) FilterOut(property string) *Query {
	if len(q.filterOut) == 0 {
		q.filterOut = make(map[string]bool)
	}
	q.filterOut[property] = true
	return q
}

func (q *Query) root() *Query {
	for q.parent != nil {
		q = q.parent
	}
	return q
}

func (q *Query) doRun(level, max int, ch chan runResult) {
	log.Debugf("Query: %+v", q)
	its, err := Get().GetEntity(q.id)
	if err != nil {
		x.LogErr(log, err).Error("While retrieving: ", q.id)
		ch <- runResult{Result: nil, Err: err}
		return
	}
	if len(its) == 0 {
		ch <- runResult{Result: new(Result), Err: nil}
		return
	}
	sort.Sort(x.Its(its))

	follow := make(map[string]*Query)
	for _, child := range q.children {
		follow[child.kind] = child
		log.WithField("kind", child.kind).WithField("child", child).Debug("Following")
	}

	result := new(Result)
	result.Columns = make(map[string]*Versions)
	it := its[0]
	result.Id = it.SubjectId
	result.Kind = it.SubjectType

	waitTimes := 0
	childChan := make(chan runResult)
	for _, it := range its {
		if it.Predicate == "_delete_" {
			// If marked as deleted, don't return this node.
			log.WithField("id", result.Id).
				WithField("kind", result.Kind).
				WithField("_delete_", true).
				Debug("Discarding due to delete bit")
			ch <- runResult{Result: new(Result), Err: nil}
			return
		}

		if it.Predicate == "_parent_" {
			log.WithField("id", result.Id).
				WithField("kind", result.Kind).
				WithField("parent", it.ObjectId).
				Debug("Not following edge back to parent")
			continue
		}

		if _, fout := q.filterOut[it.Predicate]; fout {
			log.WithField("id", result.Id).
				WithField("kind", result.Kind).
				WithField("predicate", it.Predicate).
				Debug("Discarding due to predicate filter")
			ch <- runResult{Result: new(Result), Err: nil}
			return
		}

		if len(it.ObjectId) == 0 {
			o := Object{NanoTs: it.NanoTs, Source: it.Source}
			if err := json.Unmarshal(it.Object, &o.Value); err != nil {
				x.LogErr(log, err).Error("While unmarshal")
				ch <- runResult{Result: nil, Err: err}
				return
			}

			if _, vok := result.Columns[it.Predicate]; !vok {
				result.Columns[it.Predicate] = new(Versions)
			}
			result.Columns[it.Predicate].add(o)
			continue
		}

		if childq, fw := follow[it.Predicate]; fw {
			nchildq := new(Query)
			*nchildq = *childq // This is important, otherwise id gets overwritten
			nchildq.id = it.ObjectId

			// Use child's maxDepth here, instead of parent's.
			waitTimes += 1
			log.WithField("child_id", nchildq.id).
				WithField("child_kind", nchildq.kind).Debug("Go routine for child")
			go nchildq.doRun(0, nchildq.maxDepth, childChan)
			continue
		}

		if len(it.ObjectId) > 0 && level < max {
			child := new(Query)
			child.id = it.ObjectId

			waitTimes += 1
			log.WithField("child_id", child.id).WithField("level", level+1).
				Debug("Go routine for child one level deeper")
			go child.doRun(level+1, max, childChan)
		}
	}

	// Wait for all those subroutines
	for i := 0; i < waitTimes; i++ {
		log.Debugf("Waiting for children subroutines: %v/%v", i, waitTimes-1)
		rr := <-childChan
		log.Debugf("Waiting done")
		if rr.Err != nil {
			x.LogErr(log, err).Error("While child doRun")
		} else {
			if len(rr.Result.Id) > 0 && len(rr.Result.Kind) > 0 {
				log.WithField("result", *rr.Result).Debug("Appending child")
				result.Children = append(result.Children, rr.Result)
			}
		}
	}

	ch <- runResult{Result: result, Err: nil}
	return
}

// Run finds the root from the given Query pointer, recursively executes
// the read operations, and returns back pointer to Result object.
// Any errors encountered during these stpeps is returned as well.
func (q *Query) Run() (result *Result, rerr error) {
	q = q.root()

	ch := make(chan runResult)
	go q.doRun(0, q.maxDepth, ch)
	rr := <-ch // Blocking wait
	return rr.Result, rr.Err
}

func (r *Result) Drop(pred string) {
	delete(r.Columns, pred)
}

func (r *Result) Debug(level int) {
	log.Debugf("Result level %v: %+v", level, r)
	for _, child := range r.Children {
		child.Debug(level + 1)
	}
}

func (r *Result) ToMap() (data map[string]interface{}) {
	data = make(map[string]interface{})
	data["id"] = r.Id
	data["kind"] = r.Kind
	var ts_latest int64
	ts_oldest := time.Now().UnixNano()
	for pred, versions := range r.Columns {
		// During conversion to JSON, to keep things simple,
		// we're dropping older versions of predicates, and
		// source and ts information across all the predicates,
		// keeping only the latest one.
		data[pred] = versions.Latest().Value
		if versions.Latest().NanoTs > ts_latest {
			ts_latest = versions.Latest().NanoTs
			data["modifier"] = versions.Latest().Source // Loss of information.
		}
		if versions.Oldest().NanoTs < ts_oldest {
			ts_oldest = versions.Oldest().NanoTs
			data["creator"] = versions.Oldest().Source
		}
	}
	data["creation_ms"] = int(ts_oldest / 1000000)
	data["modification_ms"] = int(ts_latest / 1000000) // Loss of information. Picking up latest mod time.
	kinds := make(map[string]bool)
	for _, child := range r.Children {
		kinds[child.Kind] = true
	}

	for kind := range kinds {
		var l []map[string]interface{}
		for _, child := range r.Children {
			if kind != child.Kind {
				continue
			}
			l = append(l, child.ToMap())
		}
		data[kind] = l
	}

	return
}

// ToJson creates the JSON for the data pointed by the Result pointer.
//
// Note that this doesn't find the "root" from the Result pointer, instead
// doing the processing only from the current Result pointer.
func (r *Result) ToJson() ([]byte, error) {
	data := r.ToMap()
	return json.Marshal(data)
}

// WriteJsonResponse does the same as ToJson. But also writes the JSON
// generated to http.ResponseWriter. In case of error, writes that error
// instead, in this format:
//
//  {
//   "code":    "E_ERROR",
//   "message": "error message"
//  }
func (r *Result) WriteJsonResponse(w http.ResponseWriter) {
	data, err := r.ToJson()
	if err != nil {
		x.SetStatus(w, x.E_ERROR, err.Error())
		return
	}
	_, err = w.Write(data)
	if err != nil {
		x.SetStatus(w, x.E_ERROR, err.Error())
		return
	}
}
<commit_msg>Add query.AllowDeleted()<commit_after>package store

import (
	"encoding/json"
	"errors"
	"net/http"
	"sort"
	"time"

	"github.com/manishrjain/gocrud/x"
)

var (
	ErrNoParent = errors.New("No parent found")
)

// Query stores the read instrutions, storing the instruction set
// for the entities Query relates to.
type Query struct {
	kind       string
	id         string
	filterOut  map[string]bool
	maxDepth   int
	children   []*Query
	parent     *Query
	getDeleted bool
}

type Object struct {
	Value  interface{}
	Source string
	NanoTs int64
}

type Versions struct {
	versions []Object
}

// Result stores the final entity state retrieved from Store upon running
// the instructions provided by Query.
type Result struct {
	Id       string
	Kind     string
	Columns  map[string]*Versions
	Children []*Result
}

type runResult struct {
	Result *Result
	Err    error
}

func (v *Versions) add(o Object) {
	if len(v.versions) > 0 {
		i := len(v.versions) - 1
		if v.versions[i].NanoTs > o.NanoTs {
			// unsorted list. Gocrud code is doing something wrong.
			log.Fatal("Appending an object with lower ts to a sorted list")
		}
	}
	v.versions = append(v.versions, o)
}

func (v Versions) Latest() Object {
	if len(v.versions) == 0 {
		return Object{}
	}
	i := len(v.versions) - 1
	return v.versions[i]
}

func (v Versions) Oldest() Object {
	if len(v.versions) == 0 {
		return Object{}
	}
	return v.versions[0]
}

func (v Versions) Count() int {
	return len(v.versions)
}

// Retrieve the parent id for given entity id. Return ErrNoParent if parent is
// not present. Otherwise, if an error occurs during retrieval, returns that.
func Parent(id string) (parentid string, rerr error) {
	its, err := Get().GetEntity(id)
	if err != nil {
		x.LogErr(log, err).WithField("id", id).Error("While retrieving entity")
		return "", err
	}
	for _, it := range its {
		if it.Predicate == "_parent_" {
			return it.ObjectId, nil
		}
	}
	return "", ErrNoParent
}

// NewQuery is the main entrypoint to data store queries. Returns back a Query
// object pointer, to run read instructions on.
func NewQuery(id string) *Query {
	q := new(Query)
	q.id = id
	return q
}

// UptoDepth specifies the number of levels of descendants that would be
// retrieved for the entity Query points to.
//
// You can think of this in terms of a tree structure, where the Entity pointed
// to by Query points to other Entities, which in turn point to more Entities,
// and so on. For e.g. Post -> Comments -> Likes.
func (q *Query) UptoDepth(level int) *Query {
	q.maxDepth = level
	return q
}

func (q *Query) AllowDeleted() *Query {
	q.getDeleted = true
	return q
}

// Collect specifies the kind of child entities to retrieve. Returns back
// a new Query pointer pointing to those children entities as a collective.
//
// Any further operations on this returned pointer would attribute to those
// children entities, and not the caller query entity.
func (q *Query) Collect(kind string) *Query {
	for _, child := range q.children {
		if child.kind == kind {
			return child
		}
	}
	child := new(Query)
	child.parent = q
	child.kind = kind
	q.children = append(q.children, child)
	return child
}

// FilterOut provides a way to well, filter out, any entities which have
// the given property.
func (q *Query) FilterOut(property string) *Query {
	if len(q.filterOut) == 0 {
		q.filterOut = make(map[string]bool)
	}
	q.filterOut[property] = true
	return q
}

func (q *Query) root() *Query {
	for q.parent != nil {
		q = q.parent
	}
	return q
}

func (q *Query) doRun(level, max int, ch chan runResult) {
	log.Debugf("Query: %+v", q)
	its, err := Get().GetEntity(q.id)
	if err != nil {
		x.LogErr(log, err).Error("While retrieving: ", q.id)
		ch <- runResult{Result: nil, Err: err}
		return
	}
	if len(its) == 0 {
		ch <- runResult{Result: new(Result), Err: nil}
		return
	}
	sort.Sort(x.Its(its))

	follow := make(map[string]*Query)
	for _, child := range q.children {
		follow[child.kind] = child
		log.WithField("kind", child.kind).WithField("child", child).Debug("Following")
	}

	result := new(Result)
	result.Columns = make(map[string]*Versions)
	it := its[0]
	result.Id = it.SubjectId
	result.Kind = it.SubjectType

	waitTimes := 0
	childChan := make(chan runResult)
	for _, it := range its {
		if it.Predicate == "_delete_" && !q.getDeleted {
			// If marked as deleted, don't return this node.
			log.WithField("id", result.Id).
				WithField("kind", result.Kind).
				WithField("_delete_", true).
				Debug("Discarding due to delete bit")
			ch <- runResult{Result: new(Result), Err: nil}
			return
		}

		if it.Predicate == "_parent_" {
			log.WithField("id", result.Id).
				WithField("kind", result.Kind).
				WithField("parent", it.ObjectId).
				Debug("Not following edge back to parent")
			continue
		}

		if _, fout := q.filterOut[it.Predicate]; fout {
			log.WithField("id", result.Id).
				WithField("kind", result.Kind).
				WithField("predicate", it.Predicate).
				Debug("Discarding due to predicate filter")
			ch <- runResult{Result: new(Result), Err: nil}
			return
		}

		if len(it.ObjectId) == 0 {
			o := Object{NanoTs: it.NanoTs, Source: it.Source}
			if err := json.Unmarshal(it.Object, &o.Value); err != nil {
				x.LogErr(log, err).Error("While unmarshal")
				ch <- runResult{Result: nil, Err: err}
				return
			}

			if _, vok := result.Columns[it.Predicate]; !vok {
				result.Columns[it.Predicate] = new(Versions)
			}
			result.Columns[it.Predicate].add(o)
			continue
		}

		if childq, fw := follow[it.Predicate]; fw {
			nchildq := new(Query)
			*nchildq = *childq // This is important, otherwise id gets overwritten
			nchildq.id = it.ObjectId

			// Use child's maxDepth here, instead of parent's.
			waitTimes += 1
			log.WithField("child_id", nchildq.id).
				WithField("child_kind", nchildq.kind).Debug("Go routine for child")
			go nchildq.doRun(0, nchildq.maxDepth, childChan)
			continue
		}

		if len(it.ObjectId) > 0 && level < max {
			child := new(Query)
			child.id = it.ObjectId

			waitTimes += 1
			log.WithField("child_id", child.id).WithField("level", level+1).
				Debug("Go routine for child one level deeper")
			go child.doRun(level+1, max, childChan)
		}
	}

	// Wait for all those subroutines
	for i := 0; i < waitTimes; i++ {
		log.Debugf("Waiting for children subroutines: %v/%v", i, waitTimes-1)
		rr := <-childChan
		log.Debugf("Waiting done")
		if rr.Err != nil {
			x.LogErr(log, err).Error("While child doRun")
		} else {
			if len(rr.Result.Id) > 0 && len(rr.Result.Kind) > 0 {
				log.WithField("result", *rr.Result).Debug("Appending child")
				result.Children = append(result.Children, rr.Result)
			}
		}
	}

	ch <- runResult{Result: result, Err: nil}
	return
}

// Run finds the root from the given Query pointer, recursively executes
// the read operations, and returns back pointer to Result object.
// Any errors encountered during these stpeps is returned as well.
func (q *Query) Run() (result *Result, rerr error) {
	q = q.root()

	ch := make(chan runResult)
	go q.doRun(0, q.maxDepth, ch)
	rr := <-ch // Blocking wait
	return rr.Result, rr.Err
}

func (r *Result) Drop(pred string) {
	delete(r.Columns, pred)
}

func (r *Result) Debug(level int) {
	log.Debugf("Result level %v: %+v", level, r)
	for _, child := range r.Children {
		child.Debug(level + 1)
	}
}

func (r *Result) ToMap() (data map[string]interface{}) {
	data = make(map[string]interface{})
	data["id"] = r.Id
	data["kind"] = r.Kind
	var ts_latest int64
	ts_oldest := time.Now().UnixNano()
	for pred, versions := range r.Columns {
		// During conversion to JSON, to keep things simple,
		// we're dropping older versions of predicates, and
		// source and ts information across all the predicates,
		// keeping only the latest one.
		data[pred] = versions.Latest().Value
		if versions.Latest().NanoTs > ts_latest {
			ts_latest = versions.Latest().NanoTs
			data["modifier"] = versions.Latest().Source // Loss of information.
		}
		if versions.Oldest().NanoTs < ts_oldest {
			ts_oldest = versions.Oldest().NanoTs
			data["creator"] = versions.Oldest().Source
		}
	}
	data["creation_ms"] = int(ts_oldest / 1000000)
	data["modification_ms"] = int(ts_latest / 1000000) // Loss of information. Picking up latest mod time.
	kinds := make(map[string]bool)
	for _, child := range r.Children {
		kinds[child.Kind] = true
	}

	for kind := range kinds {
		var l []map[string]interface{}
		for _, child := range r.Children {
			if kind != child.Kind {
				continue
			}
			l = append(l, child.ToMap())
		}
		data[kind] = l
	}

	return
}

// ToJson creates the JSON for the data pointed by the Result pointer.
//
// Note that this doesn't find the "root" from the Result pointer, instead
// doing the processing only from the current Result pointer.
func (r *Result) ToJson() ([]byte, error) {
	data := r.ToMap()
	return json.Marshal(data)
}

// WriteJsonResponse does the same as ToJson. But also writes the JSON
// generated to http.ResponseWriter. In case of error, writes that error
// instead, in this format:
//
//  {
//   "code":    "E_ERROR",
//   "message": "error message"
//  }
func (r *Result) WriteJsonResponse(w http.ResponseWriter) {
	data, err := r.ToJson()
	if err != nil {
		x.SetStatus(w, x.E_ERROR, err.Error())
		return
	}
	_, err = w.Write(data)
	if err != nil {
		x.SetStatus(w, x.E_ERROR, err.Error())
		return
	}
}
<|endoftext|>


<commit_before>// Copyright 2018 The Trickuri Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://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.

// Generates test cases for URL spoofing vulnerabilities.
package main

import (
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"flag"
	"fmt"
	"io"
	"log"
	"math/big"
	"net"
	"net/http"
	"os"
	"runtime"
	"strconv"
	"time"
)

// Command line flags
var (
	port      = flag.Int("p", 1270, "port on which to listen")
	httpsPort = flag.Int("h", 8443, "port on which the HTTPS proxy will listen")
	directory = flag.String("d", userHomeDir()+"/.trickuri", "default directory on which to save certificates")
)

// Globals
var (
	caKey       crypto.PrivateKey
	caCert      *x509.Certificate
	certMap     map[string]*tls.Certificate
	lastTunHost string
)

// Multi-platform code to get user home directory, used for the default certificate storage path.
func userHomeDir() string {
	if runtime.GOOS == "windows" {
		home := os.Getenv("HOMEDRIVE") + os.Getenv("HOMEPATH")
		if home == "" {
			home = os.Getenv("USERPROFILE")
		}
		return home
	}
	return os.Getenv("HOME")
}

// Common certificate generation code
func createCertificateTemplate() (*x509.Certificate, error) {
	notBefore := time.Now().AddDate(0, 0, -7) // Mitigate clock-skew between this server and the client.
	notAfter := notBefore.AddDate(1, 0, 7)    // Certs are valid for one year from now.
	serialNumberLimit := new(big.Int).Lsh(big.NewInt(1), 128)
	serialNumber, err := rand.Int(rand.Reader, serialNumberLimit)
	if err != nil {
		return nil, err
	}

	certTemplate := &x509.Certificate{
		SerialNumber: serialNumber,
		Subject: pkix.Name{
			Organization: []string{"TrickUri Interception Certificate"},
		},
		NotBefore:             notBefore,
		NotAfter:              notAfter,
		KeyUsage:              x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		BasicConstraintsValid: true,
	}

	return certTemplate, nil
}

// Returns a root certificate and its key.
func createRootCertificate() ([]byte, *rsa.PrivateKey, error) {
	priv, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return nil, nil, err
	}

	certTemplate, err := createCertificateTemplate()
	if err != nil {
		return nil, nil, err
	}

	var cerBytes []byte
	certTemplate.IsCA = true
	certTemplate.MaxPathLen = 0
	certTemplate.MaxPathLenZero = true
	certTemplate.Subject = pkix.Name{Organization: []string{"TrickUri Root"}}
	hash := sha256.Sum256(priv.PublicKey.N.Bytes())
	certTemplate.SubjectKeyId = hash[:]
	certTemplate.KeyUsage |= x509.KeyUsageCertSign
	cerBytes, err = x509.CreateCertificate(rand.Reader, certTemplate, certTemplate, &priv.PublicKey, priv)
	if err != nil {
		return nil, nil, err
	}

	return cerBytes, priv, nil
}

// Writes |certificate| to |filename|.cer and |key| to |filename|.pem.
func writeCertificate(certificate []byte, key *rsa.PrivateKey, filename string) error {
	certOut, err := os.Create(*directory + "/" + filename + ".cer")
	if err != nil {
		return err
	}
	err = pem.Encode(certOut, &pem.Block{Type: "CERTIFICATE", Bytes: certificate})
	certOut.Close()
	if err != nil {
		return err
	}

	keyOut, err := os.OpenFile(*directory+"/"+filename+".pem", os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0600)
	if err != nil {
		return err
	}
	err = pem.Encode(keyOut, &pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(key)})
	keyOut.Close()
	if err != nil {
		return err
	}
	return nil
}

// If a root certificate and key exist in the trickuri directory, loads them, otherwise creates them.
func loadOrCreateRootCertificate() (tls.Certificate, error) {
	rootKeys, err := tls.LoadX509KeyPair(*directory+"/root.cer", *directory+"/root.pem")
	if err != nil {
		log.Println("root.cer failed to load. Recreating root certificate.")
		rootCert, rootKey, err := createRootCertificate()
		if err != nil {
			return tls.Certificate{}, err
		}
		if err := writeCertificate(rootCert, rootKey, "root"); err != nil {
			return tls.Certificate{}, err
		}
		return tls.LoadX509KeyPair(*directory+"/root.cer", *directory+"/root.pem")
	}
	return rootKeys, err
}

// Generates a certificate for |hostname|.
func generateCertificate(hostname string) (*tls.Certificate, error) {
	priv, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return nil, err
	}
	certTemplate, err := createCertificateTemplate()
	if err != nil {
		return nil, err
	}
	if ip := net.ParseIP(hostname); ip != nil {
		certTemplate.IPAddresses = append(certTemplate.IPAddresses, ip)
	} else {
		certTemplate.DNSNames = append(certTemplate.DNSNames, hostname)
	}
	cerBytes, err := x509.CreateCertificate(rand.Reader, certTemplate, caCert, &priv.PublicKey, caKey)
	if err != nil {
		return nil, err
	}
	return &tls.Certificate{
		Certificate: [][]byte{cerBytes},
		PrivateKey:  priv,
	}, nil
}

// Retrieves the certificate for |info| if it has been created before, otherwise generates it and returns it.
func certificate(info *tls.ClientHelloInfo) (*tls.Certificate, error) {
	var hostname string
	if len(info.ServerName) > 0 {
		hostname = info.ServerName
	} else {
		// If the SNI hostname is empty (e.g., for IP addresses), create a certificate for the last tunneled hostname.
		hostname = lastTunHost
	}
	if cert := certMap[hostname]; cert != nil {
		return cert, nil
	}
	cert, err := generateCertificate(hostname)
	if err != nil {
		return nil, err
	}
	certMap[hostname] = cert
	return cert, nil
}

// Serve a proxy configuration script that directs the client to go direct for responses
// that should not be generated from this tool. This helps ensure that Chrome/Windows/etc
// are not unduly affected by our HTTP(S) interceptions.
func serveProxyConfigurationScript(w http.ResponseWriter, _ *http.Request) {
	w.Header().Set("Content-Type", "application/x-ns-proxy-autoconfig; charset=utf-8")
	w.WriteHeader(http.StatusOK)
	fmt.Fprintf(w, `function FindProxyForURL(url, host) {
			// Bypass list. See https://findproxyforurl.com/pac-functions/ for functions.
			if (shExpMatch(host, "*.google.com")) return "DIRECT";
			if (shExpMatch(host, "*.gstatic.com")) return "DIRECT";
			if (shExpMatch(host, "*.googleusercontent.com")) return "DIRECT";
			if (shExpMatch(host, "*.googleapis.com")) return "DIRECT";
			if (shExpMatch(host, "*.microsoft.com")) return "DIRECT";
			if (dnsDomainLevels(host) < 1) return "DIRECT";

			// Return a response from TrickUri.
			return "PROXY localhost:%d";
		}`,
		*port)
}

// Serve a printout of the request |r|.
func serveEcho(w http.ResponseWriter, r *http.Request) {
	err := r.Write(w)
	if err != nil {
		log.Println(err)
	}
}

// Handler for HTTP tunneling.
func tunnelHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodConnect {
		// This is a request that doesn't require tunneling, use the default handler.
		httpHandler(w, r)
		return
	}
	// All https requests will be forwarded to the HTTPS proxy running on the set HTTPS port
	httpsHost := "localhost:" + strconv.Itoa(*httpsPort)
	dst, err := net.DialTimeout("tcp", httpsHost, 3*time.Second)
	if err != nil {
		http.Error(w, "forwarding to "+httpsHost+" failed", http.StatusInternalServerError)
		return
	}
	w.WriteHeader(http.StatusOK)
	hj, ok := w.(http.Hijacker)
	if !ok {
		http.Error(w, "hijacking is not supported", http.StatusInternalServerError)
		return
	}
	hijackedConn, _, err := hj.Hijack()
	if err != nil {
		http.Error(w, "error hijacking connection", http.StatusInternalServerError)
		return
	}
	// Save a copy of the last tunneled hostname, used for certificate generation if SNI is not available.
	lastTunHost = r.URL.Hostname()
	go forward(dst, hijackedConn)
	go forward(hijackedConn, dst)
}

func forward(dst io.WriteCloser, src io.ReadCloser) {
	defer dst.Close()
	defer src.Close()
	io.Copy(dst, src)
}

func httpHandler(w http.ResponseWriter, r *http.Request) {
	if r.URL.EscapedPath() == "/proxy.pac" {
		serveProxyConfigurationScript(w, r)
		return
	}
	if r.URL.EscapedPath() == "/echo" {
		serveEcho(w, r)
		return
	}
	testcaseHandler := http.FileServer(http.Dir("testcases"))
	testcaseHandler.ServeHTTP(w, r)
}

func main() {
	flag.Parse()
	err := os.MkdirAll(*directory, os.ModePerm)
	if err != nil {
		log.Fatal(err)
	}
	certMap = make(map[string]*tls.Certificate)
	rootCert, err := loadOrCreateRootCertificate()
	if err != nil {
		log.Fatal(err)
	}
	caCert, err = x509.ParseCertificate(rootCert.Certificate[0])
	if err != nil {
		log.Fatal(err)
	}
	caKey = rootCert.PrivateKey

	log.Printf("Download the root certificate at http://localhost:%d/root.cer and import it into browser/OS certificate store.", *port)
	log.Printf("Set your proxy to \"http://localhost:%[1]d/proxy.pac\" or localhost:%[1]d. Using pac file will pass through requests to google.com and microsoft.com for common Chrome/Windows requests.", *port)
	log.Println("Ready for requests")

	cfg := &tls.Config{
		CipherSuites: []uint16{
			tls.TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256,
			tls.TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256,
		},
		MinVersion:               tls.VersionTLS12,
		PreferServerCipherSuites: true,
		GetCertificate:           certificate,
	}
	httpServer := http.Server{
		Addr:    "localhost:1270",
		Handler: http.HandlerFunc(tunnelHandler),
	}
	httpsServer := http.Server{
		Addr:      "localhost:" + strconv.Itoa(*httpsPort),
		TLSConfig: cfg,
		Handler:   http.HandlerFunc(httpHandler),
	}
	go func() {
		if err := httpsServer.ListenAndServeTLS("", ""); err != nil {
			log.Fatal(err)
		}
	}()
	if err := httpServer.ListenAndServe(); err != nil {
		log.Fatal(err)
	}
}
<commit_msg>Handle /root.cer path by serving root certificate<commit_after>// Copyright 2018 The Trickuri Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://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.

// Generates test cases for URL spoofing vulnerabilities.
package main

import (
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"flag"
	"fmt"
	"io"
	"log"
	"math/big"
	"net"
	"net/http"
	"os"
	"runtime"
	"strconv"
	"time"
)

// Command line flags
var (
	port      = flag.Int("p", 1270, "port on which to listen")
	httpsPort = flag.Int("h", 8443, "port on which the HTTPS proxy will listen")
	directory = flag.String("d", userHomeDir()+"/.trickuri", "default directory on which to save certificates")
)

// Globals
var (
	caKey       crypto.PrivateKey
	caCert      *x509.Certificate
	certMap     map[string]*tls.Certificate
	lastTunHost string
)

// Multi-platform code to get user home directory, used for the default certificate storage path.
func userHomeDir() string {
	if runtime.GOOS == "windows" {
		home := os.Getenv("HOMEDRIVE") + os.Getenv("HOMEPATH")
		if home == "" {
			home = os.Getenv("USERPROFILE")
		}
		return home
	}
	return os.Getenv("HOME")
}

// Common certificate generation code
func createCertificateTemplate() (*x509.Certificate, error) {
	notBefore := time.Now().AddDate(0, 0, -7) // Mitigate clock-skew between this server and the client.
	notAfter := notBefore.AddDate(1, 0, 7)    // Certs are valid for one year from now.
	serialNumberLimit := new(big.Int).Lsh(big.NewInt(1), 128)
	serialNumber, err := rand.Int(rand.Reader, serialNumberLimit)
	if err != nil {
		return nil, err
	}

	certTemplate := &x509.Certificate{
		SerialNumber: serialNumber,
		Subject: pkix.Name{
			Organization: []string{"TrickUri Interception Certificate"},
		},
		NotBefore:             notBefore,
		NotAfter:              notAfter,
		KeyUsage:              x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		BasicConstraintsValid: true,
	}

	return certTemplate, nil
}

// Returns a root certificate and its key.
func createRootCertificate() ([]byte, *rsa.PrivateKey, error) {
	priv, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return nil, nil, err
	}

	certTemplate, err := createCertificateTemplate()
	if err != nil {
		return nil, nil, err
	}

	var cerBytes []byte
	certTemplate.IsCA = true
	certTemplate.MaxPathLen = 0
	certTemplate.MaxPathLenZero = true
	certTemplate.Subject = pkix.Name{Organization: []string{"TrickUri Root"}}
	hash := sha256.Sum256(priv.PublicKey.N.Bytes())
	certTemplate.SubjectKeyId = hash[:]
	certTemplate.KeyUsage |= x509.KeyUsageCertSign
	cerBytes, err = x509.CreateCertificate(rand.Reader, certTemplate, certTemplate, &priv.PublicKey, priv)
	if err != nil {
		return nil, nil, err
	}

	return cerBytes, priv, nil
}

// Writes |certificate| to |filename|.cer and |key| to |filename|.pem.
func writeCertificate(certificate []byte, key *rsa.PrivateKey, filename string) error {
	certOut, err := os.Create(*directory + "/" + filename + ".cer")
	if err != nil {
		return err
	}
	err = pem.Encode(certOut, &pem.Block{Type: "CERTIFICATE", Bytes: certificate})
	certOut.Close()
	if err != nil {
		return err
	}

	keyOut, err := os.OpenFile(*directory+"/"+filename+".pem", os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0600)
	if err != nil {
		return err
	}
	err = pem.Encode(keyOut, &pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(key)})
	keyOut.Close()
	if err != nil {
		return err
	}
	return nil
}

// If a root certificate and key exist in the trickuri directory, loads them, otherwise creates them.
func loadOrCreateRootCertificate() (tls.Certificate, error) {
	rootKeys, err := tls.LoadX509KeyPair(*directory+"/root.cer", *directory+"/root.pem")
	if err != nil {
		log.Println("root.cer failed to load. Recreating root certificate.")
		rootCert, rootKey, err := createRootCertificate()
		if err != nil {
			return tls.Certificate{}, err
		}
		if err := writeCertificate(rootCert, rootKey, "root"); err != nil {
			return tls.Certificate{}, err
		}
		return tls.LoadX509KeyPair(*directory+"/root.cer", *directory+"/root.pem")
	}
	return rootKeys, err
}

// Generates a certificate for |hostname|.
func generateCertificate(hostname string) (*tls.Certificate, error) {
	priv, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return nil, err
	}
	certTemplate, err := createCertificateTemplate()
	if err != nil {
		return nil, err
	}
	if ip := net.ParseIP(hostname); ip != nil {
		certTemplate.IPAddresses = append(certTemplate.IPAddresses, ip)
	} else {
		certTemplate.DNSNames = append(certTemplate.DNSNames, hostname)
	}
	cerBytes, err := x509.CreateCertificate(rand.Reader, certTemplate, caCert, &priv.PublicKey, caKey)
	if err != nil {
		return nil, err
	}
	return &tls.Certificate{
		Certificate: [][]byte{cerBytes},
		PrivateKey:  priv,
	}, nil
}

// Retrieves the certificate for |info| if it has been created before, otherwise generates it and returns it.
func certificate(info *tls.ClientHelloInfo) (*tls.Certificate, error) {
	var hostname string
	if len(info.ServerName) > 0 {
		hostname = info.ServerName
	} else {
		// If the SNI hostname is empty (e.g., for IP addresses), create a certificate for the last tunneled hostname.
		hostname = lastTunHost
	}
	if cert := certMap[hostname]; cert != nil {
		return cert, nil
	}
	cert, err := generateCertificate(hostname)
	if err != nil {
		return nil, err
	}
	certMap[hostname] = cert
	return cert, nil
}

// Serve a proxy configuration script that directs the client to go direct for responses
// that should not be generated from this tool. This helps ensure that Chrome/Windows/etc
// are not unduly affected by our HTTP(S) interceptions.
func serveProxyConfigurationScript(w http.ResponseWriter, _ *http.Request) {
	w.Header().Set("Content-Type", "application/x-ns-proxy-autoconfig; charset=utf-8")
	w.WriteHeader(http.StatusOK)
	fmt.Fprintf(w, `function FindProxyForURL(url, host) {
			// Bypass list. See https://findproxyforurl.com/pac-functions/ for functions.
			if (shExpMatch(host, "*.google.com")) return "DIRECT";
			if (shExpMatch(host, "*.gstatic.com")) return "DIRECT";
			if (shExpMatch(host, "*.googleusercontent.com")) return "DIRECT";
			if (shExpMatch(host, "*.googleapis.com")) return "DIRECT";
			if (shExpMatch(host, "*.microsoft.com")) return "DIRECT";
			if (dnsDomainLevels(host) < 1) return "DIRECT";

			// Return a response from TrickUri.
			return "PROXY localhost:%d";
		}`,
		*port)
}

// Serve a printout of the request |r|.
func serveEcho(w http.ResponseWriter, r *http.Request) {
	err := r.Write(w)
	if err != nil {
		log.Println(err)
	}
}

func serveRootCert(w http.ResponseWriter, r *http.Request) {
	http.ServeFile(w, r, *directory+"/root.cer")
}

// Handler for HTTP tunneling.
func tunnelHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodConnect {
		// This is a request that doesn't require tunneling, use the default handler.
		httpHandler(w, r)
		return
	}
	// All https requests will be forwarded to the HTTPS proxy running on the set HTTPS port
	httpsHost := "localhost:" + strconv.Itoa(*httpsPort)
	dst, err := net.DialTimeout("tcp", httpsHost, 3*time.Second)
	if err != nil {
		http.Error(w, "forwarding to "+httpsHost+" failed", http.StatusInternalServerError)
		return
	}
	w.WriteHeader(http.StatusOK)
	hj, ok := w.(http.Hijacker)
	if !ok {
		http.Error(w, "hijacking is not supported", http.StatusInternalServerError)
		return
	}
	hijackedConn, _, err := hj.Hijack()
	if err != nil {
		http.Error(w, "error hijacking connection", http.StatusInternalServerError)
		return
	}
	// Save a copy of the last tunneled hostname, used for certificate generation if SNI is not available.
	lastTunHost = r.URL.Hostname()
	go forward(dst, hijackedConn)
	go forward(hijackedConn, dst)
}

func forward(dst io.WriteCloser, src io.ReadCloser) {
	defer dst.Close()
	defer src.Close()
	io.Copy(dst, src)
}

func httpHandler(w http.ResponseWriter, r *http.Request) {
	if r.URL.EscapedPath() == "/proxy.pac" {
		serveProxyConfigurationScript(w, r)
		return
	}
	if r.URL.EscapedPath() == "/echo" {
		serveEcho(w, r)
		return
	}
	if r.URL.EscapedPath() == "/root.cer" {
		serveRootCert(w, r)
		return
	}
	testcaseHandler := http.FileServer(http.Dir("testcases"))
	testcaseHandler.ServeHTTP(w, r)
}

func main() {
	flag.Parse()
	err := os.MkdirAll(*directory, os.ModePerm)
	if err != nil {
		log.Fatal(err)
	}
	certMap = make(map[string]*tls.Certificate)
	rootCert, err := loadOrCreateRootCertificate()
	if err != nil {
		log.Fatal(err)
	}
	caCert, err = x509.ParseCertificate(rootCert.Certificate[0])
	if err != nil {
		log.Fatal(err)
	}
	caKey = rootCert.PrivateKey

	log.Printf("Download the root certificate at http://localhost:%d/root.cer and import it into browser/OS certificate store.", *port)
	log.Printf("Set your proxy to \"http://localhost:%[1]d/proxy.pac\" or localhost:%[1]d. Using pac file will pass through requests to google.com and microsoft.com for common Chrome/Windows requests.", *port)
	log.Println("Ready for requests")

	cfg := &tls.Config{
		CipherSuites: []uint16{
			tls.TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256,
			tls.TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256,
		},
		MinVersion:               tls.VersionTLS12,
		PreferServerCipherSuites: true,
		GetCertificate:           certificate,
	}
	httpServer := http.Server{
		Addr:    "localhost:1270",
		Handler: http.HandlerFunc(tunnelHandler),
	}
	httpsServer := http.Server{
		Addr:      "localhost:" + strconv.Itoa(*httpsPort),
		TLSConfig: cfg,
		Handler:   http.HandlerFunc(httpHandler),
	}
	go func() {
		if err := httpsServer.ListenAndServeTLS("", ""); err != nil {
			log.Fatal(err)
		}
	}()
	if err := httpServer.ListenAndServe(); err != nil {
		log.Fatal(err)
	}
}
<|endoftext|>


<commit_before>package gobro

import (
	"fmt"
	"os"
	"runtime"
	"strings"
)

// If the error is not nil, exit with error code 1.
// Message is optional. Including more than one message will not have any result.
func ExitOnError(err error, message ...string) {
	if err != nil {
		errorMessage := caller() + err.Error()
		if len(message) > 0 {
			errorMessage += message[0]
		}
		fmt.Fprintf(os.Stderr, errorMessage+"\n")
		os.Exit(1)
	}
}

func caller() string {
	var stack [4096]byte
	n := runtime.Stack(stack[:], false)
	caller := strings.Split(string(stack[:n]), "\n")[6]
	caller = strings.Trim(caller, " \t")
	return strings.Split(caller, " ")[0] + ": "
}
<commit_msg>main func command parsing<commit_after>package gobro

import (
	"fmt"
	"os"
	"reflect"
	"runtime"
	"strings"
)

// If the error is not nil, exit with error code 1.
// Message is optional. Including more than one message will not have any result.
func ExitOnError(err error, message ...string) {
	if err != nil {
		var msg string
		if len(message) > 0 {
			msg = message[0] + " "
		}
		errorMessage := caller() + msg + err.Error()
		fmt.Fprintf(os.Stderr, errorMessage+"\n")
		os.Exit(1)
	}
}

func LogError(err error, message ...string) {
	if err != nil {
		var msg string
		if len(message) > 0 {
			msg = message[0] + " "
		}
		errorMessage := caller() + msg + err.Error()
		fmt.Fprintf(os.Stderr, errorMessage+"\n")
	}
}

func TrimAll(items []string) {
	for i, item := range items {
		items[i] = strings.Trim(item, " \n\r\t")
	}
}

func caller() string {
	var stack [4096]byte
	n := runtime.Stack(stack[:], false)
	caller := strings.Split(string(stack[:n]), "\n")[6]
	caller = strings.Trim(caller, " \t")
	return strings.Split(caller, " ")[0] + ": "
}

// ===== COMMAND MAPPER ======================================================

type CommandMap struct {
	// CommandMap holds a map of names to functions. Useful for handling
	// control flow in main functions writing a ton of if this else that or
	// using flag, which I find sub-optimal
	commandMap map[string]func([]string)
}

func NewCommandMap(functions ...func(args []string)) (commandMap CommandMap) {
	// Returns a new CommandMap with the functions mapped to their names.
	commandMap.commandMap = make(map[string]func([]string))

	for _, fn := range functions {
		name := runtime.FuncForPC(reflect.ValueOf(fn).Pointer()).Name()
		name = strings.Split(name, ".")[1] // main.command becomes command
		commandMap.commandMap[name] = fn
	}

	return
}

func (cm CommandMap) Run(args []string) {

	if len(os.Args) == 1 {
		fmt.Printf("Usage: %s [options] <command> [<args>]\n\n", args[0])
		fmt.Println("Available commands:")
		for k, _ := range cm.commandMap {
			fmt.Printf("   %s\n", k)
		}
		os.Exit(1)
	}

	command := os.Args[1]
	fn := cm.commandMap[command]
	if fn == nil {
		fmt.Fprintf(os.Stderr, "Command '%s' not found\n", command)
		os.Exit(1)
	}

	fn(os.Args[2:])
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	//"image"
	"github.com/nsf/termbox-go"
	"image/gif"
	"os"
	"time"
)

func exitOnError(err error) {
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}

type AttribVals struct {
	fg termbox.Attribute
	bg termbox.Attribute
}

type AttribTable [256]AttribVals

type ColourMapFunc func(idx uint8) AttribVals

func test(idx uint8) AttribVals {
	return AttribVals{fg: termbox.Attribute(idx), bg: termbox.Attribute(idx)}
}

func mapColours(g *gif.GIF, cmap ColourMapFunc) []AttribTable {
	var attribs []AttribTable

	for f := 0; f < len(g.Image); f++ {

	}

	return attribs
}

func renderFrame(g *gif.GIF, framenum int, attribs []AttribTable) error {

	width, height := termbox.Size()

	if width > g.Image[framenum].Rect.Dx() {
		width = g.Image[framenum].Rect.Dx()
	}

	if height > g.Image[framenum].Rect.Dy() {
		height = g.Image[framenum].Rect.Dy()
	}

	for y := 0; y < height; y++ {
		lineOffset := g.Image[framenum].Stride * y
		for x := 0; x < width; x++ {
			attr := attribs[framenum][x+lineOffset]
			termbox.SetCell(x, y, ' ', attr.fg, attr.bg)
		}
	}

	return nil
}

func main() {
	if len(os.Args) == 1 {
		fmt.Fprintln(os.Stderr, "Usage: gogif <filename>")
		os.Exit(1)
	}

	f, err := os.Open(os.Args[1])
	exitOnError(err)

	g, err := gif.DecodeAll(f)
	exitOnError(err)

	gc := GameCore{}
	gc.TickTime = time.Millisecond * 40

	gc.OnEvent = func(gc *GameCore, ev *termbox.Event) error {
		if ev.Type == termbox.EventKey {
			if ev.Ch == 'q' {
				gc.DoQuit = true
			}
		}
		return nil
	}

	frameNumber := 0

	attribs := mapColours(g, test)

	gc.OnTick = func(gc *GameCore) error {
		err := renderFrame(g, frameNumber, attribs)
		frameNumber++
		if len(g.Image) == frameNumber {
			frameNumber = 0
		}
		return err
	}

	gc.Run()
}
<commit_msg>contunuing development of pre calulated attribute values for colour mapping<commit_after>package main

import (
	"fmt"
	"github.com/nsf/termbox-go"
	"image/color"
	"image/gif"
	"os"
	"time"
)

func exitOnError(err error) {
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}

type AttribVals struct {
	fg termbox.Attribute
	bg termbox.Attribute
}

type AttribTable [256]AttribVals

type ColourMapFunc func(color.Color) AttribVals

func test(c color.Color) AttribVals {
	r, g, b, a := c.RGBA()

	g = g
	b = b
	a = a

	return AttribVals{fg: termbox.Attribute(r), bg: termbox.Attribute(r)}
}

func mapColours(g *gif.GIF, cmap ColourMapFunc) []AttribTable {
	var attribs []AttribTable

	for f := 0; f < len(g.Image); f++ {
		var at AttribTable
		for i := 0; i < len(g.Image[f].Palette); i++ {
			at[i] = cmap(g.Image[f].Palette[i])
		}
		attribs = append(attribs, at)
	}

	return attribs
}

func renderFrame(g *gif.GIF, framenum int, attribs []AttribTable) error {

	width, height := termbox.Size()

	if width > g.Image[framenum].Rect.Dx() {
		width = g.Image[framenum].Rect.Dx()
	}

	if height > g.Image[framenum].Rect.Dy() {
		height = g.Image[framenum].Rect.Dy()
	}

	for y := 0; y < height; y++ {
		lineOffset := g.Image[framenum].Stride * y
		for x := 0; x < width; x++ {
			attr := attribs[framenum][g.Image[framenum].Pix[x+lineOffset]]
			termbox.SetCell(x, y, ' ', attr.fg, attr.bg)
		}
	}

	return nil
}

func main() {
	if len(os.Args) == 1 {
		fmt.Fprintln(os.Stderr, "Usage: gogif <filename>")
		os.Exit(1)
	}

	f, err := os.Open(os.Args[1])
	exitOnError(err)

	g, err := gif.DecodeAll(f)
	exitOnError(err)

	gc := GameCore{}
	gc.TickTime = time.Millisecond * 40

	gc.OnEvent = func(gc *GameCore, ev *termbox.Event) error {
		if ev.Type == termbox.EventKey {
			if ev.Ch == 'q' {
				gc.DoQuit = true
			}
		}
		return nil
	}

	frameNumber := 0

	attribs := mapColours(g, test)

	gc.OnTick = func(gc *GameCore) error {
		err := renderFrame(g, frameNumber, attribs)
		frameNumber++
		if len(g.Image) == frameNumber {
			frameNumber = 0
		}
		return err
	}

	gc.Run()
}
<|endoftext|>