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

// A Left-Leaning Red-Black (LLRB) implementation of 2-3 balanced binary search trees

package llrb

// |Tree| is a Left-Leaning Red-Black (LLRB) implementation of 2-3 trees, based on:
//
//   http://www.cs.princeton.edu/~rs/talks/LLRB/08Penn.pdf
//   http://www.cs.princeton.edu/~rs/talks/LLRB/LLRB.pdf
//   http://www.cs.princeton.edu/~rs/talks/LLRB/Java/RedBlackBST.java
//
//  2-3 trees (and the run-time equivalent 2-3-4 trees) are the de facto standard BST
//  algoritms found in implementations of Python, Java, and other libraries. The LLRB
//  implementation of 2-3 trees is a recent improvement on the traditional implementation,
//  observed and documented by Robert Sedgewick.
//
// |Tree| has an associative interface, i.e. duplicate key are not allowed.
// The zero-value of a |Tree| represents a ready-for-use tree.
type Tree struct {
	count int
	root  *node
}

// An |Item| represents an object that can be inserted in |Tree|. It acts as a
// key via the method |LessThan|, which induces a full ordering on |Item|s. It is
// also a value, as the user can attach any data to it.
type Item interface {

	// |LessThan| returns true, if and only if |this| is STRICTLY less than |other|
	LessThan(other interface{}) bool
}

// Init resets (empties) the tree
func (t *Tree) Init() {
	t.root = nil
}

func (t *Tree) Len() int { return t.count }

// |Has| returns true if the tree contains an element
// whose |LessThan| order equals that of |key|.
func (t *Tree) Has(key Item) bool {
	return t.Get(key) != nil
}

// |Get| retrieves an element from the tree whose |LessThan| order
// equals that of |key|.
func (t *Tree) Get(key Item) Item {
	return get(t.root, key)
}

func get(h *node, item Item) Item {
	if h == nil {
		return nil
	}
	if item.LessThan(h.item) {
		return get(h.left, item)
	}
	if h.item.LessThan(item) {
		return get(h.right, item)
	}
	return h.item
}

// |Min| returns the minimum element in the tree.
func (t *Tree) Min() Item {
	return min(t.root)
}

func min(h *node) Item {
	if h == nil {
		return nil
	}
	if h.left == nil {
		return h.item
	}
	return min(h.left)
}

// |Max| returns the maximum element in the tree.
func (t *Tree) Max() Item {
	return max(t.root)
}

func max(h *node) Item {
	if h == nil {
		return nil
	}
	if h.right == nil {
		return h.item
	}
	return max(h.right)
}

// |Insert| inserts a new element in the tree, or replaces
// an existing one of identical |LessThan| order.
// If a replacement occurred, the replaced item is returned.
func (t *Tree) InsertOrReplace(item Item) Item {
	if item == nil {
		panic("inserting nil item")
	}
	var replaced Item
	t.root, replaced = insert(t.root, item)
	t.root.black = true
	if replaced == nil {
		t.count++
	}
	return replaced
}

func insert(h *node, item Item) (*node, Item) {
	if h == nil {
		return newNode(item), nil
	}

	// PLACEHOLDER: 2-3-4 tree (see comment below)

	var replaced Item
	if item.LessThan(h.item) {
		h.left, replaced = insert(h.left, item)
	} else if h.item.LessThan(item) {
		h.right, replaced = insert(h.right, item)
	} else {
		replaced, h.item = h.item, item
	}

	if isRed(h.right) && !isRed(h.left) {
		h = rotateLeft(h)
	}

	// PETAR: added |h.left != nil|
	if h.left != nil && isRed(h.left) && isRed(h.left.left) {
		h = rotateRight(h)
	}

	// When the next 3 lines of code are here, the LLRB behaves
	// like a 2-3 tree. If they are moved to the 2-3-4 placeholder above,
	// the LLRB tree behaves like a 2-3-4 tree.
	if isRed(h.left) && isRed(h.right) {
		flip(h)
	}

	return h, replaced
}

// |DeleteMin| deletes the minimum element in the tree and returns the
// deleted item or nil otherwise.
func (t *Tree) DeleteMin() Item {
	var deleted Item
	t.root, deleted = deleteMin(t.root)
	if t.root != nil {
		t.root.black = true
	}
	if deleted != nil {
		t.count--
	}
	return deleted
}

// deleteMin code for LLRB 2-3 trees
func deleteMin(h *node) (*node, Item) {
	if h == nil {
		return nil, nil
	}
	if h.left == nil {
		return nil, h.item
	}

	if !isRed(h.left) && !isRed(h.left.left) {
		h = moveRedLeft(h)
	}

	var deleted Item
	h.left, deleted = deleteMin(h.left)

	return fixUp(h), deleted
}

// |DeleteMax| deletes the maximum element in the tree and returns
// the deleted item or nil otherwise
func (t *Tree) DeleteMax() Item {
	var deleted Item
	t.root, deleted = deleteMax(t.root)
	if t.root != nil {
		t.root.black = true
	}
	if deleted != nil {
		t.count--
	}
	return deleted
}

func deleteMax(h *node) (*node, Item) {
	if h == nil {
		return nil, nil
	}
	if isRed(h.left) {
		h = rotateRight(h)
	}
	if h.right == nil {
		return nil, h.item
	}
	if !isRed(h.right) && !isRed(h.right.left) {
		h = moveRedRight(h)
	}
	var deleted Item
	h.right, deleted = deleteMax(h.right)

	return fixUp(h), deleted
}

// |Delete| deletes an item from the tree whose key equals |key|.
// The deleted item is return, otherwise nil is returned.
func (t *Tree) Delete(key Item) Item {
	var deleted Item
	t.root, deleted = delete(t.root, key)
	if t.root != nil {
		t.root.black = true
	}
	if deleted != nil {
		t.count--
	}
	return deleted
}

func delete(h *node, item Item) (*node, Item) {
	var deleted Item
	if h == nil {
		return nil, nil
	}
	if item.LessThan(h.item) {
		if h.left == nil { // item not present. Nothing to delete
			return h, nil
		}
		if !isRed(h.left) && !isRed(h.left.left) {
			h = moveRedLeft(h)
		}
		h.left, deleted = delete(h.left, item)
	} else {
		if isRed(h.left) {
			h = rotateRight(h)
		}
		// If |item| equals |h.item| and no right children at |h|
		if !h.item.LessThan(item) && h.right == nil {
			return nil, h.item
		}
		// PETAR: Added |h.right != nil| below
		if h.right != nil && !isRed(h.right) && !isRed(h.right.left) {
			h = moveRedRight(h)
		}
		// If |item| equals |h.item|, and (from above) |h.right != nil|
		if !h.item.LessThan(item) {
			var subDeleted Item
			h.right, subDeleted = deleteMin(h.right)
			if subDeleted == nil {
				panic("logic")
			}
			deleted, h.item = h.item, subDeleted
		} else { // Else, |item| is bigger than |h.item|
			h.right, deleted = delete(h.right, item)
		}
	}

	return fixUp(h), deleted
}

// |Iter| returns a chan that iterates through all elements in the
// tree in ascending order.
func (t *Tree) Iter() <-chan Item {
	c := make(chan Item)
	go func() {
		iterateInOrder(t.root, c)
		close(c)
	}()
	return c
}

func iterateInOrder(h *node, c chan<- Item) {
	if h == nil {
		return
	}
	iterateInOrder(h.left, c)
	c <- h.item
	iterateInOrder(h.right, c)
}

func iteratePreOrder(h *node, c chan<- Item) {
	if h == nil {
		return
	}
	c <- h.item
	iteratePreOrder(h.left, c)
	iteratePreOrder(h.right, c)
}

func iteratePostOrder(h *node, c chan<- Item) {
	if h == nil {
		return
	}
	iteratePostOrder(h.left, c)
	iteratePostOrder(h.right, c)
	c <- h.item
}

type node struct {
	item        Item
	left, right *node // Pointers to left and right child nodes
	black       bool  // If set, the color of the link (incoming from the parent) is black
	// In the LLRB, new nodes are always red, hence the zero-value for node
}

func newNode(item Item) *node { return &node{item: item} }

func isRed(h *node) bool {
	if h == nil {
		return false
	}
	return !h.black
}

func rotateLeft(h *node) *node {
	x := h.right
	if x.black {
		panic("rotating a black link")
	}
	h.right = x.left
	x.left = h
	x.black = h.black
	h.black = false
	return x
}

func rotateRight(h *node) *node {
	x := h.left
	if x.black {
		panic("rotating a black link")
	}
	h.left = x.right
	x.right = h
	x.black = h.black
	h.black = false
	return x
}

// XXX:
//      - Can flip() ever be called with |h.left == nil| or |h.right == nil|?

// Left and right children must be present
func flip(h *node) {
	h.black = !h.black
	h.left.black = !h.left.black
	h.right.black = !h.right.black
}

// Left and right children must be present
func moveRedLeft(h *node) *node {
	flip(h)
	if isRed(h.right.left) {
		h.right = rotateRight(h.right)
		h = rotateLeft(h)
		flip(h)
	}
	return h
}

// Left and right children must be present
func moveRedRight(h *node) *node {
	flip(h)
	if isRed(h.left.left) {
		h = rotateRight(h)
		flip(h)
	}
	return h
}

func fixUp(h *node) *node {
	if isRed(h.right) {
		h = rotateLeft(h)
	}

	// PETAR: added |h.left != nil|
	if h.left != nil && isRed(h.left) && isRed(h.left.left) {
		h = rotateRight(h)
	}

	if isRed(h.left) && isRed(h.right) {
		flip(h)
	}

	return h
}
<commit_msg>fixed small bug in counting after a reset<commit_after>// Copyright 2010 Petar Maymounkov. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// A Left-Leaning Red-Black (LLRB) implementation of 2-3 balanced binary search trees

package llrb

// |Tree| is a Left-Leaning Red-Black (LLRB) implementation of 2-3 trees, based on:
//
//   http://www.cs.princeton.edu/~rs/talks/LLRB/08Penn.pdf
//   http://www.cs.princeton.edu/~rs/talks/LLRB/LLRB.pdf
//   http://www.cs.princeton.edu/~rs/talks/LLRB/Java/RedBlackBST.java
//
//  2-3 trees (and the run-time equivalent 2-3-4 trees) are the de facto standard BST
//  algoritms found in implementations of Python, Java, and other libraries. The LLRB
//  implementation of 2-3 trees is a recent improvement on the traditional implementation,
//  observed and documented by Robert Sedgewick.
//
// |Tree| has an associative interface, i.e. duplicate key are not allowed.
// The zero-value of a |Tree| represents a ready-for-use tree.
type Tree struct {
	count int
	root  *node
}

// An |Item| represents an object that can be inserted in |Tree|. It acts as a
// key via the method |LessThan|, which induces a full ordering on |Item|s. It is
// also a value, as the user can attach any data to it.
type Item interface {

	// |LessThan| returns true, if and only if |this| is STRICTLY less than |other|
	LessThan(other interface{}) bool
}

// Init resets (empties) the tree
func (t *Tree) Init() {
	t.root = nil
	t.count = 0
}

func (t *Tree) Len() int { return t.count }

// |Has| returns true if the tree contains an element
// whose |LessThan| order equals that of |key|.
func (t *Tree) Has(key Item) bool {
	return t.Get(key) != nil
}

// |Get| retrieves an element from the tree whose |LessThan| order
// equals that of |key|.
func (t *Tree) Get(key Item) Item {
	return get(t.root, key)
}

func get(h *node, item Item) Item {
	if h == nil {
		return nil
	}
	if item.LessThan(h.item) {
		return get(h.left, item)
	}
	if h.item.LessThan(item) {
		return get(h.right, item)
	}
	return h.item
}

// |Min| returns the minimum element in the tree.
func (t *Tree) Min() Item {
	return min(t.root)
}

func min(h *node) Item {
	if h == nil {
		return nil
	}
	if h.left == nil {
		return h.item
	}
	return min(h.left)
}

// |Max| returns the maximum element in the tree.
func (t *Tree) Max() Item {
	return max(t.root)
}

func max(h *node) Item {
	if h == nil {
		return nil
	}
	if h.right == nil {
		return h.item
	}
	return max(h.right)
}

// |Insert| inserts a new element in the tree, or replaces
// an existing one of identical |LessThan| order.
// If a replacement occurred, the replaced item is returned.
func (t *Tree) InsertOrReplace(item Item) Item {
	if item == nil {
		panic("inserting nil item")
	}
	var replaced Item
	t.root, replaced = insert(t.root, item)
	t.root.black = true
	if replaced == nil {
		t.count++
	}
	return replaced
}

func insert(h *node, item Item) (*node, Item) {
	if h == nil {
		return newNode(item), nil
	}

	// PLACEHOLDER: 2-3-4 tree (see comment below)

	var replaced Item
	if item.LessThan(h.item) {
		h.left, replaced = insert(h.left, item)
	} else if h.item.LessThan(item) {
		h.right, replaced = insert(h.right, item)
	} else {
		replaced, h.item = h.item, item
	}

	if isRed(h.right) && !isRed(h.left) {
		h = rotateLeft(h)
	}

	// PETAR: added |h.left != nil|
	if h.left != nil && isRed(h.left) && isRed(h.left.left) {
		h = rotateRight(h)
	}

	// When the next 3 lines of code are here, the LLRB behaves
	// like a 2-3 tree. If they are moved to the 2-3-4 placeholder above,
	// the LLRB tree behaves like a 2-3-4 tree.
	if isRed(h.left) && isRed(h.right) {
		flip(h)
	}

	return h, replaced
}

// |DeleteMin| deletes the minimum element in the tree and returns the
// deleted item or nil otherwise.
func (t *Tree) DeleteMin() Item {
	var deleted Item
	t.root, deleted = deleteMin(t.root)
	if t.root != nil {
		t.root.black = true
	}
	if deleted != nil {
		t.count--
	}
	return deleted
}

// deleteMin code for LLRB 2-3 trees
func deleteMin(h *node) (*node, Item) {
	if h == nil {
		return nil, nil
	}
	if h.left == nil {
		return nil, h.item
	}

	if !isRed(h.left) && !isRed(h.left.left) {
		h = moveRedLeft(h)
	}

	var deleted Item
	h.left, deleted = deleteMin(h.left)

	return fixUp(h), deleted
}

// |DeleteMax| deletes the maximum element in the tree and returns
// the deleted item or nil otherwise
func (t *Tree) DeleteMax() Item {
	var deleted Item
	t.root, deleted = deleteMax(t.root)
	if t.root != nil {
		t.root.black = true
	}
	if deleted != nil {
		t.count--
	}
	return deleted
}

func deleteMax(h *node) (*node, Item) {
	if h == nil {
		return nil, nil
	}
	if isRed(h.left) {
		h = rotateRight(h)
	}
	if h.right == nil {
		return nil, h.item
	}
	if !isRed(h.right) && !isRed(h.right.left) {
		h = moveRedRight(h)
	}
	var deleted Item
	h.right, deleted = deleteMax(h.right)

	return fixUp(h), deleted
}

// |Delete| deletes an item from the tree whose key equals |key|.
// The deleted item is return, otherwise nil is returned.
func (t *Tree) Delete(key Item) Item {
	var deleted Item
	t.root, deleted = delete(t.root, key)
	if t.root != nil {
		t.root.black = true
	}
	if deleted != nil {
		t.count--
	}
	return deleted
}

func delete(h *node, item Item) (*node, Item) {
	var deleted Item
	if h == nil {
		return nil, nil
	}
	if item.LessThan(h.item) {
		if h.left == nil { // item not present. Nothing to delete
			return h, nil
		}
		if !isRed(h.left) && !isRed(h.left.left) {
			h = moveRedLeft(h)
		}
		h.left, deleted = delete(h.left, item)
	} else {
		if isRed(h.left) {
			h = rotateRight(h)
		}
		// If |item| equals |h.item| and no right children at |h|
		if !h.item.LessThan(item) && h.right == nil {
			return nil, h.item
		}
		// PETAR: Added |h.right != nil| below
		if h.right != nil && !isRed(h.right) && !isRed(h.right.left) {
			h = moveRedRight(h)
		}
		// If |item| equals |h.item|, and (from above) |h.right != nil|
		if !h.item.LessThan(item) {
			var subDeleted Item
			h.right, subDeleted = deleteMin(h.right)
			if subDeleted == nil {
				panic("logic")
			}
			deleted, h.item = h.item, subDeleted
		} else { // Else, |item| is bigger than |h.item|
			h.right, deleted = delete(h.right, item)
		}
	}

	return fixUp(h), deleted
}

// |Iter| returns a chan that iterates through all elements in the
// tree in ascending order.
func (t *Tree) Iter() <-chan Item {
	c := make(chan Item)
	go func() {
		iterateInOrder(t.root, c)
		close(c)
	}()
	return c
}

func iterateInOrder(h *node, c chan<- Item) {
	if h == nil {
		return
	}
	iterateInOrder(h.left, c)
	c <- h.item
	iterateInOrder(h.right, c)
}

func iteratePreOrder(h *node, c chan<- Item) {
	if h == nil {
		return
	}
	c <- h.item
	iteratePreOrder(h.left, c)
	iteratePreOrder(h.right, c)
}

func iteratePostOrder(h *node, c chan<- Item) {
	if h == nil {
		return
	}
	iteratePostOrder(h.left, c)
	iteratePostOrder(h.right, c)
	c <- h.item
}

type node struct {
	item        Item
	left, right *node // Pointers to left and right child nodes
	black       bool  // If set, the color of the link (incoming from the parent) is black
	// In the LLRB, new nodes are always red, hence the zero-value for node
}

func newNode(item Item) *node { return &node{item: item} }

func isRed(h *node) bool {
	if h == nil {
		return false
	}
	return !h.black
}

func rotateLeft(h *node) *node {
	x := h.right
	if x.black {
		panic("rotating a black link")
	}
	h.right = x.left
	x.left = h
	x.black = h.black
	h.black = false
	return x
}

func rotateRight(h *node) *node {
	x := h.left
	if x.black {
		panic("rotating a black link")
	}
	h.left = x.right
	x.right = h
	x.black = h.black
	h.black = false
	return x
}

// XXX:
//      - Can flip() ever be called with |h.left == nil| or |h.right == nil|?

// Left and right children must be present
func flip(h *node) {
	h.black = !h.black
	h.left.black = !h.left.black
	h.right.black = !h.right.black
}

// Left and right children must be present
func moveRedLeft(h *node) *node {
	flip(h)
	if isRed(h.right.left) {
		h.right = rotateRight(h.right)
		h = rotateLeft(h)
		flip(h)
	}
	return h
}

// Left and right children must be present
func moveRedRight(h *node) *node {
	flip(h)
	if isRed(h.left.left) {
		h = rotateRight(h)
		flip(h)
	}
	return h
}

func fixUp(h *node) *node {
	if isRed(h.right) {
		h = rotateLeft(h)
	}

	// PETAR: added |h.left != nil|
	if h.left != nil && isRed(h.left) && isRed(h.left.left) {
		h = rotateRight(h)
	}

	if isRed(h.left) && isRed(h.right) {
		flip(h)
	}

	return h
}
<|endoftext|>


<commit_before><commit_msg>Implemented first step<commit_after><|endoftext|>


<commit_before>package log

import (
	"fmt"
	"net/http"
	"runtime/debug"
	"time"

	"github.com/urandom/handler"
)

// PanicDateFormat is the default timestamp format for the panic messages.
const PanicDateFormat = "Jan 2, 2006 at 3:04pm (MST)"

// PanicOpts represents the various options for the Panic handler
type PanicOpts struct {
	// Logger will be used to print out detailed message whenever a panic is
	// recovered. Each message includes a stack trace and timestamp. If none is
	// provded, os.Stderr is used.
	Logger handler.Logger
	// ShowStack will print the stack in the given http.ResponseWriter if true.
	ShowStack bool
	// DateFormat is used to format the timestamp. Defaults to PanicDateFormat.
	DateFormat string
}

// Panic returns a handler that invokes the passed handler h, catching any
// panics. If one occurs, an HTTP 500 response is produced.
func Panic(h http.Handler, o PanicOpts) http.Handler {
	if o.Logger == nil {
		o.Logger = handler.ErrLogger()
	}
	if o.DateFormat == "" {
		o.DateFormat = PanicDateFormat
	}
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		defer func() {
			if rec := recover(); rec != nil {
				stack := debug.Stack()
				timestamp := time.Now().Format(o.DateFormat)
				message := fmt.Sprintf("%s - %s\n%s\n", timestamp, rec, stack)

				o.Logger.Print(message)

				w.WriteHeader(http.StatusInternalServerError)

				if !o.ShowStack {
					message = "Internal Server Error"
				}

				w.Write([]byte(message))
			}
		}()

		h.ServeHTTP(w, r)
	})
}
<commit_msg>Update panic.go<commit_after>package log

import (
	"fmt"
	"net/http"
	"runtime/debug"
	"time"

	"github.com/urandom/handler"
)

// PanicDateFormat is the default timestamp format for the panic messages.
const PanicDateFormat = "Jan 2, 2006 at 3:04pm (MST)"

// PanicOpts represents the various options for the Panic handler
type PanicOpts struct {
	// Logger will be used to print out detailed message whenever a panic is
	// recovered. Each message includes a stack trace and timestamp. If none is
	// provided, os.Stderr is used.
	Logger handler.Logger
	// ShowStack will print the stack in the given http.ResponseWriter if true.
	ShowStack bool
	// DateFormat is used to format the timestamp. Defaults to PanicDateFormat.
	DateFormat string
}

// Panic returns a handler that invokes the passed handler h, catching any
// panics. If one occurs, an HTTP 500 response is produced.
func Panic(h http.Handler, o PanicOpts) http.Handler {
	if o.Logger == nil {
		o.Logger = handler.ErrLogger()
	}
	if o.DateFormat == "" {
		o.DateFormat = PanicDateFormat
	}
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		defer func() {
			if rec := recover(); rec != nil {
				stack := debug.Stack()
				timestamp := time.Now().Format(o.DateFormat)
				message := fmt.Sprintf("%s - %s\n%s\n", timestamp, rec, stack)

				o.Logger.Print(message)

				w.WriteHeader(http.StatusInternalServerError)

				if !o.ShowStack {
					message = "Internal Server Error"
				}

				w.Write([]byte(message))
			}
		}()

		h.ServeHTTP(w, r)
	})
}
<|endoftext|>


<commit_before>package TF2RconWrapper

import (
	"bytes"
	"log"
	"math/rand"
	"net"
	"strconv"
	"sync"
	"time"
)

type Handler interface {
	PlayerConnected(PlayerData)
	PlayerDisconnected(PlayerData)

	PlayerGlobalMessage(PlayerData, string)
	PlayerTeamMessage(PlayerData, string)

	PlayerClassChange(PlayerData, string) // string -> new classes
	PlayerTeamChange(PlayerData, string)  // string -> new team

	GameOver()

	CVarChange(variable string, value string)
	LogFileClosed()
}

type Listener struct {
	mapMu    *sync.RWMutex
	sources  map[string]*Source
	channels map[string](chan string)

	listenAddr   *net.UDPAddr
	redirectAddr string
}

type Source struct {
	Secret string
	logsMu *sync.RWMutex //protects logs
	logs   *bytes.Buffer

	handlerMu *sync.Mutex //protects handler and closed
	handler   Handler
	closed    bool
}

func (s *Source) Logs() *bytes.Buffer {
	s.logsMu.RLock()
	b := s.logs.Bytes()
	var logs []byte
	copy(b, logs)
	s.logsMu.RUnlock()

	return bytes.NewBuffer(logs)
}

// NewListener returns a new Listener
func NewListener(addr string) (*Listener, error) {
	return NewListenerAddr(addr, addr)
}

func NewListenerAddr(port, redirectAddr string) (*Listener, error) {
	addr, err := net.ResolveUDPAddr("udp", ":"+port)
	if err != nil {
		return nil, err
	}

	l := &Listener{
		mapMu:    new(sync.RWMutex),
		sources:  make(map[string]*Source),
		channels: make(map[string](chan string)),

		listenAddr:   addr,
		redirectAddr: redirectAddr,
	}

	conn, err := net.ListenUDP("udp", addr)
	if err != nil {
		return nil, err
	}

	go l.start(conn)
	return l, nil
}

func (l *Listener) RemoveSource(s *Source, m *TF2RconConnection) {
	s.handlerMu.Lock()
	s.closed = true
	s.handlerMu.Unlock()

	l.mapMu.Lock()
	delete(l.sources, s.Secret)
	l.mapMu.Unlock()

	m.StopLogRedirection(l.redirectAddr)
}

func (l *Listener) start(conn *net.UDPConn) {
	for {
		buff := make([]byte, 512)
		n, err := conn.Read(buff)
		if err != nil {
			log.Println(err)
		}

		secret, pos, err := getSecret(buff[0:n])
		if err != nil {
			continue
		}

		l.mapMu.RLock()
		source, ok := l.sources[secret]
		l.mapMu.RUnlock()

		if !ok {
			continue
		}

		go func() {
			source.handlerMu.Lock()
			defer source.handlerMu.Unlock()

			if source.closed {
				return
			}

			source.logsMu.Lock()
			source.logs.WriteString("L ")
			source.logs.Write(buff[pos : n-2])
			source.logs.WriteByte('\n')
			source.logsMu.Unlock()

			handler := source.handler

			m := parse(string(buff[pos : n-2]))

			switch m.Parsed.Type {
			case PlayerGlobalMessage:
				d := m.Parsed.Data.(PlayerData)
				handler.PlayerGlobalMessage(d, d.Text)
			case PlayerTeamMessage:
				d := m.Parsed.Data.(PlayerData)
				handler.PlayerTeamMessage(d, d.Text)
			case PlayerChangedClass:
				d := m.Parsed.Data.(PlayerData)
				handler.PlayerClassChange(d, d.Class)
			case PlayerChangedTeam:
				d := m.Parsed.Data.(PlayerData)
				handler.PlayerTeamChange(d, d.NewTeam)
			case PlayerConnected:
				d := m.Parsed.Data.(PlayerData)
				handler.PlayerConnected(d)
			case PlayerDisconnected:
				d := m.Parsed.Data.(PlayerData)
				handler.PlayerDisconnected(d)
			case WorldGameOver:
				handler.GameOver()
			case ServerCvar:
				d := m.Parsed.Data.(CvarData)
				handler.CVarChange(d.Variable, d.Value)
			}
		}()
	}
}

func (l *Listener) AddSource(handler Handler, m *TF2RconConnection) *Source {
	secret := strconv.Itoa(rand.Intn(999998) + 1)
	rand.Seed(time.Now().Unix())

	l.mapMu.RLock()
	_, ok := l.sources[secret]
	for ok {
		secret = strconv.Itoa(rand.Intn(999998) + 1)
		_, ok = l.sources[secret]
	}
	l.mapMu.RUnlock()

	return l.AddSourceSecret(secret, handler, m)
}

func (l *Listener) AddSourceSecret(secret string, handler Handler, m *TF2RconConnection) *Source {
	s := newSource(secret, handler)

	l.mapMu.Lock()
	l.sources[secret] = s
	l.mapMu.Unlock()

	m.Query("sv_logsecret " + secret)
	m.RedirectLogs(l.redirectAddr)
	return s
}

func newSource(secret string, handler Handler) *Source {
	return &Source{
		Secret: secret,
		logsMu: new(sync.RWMutex),
		logs:   new(bytes.Buffer),

		handlerMu: new(sync.Mutex),
		handler:   handler,
	}
}
<commit_msg>Add handlers for new events<commit_after>package TF2RconWrapper

import (
	"bytes"
	"log"
	"math/rand"
	"net"
	"strconv"
	"sync"
	"time"
)

type Handler interface {
	PlayerConnected(PlayerData)
	PlayerDisconnected(PlayerData)

	PlayerGlobalMessage(PlayerData, string) // strings are chat message
	PlayerTeamMessage(PlayerData, string)

	PlayerSpawned(PlayerData, string)     // string is class
	PlayerClassChange(PlayerData, string) // string is new classes
	PlayerTeamChange(PlayerData, string)  // string is new team

	PlayerKilled(PlayerKill)
	PlayerDamaged(PlayerDamage)
	PlayerHealed(PlayerHeal)
	PlayerKilledMedic(PlayerTrigger)
	PlayerUberFinished(PlayerData)
	PlayerBlockedCapture(PlayerData)
	PlayerItemPickup(ItemPickup)

	TeamPointCapture(TeamData)
	TeamScoreUpdate(TeamData)
	GameOver()

	WorldRoundWin(string) // string is team which won

	CVarChange(variable string, value string)
	LogFileClosed()
}

type Listener struct {
	mapMu    *sync.RWMutex
	sources  map[string]*Source
	channels map[string](chan string)

	listenAddr   *net.UDPAddr
	redirectAddr string
}

type Source struct {
	Secret string
	logsMu *sync.RWMutex //protects logs
	logs   *bytes.Buffer

	handlerMu *sync.Mutex //protects handler and closed
	handler   Handler
	closed    bool
}

func (s *Source) Logs() *bytes.Buffer {
	s.logsMu.RLock()
	b := s.logs.Bytes()
	var logs []byte
	copy(b, logs)
	s.logsMu.RUnlock()

	return bytes.NewBuffer(logs)
}

// NewListener returns a new Listener
func NewListener(addr string) (*Listener, error) {
	return NewListenerAddr(addr, addr)
}

func NewListenerAddr(port, redirectAddr string) (*Listener, error) {
	addr, err := net.ResolveUDPAddr("udp", ":"+port)
	if err != nil {
		return nil, err
	}

	l := &Listener{
		mapMu:    new(sync.RWMutex),
		sources:  make(map[string]*Source),
		channels: make(map[string](chan string)),

		listenAddr:   addr,
		redirectAddr: redirectAddr,
	}

	conn, err := net.ListenUDP("udp", addr)
	if err != nil {
		return nil, err
	}

	go l.start(conn)
	return l, nil
}

func (l *Listener) RemoveSource(s *Source, m *TF2RconConnection) {
	s.handlerMu.Lock()
	s.closed = true
	s.handlerMu.Unlock()

	l.mapMu.Lock()
	delete(l.sources, s.Secret)
	l.mapMu.Unlock()

	m.StopLogRedirection(l.redirectAddr)
}

func (l *Listener) start(conn *net.UDPConn) {
	for {
		buff := make([]byte, 512)
		n, err := conn.Read(buff)
		if err != nil {
			log.Println(err)
		}

		secret, pos, err := getSecret(buff[0:n])
		if err != nil {
			continue
		}

		l.mapMu.RLock()
		source, ok := l.sources[secret]
		l.mapMu.RUnlock()

		if !ok {
			continue
		}

		go func() {
			source.handlerMu.Lock()
			defer source.handlerMu.Unlock()

			if source.closed {
				return
			}

			source.logsMu.Lock()
			source.logs.WriteString("L ")
			source.logs.Write(buff[pos : n-2])
			source.logs.WriteByte('\n')
			source.logsMu.Unlock()

			handler := source.handler

			m := parse(string(buff[pos : n-2]))

			m.Parsed.CallHandler(handler)
		}()
	}
}

func (l *Listener) AddSource(handler Handler, m *TF2RconConnection) *Source {
	secret := strconv.Itoa(rand.Intn(999998) + 1)
	rand.Seed(time.Now().Unix())

	l.mapMu.RLock()
	_, ok := l.sources[secret]
	for ok {
		secret = strconv.Itoa(rand.Intn(999998) + 1)
		_, ok = l.sources[secret]
	}
	l.mapMu.RUnlock()

	return l.AddSourceSecret(secret, handler, m)
}

func (l *Listener) AddSourceSecret(secret string, handler Handler, m *TF2RconConnection) *Source {
	s := newSource(secret, handler)

	l.mapMu.Lock()
	l.sources[secret] = s
	l.mapMu.Unlock()

	m.Query("sv_logsecret " + secret)
	m.RedirectLogs(l.redirectAddr)
	return s
}

func newSource(secret string, handler Handler) *Source {
	return &Source{
		Secret: secret,
		logsMu: new(sync.RWMutex),
		logs:   new(bytes.Buffer),

		handlerMu: new(sync.Mutex),
		handler:   handler,
	}
}
<|endoftext|>


<commit_before>package lora

import (
	"encoding/json"
	"log"
	"net"
)

const (
	PUSH_DATA = iota
	PUSH_ACK  = iota
	PULL_DATA = iota
	PULL_ACK  = iota
	PULL_RESP = iota
)

var buf = make([]byte, 2048)

type Message struct {
	ProtocolVersion int
	Token           []byte
	Identifier      int
	Payload         *json.RawMessage
	GatewayEUI      int64
	Address         *net.UDPAddr
}

type Conn struct {
	Raw *net.UDPConn
}

func NewConn(r *net.UDPConn) *Conn {
	return &Conn{r}
}

func (c *Conn) ReadMessage() (*Message, error) {
	n, addr, err := c.Raw.ReadFromUDP(buf)
	if err != nil {
		log.Print("Error: ", err)
		return nil, err
	}
	log.Print("Received ", string(buf[0:n]), " from ", addr)
	msg := &Message{
		Address:    addr,
		Identifier: int(buf[3]),
	}
	return msg, nil
}
<commit_msg>Extract slightly more from udp bytes and more logging<commit_after>package lora

import (
	"encoding/json"
	"log"
	"net"
)

const (
	PUSH_DATA = iota
	PUSH_ACK  = iota
	PULL_DATA = iota
	PULL_ACK  = iota
	PULL_RESP = iota
)

var buf = make([]byte, 2048)

type Message struct {
	ProtocolVersion int
	Token           []byte
	Identifier      int
	Payload         *json.RawMessage
	GatewayEUI      net.HardwareAddr
	Address         *net.UDPAddr
}

type Conn struct {
	Raw *net.UDPConn
}

func NewConn(r *net.UDPConn) *Conn {
	return &Conn{r}
}

func (c *Conn) ReadMessage() (*Message, error) {
	n, addr, err := c.Raw.ReadFromUDP(buf)
	if err != nil {
		log.Print("Error: ", err)
		return nil, err
	}
	log.Print("Received raw bytes", buf[0:n], " from ", addr)
	log.Print("Received ", string(buf[0:n]), " from ", addr)
	msg := &Message{
		Address:         addr,
		ProtocolVersion: int(buf[0]),
		Token:           buf[1:3],
		Identifier:      int(buf[3]),
	}
	return msg, nil
}

func (m *Message) Ack() error {

	return nil
}
<|endoftext|>


<commit_before>// Copyright (c) 2017 The btcsuite developers
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package wire

import (
	"fmt"
	"github.com/btcsuite/fastsha256"
	"io"
)

const (
	// MaxCFilterHeaderDataSize is the maximum byte size of a committed
	// filter header.
	MaxCFilterHeaderDataSize = fastsha256.Size
)
type MsgCFilterHeader struct {
	Data []byte
}

// BtcDecode decodes r using the bitcoin protocol encoding into the receiver.
// This is part of the Message interface implementation.
func (msg *MsgCFilterHeader) BtcDecode(r io.Reader, pver uint32) error {
	var err error
	msg.Data, err = ReadVarBytes(r, pver, MaxCFilterHeaderDataSize,
		"cf header data")
	return err
}

// BtcEncode encodes the receiver to w using the bitcoin protocol encoding.
// This is part of the Message interface implementation.
func (msg *MsgCFilterHeader) BtcEncode(w io.Writer, pver uint32) error {
	size := len(msg.Data)
	if size > MaxCFilterHeaderDataSize {
		str := fmt.Sprintf("cf header size too large for message " +
			"[size %v, max %v]", size, MaxCFilterHeaderDataSize)
		return messageError("MsgCFilterHeader.BtcEncode", str)
	}

	return WriteVarBytes(w, pver, msg.Data)
}

// Command returns the protocol command string for the message.  This is part
// of the Message interface implementation.
func (msg *MsgCFilterHeader) Command() string {
	return CmdCFilterHeader
}

// MaxPayloadLength returns the maximum length the payload can be for the
// receiver.  This is part of the Message interface implementation.
func (msg *MsgCFilterHeader) MaxPayloadLength(pver uint32) uint32 {
	return uint32(VarIntSerializeSize(MaxCFilterHeaderDataSize)) +
		MaxCFilterHeaderDataSize
}

// NewMsgFilterAdd returns a new bitcoin cfilterheader message that conforms to
// the Message interface. See MsgCFilterHeader for details.
func NewMsgCFilterHeader(data []byte) *MsgCFilterHeader {
	return &MsgCFilterHeader{
		Data: data,
	}
}
<commit_msg>define MsgCFilterHeader.Deserialize(), required by btcrpcclient<commit_after>// Copyright (c) 2017 The btcsuite developers
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package wire

import (
	"fmt"
	"github.com/btcsuite/fastsha256"
	"io"
)

const (
	// MaxCFilterHeaderDataSize is the maximum byte size of a committed
	// filter header.
	MaxCFilterHeaderDataSize = fastsha256.Size
)
type MsgCFilterHeader struct {
	Data []byte
}

// BtcDecode decodes r using the bitcoin protocol encoding into the receiver.
// This is part of the Message interface implementation.
func (msg *MsgCFilterHeader) BtcDecode(r io.Reader, pver uint32) error {
	var err error
	msg.Data, err = ReadVarBytes(r, pver, MaxCFilterHeaderDataSize,
		"cf header data")
	return err
}

// BtcEncode encodes the receiver to w using the bitcoin protocol encoding.
// This is part of the Message interface implementation.
func (msg *MsgCFilterHeader) BtcEncode(w io.Writer, pver uint32) error {
	size := len(msg.Data)
	if size > MaxCFilterHeaderDataSize {
		str := fmt.Sprintf("cf header size too large for message " +
			"[size %v, max %v]", size, MaxCFilterHeaderDataSize)
		return messageError("MsgCFilterHeader.BtcEncode", str)
	}

	return WriteVarBytes(w, pver, msg.Data)
}

// Deserialize decodes a filter header from r into the receiver using a format
// that is suitable for long-term storage such as a database. This function
// differs from BtcDecode in that BtcDecode decodes from the bitcoin wire
// protocol as it was sent across the network.  The wire encoding can
// technically differ depending on the protocol version and doesn't even really
// need to match the format of a stored filter header at all. As of the time
// this comment was written, the encoded filter header is the same in both
// instances, but there is a distinct difference and separating the two allows
// the API to be flexible enough to deal with changes.
func (msg *MsgCFilterHeader) Deserialize(r io.Reader) error {
	// At the current time, there is no difference between the wire encoding
	// and the stable long-term storage format.  As a result, make use of
	// BtcDecode.
	return msg.BtcDecode(r, 0)
}

// Command returns the protocol command string for the message.  This is part
// of the Message interface implementation.
func (msg *MsgCFilterHeader) Command() string {
	return CmdCFilterHeader
}

// MaxPayloadLength returns the maximum length the payload can be for the
// receiver.  This is part of the Message interface implementation.
func (msg *MsgCFilterHeader) MaxPayloadLength(pver uint32) uint32 {
	return uint32(VarIntSerializeSize(MaxCFilterHeaderDataSize)) +
		MaxCFilterHeaderDataSize
}

// NewMsgFilterAdd returns a new bitcoin cfilterheader message that conforms to
// the Message interface. See MsgCFilterHeader for details.
func NewMsgCFilterHeader(data []byte) *MsgCFilterHeader {
	return &MsgCFilterHeader{
		Data: data,
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2021 The logr Authors.

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

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

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

package main

import (
	"fmt"

	"github.com/go-logr/logr"
	"github.com/go-logr/logr/funcr"
)

type e struct {
	str string
}

func (e e) Error() string {
	return e.str
}

func helper(log logr.Logger, msg string) {
	helper2(log, msg)
}

func helper2(log logr.Logger, msg string) {
	log.WithCallDepth(2).Info(msg)
}

func main() {
	log := funcr.New(
		func(pfx, args string) { fmt.Println(pfx, args) },
		funcr.Options{
			LogCaller:    funcr.All,
			LogTimestamp: true,
		})
	example(log.WithValues("module", "example"))
}

func example(log logr.Logger) {
	log.Info("hello", "val1", 1, "val2", map[string]int{"k": 1})
	log.V(1).Info("you should see this")
	log.V(1).V(1).Info("you should NOT see this")
	log.Error(nil, "uh oh", "trouble", true, "reasons", []float64{0.1, 0.11, 3.14})
	log.Error(e{"an error occurred"}, "goodbye", "code", -1)
	helper(log, "thru a helper")
}
<commit_msg>funcr: fix example verbosity<commit_after>/*
Copyright 2021 The logr Authors.

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

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

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

package main

import (
	"fmt"

	"github.com/go-logr/logr"
	"github.com/go-logr/logr/funcr"
)

type e struct {
	str string
}

func (e e) Error() string {
	return e.str
}

func helper(log logr.Logger, msg string) {
	helper2(log, msg)
}

func helper2(log logr.Logger, msg string) {
	log.WithCallDepth(2).Info(msg)
}

func main() {
	log := funcr.New(
		func(pfx, args string) { fmt.Println(pfx, args) },
		funcr.Options{
			LogCaller:    funcr.All,
			LogTimestamp: true,
			Verbosity:    1,
		})
	example(log.WithValues("module", "example"))
}

func example(log logr.Logger) {
	log.Info("hello", "val1", 1, "val2", map[string]int{"k": 1})
	log.V(1).Info("you should see this")
	log.V(1).V(1).Info("you should NOT see this")
	log.Error(nil, "uh oh", "trouble", true, "reasons", []float64{0.1, 0.11, 3.14})
	log.Error(e{"an error occurred"}, "goodbye", "code", -1)
	helper(log, "thru a helper")
}
<|endoftext|>


<commit_before>package main

import (
	"os"

	log "github.com/Sirupsen/logrus"
	"github.com/bamstats"
	"github.com/codegangsta/cli"
)

func check(err error) {
	if err != nil {
		log.Fatal(err)
	}
}

var (
	bam, annotation, loglevel string
	cpu, maxBuf, reads        int
)

func run(c *cli.Context) {
	level, err := log.ParseLevel(loglevel)
	check(err)
	log.SetLevel(level)
	if bam == "" {
		log.Fatal("no file specified")
	}
	// stats := bamstats.Coverage1(bam, annotation, cpu)
	stats := bamstats.General(bam, cpu, maxBuf, reads)
	bamstats.OutputJson(stats)
}

func main() {
	app := cli.NewApp()
	app.Name = "bamstats"
	app.Usage = "Compute mapping statistics"
	app.Version = bamstats.Version
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:        "bam, b",
			Value:       "",
			Usage:       "input file",
			Destination: &bam,
		},
		cli.StringFlag{
			Name:        "annotation, a",
			Value:       "",
			Usage:       "bgzip compressed and indexed annotation file",
			Destination: &annotation,
		},
		cli.StringFlag{
			Name:        "loglevel",
			Value:       "warn",
			Usage:       "logging level",
			Destination: &loglevel,
		},
		cli.IntFlag{
			Name:        "cpu, c",
			Value:       1,
			Usage:       "number of cpus to be used",
			Destination: &cpu,
		},
		cli.IntFlag{
			Name:        "max-buf",
			Value:       1000000,
			Usage:       "maximum number of buffered records",
			Destination: &maxBuf,
		},
		cli.IntFlag{
			Name:        "n",
			Value:       -1,
			Usage:       "number of records to process",
			Destination: &reads,
		},
	}
	app.Action = run

	if len(os.Args) == 1 {
		os.Args = append(os.Args, "help")
	}
	app.Run(os.Args)
}
<commit_msg>Add output file flag<commit_after>package main

import (
	"os"

	log "github.com/Sirupsen/logrus"
	"github.com/bamstats"
	"github.com/codegangsta/cli"
)

func check(err error) {
	if err != nil {
		log.Fatal(err)
	}
}

var (
	bam, annotation, loglevel, output string
	cpu, maxBuf, reads                int
)

func run(c *cli.Context) {
	level, err := log.ParseLevel(loglevel)
	check(err)
	log.SetLevel(level)
	if bam == "" {
		log.Fatal("no file specified")
	}
	// stats := bamstats.Coverage1(bam, annotation, cpu)
	stats := bamstats.General(bam, cpu, maxBuf, reads)
	out := bamstats.NewOutput(output)
	bamstats.OutputJson(out, stats)
}

func main() {
	app := cli.NewApp()
	app.Name = "bamstats"
	app.Usage = "Compute mapping statistics"
	app.Version = bamstats.Version
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:        "bam, b",
			Value:       "",
			Usage:       "input file",
			Destination: &bam,
		},
		cli.StringFlag{
			Name:        "annotation, a",
			Value:       "",
			Usage:       "bgzip compressed and indexed annotation file",
			Destination: &annotation,
		},
		cli.StringFlag{
			Name:        "loglevel",
			Value:       "warn",
			Usage:       "logging level",
			Destination: &loglevel,
		},
		cli.IntFlag{
			Name:        "cpu, c",
			Value:       1,
			Usage:       "number of cpus to be used",
			Destination: &cpu,
		},
		cli.IntFlag{
			Name:        "max-buf",
			Value:       1000000,
			Usage:       "maximum number of buffered records",
			Destination: &maxBuf,
		},
		cli.IntFlag{
			Name:        "n",
			Value:       -1,
			Usage:       "number of records to process",
			Destination: &reads,
		},
		cli.StringFlag{
			Name:        "o",
			Value:       "-",
			Usage:       "output file",
			Destination: &output,
		},
	}
	app.Action = run

	if len(os.Args) == 1 {
		os.Args = append(os.Args, "help")
	}
	app.Run(os.Args)
}
<|endoftext|>


<commit_before>package main

import (
	log "github.com/Sirupsen/logrus"
	"github.com/bamstats"
	"github.com/codegangsta/cli"
	"os"
)

func check(err error) {
	if err != nil {
		log.Fatal(err)
	}
}

var (
	bam, annotation, loglevel string
	cpu                       int
)

func run(c *cli.Context) {
	level, err := log.ParseLevel(loglevel)
	check(err)
	log.SetLevel(level)
	if bam == "" {
		log.Fatal("no file specified")
	}
	stats := bamstats.Coverage(bam, annotation, cpu)
	bamstats.OutputJson(stats)
}

func main() {
	app := cli.NewApp()
	app.Name = "bamstats"
	app.Usage = "Compute mapping statistics"
	app.Version = bamstats.Version
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:        "bam, b",
			Value:       "",
			Usage:       "input file",
			Destination: &bam,
		},
		cli.StringFlag{
			Name:        "annotation, a",
			Value:       "",
			Usage:       "bgzip compressed and indexed annotation file",
			Destination: &annotation,
		},
		cli.StringFlag{
			Name:        "loglevel",
			Value:       "warn",
			Usage:       "logging level",
			Destination: &loglevel,
		},
		cli.IntFlag{
			Name:        "cpu, c",
			Value:       1,
			Usage:       "number of cpus to be used",
			Destination: &cpu,
		},
	}
	app.Action = run

	if len(os.Args) == 1 {
		os.Args = append(os.Args, "help")
	}
	app.Run(os.Args)
}
<commit_msg>Update main for general stats<commit_after>package main

import (
	"os"

	log "github.com/Sirupsen/logrus"
	"github.com/bamstats"
	"github.com/codegangsta/cli"
)

func check(err error) {
	if err != nil {
		log.Fatal(err)
	}
}

var (
	bam, annotation, loglevel string
	cpu, maxBuf               int
)

func run(c *cli.Context) {
	level, err := log.ParseLevel(loglevel)
	check(err)
	log.SetLevel(level)
	if bam == "" {
		log.Fatal("no file specified")
	}
	// stats := bamstats.Coverage1(bam, annotation, cpu)
	stats := bamstats.General(bam, cpu, maxBuf)
	bamstats.OutputJson(stats)
}

func main() {
	app := cli.NewApp()
	app.Name = "bamstats"
	app.Usage = "Compute mapping statistics"
	app.Version = bamstats.Version
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:        "bam, b",
			Value:       "",
			Usage:       "input file",
			Destination: &bam,
		},
		cli.StringFlag{
			Name:        "annotation, a",
			Value:       "",
			Usage:       "bgzip compressed and indexed annotation file",
			Destination: &annotation,
		},
		cli.StringFlag{
			Name:        "loglevel",
			Value:       "warn",
			Usage:       "logging level",
			Destination: &loglevel,
		},
		cli.IntFlag{
			Name:        "cpu, c",
			Value:       1,
			Usage:       "number of cpus to be used",
			Destination: &cpu,
		},
		cli.IntFlag{
			Name:        "max-buf",
			Value:       100000,
			Usage:       "maximum buffer size for reading",
			Destination: &maxBuf,
		},
	}
	app.Action = run

	if len(os.Args) == 1 {
		os.Args = append(os.Args, "help")
	}
	app.Run(os.Args)
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"os/exec"
	"os/signal"
	"runtime"
	"strings"

	"github.com/elasticHD/main/search"
	_ "github.com/elasticHD/main/statik"
	"github.com/rakyll/statik/fs"
)

var (
	Stderr io.Writer = os.Stderr // Stderr is the io.Writer to which executed commands write standard error.
	Stdout io.Writer = os.Stdout // Stdout is the io.Writer to which executed commands write standard output.
)

// Config ...
type Config struct {
	ServerHost string
}

var config *Config

func main() {

	parseFlags()
	statikFS, err := fs.New() // 静态文件编译成二进制
	if err != nil {
		log.Fatalf(err.Error())
	}
	http.FileServer(statikFS) // 加载http server file
	var mux = http.NewServeMux()
	// router
	mux.Handle("/", http.StripPrefix("/", http.FileServer(statikFS)))
	// 加载search package
	search.NewSearch(search.InitOptions{
		Mux: mux,
	})
	// start http server
	go func(serverAddr string, m *http.ServeMux) {
		if err = http.ListenAndServe(serverAddr, m); err != nil {
			fmt.Println("Cant start server:", err)
			os.Exit(1)
		}
	}(config.ServerHost, mux)
	openPage()
	// 捕捉ctrl+C 退出信号
	handleSignals()
}

func handleSignals() {
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt, os.Kill)
	<-c
}

func openPage() {
	url := fmt.Sprintf("http://%v", config.ServerHost)
	fmt.Println("To view elasticHD console open", url, "in browser")
	var err error
	switch runtime.GOOS {
	case "linux":
		err = runCmd("xdg-open", url)
	case "darwin":
		err = runCmd("open", url)
	case "windows":
		r := strings.NewReplacer("&", "^&")
		err = runCmd("cmd", "/c", "start", r.Replace(url))
	default:
		err = fmt.Errorf("unsupported platform")
	}
	if err != nil {
		fmt.Println(err)
	}
}

// runCmd run command opens a new browser window pointing to url.
func runCmd(prog string, args ...string) error {
	cmd := exec.Command(prog, args...)
	cmd.Stdout = Stdout
	cmd.Stderr = Stderr
	return cmd.Run()
}

// parseFlags parse flags of program.
func parseFlags() {
	config = &Config{}
	flag.StringVar(&config.ServerHost, "p", "0.0.0.0:9800", "local server address")
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage of %s:\n", os.Args[0])
		flag.PrintDefaults()
	}
	flag.Parse()
}
<commit_msg>Create main.go<commit_after>package main

import (
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"os/exec"
	"os/signal"
	"runtime"
	"strings"

	"github.com/farmerx/ElasticHD/main/search"
	_ "github.com/farmerx/ElasticHD/main/statik"
	"github.com/rakyll/statik/fs"
)

var (
	Stderr io.Writer = os.Stderr // Stderr is the io.Writer to which executed commands write standard error.
	Stdout io.Writer = os.Stdout // Stdout is the io.Writer to which executed commands write standard output.
)

// Config ...
type Config struct {
	ServerHost string
}

var config *Config

func main() {

	parseFlags()
	statikFS, err := fs.New() // 静态文件编译成二进制
	if err != nil {
		log.Fatalf(err.Error())
	}
	http.FileServer(statikFS) // 加载http server file
	var mux = http.NewServeMux()
	// router
	mux.Handle("/", http.StripPrefix("/", http.FileServer(statikFS)))
	// 加载search package
	search.NewSearch(search.InitOptions{
		Mux: mux,
	})
	// start http server
	go func(serverAddr string, m *http.ServeMux) {
		if err = http.ListenAndServe(serverAddr, m); err != nil {
			fmt.Println("Cant start server:", err)
			os.Exit(1)
		}
	}(config.ServerHost, mux)
	openPage()
	// 捕捉ctrl+C 退出信号
	handleSignals()
}

func handleSignals() {
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt, os.Kill)
	<-c
}

func openPage() {
	url := fmt.Sprintf("http://%v", config.ServerHost)
	fmt.Println("To view elasticHD console open", url, "in browser")
	var err error
	switch runtime.GOOS {
	case "linux":
		err = runCmd("xdg-open", url)
	case "darwin":
		err = runCmd("open", url)
	case "windows":
		r := strings.NewReplacer("&", "^&")
		err = runCmd("cmd", "/c", "start", r.Replace(url))
	default:
		err = fmt.Errorf("unsupported platform")
	}
	if err != nil {
		fmt.Println(err)
	}
}

// runCmd run command opens a new browser window pointing to url.
func runCmd(prog string, args ...string) error {
	cmd := exec.Command(prog, args...)
	cmd.Stdout = Stdout
	cmd.Stderr = Stderr
	return cmd.Run()
}

// parseFlags parse flags of program.
func parseFlags() {
	config = &Config{}
	flag.StringVar(&config.ServerHost, "p", "0.0.0.0:9800", "local server address")
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage of %s:\n", os.Args[0])
		flag.PrintDefaults()
	}
	flag.Parse()
}
<|endoftext|>


<commit_before>package main

import (
	"image/png"
	"log"
	"os"
	"strings"
	"time"

	"github.com/peterstace/grayt"
)
import "flag"

func main() {

	var (
		out string
		spp int
	)

	flag.StringVar(&out, "o", "", "output file (must end in .png)")
	flag.IntVar(&spp, "s", 10, "samples per pixel")
	flag.Parse()

	if !strings.HasSuffix(out, ".png") {
		log.Fatalf(`%q does not end in ".png"`, out)
	}

	scene := CornellBox()

	const (
		pxWide  = 300
		pxHigh  = 300
		totalPx = pxWide * pxHigh
	)
	acc := grayt.NewAccumulator(pxWide, pxHigh)
	startTime := time.Now()
	for i := 1; i <= spp; i++ {
		samplesPerSecond := float64(i*totalPx) / time.Now().Sub(startTime).Seconds()
		timeRemaining := time.Duration(((spp-i)*totalPx)/int(samplesPerSecond)) * time.Second
		log.Printf("Sample=%d/%d Sampes/sec=%.2e NRSD=%.2f%% ETA=%s\n",
			i, spp, samplesPerSecond, acc.NeighbourRelativeStdDev()*100, timeRemaining)
		grayt.TracerImage(scene, acc)
	}
	img := acc.ToImage(1.0)

	f, err := os.Create(out)
	if err != nil {
		log.Fatal(err)
	}
	defer f.Close()
	if err := png.Encode(f, img); err != nil {
		log.Fatal(err)
	}
}
<commit_msg>Add mode to stop at a given neighbour relative StdDev<commit_after>package main

import (
	"image/png"
	"log"
	"os"
	"strings"
	"time"

	"github.com/peterstace/grayt"
)
import (
	"flag"
	"fmt"
)

func main() {

	var (
		out  string
		spp  int
		nrsd float64
	)

	flag.StringVar(&out, "o", "", "output file (must end in .png)")
	flag.IntVar(&spp, "s", 0, "samples per pixel stopping point")
	flag.Float64Var(&nrsd, "d", 0.0, "neighbourhood relative std dev stopping point")
	flag.Parse()

	if !strings.HasSuffix(out, ".png") {
		log.Fatalf(`%q does not end in ".png"`, out)
	}
	if (spp == 0 && nrsd == 0) || (spp != 0 && nrsd != 0) {
		log.Fatalf(`exactly 1 of s and d must be set`)
	}

	scene := CornellBox()

	const (
		pxWide  = 300
		pxHigh  = 300
		totalPx = pxWide * pxHigh
	)
	acc := grayt.NewAccumulator(pxWide, pxHigh)

	startTime := time.Now()
	iteration := 0
	var prevRSD, currentRSD float64
	var downRSDCount int
	for {
		iteration++
		samplesPerSecond := float64(iteration*totalPx) / time.Now().Sub(startTime).Seconds()
		timeRemaining := "??"
		totalSamples := "??"
		if spp != 0 {
			timeRemaining = fmt.Sprintf("%v",
				time.Duration(((spp-iteration)*totalPx)/int(samplesPerSecond))*time.Second)
			totalSamples = fmt.Sprintf("%d", spp)
		}
		grayt.TracerImage(scene, acc)
		prevRSD, currentRSD = currentRSD, acc.NeighbourRelativeStdDev()
		log.Printf("Sample=%d/%s, Samples/sec=%.2e NRSD=%.4f ETA=%s\n",
			iteration, totalSamples, samplesPerSecond, currentRSD, timeRemaining)
		if currentRSD < prevRSD {
			downRSDCount++
		} else {
			downRSDCount = 0
		}
		if downRSDCount > 5 && currentRSD < nrsd {
			break
		}
		if spp != 0 && iteration == spp {
			break
		}
	}

	img := acc.ToImage(1.0)

	f, err := os.Create(out)
	if err != nil {
		log.Fatal(err)
	}
	defer f.Close()
	if err := png.Encode(f, img); err != nil {
		log.Fatal(err)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"testing"

	"github.com/8byt/gox/parser"
	"github.com/8byt/gox/scanner"
	"github.com/8byt/gox/token"
)

// should test for all of these
//	case token.ASSIGN, token.EQL, token.NEQ, token.DEFINE,
//token.LPAREN, token.LBRACE, token.COMMA, token.COLON,
//token.RETURN, token.IF, token.SWITCH, token.CASE:

func TestOtagParse(t *testing.T) {
	t.Log(doParseAst("return hello world"))
	t.Log(doParseAst("lol := hahaha"))
	t.Log(doParseAst("if <abc></abc> == <abc></abc> {}"))
	t.Log(doParseAst(`return <a attr="value">hello world</a>`))

	t.Log(doParseAst(`return <a attr="value" attr2={"bang"}>hello world</a>`))
	t.Log(doParseAst(`
		return
		   hello world
		`))
	t.Log(doParseAst(`
		return
		   hello world

		   	yeah that's right
		   	{<back at it="again"></back>}
		   	{func () vecty.Component {
				return <wow></wow>
			}()}

		`))
	// check to see if newlines in OTAGs okay
	t.Log(doParseAst(`
		return
		   hello world
		`))

	t.Skip()
}

func doTestExpr(strExpr string) {
	expr, err := parser.ParseExpr(strExpr)
	fmt.Printf("err: %v value: %v\n", err, expr)
}

func doParseAst(strExpr string) (result string) {
	// src is the input that we want to tokenize.
	src := []byte(strExpr)

	func () string {
		return ""
	}()
	// Initialize the scanner.
	var s scanner.Scanner
	fset := token.NewFileSet()                      // positions are relative to fset
	file := fset.AddFile("", fset.Base(), len(src)) // register input "file"
	s.Init(file, src, nil /* no error handler */, scanner.ScanComments)

	// Repeated calls to Scan yield the token sequence found in the input.
	for {
		pos, tok, lit := s.Scan()
		if tok == token.EOF {
			break
		}
		result += fmt.Sprintf("%s\t%s\t%q\n", fset.Position(pos), tok, lit)
	}

	return result
}
<commit_msg>Add better test case for newlines in bare words<commit_after>package main

import (
	"fmt"
	"testing"

	"github.com/8byt/gox/parser"
	"github.com/8byt/gox/scanner"
	"github.com/8byt/gox/token"
)

// should test for all of these
//	case token.ASSIGN, token.EQL, token.NEQ, token.DEFINE,
//token.LPAREN, token.LBRACE, token.COMMA, token.COLON,
//token.RETURN, token.IF, token.SWITCH, token.CASE:

func TestOtagParse(t *testing.T) {
	t.Log(doParseAst("return hello world"))
	t.Log(doParseAst("lol := hahaha"))
	t.Log(doParseAst("if <abc></abc> == <abc></abc> {}"))
	t.Log(doParseAst(`return <a attr="value">hello world</a>`))

	t.Log(doParseAst(`return <a attr="value" attr2={"bang"}>hello world</a>`))
	t.Log(doParseAst(`
		return
		   hello world
		`))
	t.Log(doParseAst(`
		return
		   hello world

		   	yeah that's right
		   	let's have another line
		   	and another one
		   	{<back at it="again"></back>}
		   	{func () vecty.Component {
				return <wow></wow>
			}()}

		`))
	// check to see if newlines in OTAGs okay
	t.Log(doParseAst(`
		return
		   hello world
		`))

	t.Skip()
}

func doTestExpr(strExpr string) {
	expr, err := parser.ParseExpr(strExpr)
	fmt.Printf("err: %v value: %v\n", err, expr)
}

func doParseAst(strExpr string) (result string) {
	// src is the input that we want to tokenize.
	src := []byte(strExpr)

	func () string {
		return ""
	}()
	// Initialize the scanner.
	var s scanner.Scanner
	fset := token.NewFileSet()                      // positions are relative to fset
	file := fset.AddFile("", fset.Base(), len(src)) // register input "file"
	s.Init(file, src, nil /* no error handler */, scanner.ScanComments)

	// Repeated calls to Scan yield the token sequence found in the input.
	for {
		pos, tok, lit := s.Scan()
		if tok == token.EOF {
			break
		}
		result += fmt.Sprintf("%s\t%s\t%q\n", fset.Position(pos), tok, lit)
	}

	return result
}
<|endoftext|>


<commit_before>package main_test

import (
	"fmt"
	"os/exec"
	"syscall"
	"testing"
	"time"

	"github.com/cloudfoundry-incubator/garden/backend"
	"github.com/cloudfoundry-incubator/garden/backend/fake_backend"
	"github.com/cloudfoundry/gunk/runner_support"
	"github.com/cloudfoundry/gunk/timeprovider"
	"github.com/cloudfoundry/storeadapter/storerunner/etcdstorerunner"

	GardenServer "github.com/cloudfoundry-incubator/garden/server"
	Bbs "github.com/cloudfoundry-incubator/runtime-schema/bbs"
	"github.com/cloudfoundry-incubator/runtime-schema/models/factories"
	. "github.com/onsi/ginkgo"
	"github.com/onsi/ginkgo/config"
	. "github.com/onsi/gomega"
	"github.com/vito/cmdtest"
	. "github.com/vito/cmdtest/matchers"
)

func TestExecutorMain(t *testing.T) {
	RegisterFailHandler(Fail)
	RunSpecs(t, "Executor Suite")
}

var _ = Describe("Main", func() {
	var (
		etcdRunner      *etcdstorerunner.ETCDClusterRunner
		executorSession *cmdtest.Session
		gardenServer    *GardenServer.WardenServer
		bbs             *Bbs.BBS
		fakeBackend     *fake_backend.FakeBackend
	)

	drainTimeout := 5 * time.Second
	aBit := drainTimeout / 5

	BeforeEach(func() {
		var err error

		etcdPort := 5001 + config.GinkgoConfig.ParallelNode
		wardenPort := 9001 + config.GinkgoConfig.ParallelNode

		etcdRunner = etcdstorerunner.NewETCDClusterRunner(etcdPort, 1)
		etcdRunner.Start()

		bbs = Bbs.New(etcdRunner.Adapter(), timeprovider.NewTimeProvider())

		executorPath, err := cmdtest.Build("github.com/cloudfoundry-incubator/executor", "-race")
		Ω(err).ShouldNot(HaveOccurred())

		wardenAddr := fmt.Sprintf("127.0.0.1:%d", wardenPort)

		fakeBackend = fake_backend.New()
		gardenServer = GardenServer.New("tcp", wardenAddr, 0, fakeBackend)

		err = gardenServer.Start()
		Ω(err).ShouldNot(HaveOccurred())

		executorCmd := exec.Command(
			executorPath,
			"-wardenNetwork", "tcp",
			"-wardenAddr", wardenAddr,
			"-drainTimeout", "5s",
			"-etcdCluster", fmt.Sprintf("http://127.0.0.1:%d", etcdPort),
			"-stack", "the-stack",
			"-memoryMB", "10240",
			"-diskMB", "10240",
			"-heartbeatInterval", "3s",
		)

		executorSession, err = cmdtest.StartWrapped(executorCmd, runner_support.TeeToGinkgoWriter, runner_support.TeeToGinkgoWriter)
		Ω(err).ShouldNot(HaveOccurred())

		Ω(executorSession).Should(SayWithTimeout("executor.started", 5*time.Second))
	})

	AfterEach(func() {
		executorSession.Cmd.Process.Kill()
		etcdRunner.Stop()
		gardenServer.Stop()
	})

	desireRunOnce := func(duration time.Duration) {
		exitStatus := uint32(0)

		fakeContainer := fake_backend.NewFakeContainer(backend.ContainerSpec{})
		fakeContainer.StreamDelay = duration
		fakeContainer.StreamedProcessChunks = []backend.ProcessStream{
			{ExitStatus: &exitStatus},
		}

		fakeBackend.CreateResult = fakeContainer

		err := bbs.DesireRunOnce(factories.BuildRunOnceWithRunAction("the-stack", 1024, 1024, "ls"))
		Ω(err).ShouldNot(HaveOccurred())
	}

	Describe("when the executor fails to maintain its presence", func() {
		It("stops all running tasks", func() {
			desireRunOnce(time.Hour)
			Eventually(fakeBackend.Containers).Should(HaveLen(1))

			// delete the executor's key (and everything else lol)
			etcdRunner.Reset()

			Eventually(fakeBackend.Containers, 7).Should(BeEmpty())
		})
	})

	Describe("when the executor receives the USR1 signal", func() {
		sendDrainSignal := func() {
			executorSession.Cmd.Process.Signal(syscall.SIGUSR1)
			Ω(executorSession).Should(SayWithTimeout("executor.draining", time.Second))
		}

		Context("when there are tasks running", func() {
			It("stops accepting new tasks", func() {
				sendDrainSignal()

				desireRunOnce(time.Hour)
				Consistently(bbs.GetAllPendingRunOnces, 2).Should(HaveLen(1))
			})

			Context("and USR1 is received again", func() {
				It("does not die", func() {
					desireRunOnce(time.Hour)
					Eventually(bbs.GetAllStartingRunOnces).Should(HaveLen(1))

					executorSession.Cmd.Process.Signal(syscall.SIGUSR1)
					Ω(executorSession).Should(SayWithTimeout("executor.draining", time.Second))

					executorSession.Cmd.Process.Signal(syscall.SIGUSR1)
					Ω(executorSession).Should(SayWithTimeout("executor.signal.ignored", 5*time.Second))
				})
			})

			Context("when the tasks complete before the drain timeout", func() {
				It("exits successfully", func() {
					desireRunOnce(1 * time.Second)
					Eventually(bbs.GetAllStartingRunOnces).Should(HaveLen(1))

					sendDrainSignal()

					Ω(executorSession).Should(ExitWithTimeout(0, drainTimeout-aBit))
				})
			})

			Context("when the tasks do not complete before the drain timeout", func() {
				BeforeEach(func() {
					desireRunOnce(time.Hour)
					Eventually(bbs.GetAllStartingRunOnces).Should(HaveLen(1))
				})

				It("cancels all running tasks", func() {
					Ω(fakeBackend.Containers()).Should(HaveLen(1))
					sendDrainSignal()
					Eventually(fakeBackend.Containers, (drainTimeout + aBit).Seconds()).Should(BeEmpty())
				})

				It("exits successfully", func() {
					sendDrainSignal()
					Ω(executorSession).Should(ExitWithTimeout(0, 2*drainTimeout))
				})
			})
		})

		Context("when there are no tasks running", func() {
			It("exits successfully", func() {
				sendDrainSignal()
				Ω(executorSession).Should(ExitWithTimeout(0, 1*time.Second))
			})
		})
	})
})
<commit_msg>backfill SIGTERM/SIGINT integration tests<commit_after>package main_test

import (
	"fmt"
	"os/exec"
	"syscall"
	"testing"
	"time"

	"github.com/cloudfoundry-incubator/garden/backend"
	"github.com/cloudfoundry-incubator/garden/backend/fake_backend"
	"github.com/cloudfoundry/gunk/runner_support"
	"github.com/cloudfoundry/gunk/timeprovider"
	"github.com/cloudfoundry/storeadapter/storerunner/etcdstorerunner"

	GardenServer "github.com/cloudfoundry-incubator/garden/server"
	Bbs "github.com/cloudfoundry-incubator/runtime-schema/bbs"
	"github.com/cloudfoundry-incubator/runtime-schema/models/factories"
	. "github.com/onsi/ginkgo"
	"github.com/onsi/ginkgo/config"
	. "github.com/onsi/gomega"
	"github.com/vito/cmdtest"
	. "github.com/vito/cmdtest/matchers"
)

func TestExecutorMain(t *testing.T) {
	RegisterFailHandler(Fail)
	RunSpecs(t, "Executor Suite")
}

var _ = Describe("Main", func() {
	var (
		etcdRunner      *etcdstorerunner.ETCDClusterRunner
		executorSession *cmdtest.Session
		gardenServer    *GardenServer.WardenServer
		bbs             *Bbs.BBS
		fakeBackend     *fake_backend.FakeBackend
	)

	drainTimeout := 5 * time.Second
	aBit := drainTimeout / 5

	BeforeEach(func() {
		var err error

		etcdPort := 5001 + config.GinkgoConfig.ParallelNode
		wardenPort := 9001 + config.GinkgoConfig.ParallelNode

		etcdRunner = etcdstorerunner.NewETCDClusterRunner(etcdPort, 1)
		etcdRunner.Start()

		bbs = Bbs.New(etcdRunner.Adapter(), timeprovider.NewTimeProvider())

		executorPath, err := cmdtest.Build("github.com/cloudfoundry-incubator/executor", "-race")
		Ω(err).ShouldNot(HaveOccurred())

		wardenAddr := fmt.Sprintf("127.0.0.1:%d", wardenPort)

		fakeBackend = fake_backend.New()
		gardenServer = GardenServer.New("tcp", wardenAddr, 0, fakeBackend)

		err = gardenServer.Start()
		Ω(err).ShouldNot(HaveOccurred())

		executorCmd := exec.Command(
			executorPath,
			"-wardenNetwork", "tcp",
			"-wardenAddr", wardenAddr,
			"-drainTimeout", "5s",
			"-etcdCluster", fmt.Sprintf("http://127.0.0.1:%d", etcdPort),
			"-stack", "the-stack",
			"-memoryMB", "10240",
			"-diskMB", "10240",
			"-heartbeatInterval", "3s",
		)

		executorSession, err = cmdtest.StartWrapped(executorCmd, runner_support.TeeToGinkgoWriter, runner_support.TeeToGinkgoWriter)
		Ω(err).ShouldNot(HaveOccurred())

		Ω(executorSession).Should(SayWithTimeout("executor.started", 5*time.Second))
	})

	AfterEach(func() {
		executorSession.Cmd.Process.Kill()
		etcdRunner.Stop()
		gardenServer.Stop()
	})

	desireRunOnce := func(duration time.Duration) {
		exitStatus := uint32(0)

		fakeContainer := fake_backend.NewFakeContainer(backend.ContainerSpec{})
		fakeContainer.StreamDelay = duration
		fakeContainer.StreamedProcessChunks = []backend.ProcessStream{
			{ExitStatus: &exitStatus},
		}

		fakeBackend.CreateResult = fakeContainer

		err := bbs.DesireRunOnce(factories.BuildRunOnceWithRunAction("the-stack", 1024, 1024, "ls"))
		Ω(err).ShouldNot(HaveOccurred())
	}

	Describe("when the executor fails to maintain its presence", func() {
		It("stops all running tasks", func() {
			desireRunOnce(time.Hour)
			Eventually(fakeBackend.Containers).Should(HaveLen(1))

			// delete the executor's key (and everything else lol)
			etcdRunner.Reset()

			Eventually(fakeBackend.Containers, 7).Should(BeEmpty())
		})
	})

	Describe("when the executor receives the TERM signal", func() {
		It("stops all running tasks", func() {
			desireRunOnce(time.Hour)
			Eventually(fakeBackend.Containers).Should(HaveLen(1))

			executorSession.Cmd.Process.Signal(syscall.SIGTERM)

			Eventually(fakeBackend.Containers, 7).Should(BeEmpty())
		})

		It("exits successfully", func() {
			executorSession.Cmd.Process.Signal(syscall.SIGTERM)
			Ω(executorSession).Should(ExitWithTimeout(0, 2*drainTimeout))
		})
	})

	Describe("when the executor receives the INT signal", func() {
		It("stops all running tasks", func() {
			desireRunOnce(time.Hour)
			Eventually(fakeBackend.Containers).Should(HaveLen(1))

			executorSession.Cmd.Process.Signal(syscall.SIGINT)

			Eventually(fakeBackend.Containers, 7).Should(BeEmpty())
		})

		It("exits successfully", func() {
			executorSession.Cmd.Process.Signal(syscall.SIGINT)
			Ω(executorSession).Should(ExitWithTimeout(0, 2*drainTimeout))
		})
	})

	Describe("when the executor receives the USR1 signal", func() {
		sendDrainSignal := func() {
			executorSession.Cmd.Process.Signal(syscall.SIGUSR1)
			Ω(executorSession).Should(SayWithTimeout("executor.draining", time.Second))
		}

		Context("when there are tasks running", func() {
			It("stops accepting new tasks", func() {
				sendDrainSignal()

				desireRunOnce(time.Hour)
				Consistently(bbs.GetAllPendingRunOnces, 2).Should(HaveLen(1))
			})

			Context("and USR1 is received again", func() {
				It("does not die", func() {
					desireRunOnce(time.Hour)
					Eventually(bbs.GetAllStartingRunOnces).Should(HaveLen(1))

					executorSession.Cmd.Process.Signal(syscall.SIGUSR1)
					Ω(executorSession).Should(SayWithTimeout("executor.draining", time.Second))

					executorSession.Cmd.Process.Signal(syscall.SIGUSR1)
					Ω(executorSession).Should(SayWithTimeout("executor.signal.ignored", 5*time.Second))
				})
			})

			Context("when the tasks complete before the drain timeout", func() {
				It("exits successfully", func() {
					desireRunOnce(1 * time.Second)
					Eventually(bbs.GetAllStartingRunOnces).Should(HaveLen(1))

					sendDrainSignal()

					Ω(executorSession).Should(ExitWithTimeout(0, drainTimeout-aBit))
				})
			})

			Context("when the tasks do not complete before the drain timeout", func() {
				BeforeEach(func() {
					desireRunOnce(time.Hour)
					Eventually(bbs.GetAllStartingRunOnces).Should(HaveLen(1))
				})

				It("cancels all running tasks", func() {
					Ω(fakeBackend.Containers()).Should(HaveLen(1))
					sendDrainSignal()
					Eventually(fakeBackend.Containers, (drainTimeout + aBit).Seconds()).Should(BeEmpty())
				})

				It("exits successfully", func() {
					sendDrainSignal()
					Ω(executorSession).Should(ExitWithTimeout(0, 2*drainTimeout))
				})
			})
		})

		Context("when there are no tasks running", func() {
			It("exits successfully", func() {
				sendDrainSignal()
				Ω(executorSession).Should(ExitWithTimeout(0, 1*time.Second))
			})
		})
	})
})
<|endoftext|>


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

package main

import (
	"encoding/json"
	"errors"
	"net"
	"os"
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
)

func TestIntegrationMain(t *testing.T) {
	// This function serves as an entry point for running integration tests.
	// We're wrapping it in a test case so that we can record the test coverage.
	isIntegration := os.Getenv("GHOSTUNNEL_INTEGRATION_TEST")

	// Catch panics to make sure test exits normally and writes coverage
	// even if we got a crash (we might want to test error cases)
	defer func() {
		if err := recover(); err != nil {
			t.Error(err)
		}
	}()

	if isIntegration != "true" {
		return
	}

	finished := make(chan bool, 1)
	once := &sync.Once{}

	// override exit function for test, to make sure calls to exitFunc() don't
	// actually terminate the process and kill the test w/o capturing results.
	exitFunc = func(exit int) {
		once.Do(func() {
			if exit != 0 {
				t.Errorf("exit code from ghostunnel: %d", exit)
			}
		})
		finished <- true
		select {} // block
	}

	var wrappedArgs []string
	err := json.Unmarshal([]byte(os.Getenv("GHOSTUNNEL_INTEGRATION_ARGS")), &wrappedArgs)
	panicOnError(err)

	go func() {
		err := run(wrappedArgs)
		if err != nil {
			t.Errorf("got error from run: %s", err)
		}
		finished <- true
	}()

	select {
	case <-finished:
		return
	case <-time.Tick(10 * time.Minute):
		panic("timed out")
	}
}

func TestInitLoggerSyslog(t *testing.T) {
	*useSyslog = true
	initLogger()
	assert.NotNil(t, logger, "logger should never be nil after init")
}

func TestPanicOnError(t *testing.T) {
	defer func() {
		if err := recover(); err == nil {
			t.Error("panicOnError should panic, but did not")
		}
	}()

	panicOnError(errors.New("error"))
}

func TestFlagValidation(t *testing.T) {
	*enableProf = true
	*statusAddress = ""
	err := validateFlags(nil)
	assert.NotNil(t, err, "--enable-pprof implies --status")

	*enableProf = false
	*metricsURL = "127.0.0.1"
	err = validateFlags(nil)
	assert.NotNil(t, err, "invalid --metrics-url should be rejected")
	*metricsURL = ""
}

func TestServerFlagValidation(t *testing.T) {
	*serverAllowAll = false
	*serverAllowedCNs = nil
	*serverAllowedOUs = nil
	*serverAllowedDNSs = nil
	*serverAllowedIPs = nil
	err := serverValidateFlags()
	assert.NotNil(t, err, "invalid access control flags accepted")

	*serverAllowAll = true
	*serverAllowedCNs = []string{"test"}
	err = serverValidateFlags()
	assert.NotNil(t, err, "--allow-all and --allow-cn are mutually exclusive")

	*serverAllowedCNs = nil
	*serverAllowedOUs = []string{"test"}
	err = serverValidateFlags()
	assert.NotNil(t, err, "--allow-all and --allow-ou are mutually exclusive")

	*serverAllowedOUs = nil
	*serverAllowedDNSs = []string{"test"}
	err = serverValidateFlags()
	assert.NotNil(t, err, "--allow-all and --allow-dns-san are mutually exclusive")

	*serverAllowedDNSs = nil
	*serverAllowedIPs = []net.IP{net.IPv4(0, 0, 0, 0)}
	err = serverValidateFlags()
	assert.NotNil(t, err, "--allow-all and --allow-ip-san are mutually exclusive")

	*serverAllowAll = false
	*serverUnsafeTarget = false
	*serverForwardAddress = "foo.com"
	err = serverValidateFlags()
	assert.NotNil(t, err, "unsafe target should be rejected")
}

func TestClientFlagValidation(t *testing.T) {
	*clientUnsafeListen = false
	*clientListenAddress = "0.0.0.0:8080"
	err := clientValidateFlags()
	assert.NotNil(t, err, "unsafe listen should be rejected")
}

func TestAllowsLocalhost(t *testing.T) {
	*serverUnsafeTarget = false
	assert.True(t, validateUnixOrLocalhost("localhost:1234"), "localhost should be allowed")
	assert.True(t, validateUnixOrLocalhost("127.0.0.1:1234"), "127.0.0.1 should be allowed")
	assert.True(t, validateUnixOrLocalhost("[::1]:1234"), "[::1] should be allowed")
	assert.True(t, validateUnixOrLocalhost("unix:/tmp/foo"), "unix:/tmp/foo should be allowed")
}

func TestDisallowsFooDotCom(t *testing.T) {
	*serverUnsafeTarget = false
	assert.False(t, validateUnixOrLocalhost("foo.com:1234"), "foo.com should be disallowed")
	assert.False(t, validateUnixOrLocalhost("alocalhost.com:1234"), "alocalhost.com should be disallowed")
	assert.False(t, validateUnixOrLocalhost("localhost.com.foo.com:1234"), "localhost.com.foo.com should be disallowed")
	assert.False(t, validateUnixOrLocalhost("74.122.190.83:1234"), "random ip address should be disallowed")
}

func TestServerBackendDialerError(t *testing.T) {
	*serverForwardAddress = "invalid"
	_, err := serverBackendDialer()
	assert.NotNil(t, err, "invalid forward address should not have dialer")
}

func TestInvalidCABundle(t *testing.T) {
	err := run([]string{
		"server",
		"--cacert", "/dev/null",
		"--target", "localhost:8080",
		"--keystore", "keystore.p12",
		"--listen", "localhost:8080",
	})
	assert.NotNil(t, err, "invalid CA bundle should exit with error")
}
<commit_msg>Expand server flag validation test<commit_after>/*-
 * Copyright 2015 Square Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package main

import (
	"encoding/json"
	"errors"
	"net"
	"os"
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
)

func TestIntegrationMain(t *testing.T) {
	// This function serves as an entry point for running integration tests.
	// We're wrapping it in a test case so that we can record the test coverage.
	isIntegration := os.Getenv("GHOSTUNNEL_INTEGRATION_TEST")

	// Catch panics to make sure test exits normally and writes coverage
	// even if we got a crash (we might want to test error cases)
	defer func() {
		if err := recover(); err != nil {
			t.Error(err)
		}
	}()

	if isIntegration != "true" {
		return
	}

	finished := make(chan bool, 1)
	once := &sync.Once{}

	// override exit function for test, to make sure calls to exitFunc() don't
	// actually terminate the process and kill the test w/o capturing results.
	exitFunc = func(exit int) {
		once.Do(func() {
			if exit != 0 {
				t.Errorf("exit code from ghostunnel: %d", exit)
			}
		})
		finished <- true
		select {} // block
	}

	var wrappedArgs []string
	err := json.Unmarshal([]byte(os.Getenv("GHOSTUNNEL_INTEGRATION_ARGS")), &wrappedArgs)
	panicOnError(err)

	go func() {
		err := run(wrappedArgs)
		if err != nil {
			t.Errorf("got error from run: %s", err)
		}
		finished <- true
	}()

	select {
	case <-finished:
		return
	case <-time.Tick(10 * time.Minute):
		panic("timed out")
	}
}

func TestInitLoggerSyslog(t *testing.T) {
	*useSyslog = true
	initLogger()
	assert.NotNil(t, logger, "logger should never be nil after init")
}

func TestPanicOnError(t *testing.T) {
	defer func() {
		if err := recover(); err == nil {
			t.Error("panicOnError should panic, but did not")
		}
	}()

	panicOnError(errors.New("error"))
}

func TestFlagValidation(t *testing.T) {
	*enableProf = true
	*statusAddress = ""
	err := validateFlags(nil)
	assert.NotNil(t, err, "--enable-pprof implies --status")

	*enableProf = false
	*metricsURL = "127.0.0.1"
	err = validateFlags(nil)
	assert.NotNil(t, err, "invalid --metrics-url should be rejected")
	*metricsURL = ""
}

func TestServerFlagValidation(t *testing.T) {
	*serverAllowAll = false
	*serverAllowedCNs = nil
	*serverAllowedOUs = nil
	*serverAllowedDNSs = nil
	*serverAllowedIPs = nil
	err := serverValidateFlags()
	assert.NotNil(t, err, "invalid access control flags accepted")

	*serverAllowAll = true
	*serverAllowedCNs = []string{"test"}
	err = serverValidateFlags()
	assert.NotNil(t, err, "--allow-all and --allow-cn are mutually exclusive")

	*serverAllowedCNs = nil
	*serverAllowedOUs = []string{"test"}
	err = serverValidateFlags()
	assert.NotNil(t, err, "--allow-all and --allow-ou are mutually exclusive")

	*serverAllowedOUs = nil
	*serverAllowedDNSs = []string{"test"}
	err = serverValidateFlags()
	assert.NotNil(t, err, "--allow-all and --allow-dns-san are mutually exclusive")

	*serverAllowedDNSs = nil
	*serverAllowedIPs = []net.IP{net.IPv4(0, 0, 0, 0)}
	err = serverValidateFlags()
	assert.NotNil(t, err, "--allow-all and --allow-ip-san are mutually exclusive")

	*serverAllowAll = false
	*serverUnsafeTarget = false
	*serverForwardAddress = "foo.com"
	err = serverValidateFlags()
	assert.NotNil(t, err, "unsafe target should be rejected")

	*enabledCipherSuites = "ABC"
	*serverForwardAddress = "127.0.0.1:8080"
	err = serverValidateFlags()
	assert.NotNil(t, err, "invalid cipher suite option should be rejected")
}

func TestClientFlagValidation(t *testing.T) {
	*clientUnsafeListen = false
	*clientListenAddress = "0.0.0.0:8080"
	err := clientValidateFlags()
	assert.NotNil(t, err, "unsafe listen should be rejected")
}

func TestAllowsLocalhost(t *testing.T) {
	*serverUnsafeTarget = false
	assert.True(t, validateUnixOrLocalhost("localhost:1234"), "localhost should be allowed")
	assert.True(t, validateUnixOrLocalhost("127.0.0.1:1234"), "127.0.0.1 should be allowed")
	assert.True(t, validateUnixOrLocalhost("[::1]:1234"), "[::1] should be allowed")
	assert.True(t, validateUnixOrLocalhost("unix:/tmp/foo"), "unix:/tmp/foo should be allowed")
}

func TestDisallowsFooDotCom(t *testing.T) {
	*serverUnsafeTarget = false
	assert.False(t, validateUnixOrLocalhost("foo.com:1234"), "foo.com should be disallowed")
	assert.False(t, validateUnixOrLocalhost("alocalhost.com:1234"), "alocalhost.com should be disallowed")
	assert.False(t, validateUnixOrLocalhost("localhost.com.foo.com:1234"), "localhost.com.foo.com should be disallowed")
	assert.False(t, validateUnixOrLocalhost("74.122.190.83:1234"), "random ip address should be disallowed")
}

func TestServerBackendDialerError(t *testing.T) {
	*serverForwardAddress = "invalid"
	_, err := serverBackendDialer()
	assert.NotNil(t, err, "invalid forward address should not have dialer")
}

func TestInvalidCABundle(t *testing.T) {
	err := run([]string{
		"server",
		"--cacert", "/dev/null",
		"--target", "localhost:8080",
		"--keystore", "keystore.p12",
		"--listen", "localhost:8080",
	})
	assert.NotNil(t, err, "invalid CA bundle should exit with error")
}
<|endoftext|>


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

package main

import (
	"bytes"
	"encoding/json"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"testing"
	"time"

	"github.com/realglobe-Inc/edo-auth/database/token"
	idpdb "github.com/realglobe-Inc/edo-idp-selector/database/idp"
	tadb "github.com/realglobe-Inc/edo-idp-selector/database/ta"
	"github.com/realglobe-Inc/edo-lib/jwk"
	logutil "github.com/realglobe-Inc/edo-lib/log"
	"github.com/realglobe-Inc/edo-lib/strset/strsetutil"
	"github.com/realglobe-Inc/edo-lib/test"
	"github.com/realglobe-Inc/go-lib/erro"
	"github.com/realglobe-Inc/go-lib/rglog/level"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)

func init() {
	logutil.SetupConsole(logRoot, level.OFF)
}

var monPool, _ = mgo.DialWithTimeout(test_monAddr, time.Minute)

type testIdProvider struct {
	base *test.HttpServer
	keys []jwk.Key
}

func newTestIdProvider(keys []jwk.Key) (*testIdProvider, error) {
	base, err := test.NewHttpServer(time.Minute)
	if err != nil {
		return nil, erro.Wrap(err)
	}
	return &testIdProvider{base, keys}, nil
}

func (this *testIdProvider) close() {
	this.base.Close()
}

func (this *testIdProvider) info() idpdb.Element {
	return idpdb.New(
		this.base.URL,
		nil,
		this.base.URL+"/auth",
		this.base.URL+"/token",
		this.base.URL+"/userinfo",
		this.base.URL+"/coop/from",
		this.base.URL+"/coop/to",
		this.keys,
	)
}

func (this *testIdProvider) selfId() string {
	return this.base.URL
}

func (this *testIdProvider) addResponse(status int, header http.Header, body []byte) <-chan *http.Request {
	return this.base.AddResponse(status, header, body)
}

type testTa struct {
	base *test.HttpServer
}

func newTestTa() (*testTa, error) {
	base, err := test.NewHttpServer(time.Minute)
	if err != nil {
		return nil, erro.Wrap(err)
	}
	return &testTa{base}, nil
}

func (this *testTa) close() {
	this.base.Close()
}

func (this *testTa) info() tadb.Element {
	return tadb.New(
		this.base.URL,
		nil,
		strsetutil.New(this.base.URL+"/callback"),
		nil,
		false,
		"",
	)
}

func (this *testTa) addResponse(status int, header http.Header, body []byte) <-chan *http.Request {
	return this.base.AddResponse(status, header, body)
}

// 正常系。
func TestServer(t *testing.T) {
	// ////////////////////////////////
	// logutil.SetupConsole(logRoot, level.ALL)
	// defer logutil.SetupConsole(logRoot, level.OFF)
	// ////////////////////////////////

	if monPool == nil {
		t.SkipNow()
	}
	red, err := test.NewRedisServer()
	if err != nil {
		t.Fatal(err)
	} else if red == nil {
		t.SkipNow()
	}

	// ID プロバイダの準備。
	idpServ, err := newTestIdProvider([]jwk.Key{test_idpKey})
	if err != nil {
		t.Fatal(err)
	}
	defer idpServ.close()
	idp := idpServ.info()

	// 連携先 TA の準備。
	toTaServ, err := newTestTa()
	if err != nil {
		t.Fatal(err)
	}
	defer toTaServ.close()

	// 環境設定。
	param, err := newTestParameter()
	if err != nil {
		t.Fatal(err)
	}
	param.socPort, err = test.FreePort()
	if err != nil {
		t.Fatal(err)
	}
	// 鍵の準備。
	keyDbPath, err := setupKeyDb(param.keyDbPath, []jwk.Key{test_key})
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(keyDbPath)
	param.keyDbPath = keyDbPath
	// ID プロバイダ DB の準備。
	param.idpDbTag = param.idpDbTag + strconv.FormatInt(time.Now().UnixNano(), 16)
	if err := setupIdpDb(param.idpDbTag, param.idpDbTag2, []idpdb.Element{idp}); err != nil {
		t.Fatal(err)
	}
	defer clearIdpDb(param.idpDbTag, param.idpDbTag2)
	// アクセストークン DB の準備。
	if err := token.NewRedisDb(red.Pool(), param.tokDbTag).Save(token.New(test_tokId, test_tokTag, time.Now().Add(time.Minute), idp.Id(), strsetutil.New("openid")), time.Now().Add(time.Minute)); err != nil {
		t.Fatal(err)
	}
	defer red.Close()
	param.tokDbAddr = red.Address()

	testServer(t, param, idpServ, toTaServ)
}

// データベースをちゃんと使った正常系。
func TestServerWithDb(t *testing.T) {
	// ////////////////////////////////
	// logutil.SetupConsole(logRoot, level.ALL)
	// defer logutil.SetupConsole(logRoot, level.OFF)
	// ////////////////////////////////

	if monPool == nil {
		t.SkipNow()
	}
	red, err := test.NewRedisServer()
	if err != nil {
		t.Fatal(err)
	} else if red == nil {
		t.SkipNow()
	}

	// ID プロバイダの準備。
	idpServ, err := newTestIdProvider([]jwk.Key{test_idpKey})
	if err != nil {
		t.Fatal(err)
	}
	defer idpServ.close()
	idp := idpServ.info()

	// 連携先 TA の準備。
	toTaServ, err := newTestTa()
	if err != nil {
		t.Fatal(err)
	}
	defer toTaServ.close()

	// 環境設定。
	param, err := newTestParameter()
	if err != nil {
		t.Fatal(err)
	}
	// 鍵の準備。
	param.keyDbType = "redis"
	param.keyDbAddr = red.Address()
	keyDbPath, err := setupKeyDb(param.keyDbPath, []jwk.Key{test_key})
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(keyDbPath)
	param.keyDbPath = keyDbPath
	// web DB の準備。
	param.webDbType = "redis"
	param.webDbAddr = red.Address()
	// ID プロバイダ DB の準備。
	param.idpDbTag = param.idpDbTag + strconv.FormatInt(time.Now().UnixNano(), 16)
	if err := setupIdpDb(param.idpDbTag, param.idpDbTag2, []idpdb.Element{idp}); err != nil {
		t.Fatal(err)
	}
	defer clearIdpDb(param.idpDbTag, param.idpDbTag2)
	// アクセストークン DB の準備。
	if err := token.NewRedisDb(red.Pool(), param.tokDbTag).Save(token.New(test_tokId, test_tokTag, time.Now().Add(time.Minute), idp.Id(), strsetutil.New("openid")), time.Now().Add(time.Minute)); err != nil {
		t.Fatal(err)
	}
	defer red.Close()
	param.tokDbAddr = red.Address()
	// セッション DB の準備。
	param.sessDbType = "redis"
	param.sessDbAddr = red.Address()

	testServer(t, param, idpServ, toTaServ)
}

func testServer(t *testing.T, param *parameters, idpServ *testIdProvider, toTaServ *testTa) {
	idp := idpServ.info()
	toTa := toTaServ.info()

	errCh := make(chan error, 1)
	go func() {
		errCh <- serve(param)
	}()
	defer func() { param.shutCh <- struct{}{} }()

	selfUri := "http://localhost:" + strconv.Itoa(param.socPort)
	exp := time.Now().Add(time.Minute)
	for {
		if time.Now().After(exp) {
			t.Fatal("timeout")
		}
		r, err := http.NewRequest("GET", selfUri+param.pathOk, nil)
		if err != nil {
			t.Fatal(err)
		}
		r.Header.Set("Connection", "close")
		if _, err := (&http.Client{}).Do(r); err == nil {
			break
		}

		select {
		case err := <-errCh:
			t.Fatal(err)
		default:
		}
		time.Sleep(time.Millisecond)
	}

	{
		s, h, b, err := newTestIdpResponse(param, idp, toTa.Id())
		if err != nil {
			t.Fatal(err)
		}
		idpServ.addResponse(s, h, b)
	}

	stat, head, body, err := newTestToTaResponse()
	if err != nil {
		t.Fatal(err)
	}
	toTaServ.addResponse(stat, head, body)

	reqBody := []byte("request,requester,requestest")
	r, err := newTestRequest(selfUri+param.pathProx, idp.Id(), toTa.Id(), reqBody)
	if err != nil {
		t.Fatal(err)
	}
	r.Header.Set("Connection", "close")
	resp, err := (&http.Client{}).Do(r)
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != stat {
		t.Error(resp.StatusCode)
		t.Fatal(stat)
	}
	for k, vs := range head {
		if h := resp.Header.Get(k); h != vs[0] {
			t.Error(h)
			t.Fatal(vs[0])
		}
	}
	if buff, err := ioutil.ReadAll(resp.Body); err != nil {
		t.Fatal(err)
	} else if !bytes.Equal(buff, body) {
		t.Error(string(buff))
		t.Fatal(string(body))
	}

}

func setupKeyDb(path string, keys []jwk.Key) (dir string, err error) {
	dir, err = ioutil.TempDir(filepath.Dir(path), filepath.Base(path))
	for i, key := range keys {
		if data, err := json.Marshal(key.ToMap()); err != nil {
			return "", erro.Wrap(err)
		} else if ioutil.WriteFile(filepath.Join(dir, strconv.Itoa(i)+".json"), data, 0644); err != nil {
			return "", erro.Wrap(err)
		}
	}
	return dir, nil
}

func setupIdpDb(db, coll string, idps []idpdb.Element) error {
	conn := monPool.New()
	defer conn.Close()

	for _, idp := range idps {
		keys := []map[string]interface{}{}
		for _, key := range idp.Keys() {
			keys = append(keys, key.ToMap())
		}
		m := bson.M{
			"issuer":                    idp.Id(),
			"authorization_endpoint":    idp.AuthUri(),
			"token_endpoint":            idp.TokenUri(),
			"userinfo_endpoint":         idp.AccountUri(),
			"cooperation_from_endpoint": idp.CoopFromUri(),
			"cooperation_to_endpoint":   idp.CoopToUri(),
			"jwks": keys,
		}
		for k, v := range idp.Names() {
			if k == "" {
				m["issuer_name"] = v
			} else {
				m["issuer_name#"+k] = v
			}
		}
		if err := conn.DB(db).C(coll).Insert(m); err != nil {
			return erro.Wrap(err)
		}
	}
	return nil
}

func clearIdpDb(db, coll string) error {
	conn := monPool.New()
	defer conn.Close()

	return conn.DB(db).DropDatabase()
}
<commit_msg>http.DefaultClient を使うようにした<commit_after>// Copyright 2015 realglobe, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"bytes"
	"encoding/json"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"testing"
	"time"

	"github.com/realglobe-Inc/edo-auth/database/token"
	idpdb "github.com/realglobe-Inc/edo-idp-selector/database/idp"
	tadb "github.com/realglobe-Inc/edo-idp-selector/database/ta"
	"github.com/realglobe-Inc/edo-lib/jwk"
	logutil "github.com/realglobe-Inc/edo-lib/log"
	"github.com/realglobe-Inc/edo-lib/strset/strsetutil"
	"github.com/realglobe-Inc/edo-lib/test"
	"github.com/realglobe-Inc/go-lib/erro"
	"github.com/realglobe-Inc/go-lib/rglog/level"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)

func init() {
	logutil.SetupConsole(logRoot, level.OFF)
}

var monPool, _ = mgo.DialWithTimeout(test_monAddr, time.Minute)

type testIdProvider struct {
	base *test.HttpServer
	keys []jwk.Key
}

func newTestIdProvider(keys []jwk.Key) (*testIdProvider, error) {
	base, err := test.NewHttpServer(time.Minute)
	if err != nil {
		return nil, erro.Wrap(err)
	}
	return &testIdProvider{base, keys}, nil
}

func (this *testIdProvider) close() {
	this.base.Close()
}

func (this *testIdProvider) info() idpdb.Element {
	return idpdb.New(
		this.base.URL,
		nil,
		this.base.URL+"/auth",
		this.base.URL+"/token",
		this.base.URL+"/userinfo",
		this.base.URL+"/coop/from",
		this.base.URL+"/coop/to",
		this.keys,
	)
}

func (this *testIdProvider) selfId() string {
	return this.base.URL
}

func (this *testIdProvider) addResponse(status int, header http.Header, body []byte) <-chan *http.Request {
	return this.base.AddResponse(status, header, body)
}

type testTa struct {
	base *test.HttpServer
}

func newTestTa() (*testTa, error) {
	base, err := test.NewHttpServer(time.Minute)
	if err != nil {
		return nil, erro.Wrap(err)
	}
	return &testTa{base}, nil
}

func (this *testTa) close() {
	this.base.Close()
}

func (this *testTa) info() tadb.Element {
	return tadb.New(
		this.base.URL,
		nil,
		strsetutil.New(this.base.URL+"/callback"),
		nil,
		false,
		"",
	)
}

func (this *testTa) addResponse(status int, header http.Header, body []byte) <-chan *http.Request {
	return this.base.AddResponse(status, header, body)
}

// 正常系。
func TestServer(t *testing.T) {
	// ////////////////////////////////
	// logutil.SetupConsole(logRoot, level.ALL)
	// defer logutil.SetupConsole(logRoot, level.OFF)
	// ////////////////////////////////

	if monPool == nil {
		t.SkipNow()
	}
	red, err := test.NewRedisServer()
	if err != nil {
		t.Fatal(err)
	} else if red == nil {
		t.SkipNow()
	}

	// ID プロバイダの準備。
	idpServ, err := newTestIdProvider([]jwk.Key{test_idpKey})
	if err != nil {
		t.Fatal(err)
	}
	defer idpServ.close()
	idp := idpServ.info()

	// 連携先 TA の準備。
	toTaServ, err := newTestTa()
	if err != nil {
		t.Fatal(err)
	}
	defer toTaServ.close()

	// 環境設定。
	param, err := newTestParameter()
	if err != nil {
		t.Fatal(err)
	}
	param.socPort, err = test.FreePort()
	if err != nil {
		t.Fatal(err)
	}
	// 鍵の準備。
	keyDbPath, err := setupKeyDb(param.keyDbPath, []jwk.Key{test_key})
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(keyDbPath)
	param.keyDbPath = keyDbPath
	// ID プロバイダ DB の準備。
	param.idpDbTag = param.idpDbTag + strconv.FormatInt(time.Now().UnixNano(), 16)
	if err := setupIdpDb(param.idpDbTag, param.idpDbTag2, []idpdb.Element{idp}); err != nil {
		t.Fatal(err)
	}
	defer clearIdpDb(param.idpDbTag, param.idpDbTag2)
	// アクセストークン DB の準備。
	if err := token.NewRedisDb(red.Pool(), param.tokDbTag).Save(token.New(test_tokId, test_tokTag, time.Now().Add(time.Minute), idp.Id(), strsetutil.New("openid")), time.Now().Add(time.Minute)); err != nil {
		t.Fatal(err)
	}
	defer red.Close()
	param.tokDbAddr = red.Address()

	testServer(t, param, idpServ, toTaServ)
}

// データベースをちゃんと使った正常系。
func TestServerWithDb(t *testing.T) {
	// ////////////////////////////////
	// logutil.SetupConsole(logRoot, level.ALL)
	// defer logutil.SetupConsole(logRoot, level.OFF)
	// ////////////////////////////////

	if monPool == nil {
		t.SkipNow()
	}
	red, err := test.NewRedisServer()
	if err != nil {
		t.Fatal(err)
	} else if red == nil {
		t.SkipNow()
	}

	// ID プロバイダの準備。
	idpServ, err := newTestIdProvider([]jwk.Key{test_idpKey})
	if err != nil {
		t.Fatal(err)
	}
	defer idpServ.close()
	idp := idpServ.info()

	// 連携先 TA の準備。
	toTaServ, err := newTestTa()
	if err != nil {
		t.Fatal(err)
	}
	defer toTaServ.close()

	// 環境設定。
	param, err := newTestParameter()
	if err != nil {
		t.Fatal(err)
	}
	// 鍵の準備。
	param.keyDbType = "redis"
	param.keyDbAddr = red.Address()
	keyDbPath, err := setupKeyDb(param.keyDbPath, []jwk.Key{test_key})
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(keyDbPath)
	param.keyDbPath = keyDbPath
	// web DB の準備。
	param.webDbType = "redis"
	param.webDbAddr = red.Address()
	// ID プロバイダ DB の準備。
	param.idpDbTag = param.idpDbTag + strconv.FormatInt(time.Now().UnixNano(), 16)
	if err := setupIdpDb(param.idpDbTag, param.idpDbTag2, []idpdb.Element{idp}); err != nil {
		t.Fatal(err)
	}
	defer clearIdpDb(param.idpDbTag, param.idpDbTag2)
	// アクセストークン DB の準備。
	if err := token.NewRedisDb(red.Pool(), param.tokDbTag).Save(token.New(test_tokId, test_tokTag, time.Now().Add(time.Minute), idp.Id(), strsetutil.New("openid")), time.Now().Add(time.Minute)); err != nil {
		t.Fatal(err)
	}
	defer red.Close()
	param.tokDbAddr = red.Address()
	// セッション DB の準備。
	param.sessDbType = "redis"
	param.sessDbAddr = red.Address()

	testServer(t, param, idpServ, toTaServ)
}

func testServer(t *testing.T, param *parameters, idpServ *testIdProvider, toTaServ *testTa) {
	idp := idpServ.info()
	toTa := toTaServ.info()

	errCh := make(chan error, 1)
	go func() {
		errCh <- serve(param)
	}()
	defer func() { param.shutCh <- struct{}{} }()

	selfUri := "http://localhost:" + strconv.Itoa(param.socPort)
	exp := time.Now().Add(time.Minute)
	for {
		if time.Now().After(exp) {
			t.Fatal("timeout")
		}
		r, err := http.NewRequest("GET", selfUri+param.pathOk, nil)
		if err != nil {
			t.Fatal(err)
		}
		r.Header.Set("Connection", "close")
		if _, err := http.DefaultClient.Do(r); err == nil {
			break
		}

		select {
		case err := <-errCh:
			t.Fatal(err)
		default:
		}
		time.Sleep(time.Millisecond)
	}

	{
		s, h, b, err := newTestIdpResponse(param, idp, toTa.Id())
		if err != nil {
			t.Fatal(err)
		}
		idpServ.addResponse(s, h, b)
	}

	stat, head, body, err := newTestToTaResponse()
	if err != nil {
		t.Fatal(err)
	}
	toTaServ.addResponse(stat, head, body)

	reqBody := []byte("request,requester,requestest")
	r, err := newTestRequest(selfUri+param.pathProx, idp.Id(), toTa.Id(), reqBody)
	if err != nil {
		t.Fatal(err)
	}
	r.Header.Set("Connection", "close")
	resp, err := http.DefaultClient.Do(r)
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != stat {
		t.Error(resp.StatusCode)
		t.Fatal(stat)
	}
	for k, vs := range head {
		if h := resp.Header.Get(k); h != vs[0] {
			t.Error(h)
			t.Fatal(vs[0])
		}
	}
	if buff, err := ioutil.ReadAll(resp.Body); err != nil {
		t.Fatal(err)
	} else if !bytes.Equal(buff, body) {
		t.Error(string(buff))
		t.Fatal(string(body))
	}

}

func setupKeyDb(path string, keys []jwk.Key) (dir string, err error) {
	dir, err = ioutil.TempDir(filepath.Dir(path), filepath.Base(path))
	for i, key := range keys {
		if data, err := json.Marshal(key.ToMap()); err != nil {
			return "", erro.Wrap(err)
		} else if ioutil.WriteFile(filepath.Join(dir, strconv.Itoa(i)+".json"), data, 0644); err != nil {
			return "", erro.Wrap(err)
		}
	}
	return dir, nil
}

func setupIdpDb(db, coll string, idps []idpdb.Element) error {
	conn := monPool.New()
	defer conn.Close()

	for _, idp := range idps {
		keys := []map[string]interface{}{}
		for _, key := range idp.Keys() {
			keys = append(keys, key.ToMap())
		}
		m := bson.M{
			"issuer":                    idp.Id(),
			"authorization_endpoint":    idp.AuthUri(),
			"token_endpoint":            idp.TokenUri(),
			"userinfo_endpoint":         idp.AccountUri(),
			"cooperation_from_endpoint": idp.CoopFromUri(),
			"cooperation_to_endpoint":   idp.CoopToUri(),
			"jwks": keys,
		}
		for k, v := range idp.Names() {
			if k == "" {
				m["issuer_name"] = v
			} else {
				m["issuer_name#"+k] = v
			}
		}
		if err := conn.DB(db).C(coll).Insert(m); err != nil {
			return erro.Wrap(err)
		}
	}
	return nil
}

func clearIdpDb(db, coll string) error {
	conn := monPool.New()
	defer conn.Close()

	return conn.DB(db).DropDatabase()
}
<|endoftext|>